refactor(completions): remove custom parse for policies, use tree-sitter

.github/workflows/pull_request.yml

@@ -112,6 +112,9 @@ jobs:
             sleep 1
           done
 
+      - name: Install tree-sitter-cli
+        run: cargo install tree-sitter-cli
+
       - name: Setup sqlx-cli
         run: cargo install sqlx-cli
 
@@ -167,6 +170,9 @@ jobs:
       - name: Setup Postgres
         uses: ./.github/actions/setup-postgres
 
+      - name: Install tree-sitter-cli
+        run: cargo install tree-sitter-cli
+
       - name: Run tests
         run: cargo test --workspace
 
@@ -195,6 +201,8 @@ jobs:
         uses: moonrepo/setup-rust@v1
         with:
           cache-base: main
+      - name: Install tree-sitter-cli
+        run: cargo install tree-sitter-cli
       - name: Build main binary
         run: cargo build -p pgt_cli --release
       - name: Setup Bun
@@ -236,6 +244,8 @@ jobs:
           cache-base: main
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+      - name: Install tree-sitter-cli
+        run: cargo install tree-sitter-cli
       - name: Ensure RustFMT on nightly toolchain
         run: rustup component add rustfmt --toolchain nightly
       - name: echo toolchain

.gitignore

@@ -26,3 +26,8 @@ node_modules/
 .claude-session-id
 
 squawk/
+
+# Auto generated treesitter files
+crates/pgt_treesitter_grammar/src/grammar.json
+crates/pgt_treesitter_grammar/src/node-types.json
+crates/pgt_treesitter_grammar/src/parser.c

crates/pgt_completions/src/providers/functions.rs

@@ -274,14 +274,16 @@ mod tests {
         pool.execute(setup).await.unwrap();
 
         let query = format!(
-            r#"create policy "my_pol" on public.instruments for insert with check (id = {})"#,
+            r#"create policy "my_pol" on public.coos for insert with check (id = {})"#,
             QueryWithCursorPosition::cursor_marker()
         );
 
         assert_complete_results(
             query.as_str(),
             vec![
                 CompletionAssertion::LabelNotExists("string_concat".into()),
+                CompletionAssertion::LabelAndKind("id".into(), CompletionItemKind::Column),
+                CompletionAssertion::LabelAndKind("name".into(), CompletionItemKind::Column),
                 CompletionAssertion::LabelAndKind(
                     "my_cool_foo".into(),
                     CompletionItemKind::Function,

crates/pgt_completions/src/providers/policies.rs

@@ -80,6 +80,21 @@ mod tests {
 
         pool.execute(setup).await.unwrap();
 
+        assert_complete_results(
+            format!(
+                "alter policy \"{}",
+                QueryWithCursorPosition::cursor_marker()
+            )
+            .as_str(),
+            vec![
+                CompletionAssertion::Label("read for public users disallowed".into()),
+                CompletionAssertion::Label("write for public users allowed".into()),
+            ],
+            None,
+            &pool,
+        )
+        .await;
+
         assert_complete_results(
             format!(
                 "alter policy \"{}\" on private.users;",

crates/pgt_completions/src/providers/tables.rs

@@ -656,4 +656,57 @@ mod tests {
         )
         .await;
     }
+
+    #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")]
+    async fn completes_tables_in_policies(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#"create policy "my cool pol" on {}"#,
+                QueryWithCursorPosition::cursor_marker()
+            )
+            .as_str(),
+            vec![
+                CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
+                CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema),
+            ],
+            None,
+            &pool,
+        )
+        .await;
+
+        assert_complete_results(
+            format!(
+                r#"create policy "my cool pol" on auth.{}"#,
+                QueryWithCursorPosition::cursor_marker()
+            )
+            .as_str(),
+            vec![
+                CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table),
+                CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
+            ],
+            None,
+            &pool,
+        )
+        .await;
+    }
 }

crates/pgt_completions/src/relevance/filtering.rs

@@ -1,5 +1,4 @@
 use pgt_schema_cache::ProcKind;
-
 use pgt_treesitter::context::{NodeUnderCursor, TreesitterContext, WrappingClause, WrappingNode};
 
 use super::CompletionRelevanceData;
@@ -58,8 +57,7 @@ impl CompletionFilter<'_> {
                     | Some(WrappingClause::Insert)
                     | Some(WrappingClause::DropColumn)
                     | Some(WrappingClause::AlterColumn)
-                    | Some(WrappingClause::RenameColumn)
-                    | Some(WrappingClause::PolicyCheck) => {
+                    | Some(WrappingClause::RenameColumn) => {
                         // the literal is probably a column
                     }
                     _ => return None,
@@ -123,6 +121,13 @@ impl CompletionFilter<'_> {
                                 "keyword_table",
                             ]),
 
