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

Author: juleswritescode

.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(),