arairga

juleswritescode · juleswritescode · commit bf3d96cbfea8 · 2025-05-12T17:40:18.000+02:00
diff --git a/crates/pgt_completions/src/context/context.rs b/crates/pgt_completions/src/context/context.rs
@@ -67,13 +67,6 @@ impl<'a> NodeUnderCursor<'a> {
         }
     }
 
-    pub fn range(&self) -> TextRange {
-        let start: u32 = self.start_byte().try_into().unwrap();
-        let end: u32 = self.end_byte().try_into().unwrap();
-
-        TextRange::new(start.into(), end.into())
-    }
-
     pub fn kind(&self) -> &str {
         match self {
             NodeUnderCursor::TsNode(node) => node.kind(),
@@ -216,12 +209,6 @@ impl<'a> CompletionContext<'a> {
             ctx.gather_info_from_ts_queries();
         }
 
-        tracing::warn!("SQL: {}", ctx.text);
-        tracing::warn!("Position: {}", ctx.position);
-        tracing::warn!("Node: {:#?}", ctx.node_under_cursor);
-        tracing::warn!("Relations: {:#?}", ctx.mentioned_relations);
-        tracing::warn!("Clause: {:#?}", ctx.wrapping_clause_type);
-
         ctx
     }
 
diff --git a/crates/pgt_completions/src/context/policy_parser.rs b/crates/pgt_completions/src/context/policy_parser.rs
@@ -110,10 +110,11 @@ pub(crate) struct PolicyParser {
 
 impl PolicyParser {
     pub(crate) fn get_context(sql: &str, cursor_position: usize) -> PolicyContext {
+        let trimmed = sql.trim();
         assert!(
-            sql.starts_with("create policy")
-                || sql.starts_with("drop policy")
-                || sql.starts_with("alter policy"),
+            trimmed.starts_with("create policy")
+                || trimmed.starts_with("drop policy")
+                || trimmed.starts_with("alter policy"),
             "PolicyParser should only be used for policy statements. Developer error!"
         );
 
diff --git a/crates/pgt_completions/src/providers/policies.rs b/crates/pgt_completions/src/providers/policies.rs
@@ -10,6 +10,10 @@ use super::helper::get_range_to_replace;
 pub fn complete_policies<'a>(ctx: &CompletionContext<'a>, builder: &mut CompletionBuilder<'a>) {
     let available_policies = &ctx.schema_cache.policies;
 
+    let has_quotes = ctx
+        .get_node_under_cursor_content()
+        .is_some_and(|c| c.starts_with('"') && c.ends_with('"'));
+
     for pol in available_policies {
         let relevance = CompletionRelevanceData::Policy(pol);
 
@@ -19,10 +23,14 @@ pub fn complete_policies<'a>(ctx: &CompletionContext<'a>, builder: &mut Completi
             filter: CompletionFilter::from(relevance),
             description: format!("{}", pol.table_name),
             kind: CompletionItemKind::Policy,
-            completion_text: Some(CompletionText {
-                text: format!("\"{}\"", pol.name),
-                range: get_range_to_replace(ctx),
-            }),
+            completion_text: if !has_quotes {
+                Some(CompletionText {
+                    text: format!("\"{}\"", pol.name),
+                    range: get_range_to_replace(ctx),
+                })
+            } else {
+                None
+            },
         };
 
         builder.add_item(item);
@@ -31,30 +39,69 @@ pub fn complete_policies<'a>(ctx: &CompletionContext<'a>, builder: &mut Completi
 
 #[cfg(test)]
 mod tests {
-    use crate::test_helper::{CURSOR_POS, CompletionAssertion, assert_complete_results};
+    use crate::{
+        complete,
+        test_helper::{
+            CURSOR_POS, CompletionAssertion, assert_complete_results, get_test_params,
+            test_against_connection_string,
+        },
+    };
 
     #[tokio::test]
     async fn completes_within_quotation_marks() {
         let setup = r#"
-            create table users (
+            create schema private;
+
+            create table private.users (
                 id serial primary key,
                 email text
             );
 
-            create policy "should never have access" on users
+            create policy "read for public users disallowed" on private.users
                 as restrictive
-                for all
+                for select
                 to public
                 using (false);
+
+            create policy "write for public users allowed" on private.users
+                as restrictive
+                for insert
+                to public
+                with check (true);
         "#;
 
         assert_complete_results(
-            format!("alter policy \"{}\" on users;", CURSOR_POS).as_str(),
+            format!("alter policy \"{}\" on private.users;", CURSOR_POS).as_str(),
+            vec![
+                CompletionAssertion::Label("read for public users disallowed".into()),
+                CompletionAssertion::Label("write for public users allowed".into()),
+            ],
+            setup,
+        )
+        .await;
+
+        assert_complete_results(
+            format!("alter policy \"w{}\" on private.users;", CURSOR_POS).as_str(),
             vec![CompletionAssertion::Label(
-                "should never have access".into(),
+                "write for public users allowed".into(),
             )],
             setup,
         )
         .await;
     }
+
+    #[tokio::test]
+    async fn sb_test() {
+        let input = format!("alter policy \"u{}\" on public.fcm_tokens;", CURSOR_POS);
+
+        let (tree, cache) = test_against_connection_string(
+            "postgresql://postgres:postgres@127.0.0.1:54322/postgres",
+            input.as_str().into(),
+        )
+        .await;
+
+        let result = complete(get_test_params(&tree, &cache, input.as_str().into()));
+
+        println!("{:#?}", result);
+    }
 }
diff --git a/crates/pgt_completions/src/relevance/filtering.rs b/crates/pgt_completions/src/relevance/filtering.rs
@@ -71,9 +71,8 @@ impl CompletionFilter<'_> {
 
         match self.data {
             CompletionRelevanceData::Table(_) => {
-                if in_clause(WrappingClause::Select)
-                    || in_clause(WrappingClause::Where)
-                    || in_clause(WrappingClause::PolicyName)
+                if in_clause(WrappingClause::Select) || in_clause(WrappingClause::Where)
+                // || in_clause(WrappingClause::PolicyName)
                 {
                     return None;
                 };
@@ -107,9 +106,9 @@ impl CompletionFilter<'_> {
                 }
             }
             _ => {
-                if in_clause(WrappingClause::PolicyName) {
-                    return None;
-                }
+                // if in_clause(WrappingClause::PolicyName) {
+                //     return None;
+                // }
             }
         }
 
diff --git a/crates/pgt_completions/src/relevance/scoring.rs b/crates/pgt_completions/src/relevance/scoring.rs
@@ -36,21 +36,23 @@ impl CompletionScore<'_> {
 
     fn check_matches_query_input(&mut self, ctx: &CompletionContext) {
         let content = match ctx.get_node_under_cursor_content() {
-            Some(c) => c,
+            Some(c) => c.replace('"', ""),
             None => return,
         };
 
         let name = match self.data {
-            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(),
-            CompletionRelevanceData::Policy(p) => p.name.as_str(),
+            CompletionRelevanceData::Function(f) => f.name.as_str().to_ascii_lowercase(),
+            CompletionRelevanceData::Table(t) => t.name.as_str().to_ascii_lowercase(),
+            CompletionRelevanceData::Column(c) => c.name.as_str().to_ascii_lowercase(),
+            CompletionRelevanceData::Schema(s) => s.name.as_str().to_ascii_lowercase(),
+            CompletionRelevanceData::Policy(p) => p.name.as_str().to_ascii_lowercase(),
         };
 
         let fz_matcher = SkimMatcherV2::default();
 
-        if let Some(score) = fz_matcher.fuzzy_match(name, content.as_str()) {
+        if let Some(score) =
+            fz_matcher.fuzzy_match(name.as_str(), content.to_ascii_lowercase().as_str())
+        {
             let scorei32: i32 = score
                 .try_into()
                 .expect("The length of the input exceeds i32 capacity");
diff --git a/crates/pgt_completions/src/sanitization.rs b/crates/pgt_completions/src/sanitization.rs
@@ -45,8 +45,8 @@ where
     'larger: 'smaller,
 {
     fn from(params: CompletionParams<'larger>) -> Self {
-        if cursor_inbetween_nodes(params.tree, params.position)
-            || cursor_prepared_to_write_token_after_last_node(params.tree, params.position)
+        if cursor_inbetween_nodes(&params.text, params.position)
+            || cursor_prepared_to_write_token_after_last_node(&params.text, params.position)
             || cursor_before_semicolon(params.tree, params.position)
             || cursor_on_a_dot(&params.text, params.position)
             || cursor_between_double_quotes(&params.text, params.position)
@@ -125,37 +125,17 @@ where
 /// select |from users; -- cursor "touches" from node. returns false.
 /// select | from users; -- cursor is between select and from nodes. returns true.
 /// ```
-fn cursor_inbetween_nodes(tree: &tree_sitter::Tree, position: TextSize) -> bool {
-    let mut cursor = tree.walk();
-    let mut leaf_node = tree.root_node();
-
-    let byte = position.into();
-
-    // if the cursor escapes the root node, it can't be between nodes.
-    if byte < leaf_node.start_byte() || byte >= leaf_node.end_byte() {
-        return false;
-    }
+fn cursor_inbetween_nodes(sql: &str, position: TextSize) -> bool {
+    let position: usize = position.into();
+    let mut chars = sql.chars();
 
-    /*
-     * Get closer and closer to the leaf node, until
-     *  a) there is no more child *for the node* or
-     *  b) there is no more child *under the cursor*.
-     */
-    loop {
-        let child_idx = cursor.goto_first_child_for_byte(position.into());
-        if child_idx.is_none() {
-            break;
-        }
-        leaf_node = cursor.node();
-    }
+    let previous_whitespace = chars
+        .nth(position - 1)
+        .is_some_and(|c| c.is_ascii_whitespace());
 
-    let cursor_on_leafnode = byte >= leaf_node.start_byte() && leaf_node.end_byte() >= byte;
+    let current_whitespace = chars.next().is_some_and(|c| c.is_ascii_whitespace());
 
-    /*
-     * The cursor is inbetween nodes if it is not within the range
-     * of a leaf node.
-     */
-    !cursor_on_leafnode
+    previous_whitespace && current_whitespace
 }
 
 /// Checks if the cursor is positioned after the last node,
@@ -166,12 +146,9 @@ fn cursor_inbetween_nodes(tree: &tree_sitter::Tree, position: TextSize) -> bool
 /// select * from|    -- user still needs to type a space
 /// select * from  |  -- too far off.
 /// ```
-fn cursor_prepared_to_write_token_after_last_node(
-    tree: &tree_sitter::Tree,
-    position: TextSize,
-) -> bool {
+fn cursor_prepared_to_write_token_after_last_node(sql: &str, position: TextSize) -> bool {
     let cursor_pos: usize = position.into();
-    cursor_pos == tree.root_node().end_byte() + 1
+    cursor_pos == sql.len() + 1
 }
 
 fn cursor_on_a_dot(sql: &str, position: TextSize) -> bool {
@@ -243,58 +220,44 @@ mod tests {
         // note: two spaces between select and from.
         let input = "select  from users;";
 
-        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();
-
         // select | from users; <-- just right, one space after select token, one space before from
-        assert!(cursor_inbetween_nodes(&tree, TextSize::new(7)));
+        assert!(cursor_inbetween_nodes(input, TextSize::new(7)));
 
         // select|  from users; <-- still on select token
-        assert!(!cursor_inbetween_nodes(&tree, TextSize::new(6)));
+        assert!(!cursor_inbetween_nodes(input, TextSize::new(6)));
 
         // select  |from users; <-- already on from token
-        assert!(!cursor_inbetween_nodes(&tree, TextSize::new(8)));
+        assert!(!cursor_inbetween_nodes(input, TextSize::new(8)));
 
         // select from users;|
-        assert!(!cursor_inbetween_nodes(&tree, TextSize::new(19)));
+        assert!(!cursor_inbetween_nodes(input, TextSize::new(19)));
     }
 
     #[test]
     fn test_cursor_after_nodes() {
         let input = "select * from";
 
-        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();
-
         // select * from| <-- still on previous token
         assert!(!cursor_prepared_to_write_token_after_last_node(
-            &tree,
+            input,
             TextSize::new(13)
         ));
 
         // select * from  | <-- too far off, two spaces afterward
         assert!(!cursor_prepared_to_write_token_after_last_node(
-            &tree,
+            input,
             TextSize::new(15)
         ));
 
         // select * |from  <-- it's within
         assert!(!cursor_prepared_to_write_token_after_last_node(
-            &tree,
+            input,
             TextSize::new(9)
         ));
 
         // select * from | <-- just right
         assert!(cursor_prepared_to_write_token_after_last_node(
-            &tree,
+            input,
             TextSize::new(14)
         ));
     }
@@ -353,5 +316,11 @@ mod tests {
 
         // select * from "|" <-- between quotations
         assert!(cursor_between_double_quotes(input, TextSize::new(15)));
+
+        // select * from "r|" <-- between quotations, but there's
+        // a letter inside
+        let input = "select * from \"r\"";
+
+        assert!(!cursor_between_double_quotes(input, TextSize::new(16)));
     }
 }
diff --git a/crates/pgt_lsp/src/handlers/completions.rs b/crates/pgt_lsp/src/handlers/completions.rs
@@ -54,6 +54,8 @@ pub fn get_completions(
         })
         .collect();
 
+    tracing::warn!("{:#?}", items);
+
     Ok(lsp_types::CompletionResponse::Array(items))
 }
 
@@ -65,6 +67,6 @@ fn to_lsp_types_completion_item_kind(
         pgt_completions::CompletionItemKind::Table => lsp_types::CompletionItemKind::CLASS,
         pgt_completions::CompletionItemKind::Column => lsp_types::CompletionItemKind::FIELD,
         pgt_completions::CompletionItemKind::Schema => lsp_types::CompletionItemKind::CLASS,
-        pgt_completions::CompletionItemKind::Policy => lsp_types::CompletionItemKind::VALUE,
+        pgt_completions::CompletionItemKind::Policy => lsp_types::CompletionItemKind::CONSTANT,
     }
 }
diff --git a/crates/pgt_workspace/src/workspace/server.rs b/crates/pgt_workspace/src/workspace/server.rs
@@ -481,15 +481,20 @@ impl Workspace for WorkspaceServer {
             Some(pool) => pool,
             None => {
                 tracing::debug!("No connection to database. Skipping completions.");
+                tracing::warn!("No connection to database.");
                 return Ok(CompletionsResult::default());
             }
         };
 
         let schema_cache = self.schema_cache.load(pool)?;
 
         match get_statement_for_completions(&parsed_doc, params.position) {
-            None => Ok(CompletionsResult::default()),
+            None => {
+                tracing::warn!("No statement found.");
+                Ok(CompletionsResult::default())
+            }
             Some((_id, range, content, cst)) => {
+                tracing::warn!("found matching statement, content: {}", content);
                 let position = params.position - range.start();
 
                 let items = pgt_completions::complete(pgt_completions::CompletionParams {

Original file line number	Diff line number	Diff line change
`@@ -67,13 +67,6 @@ impl<'a> NodeUnderCursor<'a> {`
`67`	`67`	`}`
`68`	`68`	`}`
`69`	`69`
`70`		`- pub fn range(&self) -> TextRange {`
`71`		`- let start: u32 = self.start_byte().try_into().unwrap();`
`72`		`- let end: u32 = self.end_byte().try_into().unwrap();`
`73`		`-`
`74`		`- TextRange::new(start.into(), end.into())`
`75`		`- }`
`76`		`-`
`77`	`70`	`pub fn kind(&self) -> &str {`
`78`	`71`	`match self {`
`79`	`72`	`NodeUnderCursor::TsNode(node) => node.kind(),`
`@@ -216,12 +209,6 @@ impl<'a> CompletionContext<'a> {`
`216`	`209`	`ctx.gather_info_from_ts_queries();`
`217`	`210`	`}`
`218`	`211`
`219`		`- tracing::warn!("SQL: {}", ctx.text);`
`220`		`- tracing::warn!("Position: {}", ctx.position);`
`221`		`- tracing::warn!("Node: {:#?}", ctx.node_under_cursor);`
`222`		`- tracing::warn!("Relations: {:#?}", ctx.mentioned_relations);`
`223`		`- tracing::warn!("Clause: {:#?}", ctx.wrapping_clause_type);`
`224`		`-`
`225`	`212`	`ctx`
`226`	`213`	`}`
`227`	`214`
Original file line number	Diff line number	Diff line change
`@@ -71,9 +71,8 @@ impl CompletionFilter<'_> {`
`71`	`71`
`72`	`72`	`match self.data {`
`73`	`73`	`CompletionRelevanceData::Table(_) => {`
`74`		`- if in_clause(WrappingClause::Select)`
`75`		`- \|\| in_clause(WrappingClause::Where)`
`76`		`- \|\| in_clause(WrappingClause::PolicyName)`
	`74`	`+ if in_clause(WrappingClause::Select) \|\| in_clause(WrappingClause::Where)`
	`75`	`+ // \|\| in_clause(WrappingClause::PolicyName)`
`77`	`76`	`{`
`78`	`77`	`return None;`
`79`	`78`	`};`
`@@ -107,9 +106,9 @@ impl CompletionFilter<'_> {`
`107`	`106`	`}`
`108`	`107`	`}`
`109`	`108`	`_ => {`
`110`		`- if in_clause(WrappingClause::PolicyName) {`
`111`		`- return None;`
`112`		`- }`
	`109`	`+ // if in_clause(WrappingClause::PolicyName) {`
	`110`	`+ // return None;`
	`111`	`+ // }`
`113`	`112`	`}`
`114`	`113`	`}`
`115`	`114`
Original file line number	Diff line number	Diff line change
`@@ -54,6 +54,8 @@ pub fn get_completions(`
`54`	`54`	`})`
`55`	`55`	`.collect();`
`56`	`56`
	`57`	`+ tracing::warn!("{:#?}", items);`
	`58`	`+`
`57`	`59`	`Ok(lsp_types::CompletionResponse::Array(items))`
`58`	`60`	`}`
`59`	`61`
`@@ -65,6 +67,6 @@ fn to_lsp_types_completion_item_kind(`
`65`	`67`	`pgt_completions::CompletionItemKind::Table => lsp_types::CompletionItemKind::CLASS,`
`66`	`68`	`pgt_completions::CompletionItemKind::Column => lsp_types::CompletionItemKind::FIELD,`
`67`	`69`	`pgt_completions::CompletionItemKind::Schema => lsp_types::CompletionItemKind::CLASS,`
`68`		`- pgt_completions::CompletionItemKind::Policy => lsp_types::CompletionItemKind::VALUE,`
	`70`	`+ pgt_completions::CompletionItemKind::Policy => lsp_types::CompletionItemKind::CONSTANT,`
`69`	`71`	`}`
`70`	`72`	`}`