From 6cba26450b5bff5d7a81a8235a9116ef130f6edc Mon Sep 17 00:00:00 2001 From: Julian Date: Thu, 11 Sep 2025 17:40:52 +0200 Subject: [PATCH 1/3] oK --- .../pgt_completions/src/providers/tables.rs | 53 +++++++++++++++++++ .../src/relevance/filtering.rs | 8 +-- 2 files changed, 57 insertions(+), 4 deletions(-) diff --git a/crates/pgt_completions/src/providers/tables.rs b/crates/pgt_completions/src/providers/tables.rs index 20100e01f..7759222bd 100644 --- a/crates/pgt_completions/src/providers/tables.rs +++ b/crates/pgt_completions/src/providers/tables.rs @@ -569,4 +569,57 @@ mod tests { ) .await; } + + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn after_quoted_schemas(pool: PgPool) { + let setup = r#" + create schema auth; + + create table auth.users ( + uid serial primary key, + name text not null, + email text unique not null + ); + + create table auth.posts ( + pid serial primary key, + user_id int not null references auth.users(uid), + title text not null, + content text, + created_at timestamp default now() + ); + "#; + + pool.execute(setup).await.unwrap(); + + assert_complete_results( + format!( + r#"select * from "auth".{}"#, + QueryWithCursorPosition::cursor_marker() + ) + .as_str(), + vec![ + CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table), + CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table), + ], + None, + &pool, + ) + .await; + + assert_complete_results( + format!( + r#"select * from "auth".{}"#, + QueryWithCursorPosition::cursor_marker() + ) + .as_str(), + vec![ + CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table), + CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table), + ], + None, + &pool, + ) + .await; + } } diff --git a/crates/pgt_completions/src/relevance/filtering.rs b/crates/pgt_completions/src/relevance/filtering.rs index b9a896c49..e386bc5ea 100644 --- a/crates/pgt_completions/src/relevance/filtering.rs +++ b/crates/pgt_completions/src/relevance/filtering.rs @@ -249,13 +249,13 @@ impl CompletionFilter<'_> { return Some(()); } - let schema_or_alias = ctx.schema_or_alias_name.as_ref().unwrap(); + let schema_or_alias = ctx.schema_or_alias_name.as_ref().unwrap().replace('"', ""); let matches = match self.data { - CompletionRelevanceData::Table(table) => &table.schema == schema_or_alias, - CompletionRelevanceData::Function(f) => &f.schema == schema_or_alias, + CompletionRelevanceData::Table(table) => table.schema == schema_or_alias, + CompletionRelevanceData::Function(f) => f.schema == schema_or_alias, CompletionRelevanceData::Column(col) => ctx - .get_mentioned_table_for_alias(schema_or_alias) + .get_mentioned_table_for_alias(&schema_or_alias) .is_some_and(|t| t == &col.table_name), // we should never allow schema suggestions if there already was one. From 77d2649fe5024bf8a47a54b8d1eda8be282a06af Mon Sep 17 00:00:00 2001 From: Julian Date: Thu, 11 Sep 2025 18:18:30 +0200 Subject: [PATCH 2/3] ack --- .../src/providers/functions.rs | 75 ++++++++++++++++++- .../pgt_completions/src/providers/helper.rs | 15 ++-- .../pgt_completions/src/providers/tables.rs | 44 +++++++++-- crates/pgt_completions/src/sanitization.rs | 56 +++++++++++--- 4 files changed, 167 insertions(+), 23 deletions(-) diff --git a/crates/pgt_completions/src/providers/functions.rs b/crates/pgt_completions/src/providers/functions.rs index f4e86509b..876b34a47 100644 --- a/crates/pgt_completions/src/providers/functions.rs +++ b/crates/pgt_completions/src/providers/functions.rs @@ -1,10 +1,11 @@ use pgt_schema_cache::{Function, SchemaCache}; +use pgt_text_size::TextSize; use pgt_treesitter::TreesitterContext; use crate::{ CompletionItemKind, CompletionText, builder::{CompletionBuilder, PossibleCompletionItem}, - providers::helper::get_range_to_replace, + providers::helper::{get_range_to_replace, node_text_surrounded_by_quotes, only_leading_quote}, relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore}, }; @@ -37,7 +38,7 @@ pub fn complete_functions<'a>( fn get_completion_text(ctx: &TreesitterContext, func: &Function) -> CompletionText { let mut text = with_schema_or_alias(ctx, func.name.as_str(), Some(func.schema.as_str())); - let range = get_range_to_replace(ctx); + let mut range = get_range_to_replace(ctx); if ctx.is_invocation { CompletionText { @@ -46,6 +47,11 @@ fn get_completion_text(ctx: &TreesitterContext, func: &Function) -> CompletionTe is_snippet: false, } } else { + if node_text_surrounded_by_quotes(ctx) && !only_leading_quote(ctx) { + text.push('"'); + range = range.checked_expand_end(1.into()).unwrap_or(range); + } + text.push('('); let num_args = func.args.args.len(); @@ -68,6 +74,7 @@ fn get_completion_text(ctx: &TreesitterContext, func: &Function) -> CompletionTe #[cfg(test)] mod tests { + use pgt_text_size::TextRange; use sqlx::{Executor, PgPool}; use crate::{ @@ -294,4 +301,68 @@ mod tests { ) .await; } + + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] + async fn autocompletes_after_schema_in_quotes(pool: PgPool) { + let setup = r#" + create schema auth; + + create or replace function auth.my_cool_foo() + returns trigger + language plpgsql + security invoker + as $$ + begin + raise exception 'dont matter'; + end; + $$; + "#; + + pool.execute(setup).await.unwrap(); + + assert_complete_results( + format!( + r#"select "auth".{}"#, + QueryWithCursorPosition::cursor_marker() + ) + .as_str(), + vec![CompletionAssertion::CompletionTextAndRange( + "my_cool_foo()".into(), + TextRange::new(14.into(), 14.into()), + )], + None, + &pool, + ) + .await; + + assert_complete_results( + format!( + r#"select "auth"."{}"#, + QueryWithCursorPosition::cursor_marker() + ) + .as_str(), + vec![CompletionAssertion::CompletionTextAndRange( + r#"my_cool_foo"()"#.into(), + TextRange::new(15.into(), 15.into()), + )], + None, + &pool, + ) + .await; + + assert_complete_results( + format!( + r#"select "auth"."{}""#, + QueryWithCursorPosition::cursor_marker() + ) + .as_str(), + vec![CompletionAssertion::CompletionTextAndRange( + r#"my_cool_foo"()"#.into(), + TextRange::new(15.into(), 16.into()), + )], + None, + &pool, + ) + .await; + } } diff --git a/crates/pgt_completions/src/providers/helper.rs b/crates/pgt_completions/src/providers/helper.rs index b6547d701..002e8649e 100644 --- a/crates/pgt_completions/src/providers/helper.rs +++ b/crates/pgt_completions/src/providers/helper.rs @@ -34,6 +34,12 @@ pub(crate) fn get_range_to_replace(ctx: &TreesitterContext) -> TextRange { } } +pub(crate) fn only_leading_quote(ctx: &TreesitterContext) -> bool { + let node_under_cursor_txt = ctx.get_node_under_cursor_content().unwrap_or("".into()); + let node_under_cursor_txt = node_under_cursor_txt.as_str(); + is_sanitized_token_with_quote(node_under_cursor_txt) +} + pub(crate) fn with_schema_or_alias( ctx: &TreesitterContext, item_name: &str, @@ -42,13 +48,10 @@ pub(crate) fn with_schema_or_alias( let is_already_prefixed_with_schema_name = ctx.schema_or_alias_name.is_some(); let with_quotes = node_text_surrounded_by_quotes(ctx); - - let node_under_cursor_txt = ctx.get_node_under_cursor_content().unwrap_or("".into()); - let node_under_cursor_txt = node_under_cursor_txt.as_str(); - let is_quote_sanitized = is_sanitized_token_with_quote(node_under_cursor_txt); + let single_leading_quote = only_leading_quote(ctx); if schema_or_alias_name.is_none_or(|s| s == "public") || is_already_prefixed_with_schema_name { - if is_quote_sanitized { + if single_leading_quote { format!(r#"{}""#, item_name) } else { item_name.to_string() @@ -56,7 +59,7 @@ pub(crate) fn with_schema_or_alias( } else { let schema_or_als = schema_or_alias_name.unwrap(); - if is_quote_sanitized { + if single_leading_quote { format!(r#"{}"."{}""#, schema_or_als.replace('"', ""), item_name) } else if with_quotes { format!(r#"{}"."{}"#, schema_or_als.replace('"', ""), item_name) diff --git a/crates/pgt_completions/src/providers/tables.rs b/crates/pgt_completions/src/providers/tables.rs index 7759222bd..6ef9a6b90 100644 --- a/crates/pgt_completions/src/providers/tables.rs +++ b/crates/pgt_completions/src/providers/tables.rs @@ -58,6 +58,7 @@ fn get_completion_text(ctx: &TreesitterContext, table: &Table) -> CompletionText #[cfg(test)] mod tests { + use pgt_text_size::TextRange; use sqlx::{Executor, PgPool}; use crate::{ @@ -599,8 +600,14 @@ mod tests { ) .as_str(), vec![ - CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table), - CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table), + CompletionAssertion::CompletionTextAndRange( + "posts".into(), + TextRange::new(21.into(), 21.into()), + ), + CompletionAssertion::CompletionTextAndRange( + "users".into(), + TextRange::new(21.into(), 21.into()), + ), ], None, &pool, @@ -609,13 +616,40 @@ mod tests { assert_complete_results( format!( - r#"select * from "auth".{}"#, + r#"select * from "auth"."{}""#, QueryWithCursorPosition::cursor_marker() ) .as_str(), vec![ - CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table), - CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table), + CompletionAssertion::CompletionTextAndRange( + "posts".into(), + TextRange::new(22.into(), 22.into()), + ), + CompletionAssertion::CompletionTextAndRange( + "users".into(), + TextRange::new(22.into(), 22.into()), + ), + ], + None, + &pool, + ) + .await; + + assert_complete_results( + format!( + r#"select * from "auth"."{}"#, + QueryWithCursorPosition::cursor_marker() + ) + .as_str(), + vec![ + CompletionAssertion::CompletionTextAndRange( + r#"posts""#.into(), + TextRange::new(22.into(), 22.into()), + ), + CompletionAssertion::CompletionTextAndRange( + r#"users""#.into(), + TextRange::new(22.into(), 22.into()), + ), ], None, &pool, diff --git a/crates/pgt_completions/src/sanitization.rs b/crates/pgt_completions/src/sanitization.rs index 5272c75e2..1d728ffca 100644 --- a/crates/pgt_completions/src/sanitization.rs +++ b/crates/pgt_completions/src/sanitization.rs @@ -98,7 +98,11 @@ where // 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); + if opened_quote && has_uneven_quotes { + sql.push_str(SANITIZED_TOKEN_WITH_QUOTE); + } else { + sql.push_str(SANITIZED_TOKEN); + } } else { sql.push(' '); } @@ -342,18 +346,50 @@ mod tests { #[test] fn should_sanitize_with_opened_quotes() { - // select "email", "| from "auth"."users"; - let input = r#"select "email", " from "auth"."users";"#; - let position = TextSize::new(17); + { + // select "email", "| from "auth"."users"; + let input = r#"select "email", " from "auth"."users";"#; + let position = TextSize::new(17); - let params = get_test_params(input, position); + let params = get_test_params(input, position); - let sanitized = SanitizedCompletionParams::from(params); + let sanitized = SanitizedCompletionParams::from(params); - assert_eq!( - sanitized.text, - r#"select "email", "REPLACED_TOKEN_WITH_QUOTE" from "auth"."users";"# - ); + assert_eq!( + sanitized.text, + r#"select "email", "REPLACED_TOKEN_WITH_QUOTE" from "auth"."users";"# + ); + } + + { + // select * from "auth"."|; <-- with semi + let input = r#"select * from "auth".";"#; + let position = TextSize::new(22); + + let params = get_test_params(input, position); + + let sanitized = SanitizedCompletionParams::from(params); + + assert_eq!( + sanitized.text, + r#"select * from "auth"."REPLACED_TOKEN_WITH_QUOTE";"# + ); + } + + { + // select * from "auth"."| <-- without semi + let input = r#"select * from "auth".""#; + let position = TextSize::new(22); + + let params = get_test_params(input, position); + + let sanitized = SanitizedCompletionParams::from(params); + + assert_eq!( + sanitized.text, + r#"select * from "auth"."REPLACED_TOKEN_WITH_QUOTE""# + ); + } } #[test] From 79cdcc402d3033a8f26f4760144f50996f459eb8 Mon Sep 17 00:00:00 2001 From: Julian Date: Thu, 11 Sep 2025 18:28:39 +0200 Subject: [PATCH 3/3] wowa wiwa --- crates/pgt_completions/src/providers/functions.rs | 1 - crates/pgt_completions/src/relevance/scoring.rs | 3 ++- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/pgt_completions/src/providers/functions.rs b/crates/pgt_completions/src/providers/functions.rs index 876b34a47..7636d3595 100644 --- a/crates/pgt_completions/src/providers/functions.rs +++ b/crates/pgt_completions/src/providers/functions.rs @@ -1,5 +1,4 @@ use pgt_schema_cache::{Function, SchemaCache}; -use pgt_text_size::TextSize; use pgt_treesitter::TreesitterContext; use crate::{ diff --git a/crates/pgt_completions/src/relevance/scoring.rs b/crates/pgt_completions/src/relevance/scoring.rs index ba45e2d0d..e4f8eb5a2 100644 --- a/crates/pgt_completions/src/relevance/scoring.rs +++ b/crates/pgt_completions/src/relevance/scoring.rs @@ -184,9 +184,10 @@ impl CompletionScore<'_> { } fn check_matches_schema(&mut self, ctx: &TreesitterContext) { + // TODO let schema_name = match ctx.schema_or_alias_name.as_ref() { None => return, - Some(n) => n, + Some(n) => n.replace('"', ""), }; let data_schema = match self.get_schema_name() {