feat(completions): correctly infer (quoted) schema for tables & functions (#509)

juleswritescode · web-flow · commit 98ae38f9a490 · 2025-09-18T09:11:00.000Z
diff --git a/crates/pgt_completions/src/providers/functions.rs b/crates/pgt_completions/src/providers/functions.rs
@@ -4,7 +4,7 @@ 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 +37,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 +46,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 +73,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 +300,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
@@ -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,21 +48,18 @@ 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()
         }
     } 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
@@ -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::{
@@ -569,4 +570,90 @@ 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::CompletionTextAndRange(
+                    "posts".into(),
+                    TextRange::new(21.into(), 21.into()),
+                ),
+                CompletionAssertion::CompletionTextAndRange(
+                    "users".into(),
+                    TextRange::new(21.into(), 21.into()),
+                ),
+            ],
+            None,
+            &pool,
+        )
+        .await;
+
+        assert_complete_results(
+            format!(
+                r#"select * from "auth"."{}""#,
+                QueryWithCursorPosition::cursor_marker()
+            )
+            .as_str(),
+            vec![
+                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,
+        )
+        .await;
+    }
 }
diff --git 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.
diff --git 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() {
diff --git 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]