+                        WrappingClause::CreatePolicy
+                        | WrappingClause::AlterPolicy
+                        | WrappingClause::DropPolicy => {
+                            ctx.matches_ancestor_history(&["object_reference"])
+                                && ctx.before_cursor_matches_kind(&["keyword_on", "."])
+                        }
+
                         _ => false,
                     },
 
@@ -162,8 +167,11 @@ impl CompletionFilter<'_> {
                                         && ctx.matches_ancestor_history(&["field"]))
                             }
 
-                            WrappingClause::PolicyCheck => {
-                                ctx.before_cursor_matches_kind(&["keyword_and", "("])
+                            WrappingClause::CheckOrUsingClause => {
+                                ctx.before_cursor_matches_kind(&["(", "keyword_and"])
+                                    || ctx.wrapping_node_kind.as_ref().is_some_and(|nk| {
+                                        matches!(nk, WrappingNode::BinaryExpression)
+                                    })
                             }
 
                             _ => false,
@@ -176,9 +184,12 @@ impl CompletionFilter<'_> {
                         | WrappingClause::Where
                         | WrappingClause::Join { .. } => true,
 
-                        WrappingClause::PolicyCheck => {
-                            ctx.before_cursor_matches_kind(&["="])
-                                && matches!(f.kind, ProcKind::Function | ProcKind::Procedure)
+                        WrappingClause::CheckOrUsingClause => {
+                            !matches!(f.kind, ProcKind::Aggregate)
+                                && (ctx.before_cursor_matches_kind(&["(", "keyword_and"])
+                                    || ctx.wrapping_node_kind.as_ref().is_some_and(|nk| {
+                                        matches!(nk, WrappingNode::BinaryExpression)
+                                    }))
                         }
 
                         _ => false,
@@ -209,11 +220,21 @@ impl CompletionFilter<'_> {
                                 && ctx.before_cursor_matches_kind(&["keyword_into"])
                         }
 
+                        WrappingClause::CreatePolicy
+                        | WrappingClause::AlterPolicy
+                        | WrappingClause::DropPolicy => {
+                            ctx.before_cursor_matches_kind(&["keyword_on"])
+                        }
+
                         _ => false,
                     },
 
                     CompletionRelevanceData::Policy(_) => {
-                        matches!(clause, WrappingClause::PolicyName)
+                        matches!(
+                            clause,
+                            // not CREATE – there can't be existing policies.
+                            WrappingClause::AlterPolicy | WrappingClause::DropPolicy
+                        ) && ctx.before_cursor_matches_kind(&["keyword_policy", "keyword_exists"])
                     }
 
                     CompletionRelevanceData::Role(_) => match clause {
@@ -224,6 +245,11 @@ impl CompletionFilter<'_> {
                         WrappingClause::SetStatement => ctx
                             .before_cursor_matches_kind(&["keyword_role", "keyword_authorization"]),
 
+                        WrappingClause::AlterPolicy | WrappingClause::CreatePolicy => {
+                            ctx.before_cursor_matches_kind(&["keyword_to"])
+                                && ctx.matches_ancestor_history(&["policy_to_role"])
+                        }
+
                         _ => false,
                     },
                 }

crates/pgt_completions/src/relevance/scoring.rs

@@ -35,6 +35,7 @@ impl CompletionScore<'_> {
         self.check_matching_wrapping_node(ctx);
         self.check_relations_in_stmt(ctx);
         self.check_columns_in_stmt(ctx);
