diff --git a/crates/pgt_completions/src/builder.rs b/crates/pgt_completions/src/builder.rs index 39439afb..45c36d3b 100644 --- a/crates/pgt_completions/src/builder.rs +++ b/crates/pgt_completions/src/builder.rs @@ -1,47 +1,80 @@ -use crate::item::CompletionItem; +use crate::{ + CompletionItemKind, + context::CompletionContext, + item::CompletionItem, + relevance::{filtering::CompletionFilter, scoring::CompletionScore}, +}; -pub(crate) struct CompletionBuilder { - items: Vec, +pub(crate) struct PossibleCompletionItem<'a> { + pub label: String, + pub description: String, + pub kind: CompletionItemKind, + pub score: CompletionScore<'a>, + pub filter: CompletionFilter<'a>, } -impl CompletionBuilder { - pub fn new() -> Self { - CompletionBuilder { items: vec![] } +pub(crate) struct CompletionBuilder<'a> { + items: Vec>, + ctx: &'a CompletionContext<'a>, +} + +impl<'a> CompletionBuilder<'a> { + pub fn new(ctx: &'a CompletionContext) -> Self { + CompletionBuilder { items: vec![], ctx } } - pub fn add_item(&mut self, item: CompletionItem) { + pub fn add_item(&mut self, item: PossibleCompletionItem<'a>) { self.items.push(item); } - pub fn finish(mut self) -> Vec { - self.items - .sort_by(|a, b| b.score.cmp(&a.score).then_with(|| a.label.cmp(&b.label))); + pub fn finish(self) -> Vec { + let mut items: Vec = self + .items + .into_iter() + .filter(|i| i.filter.is_relevant(self.ctx).is_some()) + .collect(); + + for item in items.iter_mut() { + item.score.calc_score(self.ctx); + } - self.items.dedup_by(|a, b| a.label == b.label); - self.items.truncate(crate::LIMIT); + items.sort_by(|a, b| { + b.score + .get_score() + .cmp(&a.score.get_score()) + .then_with(|| a.label.cmp(&b.label)) + }); - let should_preselect_first_item = self.should_preselect_first_item(); + items.dedup_by(|a, b| a.label == b.label); + items.truncate(crate::LIMIT); - self.items + let should_preselect_first_item = should_preselect_first_item(&items); + + items .into_iter() .enumerate() - .map(|(idx, mut item)| { - if idx == 0 { - item.preselected = should_preselect_first_item; + .map(|(idx, item)| { + let preselected = idx == 0 && should_preselect_first_item; + + CompletionItem { + description: item.description, + kind: item.kind, + label: item.label, + preselected, + score: item.score.get_score(), } - item }) .collect() } +} - fn should_preselect_first_item(&mut self) -> bool { - let mut items_iter = self.items.iter(); - let first = items_iter.next(); - let second = items_iter.next(); +fn should_preselect_first_item(items: &Vec) -> bool { + let mut items_iter = items.iter(); + let first = items_iter.next(); + let second = items_iter.next(); - first.is_some_and(|f| match second { - Some(s) => (f.score - s.score) > 10, - None => true, - }) - } + first.is_some_and(|f| match second { + Some(s) => (f.score.get_score() - s.score.get_score()) > 10, + None => true, + }) } diff --git a/crates/pgt_completions/src/complete.rs b/crates/pgt_completions/src/complete.rs index 89d25738..442ee546 100644 --- a/crates/pgt_completions/src/complete.rs +++ b/crates/pgt_completions/src/complete.rs @@ -27,7 +27,7 @@ pub fn complete(params: CompletionParams) -> Vec { let ctx = CompletionContext::new(&sanitized_params); - let mut builder = CompletionBuilder::new(); + let mut builder = CompletionBuilder::new(&ctx); complete_tables(&ctx, &mut builder); complete_functions(&ctx, &mut builder); diff --git a/crates/pgt_completions/src/providers/columns.rs b/crates/pgt_completions/src/providers/columns.rs index d1c3e110..e8a51e48 100644 --- a/crates/pgt_completions/src/providers/columns.rs +++ b/crates/pgt_completions/src/providers/columns.rs @@ -1,17 +1,21 @@ use crate::{ - CompletionItem, CompletionItemKind, builder::CompletionBuilder, context::CompletionContext, - relevance::CompletionRelevanceData, + CompletionItemKind, + builder::{CompletionBuilder, PossibleCompletionItem}, + context::CompletionContext, + relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore}, }; -pub fn complete_columns(ctx: &CompletionContext, builder: &mut CompletionBuilder) { +pub fn complete_columns<'a>(ctx: &CompletionContext<'a>, builder: &mut CompletionBuilder<'a>) { let available_columns = &ctx.schema_cache.columns; for col in available_columns { - let item = CompletionItem { + let relevance = CompletionRelevanceData::Column(col); + + let item = PossibleCompletionItem { label: col.name.clone(), - score: CompletionRelevanceData::Column(col).get_score(ctx), + score: CompletionScore::from(relevance.clone()), + filter: CompletionFilter::from(relevance), description: format!("Table: {}.{}", col.schema_name, col.table_name), - preselected: false, kind: CompletionItemKind::Column, }; @@ -22,7 +26,7 @@ pub fn complete_columns(ctx: &CompletionContext, builder: &mut CompletionBuilder #[cfg(test)] mod tests { use crate::{ - CompletionItem, complete, + CompletionItem, CompletionItemKind, complete, test_helper::{CURSOR_POS, InputQuery, get_test_deps, get_test_params}, }; @@ -225,6 +229,37 @@ mod tests { ); } + #[tokio::test] + async fn ignores_cols_in_from_clause() { + let setup = r#" + create schema private; + + create table private.users ( + id serial primary key, + name text, + address text, + email text + ); + "#; + + let test_case = TestCase { + message: "suggests user created tables first", + query: format!(r#"select * from private.{}"#, CURSOR_POS), + label: "", + description: "", + }; + + let (tree, cache) = get_test_deps(setup, test_case.get_input_query()).await; + let params = get_test_params(&tree, &cache, test_case.get_input_query()); + let results = complete(params); + + assert!( + !results + .into_iter() + .any(|item| item.kind == CompletionItemKind::Column) + ); + } + #[tokio::test] async fn prefers_columns_of_mentioned_tables() { let setup = r#" diff --git a/crates/pgt_completions/src/providers/functions.rs b/crates/pgt_completions/src/providers/functions.rs index b4a9c35a..b44a5ef5 100644 --- a/crates/pgt_completions/src/providers/functions.rs +++ b/crates/pgt_completions/src/providers/functions.rs @@ -1,17 +1,21 @@ use crate::{ - CompletionItem, CompletionItemKind, builder::CompletionBuilder, context::CompletionContext, - relevance::CompletionRelevanceData, + CompletionItemKind, + builder::{CompletionBuilder, PossibleCompletionItem}, + context::CompletionContext, + relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore}, }; -pub fn complete_functions(ctx: &CompletionContext, builder: &mut CompletionBuilder) { +pub fn complete_functions<'a>(ctx: &'a CompletionContext, builder: &mut CompletionBuilder<'a>) { let available_functions = &ctx.schema_cache.functions; for func in available_functions { - let item = CompletionItem { + let relevance = CompletionRelevanceData::Function(func); + + let item = PossibleCompletionItem { label: func.name.clone(), - score: CompletionRelevanceData::Function(func).get_score(ctx), + score: CompletionScore::from(relevance.clone()), + filter: CompletionFilter::from(relevance), description: format!("Schema: {}", func.schema), - preselected: false, kind: CompletionItemKind::Function, }; diff --git a/crates/pgt_completions/src/providers/schemas.rs b/crates/pgt_completions/src/providers/schemas.rs index 6e86ab56..3d8f622e 100644 --- a/crates/pgt_completions/src/providers/schemas.rs +++ b/crates/pgt_completions/src/providers/schemas.rs @@ -1,20 +1,21 @@ use crate::{ - CompletionItem, builder::CompletionBuilder, context::CompletionContext, - relevance::CompletionRelevanceData, + builder::{CompletionBuilder, PossibleCompletionItem}, + context::CompletionContext, + relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore}, }; -pub fn complete_schemas(ctx: &CompletionContext, builder: &mut CompletionBuilder) { +pub fn complete_schemas<'a>(ctx: &'a CompletionContext, builder: &mut CompletionBuilder<'a>) { let available_schemas = &ctx.schema_cache.schemas; for schema in available_schemas { let relevance = CompletionRelevanceData::Schema(schema); - let item = CompletionItem { + let item = PossibleCompletionItem { label: schema.name.clone(), description: "Schema".into(), - preselected: false, kind: crate::CompletionItemKind::Schema, - score: relevance.get_score(ctx), + score: CompletionScore::from(relevance.clone()), + filter: CompletionFilter::from(relevance), }; builder.add_item(item); diff --git a/crates/pgt_completions/src/providers/tables.rs b/crates/pgt_completions/src/providers/tables.rs index 2074a4f1..fcc8fa00 100644 --- a/crates/pgt_completions/src/providers/tables.rs +++ b/crates/pgt_completions/src/providers/tables.rs @@ -1,19 +1,21 @@ use crate::{ - builder::CompletionBuilder, + builder::{CompletionBuilder, PossibleCompletionItem}, context::CompletionContext, - item::{CompletionItem, CompletionItemKind}, - relevance::CompletionRelevanceData, + item::CompletionItemKind, + relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore}, }; -pub fn complete_tables(ctx: &CompletionContext, builder: &mut CompletionBuilder) { +pub fn complete_tables<'a>(ctx: &'a CompletionContext, builder: &mut CompletionBuilder<'a>) { let available_tables = &ctx.schema_cache.tables; for table in available_tables { - let item = CompletionItem { + let relevance = CompletionRelevanceData::Table(table); + + let item = PossibleCompletionItem { label: table.name.clone(), - score: CompletionRelevanceData::Table(table).get_score(ctx), + score: CompletionScore::from(relevance.clone()), + filter: CompletionFilter::from(relevance), description: format!("Schema: {}", table.schema), - preselected: false, kind: CompletionItemKind::Table, }; diff --git a/crates/pgt_completions/src/relevance.rs b/crates/pgt_completions/src/relevance.rs index 2abb9f2c..911a6433 100644 --- a/crates/pgt_completions/src/relevance.rs +++ b/crates/pgt_completions/src/relevance.rs @@ -1,192 +1,10 @@ -use crate::context::{ClauseType, CompletionContext, NodeText}; +pub(crate) mod filtering; +pub(crate) mod scoring; -#[derive(Debug)] +#[derive(Debug, Clone)] pub(crate) enum CompletionRelevanceData<'a> { Table(&'a pgt_schema_cache::Table), Function(&'a pgt_schema_cache::Function), Column(&'a pgt_schema_cache::Column), Schema(&'a pgt_schema_cache::Schema), } - -impl CompletionRelevanceData<'_> { - pub fn get_score(self, ctx: &CompletionContext) -> i32 { - CompletionRelevance::from(self).into_score(ctx) - } -} - -impl<'a> From> for CompletionRelevance<'a> { - fn from(value: CompletionRelevanceData<'a>) -> Self { - Self { - score: 0, - data: value, - } - } -} - -#[derive(Debug)] -pub(crate) struct CompletionRelevance<'a> { - score: i32, - data: CompletionRelevanceData<'a>, -} - -impl CompletionRelevance<'_> { - pub fn into_score(mut self, ctx: &CompletionContext) -> i32 { - self.check_is_user_defined(); - self.check_matches_schema(ctx); - self.check_matches_query_input(ctx); - self.check_is_invocation(ctx); - self.check_matching_clause_type(ctx); - self.check_relations_in_stmt(ctx); - - self.score - } - - fn check_matches_query_input(&mut self, ctx: &CompletionContext) { - let node = match ctx.node_under_cursor { - Some(node) => node, - None => return, - }; - - let content = match ctx.get_ts_node_content(node) { - Some(c) => match c { - NodeText::Original(s) => s, - NodeText::Replaced => return, - }, - None => return, - }; - - let name = match self.data { - CompletionRelevanceData::Function(f) => f.name.as_str(), - CompletionRelevanceData::Table(t) => t.name.as_str(), - CompletionRelevanceData::Column(c) => c.name.as_str(), - CompletionRelevanceData::Schema(s) => s.name.as_str(), - }; - - if name.starts_with(content) { - let len: i32 = content - .len() - .try_into() - .expect("The length of the input exceeds i32 capacity"); - - self.score += len * 10; - }; - } - - fn check_matching_clause_type(&mut self, ctx: &CompletionContext) { - let clause_type = match ctx.wrapping_clause_type.as_ref() { - None => return, - Some(ct) => ct, - }; - - let has_mentioned_tables = !ctx.mentioned_relations.is_empty(); - - self.score += match self.data { - CompletionRelevanceData::Table(_) => match clause_type { - ClauseType::From => 5, - ClauseType::Update => 15, - ClauseType::Delete => 15, - _ => -50, - }, - CompletionRelevanceData::Function(_) => match clause_type { - ClauseType::Select if !has_mentioned_tables => 15, - ClauseType::Select if has_mentioned_tables => 0, - ClauseType::From => 0, - _ => -50, - }, - CompletionRelevanceData::Column(_) => match clause_type { - ClauseType::Select if has_mentioned_tables => 10, - ClauseType::Select if !has_mentioned_tables => 0, - ClauseType::Where => 10, - _ => -15, - }, - CompletionRelevanceData::Schema(_) => match clause_type { - ClauseType::From => 10, - _ => -50, - }, - } - } - - fn check_is_invocation(&mut self, ctx: &CompletionContext) { - self.score += match self.data { - CompletionRelevanceData::Function(_) if ctx.is_invocation => 30, - CompletionRelevanceData::Function(_) if !ctx.is_invocation => -10, - _ if ctx.is_invocation => -10, - _ => 0, - }; - } - - fn check_matches_schema(&mut self, ctx: &CompletionContext) { - let schema_name = match ctx.schema_name.as_ref() { - None => return, - Some(n) => n, - }; - - let data_schema = self.get_schema_name(); - - if schema_name == data_schema { - self.score += 25; - } else { - self.score -= 10; - } - } - - fn get_schema_name(&self) -> &str { - match self.data { - CompletionRelevanceData::Function(f) => f.schema.as_str(), - CompletionRelevanceData::Table(t) => t.schema.as_str(), - CompletionRelevanceData::Column(c) => c.schema_name.as_str(), - CompletionRelevanceData::Schema(s) => s.name.as_str(), - } - } - - fn get_table_name(&self) -> Option<&str> { - match self.data { - CompletionRelevanceData::Column(c) => Some(c.table_name.as_str()), - CompletionRelevanceData::Table(t) => Some(t.name.as_str()), - _ => None, - } - } - - fn check_relations_in_stmt(&mut self, ctx: &CompletionContext) { - match self.data { - CompletionRelevanceData::Table(_) | CompletionRelevanceData::Function(_) => return, - _ => {} - } - - let schema = self.get_schema_name().to_string(); - let table_name = match self.get_table_name() { - Some(t) => t, - None => return, - }; - - if ctx - .mentioned_relations - .get(&Some(schema.to_string())) - .is_some_and(|tables| tables.contains(table_name)) - { - self.score += 45; - } else if ctx - .mentioned_relations - .get(&None) - .is_some_and(|tables| tables.contains(table_name)) - { - self.score += 30; - } - } - - fn check_is_user_defined(&mut self) { - let schema = self.get_schema_name().to_string(); - - let system_schemas = ["pg_catalog", "information_schema", "pg_toast"]; - - if system_schemas.contains(&schema.as_str()) { - self.score -= 10; - } - - // "public" is the default postgres schema where users - // create objects. Prefer it by a slight bit. - if schema.as_str() == "public" { - self.score += 2; - } - } -} diff --git a/crates/pgt_completions/src/relevance/filtering.rs b/crates/pgt_completions/src/relevance/filtering.rs new file mode 100644 index 00000000..214fda56 --- /dev/null +++ b/crates/pgt_completions/src/relevance/filtering.rs @@ -0,0 +1,105 @@ +use crate::context::{ClauseType, CompletionContext}; + +use super::CompletionRelevanceData; + +#[derive(Debug)] +pub(crate) struct CompletionFilter<'a> { + data: CompletionRelevanceData<'a>, +} + +impl<'a> From> for CompletionFilter<'a> { + fn from(value: CompletionRelevanceData<'a>) -> Self { + Self { data: value } + } +} + +impl CompletionFilter<'_> { + pub fn is_relevant(&self, ctx: &CompletionContext) -> Option<()> { + self.completable_context(ctx)?; + self.check_clause(ctx)?; + self.check_invocation(ctx)?; + self.check_mentioned_schema(ctx)?; + + Some(()) + } + + fn completable_context(&self, ctx: &CompletionContext) -> Option<()> { + let current_node_kind = ctx.node_under_cursor.map(|n| n.kind()).unwrap_or(""); + + if current_node_kind.starts_with("keyword_") + || current_node_kind == "=" + || current_node_kind == "," + || current_node_kind == "literal" + || current_node_kind == "ERROR" + { + return None; + } + + Some(()) + } + + fn check_clause(&self, ctx: &CompletionContext) -> Option<()> { + let clause = ctx.wrapping_clause_type.as_ref(); + + match self.data { + CompletionRelevanceData::Table(_) => { + let in_select_clause = clause.is_some_and(|c| c == &ClauseType::Select); + let in_where_clause = clause.is_some_and(|c| c == &ClauseType::Where); + + if in_select_clause || in_where_clause { + return None; + }; + } + CompletionRelevanceData::Column(_) => { + let in_from_clause = clause.is_some_and(|c| c == &ClauseType::From); + + if in_from_clause { + return None; + } + } + _ => {} + } + + Some(()) + } + + fn check_invocation(&self, ctx: &CompletionContext) -> Option<()> { + if !ctx.is_invocation { + return Some(()); + } + + match self.data { + CompletionRelevanceData::Table(_) | CompletionRelevanceData::Column(_) => return None, + _ => {} + } + + Some(()) + } + + fn check_mentioned_schema(&self, ctx: &CompletionContext) -> Option<()> { + if ctx.schema_name.is_none() { + return Some(()); + } + + let name = ctx.schema_name.as_ref().unwrap(); + + let does_not_match = match self.data { + CompletionRelevanceData::Table(table) => &table.schema != name, + CompletionRelevanceData::Function(f) => &f.schema != name, + CompletionRelevanceData::Column(_) => { + // columns belong to tables, not schemas + true + } + CompletionRelevanceData::Schema(_) => { + // we should never allow schema suggestions if there already was one. + true + } + }; + + if does_not_match { + return None; + } + + Some(()) + } +} diff --git a/crates/pgt_completions/src/relevance/scoring.rs b/crates/pgt_completions/src/relevance/scoring.rs new file mode 100644 index 00000000..7c3f3a06 --- /dev/null +++ b/crates/pgt_completions/src/relevance/scoring.rs @@ -0,0 +1,182 @@ +use crate::context::{ClauseType, CompletionContext, NodeText}; + +use super::CompletionRelevanceData; + +#[derive(Debug)] +pub(crate) struct CompletionScore<'a> { + score: i32, + data: CompletionRelevanceData<'a>, +} + +impl<'a> From> for CompletionScore<'a> { + fn from(value: CompletionRelevanceData<'a>) -> Self { + Self { + score: 0, + data: value, + } + } +} + +impl CompletionScore<'_> { + pub fn get_score(&self) -> i32 { + self.score + } + + pub fn calc_score(&mut self, ctx: &CompletionContext) { + self.check_is_user_defined(); + self.check_matches_schema(ctx); + self.check_matches_query_input(ctx); + self.check_is_invocation(ctx); + self.check_matching_clause_type(ctx); + self.check_relations_in_stmt(ctx); + } + + fn check_matches_query_input(&mut self, ctx: &CompletionContext) { + let node = match ctx.node_under_cursor { + Some(node) => node, + None => return, + }; + + let content = match ctx.get_ts_node_content(node) { + Some(c) => match c { + NodeText::Original(s) => s, + NodeText::Replaced => return, + }, + None => return, + }; + + let name = match self.data { + CompletionRelevanceData::Function(f) => f.name.as_str(), + CompletionRelevanceData::Table(t) => t.name.as_str(), + CompletionRelevanceData::Column(c) => c.name.as_str(), + CompletionRelevanceData::Schema(s) => s.name.as_str(), + }; + + if name.starts_with(content) { + let len: i32 = content + .len() + .try_into() + .expect("The length of the input exceeds i32 capacity"); + + self.score += len * 10; + }; + } + + fn check_matching_clause_type(&mut self, ctx: &CompletionContext) { + let clause_type = match ctx.wrapping_clause_type.as_ref() { + None => return, + Some(ct) => ct, + }; + + let has_mentioned_tables = !ctx.mentioned_relations.is_empty(); + + self.score += match self.data { + CompletionRelevanceData::Table(_) => match clause_type { + ClauseType::From => 5, + ClauseType::Update => 15, + ClauseType::Delete => 15, + _ => -50, + }, + CompletionRelevanceData::Function(_) => match clause_type { + ClauseType::Select if !has_mentioned_tables => 15, + ClauseType::Select if has_mentioned_tables => 0, + ClauseType::From => 0, + _ => -50, + }, + CompletionRelevanceData::Column(_) => match clause_type { + ClauseType::Select if has_mentioned_tables => 10, + ClauseType::Select if !has_mentioned_tables => 0, + ClauseType::Where => 10, + _ => -15, + }, + CompletionRelevanceData::Schema(_) => match clause_type { + ClauseType::From => 10, + _ => -50, + }, + } + } + + fn check_is_invocation(&mut self, ctx: &CompletionContext) { + self.score += match self.data { + CompletionRelevanceData::Function(_) if ctx.is_invocation => 30, + CompletionRelevanceData::Function(_) if !ctx.is_invocation => -10, + _ if ctx.is_invocation => -10, + _ => 0, + }; + } + + fn check_matches_schema(&mut self, ctx: &CompletionContext) { + let schema_name = match ctx.schema_name.as_ref() { + None => return, + Some(n) => n, + }; + + let data_schema = self.get_schema_name(); + + if schema_name == data_schema { + self.score += 25; + } else { + self.score -= 10; + } + } + + fn get_schema_name(&self) -> &str { + match self.data { + CompletionRelevanceData::Function(f) => f.schema.as_str(), + CompletionRelevanceData::Table(t) => t.schema.as_str(), + CompletionRelevanceData::Column(c) => c.schema_name.as_str(), + CompletionRelevanceData::Schema(s) => s.name.as_str(), + } + } + + fn get_table_name(&self) -> Option<&str> { + match self.data { + CompletionRelevanceData::Column(c) => Some(c.table_name.as_str()), + CompletionRelevanceData::Table(t) => Some(t.name.as_str()), + _ => None, + } + } + + fn check_relations_in_stmt(&mut self, ctx: &CompletionContext) { + match self.data { + CompletionRelevanceData::Table(_) | CompletionRelevanceData::Function(_) => return, + _ => {} + } + + let schema = self.get_schema_name().to_string(); + let table_name = match self.get_table_name() { + Some(t) => t, + None => return, + }; + + if ctx + .mentioned_relations + .get(&Some(schema.to_string())) + .is_some_and(|tables| tables.contains(table_name)) + { + self.score += 45; + } else if ctx + .mentioned_relations + .get(&None) + .is_some_and(|tables| tables.contains(table_name)) + { + self.score += 30; + } + } + + fn check_is_user_defined(&mut self) { + let schema = self.get_schema_name().to_string(); + + let system_schemas = ["pg_catalog", "information_schema", "pg_toast"]; + + if system_schemas.contains(&schema.as_str()) { + self.score -= 10; + } + + // "public" is the default postgres schema where users + // create objects. Prefer it by a slight bit. + if schema.as_str() == "public" { + self.score += 2; + } + } +} diff --git a/crates/pgt_completions/src/test_helper.rs b/crates/pgt_completions/src/test_helper.rs index 4edf486f..fc2cf403 100644 --- a/crates/pgt_completions/src/test_helper.rs +++ b/crates/pgt_completions/src/test_helper.rs @@ -6,6 +6,7 @@ use crate::CompletionParams; pub static CURSOR_POS: char = '€'; +#[derive(Clone)] pub struct InputQuery { sql: String, position: usize, @@ -55,6 +56,31 @@ pub(crate) async fn get_test_deps( (tree, schema_cache) } +/// Careful: This will connect against the passed database. +/// Use this only to debug issues. Do not commit to version control. +#[allow(dead_code)] +pub(crate) async fn test_against_connection_string( + conn_str: &str, + input: InputQuery, +) -> (tree_sitter::Tree, pgt_schema_cache::SchemaCache) { + let pool = sqlx::PgPool::connect(conn_str) + .await + .expect("Unable to connect to database."); + + let schema_cache = SchemaCache::load(&pool) + .await + .expect("Failed to load Schema Cache"); + + let mut parser = tree_sitter::Parser::new(); + parser + .set_language(tree_sitter_sql::language()) + .expect("Error loading sql language"); + + let tree = parser.parse(input.to_string(), None).unwrap(); + + (tree, schema_cache) +} + pub(crate) fn get_text_and_position(q: InputQuery) -> (usize, String) { (q.position, q.sql) }