6/11: fix(completions): sanitization edge cases (#659)

Author: juleswritescode

crates/pgls_completions/src/builder.rs

@@ -32,10 +32,12 @@ impl<'a> CompletionBuilder<'a> {
     }
 
     pub fn finish(self) -> Vec<CompletionItem> {
+        let mut shared_tree = self.ctx.tree.clone();
+
         let mut items: Vec<PossibleCompletionItem> = self
             .items
             .into_iter()
-            .filter(|i| i.filter.is_relevant(self.ctx).is_some())
+            .filter(|i| i.filter.is_relevant(self.ctx, &mut shared_tree).is_some())
             .collect();
 
         for item in items.iter_mut() {

crates/pgls_completions/src/complete.rs

@@ -6,8 +6,8 @@ use crate::{
     builder::CompletionBuilder,
     item::CompletionItem,
     providers::{
-        complete_columns, complete_functions, complete_policies, complete_roles, complete_schemas,
-        complete_tables,
+        complete_columns, complete_functions, complete_keywords, complete_policies, complete_roles,
+        complete_schemas, complete_tables,
     },
     sanitization::SanitizedCompletionParams,
 };
@@ -27,6 +27,13 @@ pub struct CompletionParams<'a> {
     position = params.position.to_string()
 ))]
 pub fn complete(params: CompletionParams) -> Vec<CompletionItem> {
+    let uses_upper_case = params
+        .text
+        .split_ascii_whitespace()
+        // filter out special chars and numbers
+        .filter(|word| word.chars().all(|c| c.is_alphabetic()))
+        .any(|t| t == t.to_ascii_uppercase());
+
     let sanitized_params = SanitizedCompletionParams::from(params);
 
     let ctx = TreesitterContext::new(TreeSitterContextParams {
@@ -43,6 +50,7 @@ pub fn complete(params: CompletionParams) -> Vec<CompletionItem> {
     complete_schemas(&ctx, sanitized_params.schema, &mut builder);
     complete_policies(&ctx, sanitized_params.schema, &mut builder);
     complete_roles(&ctx, sanitized_params.schema, &mut builder);
+    complete_keywords(&ctx, &mut builder, uses_upper_case);
 
     builder.finish()
 }

crates/pgls_completions/src/item.rs

@@ -13,6 +13,7 @@ pub enum CompletionItemKind {
     Schema,
     Policy,
     Role,
+    Keyword,
 }
 
 impl Display for CompletionItemKind {
@@ -24,6 +25,7 @@ impl Display for CompletionItemKind {
             CompletionItemKind::Schema => "Schema",
             CompletionItemKind::Policy => "Policy",
             CompletionItemKind::Role => "Role",
+            CompletionItemKind::Keyword => "Keyword",
         };
 
         write!(f, "{txt}")

crates/pgls_completions/src/providers/columns.rs

@@ -417,7 +417,6 @@ mod tests {
     }
 
     #[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")]
-    #[ignore = "will be reintroduced after stacked keyword-completion PRs merge"]
     async fn suggests_columns_in_where_clause(pool: PgPool) {
         let setup = r#"
             create table instruments (
@@ -441,7 +440,7 @@ mod tests {
                         "select name from instruments i join others o on i.z = o.a <sql>",
                     )
                     .type_sql("where o.<1>a = <2>i.z and <3>i.id > 5;")
-                    .comment("should respect alias speciifcation")
+                    .comment("should respect alias specification")
                     .comment("should not prioritize suggest columns or schemas (right side of binary expression)")
                     .comment("should prioritize columns that aren't already mentioned"),
             )