+        self.check_is_not_wellknown_migration(ctx);
     }
 
     fn check_matches_query_input(&mut self, ctx: &TreesitterContext) {
@@ -100,12 +101,14 @@ impl CompletionScore<'_> {
                 WrappingClause::Select if !has_mentioned_tables => 15,
                 WrappingClause::Select if has_mentioned_tables => 0,
                 WrappingClause::From => 0,
+                WrappingClause::CheckOrUsingClause => 0,
                 _ => -50,
             },
             CompletionRelevanceData::Column(col) => match clause_type {
                 WrappingClause::Select if has_mentioned_tables => 10,
                 WrappingClause::Select if !has_mentioned_tables => 0,
                 WrappingClause::Where => 10,
+                WrappingClause::CheckOrUsingClause => 0,
                 WrappingClause::Join { on_node }
                     if on_node.is_some_and(|on| {
                         ctx.node_under_cursor
@@ -123,10 +126,13 @@ impl CompletionScore<'_> {
                 WrappingClause::Join { .. } if !has_mentioned_schema => 15,
                 WrappingClause::Update if !has_mentioned_schema => 15,
                 WrappingClause::Delete if !has_mentioned_schema => 15,
+                WrappingClause::AlterPolicy if !has_mentioned_schema => 15,
+                WrappingClause::DropPolicy if !has_mentioned_schema => 15,
+                WrappingClause::CreatePolicy if !has_mentioned_schema => 15,
                 _ => -50,
             },
             CompletionRelevanceData::Policy(_) => match clause_type {
-                WrappingClause::PolicyName => 25,
+                WrappingClause::AlterPolicy | WrappingClause::DropPolicy => 25,
                 _ => -50,
             },
 
@@ -156,6 +162,7 @@ impl CompletionScore<'_> {
                 _ => -50,
             },
             CompletionRelevanceData::Function(_) => match wrapping_node {
+                WrappingNode::BinaryExpression => 15,
                 WrappingNode::Relation => 10,
                 _ => -50,
             },
@@ -184,7 +191,6 @@ 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.replace('"', ""),
@@ -349,4 +355,18 @@ impl CompletionScore<'_> {
             }
         }
     }
+
+    fn check_is_not_wellknown_migration(&mut self, _ctx: &TreesitterContext) {
+        if let Some(table_name) = self.get_table_name() {
+            if ["_sqlx_migrations"].contains(&table_name) {
+                self.score -= 10;
+            }
+        }
+
+        if let Some(schema_name) = self.get_schema_name() {
+            if ["supabase_migrations"].contains(&schema_name) {
+                self.score -= 10;
+            }
+        }
+    }
 }

crates/pgt_hover/src/hovered_node.rs

@@ -1,3 +1,5 @@
+use pgt_treesitter::WrappingClause;
+
 #[derive(Debug)]
 pub(crate) enum NodeIdentification {
     Name(String),
@@ -40,6 +42,28 @@ impl HoveredNode {
                     Some(HoveredNode::Table(NodeIdentification::Name(node_content)))
                 }
             }
+
+            "identifier"
+                if ctx.matches_ancestor_history(&["object_reference"])
+                    && ctx.wrapping_clause_type.as_ref().is_some_and(|clause| {
+                        matches!(
+                            clause,
+                            WrappingClause::AlterPolicy
+                                | WrappingClause::CreatePolicy
+                                | WrappingClause::DropPolicy
+                        )
+                    }) =>
+            {
+                if let Some(schema) = ctx.schema_or_alias_name.as_ref() {
+                    Some(HoveredNode::Table(NodeIdentification::SchemaAndName((
+                        schema.clone(),
+                        node_content,
+                    ))))
+                } else {
+                    Some(HoveredNode::Table(NodeIdentification::Name(node_content)))
+                }
+            }
+
             "identifier" if ctx.matches_ancestor_history(&["field"]) => {
                 if let Some(table_or_alias) = ctx.schema_or_alias_name.as_ref() {
                     Some(HoveredNode::Column(NodeIdentification::SchemaAndName((
@@ -62,7 +86,7 @@ impl HoveredNode {
                     )))
                 }
             }
-            "identifier" if ctx.matches_ancestor_history(&["alter_role"]) => {
+            "identifier" if ctx.matches_one_of_ancestors(&["alter_role", "policy_to_role"]) => {
                 Some(HoveredNode::Role(NodeIdentification::Name(node_content)))
             }
 
@@ -75,7 +99,7 @@ impl HoveredNode {
                 Some(HoveredNode::Column(NodeIdentification::Name(node_content)))
             }
 
-            "policy_table" | "revoke_table" | "grant_table" => {
+            "revoke_table" | "grant_table" => {
                 if let Some(schema) = ctx.schema_or_alias_name.as_ref() {
                     Some(HoveredNode::Table(NodeIdentification::SchemaAndName((
                         schema.clone(),