From ad327d20223c15aa6e2242ee027d486a8fe2e4da Mon Sep 17 00:00:00 2001 From: Julian Date: Tue, 15 Apr 2025 09:24:36 +0200 Subject: [PATCH 1/2] feat: autocomplete schemas --- .../pgt_completions/benches/sanitization.rs | 2 +- crates/pgt_completions/src/complete.rs | 3 +- crates/pgt_completions/src/item.rs | 1 + crates/pgt_completions/src/providers/mod.rs | 2 + .../pgt_completions/src/providers/schemas.rs | 69 +++++++++++++++++++ crates/pgt_completions/src/relevance.rs | 13 ++-- crates/pgt_completions/src/sanitization.rs | 2 +- crates/pgt_lsp/src/handlers/completions.rs | 1 + crates/pgt_schema_cache/src/lib.rs | 1 + crates/pgt_schema_cache/src/schemas.rs | 6 +- 10 files changed, 89 insertions(+), 11 deletions(-) create mode 100644 crates/pgt_completions/src/providers/schemas.rs diff --git a/crates/pgt_completions/benches/sanitization.rs b/crates/pgt_completions/benches/sanitization.rs index c21538de..1e5333ff 100644 --- a/crates/pgt_completions/benches/sanitization.rs +++ b/crates/pgt_completions/benches/sanitization.rs @@ -27,7 +27,7 @@ fn to_params<'a>( position: TextSize::new(pos), schema: &cache, text, - tree: Some(tree), + tree: tree, } } diff --git a/crates/pgt_completions/src/complete.rs b/crates/pgt_completions/src/complete.rs index ec1232a5..89d25738 100644 --- a/crates/pgt_completions/src/complete.rs +++ b/crates/pgt_completions/src/complete.rs @@ -4,7 +4,7 @@ use crate::{ builder::CompletionBuilder, context::CompletionContext, item::CompletionItem, - providers::{complete_columns, complete_functions, complete_tables}, + providers::{complete_columns, complete_functions, complete_schemas, complete_tables}, sanitization::SanitizedCompletionParams, }; @@ -32,6 +32,7 @@ pub fn complete(params: CompletionParams) -> Vec { complete_tables(&ctx, &mut builder); complete_functions(&ctx, &mut builder); complete_columns(&ctx, &mut builder); + complete_schemas(&ctx, &mut builder); builder.finish() } diff --git a/crates/pgt_completions/src/item.rs b/crates/pgt_completions/src/item.rs index 8f0e3b95..1f306d78 100644 --- a/crates/pgt_completions/src/item.rs +++ b/crates/pgt_completions/src/item.rs @@ -7,6 +7,7 @@ pub enum CompletionItemKind { Table, Function, Column, + Schema, } #[derive(Debug, Serialize, Deserialize)] diff --git a/crates/pgt_completions/src/providers/mod.rs b/crates/pgt_completions/src/providers/mod.rs index 93055129..d760fea0 100644 --- a/crates/pgt_completions/src/providers/mod.rs +++ b/crates/pgt_completions/src/providers/mod.rs @@ -1,7 +1,9 @@ mod columns; mod functions; +mod schemas; mod tables; pub use columns::*; pub use functions::*; +pub use schemas::*; pub use tables::*; diff --git a/crates/pgt_completions/src/providers/schemas.rs b/crates/pgt_completions/src/providers/schemas.rs new file mode 100644 index 00000000..f379dcbb --- /dev/null +++ b/crates/pgt_completions/src/providers/schemas.rs @@ -0,0 +1,69 @@ +use crate::{ + CompletionItem, builder::CompletionBuilder, context::CompletionContext, + relevance::CompletionRelevanceData, +}; + +pub fn complete_schemas(ctx: &CompletionContext, builder: &mut CompletionBuilder) { + let available_schemas = &ctx.schema_cache.schemas; + + for schema in available_schemas { + let relevance = CompletionRelevanceData::Schema(&schema); + + let item = CompletionItem { + label: schema.name.clone(), + description: "Schema".into(), + preselected: false, + kind: crate::CompletionItemKind::Schema, + score: relevance.get_score(ctx), + }; + + builder.add_item(item); + } +} + +#[cfg(test)] +mod tests { + + use crate::{ + CompletionItem, CompletionItemKind, complete, + test_helper::{CURSOR_POS, get_test_deps, get_test_params}, + }; + + #[tokio::test] + async fn autocompletes_schemas() { + let setup = r#" + create schema private; + create schema auth; + create schema internal; + "#; + + let query = format!("select * from {}", CURSOR_POS); + + let (tree, cache) = get_test_deps(setup, query.as_str().into()).await; + let params = get_test_params(&tree, &cache, query.as_str().into()); + let items = complete(params); + + assert!(!items.is_empty()); + + for item in items.iter().take(10) { + println!( + r#""{}", score: {}, kind: {:?}"#, + item.label, item.score, item.kind + ); + } + + assert_eq!( + items + .into_iter() + .take(4) + .map(|i| i.label) + .collect::>(), + vec![ + "public".to_string(), // public always preferred + "auth".to_string(), + "internal".to_string(), + "private".to_string() + ] + ); + } +} diff --git a/crates/pgt_completions/src/relevance.rs b/crates/pgt_completions/src/relevance.rs index 9650a94d..307de0f9 100644 --- a/crates/pgt_completions/src/relevance.rs +++ b/crates/pgt_completions/src/relevance.rs @@ -5,6 +5,7 @@ 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<'_> { @@ -58,6 +59,7 @@ impl CompletionRelevance<'_> { 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) { @@ -97,6 +99,10 @@ impl CompletionRelevance<'_> { ClauseType::Where => 10, _ => -15, }, + CompletionRelevanceData::Schema(_) => match clause_type { + ClauseType::From => 15, + _ => -50, + }, } } @@ -129,6 +135,7 @@ impl CompletionRelevance<'_> { 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(), } } @@ -168,11 +175,7 @@ impl CompletionRelevance<'_> { } fn check_is_user_defined(&mut self) { - let schema = match self.data { - CompletionRelevanceData::Column(c) => &c.schema_name, - CompletionRelevanceData::Function(f) => &f.schema, - CompletionRelevanceData::Table(t) => &t.schema, - }; + let schema = self.get_schema_name().to_string(); let system_schemas = ["pg_catalog", "information_schema", "pg_toast"]; diff --git a/crates/pgt_completions/src/sanitization.rs b/crates/pgt_completions/src/sanitization.rs index 5ad8ba0e..0e059e1d 100644 --- a/crates/pgt_completions/src/sanitization.rs +++ b/crates/pgt_completions/src/sanitization.rs @@ -213,7 +213,7 @@ mod tests { #[test] fn test_cursor_after_nodes() { - let input = "select * from"; + let input = "select * from "; let mut parser = tree_sitter::Parser::new(); parser diff --git a/crates/pgt_lsp/src/handlers/completions.rs b/crates/pgt_lsp/src/handlers/completions.rs index e9a18a6e..f9b68e7d 100644 --- a/crates/pgt_lsp/src/handlers/completions.rs +++ b/crates/pgt_lsp/src/handlers/completions.rs @@ -50,5 +50,6 @@ fn to_lsp_types_completion_item_kind( pgt_completions::CompletionItemKind::Function => lsp_types::CompletionItemKind::FUNCTION, pgt_completions::CompletionItemKind::Table => lsp_types::CompletionItemKind::CLASS, pgt_completions::CompletionItemKind::Column => lsp_types::CompletionItemKind::FIELD, + pgt_completions::CompletionItemKind::Schema => lsp_types::CompletionItemKind::CLASS, } } diff --git a/crates/pgt_schema_cache/src/lib.rs b/crates/pgt_schema_cache/src/lib.rs index c6dad0b7..28c5b641 100644 --- a/crates/pgt_schema_cache/src/lib.rs +++ b/crates/pgt_schema_cache/src/lib.rs @@ -13,4 +13,5 @@ mod versions; pub use columns::*; pub use functions::{Behavior, Function, FunctionArg, FunctionArgs}; pub use schema_cache::SchemaCache; +pub use schemas::Schema; pub use tables::{ReplicaIdentity, Table}; diff --git a/crates/pgt_schema_cache/src/schemas.rs b/crates/pgt_schema_cache/src/schemas.rs index 51eb8ea3..41747194 100644 --- a/crates/pgt_schema_cache/src/schemas.rs +++ b/crates/pgt_schema_cache/src/schemas.rs @@ -4,9 +4,9 @@ use crate::schema_cache::SchemaCacheItem; #[derive(Debug, Clone, Default)] pub struct Schema { - id: i64, - name: String, - owner: String, + pub id: i64, + pub name: String, + pub owner: String, } impl SchemaCacheItem for Schema { From d94bd1f79e33277c16ebe5d131d1bbd251fe9583 Mon Sep 17 00:00:00 2001 From: Julian Date: Tue, 15 Apr 2025 10:06:12 +0200 Subject: [PATCH 2/2] ok --- .../pgt_completions/src/providers/schemas.rs | 31 ++++++++++--------- crates/pgt_completions/src/relevance.rs | 2 +- crates/pgt_completions/src/sanitization.rs | 28 +++++++++++++---- crates/pgt_completions/src/test_helper.rs | 2 +- 4 files changed, 40 insertions(+), 23 deletions(-) diff --git a/crates/pgt_completions/src/providers/schemas.rs b/crates/pgt_completions/src/providers/schemas.rs index f379dcbb..2f41e8c3 100644 --- a/crates/pgt_completions/src/providers/schemas.rs +++ b/crates/pgt_completions/src/providers/schemas.rs @@ -25,7 +25,7 @@ pub fn complete_schemas(ctx: &CompletionContext, builder: &mut CompletionBuilder mod tests { use crate::{ - CompletionItem, CompletionItemKind, complete, + CompletionItemKind, complete, test_helper::{CURSOR_POS, get_test_deps, get_test_params}, }; @@ -35,6 +35,13 @@ mod tests { create schema private; create schema auth; create schema internal; + + -- add a table to compete against schemas + create table users ( + id serial primary key, + name text, + password text + ); "#; let query = format!("select * from {}", CURSOR_POS); @@ -45,24 +52,18 @@ mod tests { assert!(!items.is_empty()); - for item in items.iter().take(10) { - println!( - r#""{}", score: {}, kind: {:?}"#, - item.label, item.score, item.kind - ); - } - assert_eq!( items .into_iter() - .take(4) - .map(|i| i.label) - .collect::>(), + .take(5) + .map(|i| (i.label, i.kind)) + .collect::>(), vec![ - "public".to_string(), // public always preferred - "auth".to_string(), - "internal".to_string(), - "private".to_string() + ("public".to_string(), CompletionItemKind::Schema), + ("auth".to_string(), CompletionItemKind::Schema), + ("internal".to_string(), CompletionItemKind::Schema), + ("private".to_string(), CompletionItemKind::Schema), + ("users".to_string(), CompletionItemKind::Table), ] ); } diff --git a/crates/pgt_completions/src/relevance.rs b/crates/pgt_completions/src/relevance.rs index 307de0f9..2abb9f2c 100644 --- a/crates/pgt_completions/src/relevance.rs +++ b/crates/pgt_completions/src/relevance.rs @@ -100,7 +100,7 @@ impl CompletionRelevance<'_> { _ => -15, }, CompletionRelevanceData::Schema(_) => match clause_type { - ClauseType::From => 15, + ClauseType::From => 10, _ => -50, }, } diff --git a/crates/pgt_completions/src/sanitization.rs b/crates/pgt_completions/src/sanitization.rs index 0e059e1d..dc093847 100644 --- a/crates/pgt_completions/src/sanitization.rs +++ b/crates/pgt_completions/src/sanitization.rs @@ -42,12 +42,28 @@ where let cursor_pos: usize = params.position.into(); let mut sql = String::new(); - for (idx, c) in params.text.chars().enumerate() { - if idx == cursor_pos { - sql.push_str(SANITIZED_TOKEN); - sql.push(' '); + let mut sql_iter = params.text.chars(); + + for idx in 0..cursor_pos + 1 { + match sql_iter.next() { + Some(c) => { + if idx == cursor_pos { + sql.push_str(SANITIZED_TOKEN); + sql.push(' '); + } + sql.push(c); + } + None => { + // the cursor is outside the statement, + // we want to push spaces until we arrive at the cursor position. + // we'll then add the SANITIZED_TOKEN + if idx == cursor_pos { + sql.push_str(SANITIZED_TOKEN); + } else { + sql.push(' '); + } + } } - sql.push(c); } let mut parser = tree_sitter::Parser::new(); @@ -213,7 +229,7 @@ mod tests { #[test] fn test_cursor_after_nodes() { - let input = "select * from "; + let input = "select * from"; let mut parser = tree_sitter::Parser::new(); parser diff --git a/crates/pgt_completions/src/test_helper.rs b/crates/pgt_completions/src/test_helper.rs index 4339688e..4edf486f 100644 --- a/crates/pgt_completions/src/test_helper.rs +++ b/crates/pgt_completions/src/test_helper.rs @@ -18,7 +18,7 @@ impl From<&str> for InputQuery { .expect("Insert Cursor Position into your Query."); InputQuery { - sql: value.replace(CURSOR_POS, ""), + sql: value.replace(CURSOR_POS, "").trim().to_string(), position, } }