From 337687f8f48d0423761ffde5949d08426d2a650d Mon Sep 17 00:00:00 2001 From: Julian Date: Thu, 28 Nov 2024 08:57:28 +0100 Subject: [PATCH 01/11] impl From for PR --- crates/pg_completions/src/builder.rs | 12 +- crates/pg_completions/src/complete.rs | 128 +++++++++++++++++++ crates/pg_completions/src/context.rs | 9 ++ crates/pg_completions/src/item.rs | 33 +++++ crates/pg_completions/src/lib.rs | 168 +------------------------ crates/pg_completions/src/providers.rs | 47 ------- crates/pg_completions/src/relevance.rs | 1 + 7 files changed, 182 insertions(+), 216 deletions(-) create mode 100644 crates/pg_completions/src/complete.rs create mode 100644 crates/pg_completions/src/context.rs create mode 100644 crates/pg_completions/src/item.rs delete mode 100644 crates/pg_completions/src/providers.rs create mode 100644 crates/pg_completions/src/relevance.rs diff --git a/crates/pg_completions/src/builder.rs b/crates/pg_completions/src/builder.rs index 4075050c..f43768db 100644 --- a/crates/pg_completions/src/builder.rs +++ b/crates/pg_completions/src/builder.rs @@ -1,19 +1,19 @@ -use crate::{CompletionItem, CompletionResult}; +use crate::{item::CompletionItem, CompletionResult}; -pub struct CompletionBuilder<'a> { - pub items: Vec>, +pub struct CompletionBuilder { + pub items: Vec, } pub struct CompletionConfig {} -impl<'a> From<&'a CompletionConfig> for CompletionBuilder<'a> { +impl From<&CompletionConfig> for CompletionBuilder { fn from(_config: &CompletionConfig) -> Self { Self { items: Vec::new() } } } -impl<'a> CompletionBuilder<'a> { - pub fn finish(mut self) -> CompletionResult<'a> { +impl CompletionBuilder { + pub fn finish(mut self) -> CompletionResult { self.items.sort_by(|a, b| { b.preselect .cmp(&a.preselect) diff --git a/crates/pg_completions/src/complete.rs b/crates/pg_completions/src/complete.rs new file mode 100644 index 00000000..3f26da7d --- /dev/null +++ b/crates/pg_completions/src/complete.rs @@ -0,0 +1,128 @@ +use text_size::TextSize; + +use crate::{builder, context::CompletionContext, item::CompletionItem}; + +pub const LIMIT: usize = 50; + +#[derive(Debug)] +pub struct CompletionParams<'a> { + pub position: TextSize, + pub schema: &'a pg_schema_cache::SchemaCache, + pub text: &'a str, + pub tree: Option<&'a tree_sitter::Tree>, +} + +#[derive(Debug, Default)] +pub struct CompletionResult { + pub items: Vec, +} + +pub fn complete(params: &CompletionParams) -> CompletionResult { + let ctx = CompletionContext::new(params); + let mut builder = builder::CompletionBuilder::from(&builder::CompletionConfig {}); + + builder.finish() +} + +#[cfg(test)] +mod tests { + use pg_schema_cache::SchemaCache; + use pg_test_utils::test_database::*; + + use sqlx::Executor; + + use crate::{complete, CompletionParams}; + + #[tokio::test] + async fn test_complete() { + let pool = get_new_test_db().await; + + let input = "select id from c;"; + + let mut parser = tree_sitter::Parser::new(); + parser + .set_language(tree_sitter_sql::language()) + .expect("Error loading sql language"); + + let tree = parser.parse(input, None).unwrap(); + + let schema_cache = SchemaCache::load(&pool).await; + + let p = CompletionParams { + position: 15.into(), + schema: &schema_cache, + text: input, + tree: Some(&tree), + }; + + let result = complete(&p); + + assert!(result.items.len() > 0); + } + + #[tokio::test] + async fn test_complete_two() { + let pool = get_new_test_db().await; + + let input = "select id, name, test1231234123, unknown from co;"; + + let mut parser = tree_sitter::Parser::new(); + parser + .set_language(tree_sitter_sql::language()) + .expect("Error loading sql language"); + + let tree = parser.parse(input, None).unwrap(); + let schema_cache = SchemaCache::load(&pool).await; + + let p = CompletionParams { + position: 47.into(), + schema: &schema_cache, + text: input, + tree: Some(&tree), + }; + + let result = complete(&p); + + assert!(result.items.len() > 0); + } + + #[tokio::test] + async fn test_complete_three() { + let test_db = get_new_test_db().await; + + let setup = r#" + create table users ( + id serial primary key, + name text, + password text + ); + "#; + + test_db + .execute(setup) + .await + .expect("Failed to execute setup query"); + + let input = "select * from u"; + + let mut parser = tree_sitter::Parser::new(); + parser + .set_language(tree_sitter_sql::language()) + .expect("Error loading sql language"); + + let tree = parser.parse(input, None).unwrap(); + let schema_cache = SchemaCache::load(&test_db).await; + + let p = CompletionParams { + position: ((input.len() - 1) as u32).into(), + schema: &schema_cache, + text: input, + tree: Some(&tree), + }; + + let result = complete(&p); + + // TODO: actually assert that we get good autocompletion suggestions + assert!(result.items.len() > 0); + } +} diff --git a/crates/pg_completions/src/context.rs b/crates/pg_completions/src/context.rs new file mode 100644 index 00000000..abc9d7d9 --- /dev/null +++ b/crates/pg_completions/src/context.rs @@ -0,0 +1,9 @@ +use crate::CompletionParams; + +pub struct CompletionContext {} + +impl CompletionContext { + pub fn new(params: &CompletionParams) -> Self { + todo!() + } +} diff --git a/crates/pg_completions/src/item.rs b/crates/pg_completions/src/item.rs new file mode 100644 index 00000000..b8239e42 --- /dev/null +++ b/crates/pg_completions/src/item.rs @@ -0,0 +1,33 @@ +use text_size::TextRange; + +#[derive(Debug, PartialEq, Eq)] +pub enum CompletionItemData { + Table(pg_schema_cache::Table), +} + +impl CompletionItemData { + pub fn label(&self) -> &str { + match self { + CompletionItemData::Table(t) => t.name.as_str(), + } + } +} + +#[derive(Debug, PartialEq, Eq)] +pub struct CompletionItem { + pub score: i32, + pub range: TextRange, + pub preselect: bool, + pub data: CompletionItemData, +} + +impl CompletionItem { + pub fn new_simple(score: i32, range: TextRange, data: CompletionItemData) -> Self { + Self { + score, + range, + preselect: false, + data, + } + } +} diff --git a/crates/pg_completions/src/lib.rs b/crates/pg_completions/src/lib.rs index 382e1fc4..e6d2a6d9 100644 --- a/crates/pg_completions/src/lib.rs +++ b/crates/pg_completions/src/lib.rs @@ -1,165 +1,7 @@ mod builder; -mod providers; +mod complete; +mod context; +mod item; +mod relevance; -pub use providers::CompletionProviderParams; -use text_size::{TextRange, TextSize}; - -pub const LIMIT: usize = 50; - -#[derive(Debug)] -pub struct CompletionParams<'a> { - pub position: TextSize, - pub schema: &'a pg_schema_cache::SchemaCache, - pub text: &'a str, - pub tree: Option<&'a tree_sitter::Tree>, -} - -#[derive(Debug, Default)] -pub struct CompletionResult<'a> { - pub items: Vec>, -} - -#[derive(Debug, PartialEq, Eq)] -pub struct CompletionItem<'a> { - pub score: i32, - pub range: TextRange, - pub preselect: bool, - pub data: CompletionItemData<'a>, -} - -#[derive(Debug, PartialEq, Eq)] -pub enum CompletionItemData<'a> { - Table(&'a pg_schema_cache::Table), -} - -impl<'a> CompletionItemData<'a> { - pub fn label(&self) -> &'a str { - match self { - CompletionItemData::Table(t) => t.name.as_str(), - } - } -} - -impl<'a> CompletionItem<'a> { - pub fn new_simple(score: i32, range: TextRange, data: CompletionItemData<'a>) -> Self { - Self { - score, - range, - preselect: false, - data, - } - } -} - -pub fn complete<'a>(params: &'a CompletionParams<'a>) -> CompletionResult<'a> { - let mut builder = builder::CompletionBuilder::from(&builder::CompletionConfig {}); - - let params = CompletionProviderParams::from(params); - - providers::complete_tables(params, &mut builder); - - builder.finish() -} - -#[cfg(test)] -mod tests { - use pg_schema_cache::SchemaCache; - use pg_test_utils::test_database::*; - - use sqlx::Executor; - - use crate::{complete, CompletionParams}; - - #[tokio::test] - async fn test_complete() { - let pool = get_new_test_db().await; - - let input = "select id from c;"; - - let mut parser = tree_sitter::Parser::new(); - parser - .set_language(tree_sitter_sql::language()) - .expect("Error loading sql language"); - - let tree = parser.parse(input, None).unwrap(); - - let schema_cache = SchemaCache::load(&pool).await; - - let p = CompletionParams { - position: 15.into(), - schema: &schema_cache, - text: input, - tree: Some(&tree), - }; - - let result = complete(&p); - - assert!(result.items.len() > 0); - } - - #[tokio::test] - async fn test_complete_two() { - let pool = get_new_test_db().await; - - let input = "select id, name, test1231234123, unknown from co;"; - - let mut parser = tree_sitter::Parser::new(); - parser - .set_language(tree_sitter_sql::language()) - .expect("Error loading sql language"); - - let tree = parser.parse(input, None).unwrap(); - let schema_cache = SchemaCache::load(&pool).await; - - let p = CompletionParams { - position: 47.into(), - schema: &schema_cache, - text: input, - tree: Some(&tree), - }; - - let result = complete(&p); - - assert!(result.items.len() > 0); - } - - #[tokio::test] - async fn test_complete_three() { - let test_db = get_new_test_db().await; - - let setup = r#" - create table users ( - id serial primary key, - name text, - password text - ); - "#; - - test_db - .execute(setup) - .await - .expect("Failed to execute setup query"); - - let input = "select * from u"; - - let mut parser = tree_sitter::Parser::new(); - parser - .set_language(tree_sitter_sql::language()) - .expect("Error loading sql language"); - - let tree = parser.parse(input, None).unwrap(); - let schema_cache = SchemaCache::load(&test_db).await; - - let p = CompletionParams { - position: ((input.len() - 1) as u32).into(), - schema: &schema_cache, - text: input, - tree: Some(&tree), - }; - - let result = complete(&p); - - // TODO: actually assert that we get good autocompletion suggestions - assert!(result.items.len() > 0); - } -} +pub use complete::*; diff --git a/crates/pg_completions/src/providers.rs b/crates/pg_completions/src/providers.rs deleted file mode 100644 index 355e0ea0..00000000 --- a/crates/pg_completions/src/providers.rs +++ /dev/null @@ -1,47 +0,0 @@ -mod tables; - -pub use tables::complete_tables; - -use crate::CompletionParams; - -#[derive(Debug)] -pub struct CompletionProviderParams<'a> { - pub ts_node: Option>, - pub schema: &'a pg_schema_cache::SchemaCache, - pub source: &'a str, -} - -impl<'a> From<&'a CompletionParams<'a>> for CompletionProviderParams<'a> { - fn from(params: &'a CompletionParams) -> Self { - let ts_node = if let Some(tree) = params.tree { - let node = tree.root_node().named_descendant_for_byte_range( - usize::from(params.position), - usize::from(params.position), - ); - - if let Some(mut n) = node { - let node_range = n.range(); - - while let Some(parent) = n.parent() { - if parent.range() != node_range { - break; - } - - n = parent; - } - - Some(n) - } else { - None - } - } else { - None - }; - - Self { - ts_node, - schema: params.schema, - source: params.text, - } - } -} diff --git a/crates/pg_completions/src/relevance.rs b/crates/pg_completions/src/relevance.rs new file mode 100644 index 00000000..5768ed0c --- /dev/null +++ b/crates/pg_completions/src/relevance.rs @@ -0,0 +1 @@ +struct CompletionItemScore {} From 7c279d508dbbe3ee866ee3650efbe900b8a2e23f Mon Sep 17 00:00:00 2001 From: Julian Date: Thu, 28 Nov 2024 09:23:22 +0100 Subject: [PATCH 02/11] completion context getting ahead --- crates/pg_completions/src/context.rs | 47 +++++++++++++++++++++++++--- 1 file changed, 43 insertions(+), 4 deletions(-) diff --git a/crates/pg_completions/src/context.rs b/crates/pg_completions/src/context.rs index abc9d7d9..ca5cc9ec 100644 --- a/crates/pg_completions/src/context.rs +++ b/crates/pg_completions/src/context.rs @@ -1,9 +1,48 @@ +use pg_schema_cache::SchemaCache; + use crate::CompletionParams; -pub struct CompletionContext {} +pub struct CompletionContext<'a> { + pub ts_node: Option>, + pub tree: Option<&'a tree_sitter::Tree>, + pub text: &'a str, + pub schema_cache: &'a SchemaCache, + pub original_token: Option, +} -impl CompletionContext { - pub fn new(params: &CompletionParams) -> Self { - todo!() +impl<'a> CompletionContext<'a> { + pub fn new(params: &'a CompletionParams) -> Self { + Self { + ts_node: find_ts_node(params), + tree: params.tree, + text: params.text, + schema_cache: params.schema, + original_token: find_original_token(params), + } } } + +fn find_original_token<'a>(params: &'a CompletionParams) -> Option { + let idx = usize::from(params.position); + params.text.chars().nth(idx) +} + +fn find_ts_node<'a>(params: &'a CompletionParams) -> Option> { + let tree = params.tree?; + + let mut node = tree.root_node().named_descendant_for_byte_range( + usize::from(params.position), + usize::from(params.position), + )?; + + let node_range = node.range(); + while let Some(parent) = node.parent() { + if parent.range() != node_range { + break; + } + + node = parent; + } + + Some(node) +} From 4ccda40f383b4f37d9a83d2ded32739297f7243e Mon Sep 17 00:00:00 2001 From: Julian Date: Sun, 1 Dec 2024 19:44:30 +0100 Subject: [PATCH 03/11] nice, comes out as expected --- crates/pg_completions/src/builder.rs | 27 ++++---- crates/pg_completions/src/complete.rs | 25 +++++--- crates/pg_completions/src/context.rs | 23 +++++-- crates/pg_completions/src/item.rs | 36 ++++++----- crates/pg_completions/src/lib.rs | 1 + crates/pg_completions/src/providers/mod.rs | 3 + crates/pg_completions/src/providers/tables.rs | 61 +++++++++++-------- crates/pg_completions/src/relevance.rs | 55 ++++++++++++++++- crates/pg_lsp/src/session.rs | 12 ++-- 9 files changed, 172 insertions(+), 71 deletions(-) create mode 100644 crates/pg_completions/src/providers/mod.rs diff --git a/crates/pg_completions/src/builder.rs b/crates/pg_completions/src/builder.rs index f43768db..35ddb5dd 100644 --- a/crates/pg_completions/src/builder.rs +++ b/crates/pg_completions/src/builder.rs @@ -1,29 +1,30 @@ use crate::{item::CompletionItem, CompletionResult}; -pub struct CompletionBuilder { - pub items: Vec, +pub(crate) struct CompletionBuilder { + items: Vec, } -pub struct CompletionConfig {} +impl CompletionBuilder { + pub fn new() -> Self { + CompletionBuilder { items: vec![] } + } -impl From<&CompletionConfig> for CompletionBuilder { - fn from(_config: &CompletionConfig) -> Self { - Self { items: Vec::new() } + pub fn add_item(&mut self, item: CompletionItem) { + self.items.push(item) } -} -impl CompletionBuilder { pub fn finish(mut self) -> CompletionResult { self.items.sort_by(|a, b| { - b.preselect - .cmp(&a.preselect) - .then_with(|| b.score.cmp(&a.score)) - .then_with(|| a.data.label().cmp(b.data.label())) + b.score() + .cmp(&a.score()) + .then_with(|| a.label.cmp(&b.label)) }); - self.items.dedup_by(|a, b| a.data.label() == b.data.label()); + self.items.dedup_by(|a, b| a.label == b.label); self.items.truncate(crate::LIMIT); + let Self { items, .. } = self; + CompletionResult { items } } } diff --git a/crates/pg_completions/src/complete.rs b/crates/pg_completions/src/complete.rs index 3f26da7d..18c4738d 100644 --- a/crates/pg_completions/src/complete.rs +++ b/crates/pg_completions/src/complete.rs @@ -1,6 +1,8 @@ use text_size::TextSize; -use crate::{builder, context::CompletionContext, item::CompletionItem}; +use crate::{ + builder::CompletionBuilder, context::CompletionContext, item::CompletionItem, providers, +}; pub const LIMIT: usize = 50; @@ -17,9 +19,18 @@ pub struct CompletionResult { pub items: Vec, } -pub fn complete(params: &CompletionParams) -> CompletionResult { - let ctx = CompletionContext::new(params); - let mut builder = builder::CompletionBuilder::from(&builder::CompletionConfig {}); +pub fn complete(params: CompletionParams) -> CompletionResult { + let ctx = CompletionContext::new(¶ms); + let mut builder = CompletionBuilder::new(); + + if let Some(node) = ctx.ts_node { + match node.kind() { + "relation" => providers::complete_tables(&ctx, &mut builder), + _ => {} + } + } else { + // if query emtpy, autocomplete select keywords etc? + } builder.finish() } @@ -55,7 +66,7 @@ mod tests { tree: Some(&tree), }; - let result = complete(&p); + let result = complete(p); assert!(result.items.len() > 0); } @@ -81,7 +92,7 @@ mod tests { tree: Some(&tree), }; - let result = complete(&p); + let result = complete(p); assert!(result.items.len() > 0); } @@ -120,7 +131,7 @@ mod tests { tree: Some(&tree), }; - let result = complete(&p); + let result = complete(p); // TODO: actually assert that we get good autocompletion suggestions assert!(result.items.len() > 0); diff --git a/crates/pg_completions/src/context.rs b/crates/pg_completions/src/context.rs index ca5cc9ec..3943ec6c 100644 --- a/crates/pg_completions/src/context.rs +++ b/crates/pg_completions/src/context.rs @@ -1,34 +1,47 @@ use pg_schema_cache::SchemaCache; +use text_size::TextSize; use crate::CompletionParams; -pub struct CompletionContext<'a> { +pub(crate) struct CompletionContext<'a> { pub ts_node: Option>, pub tree: Option<&'a tree_sitter::Tree>, pub text: &'a str, pub schema_cache: &'a SchemaCache, pub original_token: Option, + pub position: TextSize, } impl<'a> CompletionContext<'a> { pub fn new(params: &'a CompletionParams) -> Self { + let ts_node = find_ts_node(¶ms); + Self { - ts_node: find_ts_node(params), + ts_node, tree: params.tree, - text: params.text, + text: ¶ms.text, schema_cache: params.schema, original_token: find_original_token(params), + position: params.position, + } + } + + pub fn get_ts_node_content(&self, ts_node: tree_sitter::Node<'a>) -> Option<&'a str> { + let source = self.text; + match ts_node.utf8_text(source.as_bytes()) { + Ok(content) => Some(content), + Err(_) => None, } } } -fn find_original_token<'a>(params: &'a CompletionParams) -> Option { +fn find_original_token<'a>(params: &CompletionParams) -> Option { let idx = usize::from(params.position); params.text.chars().nth(idx) } fn find_ts_node<'a>(params: &'a CompletionParams) -> Option> { - let tree = params.tree?; + let tree = params.tree.as_ref()?; let mut node = tree.root_node().named_descendant_for_byte_range( usize::from(params.position), diff --git a/crates/pg_completions/src/item.rs b/crates/pg_completions/src/item.rs index b8239e42..45e1e5e6 100644 --- a/crates/pg_completions/src/item.rs +++ b/crates/pg_completions/src/item.rs @@ -1,33 +1,41 @@ use text_size::TextRange; -#[derive(Debug, PartialEq, Eq)] -pub enum CompletionItemData { - Table(pg_schema_cache::Table), +use crate::relevance::CompletionRelevance; + +#[derive(Debug)] +pub enum CompletionItemData<'a> { + Table(&'a pg_schema_cache::Table), } -impl CompletionItemData { - pub fn label(&self) -> &str { +impl<'a> CompletionItemData<'a> { + pub fn label(&self) -> String { match self { - CompletionItemData::Table(t) => t.name.as_str(), + CompletionItemData::Table(t) => t.name.clone(), } } } -#[derive(Debug, PartialEq, Eq)] +#[derive(Debug)] pub struct CompletionItem { - pub score: i32, pub range: TextRange, - pub preselect: bool, - pub data: CompletionItemData, + pub label: String, + relevance: CompletionRelevance, } impl CompletionItem { - pub fn new_simple(score: i32, range: TextRange, data: CompletionItemData) -> Self { + pub(crate) fn new( + range: TextRange, + data: CompletionItemData, + relevance: CompletionRelevance, + ) -> Self { Self { - score, range, - preselect: false, - data, + label: data.label(), + relevance, } } + + pub(crate) fn score(&self) -> i32 { + self.relevance.score() + } } diff --git a/crates/pg_completions/src/lib.rs b/crates/pg_completions/src/lib.rs index e6d2a6d9..82adf0b8 100644 --- a/crates/pg_completions/src/lib.rs +++ b/crates/pg_completions/src/lib.rs @@ -2,6 +2,7 @@ mod builder; mod complete; mod context; mod item; +mod providers; mod relevance; pub use complete::*; diff --git a/crates/pg_completions/src/providers/mod.rs b/crates/pg_completions/src/providers/mod.rs new file mode 100644 index 00000000..81043e5f --- /dev/null +++ b/crates/pg_completions/src/providers/mod.rs @@ -0,0 +1,3 @@ +mod tables; + +pub use tables::*; diff --git a/crates/pg_completions/src/providers/tables.rs b/crates/pg_completions/src/providers/tables.rs index 74a52c7e..6f50ed08 100644 --- a/crates/pg_completions/src/providers/tables.rs +++ b/crates/pg_completions/src/providers/tables.rs @@ -1,31 +1,38 @@ +use pg_schema_cache::Table; use text_size::{TextRange, TextSize}; -use crate::{builder::CompletionBuilder, CompletionItem, CompletionItemData}; - -use super::CompletionProviderParams; - -// todo unify this in a type resolver crate -pub fn complete_tables<'a>( - params: CompletionProviderParams<'a>, - builder: &mut CompletionBuilder<'a>, -) { - if let Some(ts) = params.ts_node { - let range = TextRange::new( - TextSize::try_from(ts.start_byte()).unwrap(), - TextSize::try_from(ts.end_byte()).unwrap(), - ); - match ts.kind() { - "relation" => { - // todo better search - params.schema.tables.iter().for_each(|table| { - builder.items.push(CompletionItem::new_simple( - 1, - range, - CompletionItemData::Table(table), - )); - }); - } - _ => {} - } +use crate::{ + builder::CompletionBuilder, + context::CompletionContext, + item::{CompletionItem, CompletionItemData}, + relevance::CompletionRelevance, +}; + +pub fn complete_tables(ctx: &CompletionContext, builder: &mut CompletionBuilder) { + let available_tables = &ctx.schema_cache.tables; + + let completion_items: Vec = available_tables + .iter() + .map(|table| to_completion_item(ctx, table)) + .collect(); + + for item in completion_items { + builder.add_item(item); } } + +fn to_completion_item(ctx: &CompletionContext, table: &Table) -> CompletionItem { + let data = CompletionItemData::Table(table); + + let start = ctx.position; + let end = start + TextSize::from(table.name.len() as u32); + let range = TextRange::new(start, end); + + let mut relevance = CompletionRelevance::default(); + + relevance.set_is_catalog(&table.schema); + relevance.set_matches_prefix(ctx, &table.name); + relevance.set_matches_schema(ctx, &table.schema); + + CompletionItem::new(range, data, relevance) +} diff --git a/crates/pg_completions/src/relevance.rs b/crates/pg_completions/src/relevance.rs index 5768ed0c..b468c24e 100644 --- a/crates/pg_completions/src/relevance.rs +++ b/crates/pg_completions/src/relevance.rs @@ -1 +1,54 @@ -struct CompletionItemScore {} +use crate::context::CompletionContext; + +#[derive(Debug, Default)] +pub(crate) struct CompletionRelevance { + /// does the underlying data match the expected schema we can determine from the query? + matches_schema: bool, + + /// Is the underlying item from the pg_catalog schema? + is_catalog: bool, + + /// Do the characters the users typed match at least the first 3 characters + /// of the underlying data's name? + matches_prefix: usize, +} + +impl CompletionRelevance { + pub fn score(&self) -> i32 { + let mut score: i32 = 0; + + if self.matches_schema { + score += 5; + } else if self.is_catalog { + score -= 1; + } + + score += (self.matches_prefix * 5) as i32; + + score + } + + pub fn set_matches_schema(&mut self, ctx: &CompletionContext, schema: &str) { + let node = ctx.ts_node.unwrap(); + self.matches_schema = node + .prev_named_sibling() + .is_some_and(|n| ctx.get_ts_node_content(n).is_some_and(|c| c == schema)); + } + + pub fn set_is_catalog(&mut self, schema: &str) { + self.is_catalog = schema == "pg_catalog" + } + + pub fn set_matches_prefix(&mut self, ctx: &CompletionContext, name: &str) { + let node = ctx.ts_node.unwrap(); + + let content = match ctx.get_ts_node_content(node) { + Some(c) => c, + None => return, + }; + + if name.starts_with(content) { + self.matches_prefix = content.len(); + }; + } +} diff --git a/crates/pg_lsp/src/session.rs b/crates/pg_lsp/src/session.rs index 9d7ef641..fad8c263 100644 --- a/crates/pg_lsp/src/session.rs +++ b/crates/pg_lsp/src/session.rs @@ -235,17 +235,21 @@ impl Session { let schema_cache = ide.schema_cache.read().expect("No Schema Cache"); - let completion_items = pg_completions::complete(&CompletionParams { + let completion_items = pg_completions::complete(CompletionParams { position: offset - range.start() - TextSize::from(1), - text: stmt.text.as_str(), - tree: ide.tree_sitter.tree(&stmt).as_ref().map(|x| x.as_ref()), + text: &stmt.text, + tree: ide + .tree_sitter + .tree(&stmt) + .as_ref() + .and_then(|t| Some(t.as_ref())), schema: &schema_cache, }) .items .into_iter() .map(|i| CompletionItem { // TODO: add more data - label: i.data.label().to_string(), + label: i.label, label_details: None, kind: Some(CompletionItemKind::CLASS), detail: None, From 2afc5e0977e2dfb8bb0fa9ef935cb827e0ec5b80 Mon Sep 17 00:00:00 2001 From: Julian Date: Sun, 1 Dec 2024 19:47:04 +0100 Subject: [PATCH 04/11] feat(completions): basic scoring algorithm for tables --- crates/pg_completions/src/complete.rs | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/crates/pg_completions/src/complete.rs b/crates/pg_completions/src/complete.rs index 18c4738d..2c425137 100644 --- a/crates/pg_completions/src/complete.rs +++ b/crates/pg_completions/src/complete.rs @@ -133,7 +133,14 @@ mod tests { let result = complete(p); - // TODO: actually assert that we get good autocompletion suggestions assert!(result.items.len() > 0); + + let best_match = &result.items[0]; + + assert_eq!( + best_match.label, "users", + "Does not return the expected table to autocomplete: {}", + best_match.label + ) } } From 624ea8e758f8e968e02c5429a7a0979cb9959899 Mon Sep 17 00:00:00 2001 From: Julian Date: Wed, 4 Dec 2024 09:09:14 +0100 Subject: [PATCH 05/11] refactorings --- crates/pg_completions/Cargo.toml | 1 + crates/pg_completions/src/builder.rs | 10 +- crates/pg_completions/src/complete.rs | 110 +++++++++--------- crates/pg_completions/src/data.rs | 36 ++++++ crates/pg_completions/src/item.rs | 37 ++---- crates/pg_completions/src/lib.rs | 1 + crates/pg_completions/src/providers/tables.rs | 25 +--- crates/pg_completions/src/relevance.rs | 16 ++- crates/pg_lsp/Cargo.toml | 2 +- 9 files changed, 134 insertions(+), 104 deletions(-) create mode 100644 crates/pg_completions/src/data.rs diff --git a/crates/pg_completions/Cargo.toml b/crates/pg_completions/Cargo.toml index 7ff13e50..d796478d 100644 --- a/crates/pg_completions/Cargo.toml +++ b/crates/pg_completions/Cargo.toml @@ -12,6 +12,7 @@ tree-sitter.workspace = true tree_sitter_sql.workspace = true pg_schema_cache.workspace = true pg_test_utils.workspace = true +tower-lsp.workspace = true sqlx.workspace = true diff --git a/crates/pg_completions/src/builder.rs b/crates/pg_completions/src/builder.rs index 35ddb5dd..fdfc1282 100644 --- a/crates/pg_completions/src/builder.rs +++ b/crates/pg_completions/src/builder.rs @@ -1,7 +1,7 @@ -use crate::{item::CompletionItem, CompletionResult}; +use crate::{item::CompletionItemWithRelevance, CompletionResult}; pub(crate) struct CompletionBuilder { - items: Vec, + items: Vec, } impl CompletionBuilder { @@ -9,7 +9,7 @@ impl CompletionBuilder { CompletionBuilder { items: vec![] } } - pub fn add_item(&mut self, item: CompletionItem) { + pub fn add_item(&mut self, item: CompletionItemWithRelevance) { self.items.push(item) } @@ -17,10 +17,10 @@ impl CompletionBuilder { self.items.sort_by(|a, b| { b.score() .cmp(&a.score()) - .then_with(|| a.label.cmp(&b.label)) + .then_with(|| a.label().cmp(&b.label())) }); - self.items.dedup_by(|a, b| a.label == b.label); + self.items.dedup_by(|a, b| a.label() == b.label()); self.items.truncate(crate::LIMIT); let Self { items, .. } = self; diff --git a/crates/pg_completions/src/complete.rs b/crates/pg_completions/src/complete.rs index 2c425137..78319a31 100644 --- a/crates/pg_completions/src/complete.rs +++ b/crates/pg_completions/src/complete.rs @@ -1,7 +1,8 @@ use text_size::TextSize; use crate::{ - builder::CompletionBuilder, context::CompletionContext, item::CompletionItem, providers, + builder::CompletionBuilder, context::CompletionContext, item::CompletionItemWithRelevance, + providers, }; pub const LIMIT: usize = 50; @@ -16,7 +17,7 @@ pub struct CompletionParams<'a> { #[derive(Debug, Default)] pub struct CompletionResult { - pub items: Vec, + pub items: Vec, } pub fn complete(params: CompletionParams) -> CompletionResult { @@ -45,37 +46,23 @@ mod tests { use crate::{complete, CompletionParams}; #[tokio::test] - async fn test_complete() { - let pool = get_new_test_db().await; - - let input = "select id from c;"; - - let mut parser = tree_sitter::Parser::new(); - parser - .set_language(tree_sitter_sql::language()) - .expect("Error loading sql language"); - - let tree = parser.parse(input, None).unwrap(); - - let schema_cache = SchemaCache::load(&pool).await; - - let p = CompletionParams { - position: 15.into(), - schema: &schema_cache, - text: input, - tree: Some(&tree), - }; - - let result = complete(p); + async fn autocompletes_simple_table() { + let test_db = get_new_test_db().await; - assert!(result.items.len() > 0); - } + let setup = r#" + create table users ( + id serial primary key, + name text, + password text + ); + "#; - #[tokio::test] - async fn test_complete_two() { - let pool = get_new_test_db().await; + test_db + .execute(setup) + .await + .expect("Failed to execute setup query"); - let input = "select id, name, test1231234123, unknown from co;"; + let input = "select * from u"; let mut parser = tree_sitter::Parser::new(); parser @@ -83,10 +70,10 @@ mod tests { .expect("Error loading sql language"); let tree = parser.parse(input, None).unwrap(); - let schema_cache = SchemaCache::load(&pool).await; + let schema_cache = SchemaCache::load(&test_db).await; let p = CompletionParams { - position: 47.into(), + position: ((input.len() - 1) as u32).into(), schema: &schema_cache, text: input, tree: Some(&tree), @@ -95,18 +82,34 @@ mod tests { let result = complete(p); assert!(result.items.len() > 0); + + let best_match = &result.items[0]; + + assert_eq!( + best_match.label, "users", + "Does not return the expected table to autocomplete: {}", + best_match.label + ) } - #[tokio::test] - async fn test_complete_three() { + async fn autocompletes_table_with_schema() { let test_db = get_new_test_db().await; let setup = r#" - create table users ( + create schema public; + create schema private; + + create table private.users ( id serial primary key, name text, password text ); + + create table public.user_requests ( + id serial primary key, + request text, + send_at timestamp with time zone + ); "#; test_db @@ -114,33 +117,36 @@ mod tests { .await .expect("Failed to execute setup query"); - let input = "select * from u"; + let schema_cache = SchemaCache::load(&test_db).await; let mut parser = tree_sitter::Parser::new(); parser .set_language(tree_sitter_sql::language()) .expect("Error loading sql language"); - let tree = parser.parse(input, None).unwrap(); - let schema_cache = SchemaCache::load(&test_db).await; + // testing the private schema + { + let input = "select * from private.u"; + let tree = parser.parse(input, None).unwrap(); - let p = CompletionParams { - position: ((input.len() - 1) as u32).into(), - schema: &schema_cache, - text: input, - tree: Some(&tree), - }; + let p = CompletionParams { + position: ((input.len() - 1) as u32).into(), + schema: &schema_cache, + text: input, + tree: Some(&tree), + }; - let result = complete(p); + let result = complete(p); - assert!(result.items.len() > 0); + assert!(result.items.len() > 0); - let best_match = &result.items[0]; + let best_match = &result.items[0]; - assert_eq!( - best_match.label, "users", - "Does not return the expected table to autocomplete: {}", - best_match.label - ) + assert_eq!( + best_match.label, "users", + "Does not return the expected table to autocomplete: {}", + best_match.label + ) + } } } diff --git a/crates/pg_completions/src/data.rs b/crates/pg_completions/src/data.rs new file mode 100644 index 00000000..32ce2c07 --- /dev/null +++ b/crates/pg_completions/src/data.rs @@ -0,0 +1,36 @@ +use tower_lsp::lsp_types::{CompletionItem, CompletionItemKind, CompletionItemLabelDetails}; + +#[derive(Debug)] +pub(crate) enum CompletionItemData<'a> { + Table(&'a pg_schema_cache::Table), +} + +impl<'a> Into for CompletionItemData<'a> { + fn into(self) -> CompletionItem { + match self { + Self::Table(tb) => CompletionItem { + label: tb.name.clone(), + label_details: Some(CompletionItemLabelDetails { + description: Some(format!("Schema: {}", tb.schema)), + detail: None, + }), + kind: Some(CompletionItemKind::CLASS), + detail: None, + documentation: None, + deprecated: None, + preselect: None, + sort_text: None, + filter_text: None, + insert_text: None, + insert_text_format: None, + insert_text_mode: None, + text_edit: None, + additional_text_edits: None, + commit_characters: None, + data: None, + tags: None, + command: None, + }, + } + } +} diff --git a/crates/pg_completions/src/item.rs b/crates/pg_completions/src/item.rs index 45e1e5e6..cf3237cb 100644 --- a/crates/pg_completions/src/item.rs +++ b/crates/pg_completions/src/item.rs @@ -1,36 +1,17 @@ -use text_size::TextRange; +use tower_lsp::lsp_types; -use crate::relevance::CompletionRelevance; +use crate::{data::CompletionItemData, relevance::CompletionRelevance}; #[derive(Debug)] -pub enum CompletionItemData<'a> { - Table(&'a pg_schema_cache::Table), -} - -impl<'a> CompletionItemData<'a> { - pub fn label(&self) -> String { - match self { - CompletionItemData::Table(t) => t.name.clone(), - } - } -} - -#[derive(Debug)] -pub struct CompletionItem { - pub range: TextRange, - pub label: String, +pub struct CompletionItemWithRelevance { + item: lsp_types::CompletionItem, relevance: CompletionRelevance, } -impl CompletionItem { - pub(crate) fn new( - range: TextRange, - data: CompletionItemData, - relevance: CompletionRelevance, - ) -> Self { +impl CompletionItemWithRelevance { + pub(crate) fn new(data: CompletionItemData, relevance: CompletionRelevance) -> Self { Self { - range, - label: data.label(), + item: data.into(), relevance, } } @@ -38,4 +19,8 @@ impl CompletionItem { pub(crate) fn score(&self) -> i32 { self.relevance.score() } + + pub(crate) fn label(&self) -> &str { + &self.item.label + } } diff --git a/crates/pg_completions/src/lib.rs b/crates/pg_completions/src/lib.rs index 82adf0b8..db24a61e 100644 --- a/crates/pg_completions/src/lib.rs +++ b/crates/pg_completions/src/lib.rs @@ -1,6 +1,7 @@ mod builder; mod complete; mod context; +mod data; mod item; mod providers; mod relevance; diff --git a/crates/pg_completions/src/providers/tables.rs b/crates/pg_completions/src/providers/tables.rs index 6f50ed08..f5f907b7 100644 --- a/crates/pg_completions/src/providers/tables.rs +++ b/crates/pg_completions/src/providers/tables.rs @@ -1,17 +1,14 @@ use pg_schema_cache::Table; -use text_size::{TextRange, TextSize}; use crate::{ - builder::CompletionBuilder, - context::CompletionContext, - item::{CompletionItem, CompletionItemData}, - relevance::CompletionRelevance, + builder::CompletionBuilder, context::CompletionContext, data::CompletionItemData, + item::CompletionItemWithRelevance, relevance::CompletionRelevance, }; pub fn complete_tables(ctx: &CompletionContext, builder: &mut CompletionBuilder) { let available_tables = &ctx.schema_cache.tables; - let completion_items: Vec = available_tables + let completion_items: Vec = available_tables .iter() .map(|table| to_completion_item(ctx, table)) .collect(); @@ -21,18 +18,8 @@ pub fn complete_tables(ctx: &CompletionContext, builder: &mut CompletionBuilder) } } -fn to_completion_item(ctx: &CompletionContext, table: &Table) -> CompletionItem { +fn to_completion_item(ctx: &CompletionContext, table: &Table) -> CompletionItemWithRelevance { let data = CompletionItemData::Table(table); - - let start = ctx.position; - let end = start + TextSize::from(table.name.len() as u32); - let range = TextRange::new(start, end); - - let mut relevance = CompletionRelevance::default(); - - relevance.set_is_catalog(&table.schema); - relevance.set_matches_prefix(ctx, &table.name); - relevance.set_matches_schema(ctx, &table.schema); - - CompletionItem::new(range, data, relevance) + let relevance = CompletionRelevance::from_data_and_ctx(&data, ctx); + CompletionItemWithRelevance::new(data, relevance) } diff --git a/crates/pg_completions/src/relevance.rs b/crates/pg_completions/src/relevance.rs index b468c24e..c2e5ec9a 100644 --- a/crates/pg_completions/src/relevance.rs +++ b/crates/pg_completions/src/relevance.rs @@ -1,4 +1,4 @@ -use crate::context::CompletionContext; +use crate::{context::CompletionContext, data::CompletionItemData}; #[derive(Debug, Default)] pub(crate) struct CompletionRelevance { @@ -14,6 +14,20 @@ pub(crate) struct CompletionRelevance { } impl CompletionRelevance { + pub fn from_data_and_ctx(data: &CompletionItemData, ctx: &CompletionContext) -> Self { + let mut relevance = CompletionRelevance::default(); + + match data { + CompletionItemData::Table(tb) => { + relevance.set_is_catalog(&tb.schema); + relevance.set_matches_schema(ctx, &tb.schema); + relevance.set_matches_prefix(ctx, &tb.name); + } + } + + relevance + } + pub fn score(&self) -> i32 { let mut score: i32 = 0; diff --git a/crates/pg_lsp/Cargo.toml b/crates/pg_lsp/Cargo.toml index 69560b55..cfe2877c 100644 --- a/crates/pg_lsp/Cargo.toml +++ b/crates/pg_lsp/Cargo.toml @@ -23,6 +23,7 @@ text-size = "1.1.1" line_index.workspace = true sqlx.workspace = true +tower-lsp.workspace = true pg_hover.workspace = true pg_fs.workspace = true @@ -35,7 +36,6 @@ pg_workspace.workspace = true pg_diagnostics.workspace = true tokio = { version = "1.40.0", features = ["io-std", "macros", "rt-multi-thread", "sync", "time"] } tokio-util = "0.7.12" -tower-lsp = "0.20.0" tracing = "0.1.40" tracing-subscriber = "0.3.18" From 99c14632aca04c5e21af4225e6adccd05f4ae0da Mon Sep 17 00:00:00 2001 From: Julian Date: Wed, 4 Dec 2024 09:09:27 +0100 Subject: [PATCH 06/11] pin tower_lsp --- Cargo.lock | 1 + Cargo.toml | 1 + 2 files changed, 2 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index 6519bb12..6d7325d7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1645,6 +1645,7 @@ dependencies = [ "sqlx", "text-size", "tokio", + "tower-lsp", "tree-sitter", "tree_sitter_sql", ] diff --git a/Cargo.toml b/Cargo.toml index f4a063ea..afaef7ec 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,6 +17,7 @@ line_index = { path = "./lib/line_index", version = "0.0.0" } tree_sitter_sql = { path = "./lib/tree_sitter_sql", version = "0.0.0" } tree-sitter = "0.20.10" tracing = "0.1.40" +tower-lsp = "0.20.0" sqlx = { version = "0.8.2", features = [ "runtime-async-std", "tls-rustls", "postgres", "json" ] } # postgres specific crates From 1ecb7f0e7d9fae0fce47d7cb1316c0b3b2867e32 Mon Sep 17 00:00:00 2001 From: Julian Date: Wed, 4 Dec 2024 09:44:11 +0100 Subject: [PATCH 07/11] ok --- crates/pg_completions/src/builder.rs | 37 ++++++++++++++++--- crates/pg_completions/src/complete.rs | 14 +++---- crates/pg_completions/src/item.rs | 26 ++++++++----- crates/pg_completions/src/providers/tables.rs | 18 ++++----- 4 files changed, 61 insertions(+), 34 deletions(-) diff --git a/crates/pg_completions/src/builder.rs b/crates/pg_completions/src/builder.rs index fdfc1282..481900fa 100644 --- a/crates/pg_completions/src/builder.rs +++ b/crates/pg_completions/src/builder.rs @@ -1,7 +1,9 @@ -use crate::{item::CompletionItemWithRelevance, CompletionResult}; +use tower_lsp::lsp_types::CompletionItem; + +use crate::{item::CompletionItemWithScore, CompletionResult}; pub(crate) struct CompletionBuilder { - items: Vec, + items: Vec, } impl CompletionBuilder { @@ -9,22 +11,45 @@ impl CompletionBuilder { CompletionBuilder { items: vec![] } } - pub fn add_item(&mut self, item: CompletionItemWithRelevance) { + pub fn add_item(&mut self, item: CompletionItemWithScore) { self.items.push(item) } pub fn finish(mut self) -> CompletionResult { self.items.sort_by(|a, b| { - b.score() - .cmp(&a.score()) + b.score + .cmp(&a.score) .then_with(|| a.label().cmp(&b.label())) }); self.items.dedup_by(|a, b| a.label() == b.label()); self.items.truncate(crate::LIMIT); - let Self { items, .. } = self; + let should_preselect_first_item = self.should_preselect_first_item(); + + let items: Vec = self + .items + .into_iter() + .enumerate() + .map(|(idx, mut item)| { + if idx == 0 { + item.set_preselected(should_preselect_first_item); + } + item.into() + }) + .collect(); CompletionResult { items } } + + 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(); + + first.is_some_and(|f| match second { + Some(s) => (f.score - s.score) > 10, + None => true, + }) + } } diff --git a/crates/pg_completions/src/complete.rs b/crates/pg_completions/src/complete.rs index 78319a31..15681da9 100644 --- a/crates/pg_completions/src/complete.rs +++ b/crates/pg_completions/src/complete.rs @@ -1,9 +1,7 @@ use text_size::TextSize; +use tower_lsp::lsp_types::CompletionItem; -use crate::{ - builder::CompletionBuilder, context::CompletionContext, item::CompletionItemWithRelevance, - providers, -}; +use crate::{builder::CompletionBuilder, context::CompletionContext, providers}; pub const LIMIT: usize = 50; @@ -17,7 +15,7 @@ pub struct CompletionParams<'a> { #[derive(Debug, Default)] pub struct CompletionResult { - pub items: Vec, + pub items: Vec, } pub fn complete(params: CompletionParams) -> CompletionResult { @@ -25,6 +23,7 @@ pub fn complete(params: CompletionParams) -> CompletionResult { let mut builder = CompletionBuilder::new(); if let Some(node) = ctx.ts_node { + println!("{}", node.kind()); match node.kind() { "relation" => providers::complete_tables(&ctx, &mut builder), _ => {} @@ -92,11 +91,12 @@ mod tests { ) } + #[tokio::test] async fn autocompletes_table_with_schema() { let test_db = get_new_test_db().await; let setup = r#" - create schema public; + create schema open; create schema private; create table private.users ( @@ -105,7 +105,7 @@ mod tests { password text ); - create table public.user_requests ( + create table open.user_requests ( id serial primary key, request text, send_at timestamp with time zone diff --git a/crates/pg_completions/src/item.rs b/crates/pg_completions/src/item.rs index cf3237cb..7ade6af8 100644 --- a/crates/pg_completions/src/item.rs +++ b/crates/pg_completions/src/item.rs @@ -1,26 +1,32 @@ -use tower_lsp::lsp_types; +use tower_lsp::lsp_types::{self, CompletionItem}; use crate::{data::CompletionItemData, relevance::CompletionRelevance}; #[derive(Debug)] -pub struct CompletionItemWithRelevance { - item: lsp_types::CompletionItem, - relevance: CompletionRelevance, +pub struct CompletionItemWithScore { + pub item: lsp_types::CompletionItem, + pub score: i32, } -impl CompletionItemWithRelevance { +impl CompletionItemWithScore { pub(crate) fn new(data: CompletionItemData, relevance: CompletionRelevance) -> Self { Self { item: data.into(), - relevance, + score: relevance.score(), } } - pub(crate) fn score(&self) -> i32 { - self.relevance.score() - } - pub(crate) fn label(&self) -> &str { &self.item.label } + + pub(crate) fn set_preselected(&mut self, is_preselected: bool) { + self.item.preselect = Some(is_preselected) + } +} + +impl Into for CompletionItemWithScore { + fn into(self) -> CompletionItem { + self.item + } } diff --git a/crates/pg_completions/src/providers/tables.rs b/crates/pg_completions/src/providers/tables.rs index f5f907b7..1a0e4d83 100644 --- a/crates/pg_completions/src/providers/tables.rs +++ b/crates/pg_completions/src/providers/tables.rs @@ -1,25 +1,21 @@ -use pg_schema_cache::Table; - use crate::{ builder::CompletionBuilder, context::CompletionContext, data::CompletionItemData, - item::CompletionItemWithRelevance, relevance::CompletionRelevance, + item::CompletionItemWithScore, relevance::CompletionRelevance, }; pub fn complete_tables(ctx: &CompletionContext, builder: &mut CompletionBuilder) { let available_tables = &ctx.schema_cache.tables; - let completion_items: Vec = available_tables + let completion_items: Vec = available_tables .iter() - .map(|table| to_completion_item(ctx, table)) + .map(|table| { + let data = CompletionItemData::Table(table); + let relevance = CompletionRelevance::from_data_and_ctx(&data, ctx); + CompletionItemWithScore::new(data, relevance) + }) .collect(); for item in completion_items { builder.add_item(item); } } - -fn to_completion_item(ctx: &CompletionContext, table: &Table) -> CompletionItemWithRelevance { - let data = CompletionItemData::Table(table); - let relevance = CompletionRelevance::from_data_and_ctx(&data, ctx); - CompletionItemWithRelevance::new(data, relevance) -} From aaae3faeff515efe9f322f3be85e90d9f403b826 Mon Sep 17 00:00:00 2001 From: Julian Date: Sun, 8 Dec 2024 16:39:58 +0100 Subject: [PATCH 08/11] feat: better autocompletions --- Cargo.lock | 116 +++++++++ crates/pg_completions/README.md | 15 ++ crates/pg_completions/src/builder.rs | 2 +- crates/pg_completions/src/complete.rs | 31 ++- crates/pg_completions/src/context.rs | 223 ++++++++++++++++-- crates/pg_completions/src/providers/tables.rs | 2 +- crates/pg_completions/src/relevance.rs | 78 +++--- crates/pg_test_utils/Cargo.toml | 10 +- crates/pg_test_utils/src/bin/tree_print.rs | 47 ++++ 9 files changed, 438 insertions(+), 86 deletions(-) create mode 100644 crates/pg_completions/README.md create mode 100644 crates/pg_test_utils/src/bin/tree_print.rs diff --git a/Cargo.lock b/Cargo.lock index 6d7325d7..83d332e6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -44,6 +44,55 @@ version = "0.2.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "45862d1c77f2228b9e10bc609d5bc203d86ebc9b87ad8d5d5167a6c9abf739d9" +[[package]] +name = "anstream" +version = "0.6.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8acc5369981196006228e28809f761875c0327210a891e941f4c683b3a99529b" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "55cc3b69f167a1ef2e161439aa98aed94e6028e5f9a59be9a6ffb47aef1651f9" + +[[package]] +name = "anstyle-parse" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b2d16507662817a6a20a9ea92df6652ee4f94f914589377d69f3b21bc5798a9" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79947af37f4177cfead1110013d678905c37501914fba0efea834c3fe9a8d60c" +dependencies = [ + "windows-sys 0.59.0", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2109dbce0e72be3ec00bed26e6a7479ca384ad226efdd66db8fa2e3a38c83125" +dependencies = [ + "anstyle", + "windows-sys 0.59.0", +] + [[package]] name = "anyhow" version = "1.0.93" @@ -375,6 +424,46 @@ dependencies = [ "libloading", ] +[[package]] +name = "clap" +version = "4.5.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3135e7ec2ef7b10c6ed8950f0f792ed96ee093fa088608f1c76e569722700c84" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.5.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30582fc632330df2bd26877bde0c1f4470d57c582bbc070376afcd04d8cb4838" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ac6a0c7b1a9e9a5186361f67dfa1b88213572f427fb9ab038efb2bd8c582dab" +dependencies = [ + "heck 0.5.0", + "proc-macro2", + "quote", + "syn 2.0.89", +] + +[[package]] +name = "clap_lex" +version = "0.7.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f46ad14479a25103f283c0f10005961cf086d8dc42205bb44c46ac563475dca6" + [[package]] name = "cmake" version = "0.1.52" @@ -384,6 +473,12 @@ dependencies = [ "cc", ] +[[package]] +name = "colorchoice" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b63caa9aa9397e2d9480a9b13673856c78d8ac123288526c37d7839f2a86990" + [[package]] name = "concurrent-queue" version = "2.5.0" @@ -1178,6 +1273,12 @@ dependencies = [ "windows-sys 0.48.0", ] +[[package]] +name = "is_terminal_polyfill" +version = "1.70.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7943c866cc5cd64cbc25b2e01621d07fa8eb2a1a23160ee81ce38704e97b8ecf" + [[package]] name = "itertools" version = "0.10.5" @@ -1842,8 +1943,11 @@ name = "pg_test_utils" version = "0.0.0" dependencies = [ "anyhow", + "clap", "dotenv", "sqlx", + "tree-sitter", + "tree_sitter_sql", "uuid", ] @@ -2751,6 +2855,12 @@ dependencies = [ "unicode-properties", ] +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + [[package]] name = "subtle" version = "2.6.1" @@ -3171,6 +3281,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + [[package]] name = "uuid" version = "1.11.0" diff --git a/crates/pg_completions/README.md b/crates/pg_completions/README.md new file mode 100644 index 00000000..1d3218ac --- /dev/null +++ b/crates/pg_completions/README.md @@ -0,0 +1,15 @@ +# Auto-Completions + +## What does this crate do? + +The `pg_completions` identifies and ranks autocompletion items that can be displayed in your code editor. +Its main export is the `complete` function. The function takes a PostgreSQL statement, a cursor position, and a datastructure representing the underlying databases schema. It returns a list of completion items. + +Postgres's statement-parsing-engine, `libpg_query`, which is used in other parts of this LSP, is only capable of parsing _complete and valid_ statements. Since autocompletion should work for incomplete statements, we rely heavily on tree-sitter – an incremental parsing library. + +### Working with TreeSitter + +In the `pg_test_utils` crate, there's a binary that parses an SQL file and prints out the matching tree-sitter tree. +This makes writing tree-sitter queries for this crate easy. + +To print a tree, run `cargo run --bin tree_print -- -f `. diff --git a/crates/pg_completions/src/builder.rs b/crates/pg_completions/src/builder.rs index 481900fa..8217435c 100644 --- a/crates/pg_completions/src/builder.rs +++ b/crates/pg_completions/src/builder.rs @@ -12,7 +12,7 @@ impl CompletionBuilder { } pub fn add_item(&mut self, item: CompletionItemWithScore) { - self.items.push(item) + self.items.push(item); } pub fn finish(mut self) -> CompletionResult { diff --git a/crates/pg_completions/src/complete.rs b/crates/pg_completions/src/complete.rs index 15681da9..cb750271 100644 --- a/crates/pg_completions/src/complete.rs +++ b/crates/pg_completions/src/complete.rs @@ -1,7 +1,7 @@ use text_size::TextSize; use tower_lsp::lsp_types::CompletionItem; -use crate::{builder::CompletionBuilder, context::CompletionContext, providers}; +use crate::{builder::CompletionBuilder, context::CompletionContext, providers::complete_tables}; pub const LIMIT: usize = 50; @@ -20,17 +20,10 @@ pub struct CompletionResult { pub fn complete(params: CompletionParams) -> CompletionResult { let ctx = CompletionContext::new(¶ms); + let mut builder = CompletionBuilder::new(); - if let Some(node) = ctx.ts_node { - println!("{}", node.kind()); - match node.kind() { - "relation" => providers::complete_tables(&ctx, &mut builder), - _ => {} - } - } else { - // if query emtpy, autocomplete select keywords etc? - } + complete_tables(&ctx, &mut builder); builder.finish() } @@ -96,16 +89,16 @@ mod tests { let test_db = get_new_test_db().await; let setup = r#" - create schema open; + create schema customer_support; create schema private; - create table private.users ( + create table private.user_z ( id serial primary key, name text, password text ); - create table open.user_requests ( + create table customer_support.user_y ( id serial primary key, request text, send_at timestamp with time zone @@ -124,9 +117,13 @@ mod tests { .set_language(tree_sitter_sql::language()) .expect("Error loading sql language"); - // testing the private schema - { - let input = "select * from private.u"; + let test_cases = vec![ + ("select * from u", "user_y"), // user_y is preferred alphanumerically + ("select * from private.u", "user_z"), + ("select * from customer_support.u", "user_y"), + ]; + + for (input, expected_label) in test_cases { let tree = parser.parse(input, None).unwrap(); let p = CompletionParams { @@ -143,7 +140,7 @@ mod tests { let best_match = &result.items[0]; assert_eq!( - best_match.label, "users", + best_match.label, expected_label, "Does not return the expected table to autocomplete: {}", best_match.label ) diff --git a/crates/pg_completions/src/context.rs b/crates/pg_completions/src/context.rs index 3943ec6c..92cfc559 100644 --- a/crates/pg_completions/src/context.rs +++ b/crates/pg_completions/src/context.rs @@ -1,5 +1,4 @@ use pg_schema_cache::SchemaCache; -use text_size::TextSize; use crate::CompletionParams; @@ -8,22 +7,30 @@ pub(crate) struct CompletionContext<'a> { pub tree: Option<&'a tree_sitter::Tree>, pub text: &'a str, pub schema_cache: &'a SchemaCache, - pub original_token: Option, - pub position: TextSize, + pub position: usize, + + pub schema_name: Option, + pub wrapping_clause_type: Option, + pub is_invocation: bool, } impl<'a> CompletionContext<'a> { pub fn new(params: &'a CompletionParams) -> Self { - let ts_node = find_ts_node(¶ms); - - Self { - ts_node, + let mut tree = Self { tree: params.tree, text: ¶ms.text, schema_cache: params.schema, - original_token: find_original_token(params), - position: params.position, - } + position: usize::from(params.position), + + ts_node: None, + schema_name: None, + wrapping_clause_type: None, + is_invocation: false, + }; + + tree.gather_tree_context(); + + tree } pub fn get_ts_node_content(&self, ts_node: tree_sitter::Node<'a>) -> Option<&'a str> { @@ -33,29 +40,193 @@ impl<'a> CompletionContext<'a> { Err(_) => None, } } -} -fn find_original_token<'a>(params: &CompletionParams) -> Option { - let idx = usize::from(params.position); - params.text.chars().nth(idx) + fn gather_tree_context(&mut self) { + if self.tree.is_none() { + return; + } + + let mut cursor = self.tree.as_ref().unwrap().root_node().walk(); + + // go to the statement node that matches the position + let current_node_kind = cursor.node().kind(); + + cursor.goto_first_child_for_byte(self.position); + + self.gather_context_from_node(cursor, current_node_kind); + } + + fn gather_context_from_node( + &mut self, + mut cursor: tree_sitter::TreeCursor<'a>, + previous_node_kind: &str, + ) { + let current_node = cursor.node(); + let current_node_kind = current_node.kind(); + + match previous_node_kind { + "statement" => self.wrapping_clause_type = Some(current_node_kind.to_string()), + "invocation" => self.is_invocation = true, + + _ => {} + } + + match current_node_kind { + "object_reference" => { + let txt = self.get_ts_node_content(current_node); + if let Some(txt) = txt { + let parts: Vec<&str> = txt.split('.').collect(); + if parts.len() == 2 { + self.schema_name = Some(parts[0].to_string()); + } + } + } + + // in Treesitter, the Where clause is nested inside other clauses + "where" => { + self.wrapping_clause_type = Some("where".to_string()); + } + + _ => {} + } + + if current_node.child_count() == 0 { + self.ts_node = Some(current_node); + return; + } + + cursor.goto_first_child_for_byte(self.position); + self.gather_context_from_node(cursor, current_node_kind); + } } -fn find_ts_node<'a>(params: &'a CompletionParams) -> Option> { - let tree = params.tree.as_ref()?; +#[cfg(test)] +mod tests { + use crate::context::CompletionContext; + + fn get_tree(input: &str) -> tree_sitter::Tree { + let mut parser = tree_sitter::Parser::new(); + parser + .set_language(tree_sitter_sql::language()) + .expect("Couldn't set language"); + + parser.parse(input, None).expect("Unable to parse tree") + } - let mut node = tree.root_node().named_descendant_for_byte_range( - usize::from(params.position), - usize::from(params.position), - )?; + static CURSOR_POS: &str = "XXX"; - let node_range = node.range(); - while let Some(parent) = node.parent() { - if parent.range() != node_range { - break; + #[test] + fn identifies_clauses() { + let test_cases = vec![ + (format!("Select {}* from users;", CURSOR_POS), "select"), + (format!("Select * from u{};", CURSOR_POS), "from"), + ( + format!("Select {}* from users where n = 1;", CURSOR_POS), + "select", + ), + ( + format!("Select * from users where {}n = 1;", CURSOR_POS), + "where", + ), + ( + format!("update users set u{} = 1 where n = 2;", CURSOR_POS), + "update", + ), + ( + format!("update users set u = 1 where n{} = 2;", CURSOR_POS), + "where", + ), + (format!("delete{} from users;", CURSOR_POS), "delete"), + (format!("delete from {}users;", CURSOR_POS), "from"), + ( + format!("select name, age, location from public.u{}sers", CURSOR_POS), + "from", + ), + ]; + + for (text, expected_clause) in test_cases { + let position = text.find(CURSOR_POS).unwrap(); + let text = text.replace(CURSOR_POS, ""); + + let tree = get_tree(text.as_str()); + let params = crate::CompletionParams { + position: (position as u32).into(), + text: text.as_str(), + tree: Some(&tree), + schema: &pg_schema_cache::SchemaCache::new(), + }; + + let ctx = CompletionContext::new(¶ms); + + assert_eq!(ctx.wrapping_clause_type, Some(expected_clause.to_string())); } + } - node = parent; + #[test] + fn identifies_schema() { + let test_cases = vec![ + ( + format!("Select * from private.u{}", CURSOR_POS), + Some("private"), + ), + ( + format!("Select * from private.u{}sers()", CURSOR_POS), + Some("private"), + ), + (format!("Select * from u{}sers", CURSOR_POS), None), + (format!("Select * from u{}sers()", CURSOR_POS), None), + ]; + + for (text, expected_schema) in test_cases { + let position = text.find(CURSOR_POS).unwrap(); + let text = text.replace(CURSOR_POS, ""); + + let tree = get_tree(text.as_str()); + let params = crate::CompletionParams { + position: (position as u32).into(), + text: text.as_str(), + tree: Some(&tree), + schema: &pg_schema_cache::SchemaCache::new(), + }; + + let ctx = CompletionContext::new(¶ms); + + assert_eq!(ctx.schema_name, expected_schema.map(|f| f.to_string())); + } } - Some(node) + #[test] + fn identifies_invocation() { + let test_cases = vec![ + (format!("Select * from u{}sers", CURSOR_POS), false), + (format!("Select * from u{}sers()", CURSOR_POS), true), + (format!("Select cool{};", CURSOR_POS), false), + (format!("Select cool{}();", CURSOR_POS), true), + ( + format!("Select upp{}ercase as title from users;", CURSOR_POS), + false, + ), + ( + format!("Select upp{}ercase(name) as title from users;", CURSOR_POS), + true, + ), + ]; + + for (text, is_invocation) in test_cases { + let position = text.find(CURSOR_POS).unwrap(); + let text = text.replace(CURSOR_POS, ""); + + let tree = get_tree(text.as_str()); + let params = crate::CompletionParams { + position: (position as u32).into(), + text: text.as_str(), + tree: Some(&tree), + schema: &pg_schema_cache::SchemaCache::new(), + }; + + let ctx = CompletionContext::new(¶ms); + + assert_eq!(ctx.is_invocation, is_invocation); + } + } } diff --git a/crates/pg_completions/src/providers/tables.rs b/crates/pg_completions/src/providers/tables.rs index 1a0e4d83..30d1f905 100644 --- a/crates/pg_completions/src/providers/tables.rs +++ b/crates/pg_completions/src/providers/tables.rs @@ -10,7 +10,7 @@ pub fn complete_tables(ctx: &CompletionContext, builder: &mut CompletionBuilder) .iter() .map(|table| { let data = CompletionItemData::Table(table); - let relevance = CompletionRelevance::from_data_and_ctx(&data, ctx); + let relevance = CompletionRelevance::new(&data, ctx); CompletionItemWithScore::new(data, relevance) }) .collect(); diff --git a/crates/pg_completions/src/relevance.rs b/crates/pg_completions/src/relevance.rs index c2e5ec9a..038aaad5 100644 --- a/crates/pg_completions/src/relevance.rs +++ b/crates/pg_completions/src/relevance.rs @@ -2,58 +2,29 @@ use crate::{context::CompletionContext, data::CompletionItemData}; #[derive(Debug, Default)] pub(crate) struct CompletionRelevance { - /// does the underlying data match the expected schema we can determine from the query? - matches_schema: bool, - - /// Is the underlying item from the pg_catalog schema? - is_catalog: bool, - - /// Do the characters the users typed match at least the first 3 characters - /// of the underlying data's name? - matches_prefix: usize, + score: i32, } impl CompletionRelevance { - pub fn from_data_and_ctx(data: &CompletionItemData, ctx: &CompletionContext) -> Self { + pub fn score(&self) -> i32 { + self.score + } + + pub fn new(data: &CompletionItemData, ctx: &CompletionContext) -> Self { let mut relevance = CompletionRelevance::default(); match data { CompletionItemData::Table(tb) => { - relevance.set_is_catalog(&tb.schema); - relevance.set_matches_schema(ctx, &tb.schema); - relevance.set_matches_prefix(ctx, &tb.name); + relevance.check_if_catalog(ctx); + relevance.check_matches_schema(ctx, &tb.schema); + relevance.check_matches_query_input(ctx, &tb.name); } } relevance } - pub fn score(&self) -> i32 { - let mut score: i32 = 0; - - if self.matches_schema { - score += 5; - } else if self.is_catalog { - score -= 1; - } - - score += (self.matches_prefix * 5) as i32; - - score - } - - pub fn set_matches_schema(&mut self, ctx: &CompletionContext, schema: &str) { - let node = ctx.ts_node.unwrap(); - self.matches_schema = node - .prev_named_sibling() - .is_some_and(|n| ctx.get_ts_node_content(n).is_some_and(|c| c == schema)); - } - - pub fn set_is_catalog(&mut self, schema: &str) { - self.is_catalog = schema == "pg_catalog" - } - - pub fn set_matches_prefix(&mut self, ctx: &CompletionContext, name: &str) { + fn check_matches_query_input(&mut self, ctx: &CompletionContext, name: &str) { let node = ctx.ts_node.unwrap(); let content = match ctx.get_ts_node_content(node) { @@ -62,7 +33,34 @@ impl CompletionRelevance { }; if name.starts_with(content) { - self.matches_prefix = content.len(); + let len: i32 = content + .len() + .try_into() + .expect("The length of the input exceeds i32 capacity"); + + self.score += len * 5; }; } + + fn check_matches_schema(&mut self, ctx: &CompletionContext, schema: &str) { + if ctx.schema_name.is_none() { + return; + } + + let name = ctx.schema_name.as_ref().unwrap(); + + if name == schema { + self.score += 25; + } else { + self.score -= 10; + } + } + + fn check_if_catalog(&mut self, ctx: &CompletionContext) { + if ctx.schema_name.as_ref().is_some_and(|n| n == "pg_catalog") { + return; + } + + self.score -= 5; // unlikely that the user wants schema data + } } diff --git a/crates/pg_test_utils/Cargo.toml b/crates/pg_test_utils/Cargo.toml index a61688ee..ce4eb139 100644 --- a/crates/pg_test_utils/Cargo.toml +++ b/crates/pg_test_utils/Cargo.toml @@ -1,3 +1,7 @@ +[[bin]] +name = "tree_print" +path = "src/bin/tree_print.rs" + [package] name = "pg_test_utils" version = "0.0.0" @@ -6,6 +10,10 @@ edition = "2021" [dependencies] anyhow = "1.0.81" uuid = { version = "1.11.0", features = ["v4"] } +dotenv = "0.15.0" +clap = { version = "4.5.23", features = ["derive"] } sqlx.workspace = true -dotenv = "0.15.0" +tree-sitter.workspace = true +tree_sitter_sql.workspace = true + diff --git a/crates/pg_test_utils/src/bin/tree_print.rs b/crates/pg_test_utils/src/bin/tree_print.rs new file mode 100644 index 00000000..8a04365e --- /dev/null +++ b/crates/pg_test_utils/src/bin/tree_print.rs @@ -0,0 +1,47 @@ +use clap::*; + +#[derive(Parser)] +#[command( + name = "tree-printer", + about = "Prints the TreeSitter tree of the given file." +)] +struct Args { + #[arg(long = "file", short = 'f')] + file: String, +} + +fn main() { + let args = Args::parse(); + + let query = std::fs::read_to_string(&args.file).expect("Failed to read file."); + + let mut parser = tree_sitter::Parser::new(); + let lang = tree_sitter_sql::language(); + + parser.set_language(lang).expect("Setting Language failed."); + + let tree = parser + .parse(query.clone(), None) + .expect("Failed to parse query."); + + print_tree(&tree.root_node(), &query, 0); +} + +fn print_tree(node: &tree_sitter::Node, source: &str, level: usize) { + let indent = " ".repeat(level); + let node_text = node.utf8_text(source.as_bytes()).unwrap_or("NO_NAME"); + + println!( + "{}{} [{}..{}] '{}'", + indent, + node.kind(), + node.start_position().column, + node.end_position().column, + node_text + ); + + let mut cursor = node.walk(); + for child in node.children(&mut cursor) { + print_tree(&child, source, level + 1); + } +} From 1fd51cc3af23171e92f66ee5d84aa200fbc941c2 Mon Sep 17 00:00:00 2001 From: Julian Date: Sun, 8 Dec 2024 16:58:46 +0100 Subject: [PATCH 09/11] add alphanumeric test --- crates/pg_completions/src/complete.rs | 60 +++++++++++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/crates/pg_completions/src/complete.rs b/crates/pg_completions/src/complete.rs index cb750271..15785442 100644 --- a/crates/pg_completions/src/complete.rs +++ b/crates/pg_completions/src/complete.rs @@ -84,6 +84,66 @@ mod tests { ) } + #[tokio::test] + async fn autocompletes_table_alphanumerically() { + let test_db = get_new_test_db().await; + + let setup = r#" + create table addresses ( + id serial primary key + ); + + create table users ( + id serial primary key + ); + + create table emails ( + id serial primary key + ); + "#; + + test_db + .execute(setup) + .await + .expect("Failed to execute setup query"); + + let schema_cache = SchemaCache::load(&test_db).await; + + let mut parser = tree_sitter::Parser::new(); + parser + .set_language(tree_sitter_sql::language()) + .expect("Error loading sql language"); + + let test_cases = vec![ + ("select * from us", "users"), + ("select * from em", "emails"), + ("select * from ", "addresses"), + ]; + + for (input, expected_label) in test_cases { + let tree = parser.parse(input, None).unwrap(); + + let p = CompletionParams { + position: ((input.len() - 1) as u32).into(), + schema: &schema_cache, + text: input, + tree: Some(&tree), + }; + + let result = complete(p); + + assert!(result.items.len() > 0); + + let best_match = &result.items[0]; + + assert_eq!( + best_match.label, expected_label, + "Does not return the expected table to autocomplete: {}", + best_match.label + ) + } + } + #[tokio::test] async fn autocompletes_table_with_schema() { let test_db = get_new_test_db().await; From 7d9f8e270ede866b9667f5614f741204a4f868fd Mon Sep 17 00:00:00 2001 From: Julian Date: Tue, 10 Dec 2024 08:55:34 +0100 Subject: [PATCH 10/11] simplify, keep api boundary, add traits --- crates/pg_completions/src/builder.rs | 19 +++--- crates/pg_completions/src/complete.rs | 14 ++++- crates/pg_completions/src/data.rs | 36 ----------- crates/pg_completions/src/item.rs | 61 ++++++++++++------- crates/pg_completions/src/lib.rs | 1 - crates/pg_completions/src/providers/tables.rs | 30 ++++++--- crates/pg_completions/src/relevance.rs | 22 ++----- crates/pg_lsp/src/session.rs | 29 ++------- 8 files changed, 89 insertions(+), 123 deletions(-) delete mode 100644 crates/pg_completions/src/data.rs diff --git a/crates/pg_completions/src/builder.rs b/crates/pg_completions/src/builder.rs index 8217435c..c5a89889 100644 --- a/crates/pg_completions/src/builder.rs +++ b/crates/pg_completions/src/builder.rs @@ -1,9 +1,7 @@ -use tower_lsp::lsp_types::CompletionItem; - -use crate::{item::CompletionItemWithScore, CompletionResult}; +use crate::{item::CompletionItem, CompletionResult}; pub(crate) struct CompletionBuilder { - items: Vec, + items: Vec, } impl CompletionBuilder { @@ -11,18 +9,15 @@ impl CompletionBuilder { CompletionBuilder { items: vec![] } } - pub fn add_item(&mut self, item: CompletionItemWithScore) { + pub fn add_item(&mut self, item: CompletionItem) { self.items.push(item); } pub fn finish(mut self) -> CompletionResult { - self.items.sort_by(|a, b| { - b.score - .cmp(&a.score) - .then_with(|| a.label().cmp(&b.label())) - }); + self.items + .sort_by(|a, b| b.score.cmp(&a.score).then_with(|| a.label.cmp(&b.label))); - self.items.dedup_by(|a, b| a.label() == b.label()); + self.items.dedup_by(|a, b| a.label == b.label); self.items.truncate(crate::LIMIT); let should_preselect_first_item = self.should_preselect_first_item(); @@ -33,7 +28,7 @@ impl CompletionBuilder { .enumerate() .map(|(idx, mut item)| { if idx == 0 { - item.set_preselected(should_preselect_first_item); + item.preselected = Some(should_preselect_first_item); } item.into() }) diff --git a/crates/pg_completions/src/complete.rs b/crates/pg_completions/src/complete.rs index 15785442..58c08897 100644 --- a/crates/pg_completions/src/complete.rs +++ b/crates/pg_completions/src/complete.rs @@ -1,7 +1,9 @@ use text_size::TextSize; -use tower_lsp::lsp_types::CompletionItem; -use crate::{builder::CompletionBuilder, context::CompletionContext, providers::complete_tables}; +use crate::{ + builder::CompletionBuilder, context::CompletionContext, item::CompletionItem, + providers::complete_tables, +}; pub const LIMIT: usize = 50; @@ -18,6 +20,14 @@ pub struct CompletionResult { pub items: Vec, } +impl IntoIterator for CompletionResult { + type Item = CompletionItem; + type IntoIter = as IntoIterator>::IntoIter; + fn into_iter(self) -> Self::IntoIter { + self.items.into_iter() + } +} + pub fn complete(params: CompletionParams) -> CompletionResult { let ctx = CompletionContext::new(¶ms); diff --git a/crates/pg_completions/src/data.rs b/crates/pg_completions/src/data.rs deleted file mode 100644 index 32ce2c07..00000000 --- a/crates/pg_completions/src/data.rs +++ /dev/null @@ -1,36 +0,0 @@ -use tower_lsp::lsp_types::{CompletionItem, CompletionItemKind, CompletionItemLabelDetails}; - -#[derive(Debug)] -pub(crate) enum CompletionItemData<'a> { - Table(&'a pg_schema_cache::Table), -} - -impl<'a> Into for CompletionItemData<'a> { - fn into(self) -> CompletionItem { - match self { - Self::Table(tb) => CompletionItem { - label: tb.name.clone(), - label_details: Some(CompletionItemLabelDetails { - description: Some(format!("Schema: {}", tb.schema)), - detail: None, - }), - kind: Some(CompletionItemKind::CLASS), - detail: None, - documentation: None, - deprecated: None, - preselect: None, - sort_text: None, - filter_text: None, - insert_text: None, - insert_text_format: None, - insert_text_mode: None, - text_edit: None, - additional_text_edits: None, - commit_characters: None, - data: None, - tags: None, - command: None, - }, - } - } -} diff --git a/crates/pg_completions/src/item.rs b/crates/pg_completions/src/item.rs index 7ade6af8..49a50cc6 100644 --- a/crates/pg_completions/src/item.rs +++ b/crates/pg_completions/src/item.rs @@ -1,32 +1,49 @@ -use tower_lsp::lsp_types::{self, CompletionItem}; - -use crate::{data::CompletionItemData, relevance::CompletionRelevance}; +#[derive(Debug)] +pub enum CompletionItemKind { + Table, +} #[derive(Debug)] -pub struct CompletionItemWithScore { - pub item: lsp_types::CompletionItem, - pub score: i32, +pub struct CompletionItem { + pub label: String, + pub(crate) score: i32, + pub description: String, + pub preselected: Option, + pub kind: CompletionItemKind, } -impl CompletionItemWithScore { - pub(crate) fn new(data: CompletionItemData, relevance: CompletionRelevance) -> Self { - Self { - item: data.into(), - score: relevance.score(), +impl From for tower_lsp::lsp_types::CompletionItem { + fn from(i: CompletionItem) -> Self { + tower_lsp::lsp_types::CompletionItem { + label: i.label, + label_details: Some(tower_lsp::lsp_types::CompletionItemLabelDetails { + description: Some(i.description), + detail: None, + }), + kind: Some(i.kind.into()), + detail: None, + documentation: None, + deprecated: None, + preselect: None, + sort_text: None, + filter_text: None, + insert_text: None, + insert_text_format: None, + insert_text_mode: None, + text_edit: None, + additional_text_edits: None, + commit_characters: None, + data: None, + tags: None, + command: None, } } - - pub(crate) fn label(&self) -> &str { - &self.item.label - } - - pub(crate) fn set_preselected(&mut self, is_preselected: bool) { - self.item.preselect = Some(is_preselected) - } } -impl Into for CompletionItemWithScore { - fn into(self) -> CompletionItem { - self.item +impl From for tower_lsp::lsp_types::CompletionItemKind { + fn from(value: CompletionItemKind) -> Self { + match value { + CompletionItemKind::Table => tower_lsp::lsp_types::CompletionItemKind::CLASS, + } } } diff --git a/crates/pg_completions/src/lib.rs b/crates/pg_completions/src/lib.rs index db24a61e..82adf0b8 100644 --- a/crates/pg_completions/src/lib.rs +++ b/crates/pg_completions/src/lib.rs @@ -1,7 +1,6 @@ mod builder; mod complete; mod context; -mod data; mod item; mod providers; mod relevance; diff --git a/crates/pg_completions/src/providers/tables.rs b/crates/pg_completions/src/providers/tables.rs index 30d1f905..ea78deef 100644 --- a/crates/pg_completions/src/providers/tables.rs +++ b/crates/pg_completions/src/providers/tables.rs @@ -1,17 +1,23 @@ +use pg_schema_cache::Table; + use crate::{ - builder::CompletionBuilder, context::CompletionContext, data::CompletionItemData, - item::CompletionItemWithScore, relevance::CompletionRelevance, + builder::CompletionBuilder, + context::CompletionContext, + item::{CompletionItem, CompletionItemKind}, + relevance::CompletionRelevance, }; pub fn complete_tables(ctx: &CompletionContext, builder: &mut CompletionBuilder) { let available_tables = &ctx.schema_cache.tables; - let completion_items: Vec = available_tables + let completion_items: Vec = available_tables .iter() - .map(|table| { - let data = CompletionItemData::Table(table); - let relevance = CompletionRelevance::new(&data, ctx); - CompletionItemWithScore::new(data, relevance) + .map(|table| CompletionItem { + label: table.name.clone(), + score: get_score(ctx, table), + description: format!("Schema: {}", table.schema), + preselected: None, + kind: CompletionItemKind::Table, }) .collect(); @@ -19,3 +25,13 @@ pub fn complete_tables(ctx: &CompletionContext, builder: &mut CompletionBuilder) builder.add_item(item); } } + +fn get_score(ctx: &CompletionContext, table: &Table) -> i32 { + let mut relevance = CompletionRelevance::default(); + + relevance.check_matches_query_input(ctx, &table.name); + relevance.check_matches_schema(ctx, &table.schema); + relevance.check_if_catalog(ctx); + + relevance.score() +} diff --git a/crates/pg_completions/src/relevance.rs b/crates/pg_completions/src/relevance.rs index 038aaad5..ddf52ae4 100644 --- a/crates/pg_completions/src/relevance.rs +++ b/crates/pg_completions/src/relevance.rs @@ -1,4 +1,4 @@ -use crate::{context::CompletionContext, data::CompletionItemData}; +use crate::context::CompletionContext; #[derive(Debug, Default)] pub(crate) struct CompletionRelevance { @@ -10,21 +10,7 @@ impl CompletionRelevance { self.score } - pub fn new(data: &CompletionItemData, ctx: &CompletionContext) -> Self { - let mut relevance = CompletionRelevance::default(); - - match data { - CompletionItemData::Table(tb) => { - relevance.check_if_catalog(ctx); - relevance.check_matches_schema(ctx, &tb.schema); - relevance.check_matches_query_input(ctx, &tb.name); - } - } - - relevance - } - - fn check_matches_query_input(&mut self, ctx: &CompletionContext, name: &str) { + pub fn check_matches_query_input(&mut self, ctx: &CompletionContext, name: &str) { let node = ctx.ts_node.unwrap(); let content = match ctx.get_ts_node_content(node) { @@ -42,7 +28,7 @@ impl CompletionRelevance { }; } - fn check_matches_schema(&mut self, ctx: &CompletionContext, schema: &str) { + pub fn check_matches_schema(&mut self, ctx: &CompletionContext, schema: &str) { if ctx.schema_name.is_none() { return; } @@ -56,7 +42,7 @@ impl CompletionRelevance { } } - fn check_if_catalog(&mut self, ctx: &CompletionContext) { + pub fn check_if_catalog(&mut self, ctx: &CompletionContext) { if ctx.schema_name.as_ref().is_some_and(|n| n == "pg_catalog") { return; } diff --git a/crates/pg_lsp/src/session.rs b/crates/pg_lsp/src/session.rs index fad8c263..b9a58f37 100644 --- a/crates/pg_lsp/src/session.rs +++ b/crates/pg_lsp/src/session.rs @@ -10,8 +10,8 @@ use pg_workspace::Workspace; use text_size::TextSize; use tokio::sync::RwLock; use tower_lsp::lsp_types::{ - CodeActionOrCommand, CompletionItem, CompletionItemKind, CompletionList, Hover, HoverContents, - InlayHint, InlayHintKind, InlayHintLabel, MarkedString, Position, Range, + CodeActionOrCommand, CompletionItem, CompletionList, Hover, HoverContents, InlayHint, + InlayHintKind, InlayHintLabel, MarkedString, Position, Range, }; use crate::{db_connection::DbConnection, utils::line_index_ext::LineIndexExt}; @@ -235,7 +235,7 @@ impl Session { let schema_cache = ide.schema_cache.read().expect("No Schema Cache"); - let completion_items = pg_completions::complete(CompletionParams { + let completion_items: Vec = pg_completions::complete(CompletionParams { position: offset - range.start() - TextSize::from(1), text: &stmt.text, tree: ide @@ -245,29 +245,8 @@ impl Session { .and_then(|t| Some(t.as_ref())), schema: &schema_cache, }) - .items .into_iter() - .map(|i| CompletionItem { - // TODO: add more data - label: i.label, - label_details: None, - kind: Some(CompletionItemKind::CLASS), - detail: None, - documentation: None, - deprecated: None, - preselect: None, - sort_text: None, - filter_text: None, - insert_text: None, - insert_text_format: None, - insert_text_mode: None, - text_edit: None, - additional_text_edits: None, - commit_characters: None, - data: None, - tags: None, - command: None, - }) + .map(|i| i.into()) .collect(); Some(CompletionList { From e37202f3a577d4f76ec41c49fb8305f7c3d8da0e Mon Sep 17 00:00:00 2001 From: Julian Date: Wed, 11 Dec 2024 08:08:39 +0100 Subject: [PATCH 11/11] remove dep --- Cargo.lock | 1 - crates/pg_completions/Cargo.toml | 1 - crates/pg_completions/src/item.rs | 36 ------------------------- crates/pg_completions/src/lib.rs | 1 + crates/pg_lsp/src/session.rs | 29 ++++++++++++++++++-- crates/pg_lsp/src/utils.rs | 1 + crates/pg_lsp/src/utils/to_lsp_types.rs | 9 +++++++ 7 files changed, 38 insertions(+), 40 deletions(-) create mode 100644 crates/pg_lsp/src/utils/to_lsp_types.rs diff --git a/Cargo.lock b/Cargo.lock index 83d332e6..67959947 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1746,7 +1746,6 @@ dependencies = [ "sqlx", "text-size", "tokio", - "tower-lsp", "tree-sitter", "tree_sitter_sql", ] diff --git a/crates/pg_completions/Cargo.toml b/crates/pg_completions/Cargo.toml index d796478d..7ff13e50 100644 --- a/crates/pg_completions/Cargo.toml +++ b/crates/pg_completions/Cargo.toml @@ -12,7 +12,6 @@ tree-sitter.workspace = true tree_sitter_sql.workspace = true pg_schema_cache.workspace = true pg_test_utils.workspace = true -tower-lsp.workspace = true sqlx.workspace = true diff --git a/crates/pg_completions/src/item.rs b/crates/pg_completions/src/item.rs index 49a50cc6..7a015e72 100644 --- a/crates/pg_completions/src/item.rs +++ b/crates/pg_completions/src/item.rs @@ -11,39 +11,3 @@ pub struct CompletionItem { pub preselected: Option, pub kind: CompletionItemKind, } - -impl From for tower_lsp::lsp_types::CompletionItem { - fn from(i: CompletionItem) -> Self { - tower_lsp::lsp_types::CompletionItem { - label: i.label, - label_details: Some(tower_lsp::lsp_types::CompletionItemLabelDetails { - description: Some(i.description), - detail: None, - }), - kind: Some(i.kind.into()), - detail: None, - documentation: None, - deprecated: None, - preselect: None, - sort_text: None, - filter_text: None, - insert_text: None, - insert_text_format: None, - insert_text_mode: None, - text_edit: None, - additional_text_edits: None, - commit_characters: None, - data: None, - tags: None, - command: None, - } - } -} - -impl From for tower_lsp::lsp_types::CompletionItemKind { - fn from(value: CompletionItemKind) -> Self { - match value { - CompletionItemKind::Table => tower_lsp::lsp_types::CompletionItemKind::CLASS, - } - } -} diff --git a/crates/pg_completions/src/lib.rs b/crates/pg_completions/src/lib.rs index 82adf0b8..c31e9337 100644 --- a/crates/pg_completions/src/lib.rs +++ b/crates/pg_completions/src/lib.rs @@ -6,3 +6,4 @@ mod providers; mod relevance; pub use complete::*; +pub use item::*; diff --git a/crates/pg_lsp/src/session.rs b/crates/pg_lsp/src/session.rs index b9a58f37..215e54b1 100644 --- a/crates/pg_lsp/src/session.rs +++ b/crates/pg_lsp/src/session.rs @@ -14,7 +14,10 @@ use tower_lsp::lsp_types::{ InlayHintKind, InlayHintLabel, MarkedString, Position, Range, }; -use crate::{db_connection::DbConnection, utils::line_index_ext::LineIndexExt}; +use crate::{ + db_connection::DbConnection, + utils::{line_index_ext::LineIndexExt, to_lsp_types::to_completion_kind}, +}; pub struct Session { db: RwLock>, @@ -246,7 +249,29 @@ impl Session { schema: &schema_cache, }) .into_iter() - .map(|i| i.into()) + .map(|item| CompletionItem { + label: item.label, + label_details: Some(tower_lsp::lsp_types::CompletionItemLabelDetails { + description: Some(item.description), + detail: None, + }), + kind: Some(to_completion_kind(item.kind)), + detail: None, + documentation: None, + deprecated: None, + preselect: None, + sort_text: None, + filter_text: None, + insert_text: None, + insert_text_format: None, + insert_text_mode: None, + text_edit: None, + additional_text_edits: None, + commit_characters: None, + data: None, + tags: None, + command: None, + }) .collect(); Some(CompletionList { diff --git a/crates/pg_lsp/src/utils.rs b/crates/pg_lsp/src/utils.rs index bfbd57d1..50e9edd8 100644 --- a/crates/pg_lsp/src/utils.rs +++ b/crates/pg_lsp/src/utils.rs @@ -1,4 +1,5 @@ pub mod line_index_ext; +pub mod to_lsp_types; pub mod to_proto; use std::path::PathBuf; diff --git a/crates/pg_lsp/src/utils/to_lsp_types.rs b/crates/pg_lsp/src/utils/to_lsp_types.rs new file mode 100644 index 00000000..d090386b --- /dev/null +++ b/crates/pg_lsp/src/utils/to_lsp_types.rs @@ -0,0 +1,9 @@ +use tower_lsp::lsp_types; + +pub fn to_completion_kind( + kind: pg_completions::CompletionItemKind, +) -> lsp_types::CompletionItemKind { + match kind { + pg_completions::CompletionItemKind::Table => lsp_types::CompletionItemKind::CLASS, + } +}