ok?

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/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,
@@ -162,8 +160,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 +177,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,
@@ -213,7 +217,11 @@ impl CompletionFilter<'_> {
  },
 
  CompletionRelevanceData::Policy(_) => {
- matches!(clause, WrappingClause::PolicyName)
+ matches!(
+ clause,
+ // not CREATE – there can't be existing policies.
+ WrappingClause::DropPolicy | WrappingClause::AlterPolicy
+ )
  }
 
  CompletionRelevanceData::Role(_) => match clause {
@@ -224,6 +232,10 @@ impl CompletionFilter<'_> {
  WrappingClause::SetStatement => ctx
  .before_cursor_matches_kind(&["keyword_role", "keyword_authorization"]),
 
+ WrappingClause::AlterPolicy | WrappingClause::CreatePolicy => {
+ ctx.before_cursor_matches_kind(&["keyword_to"])
+ }
+
  _ => 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
@@ -126,7 +129,7 @@ impl CompletionScore<'_> {
  _ => -50,
  },
  CompletionRelevanceData::Policy(_) => match clause_type {
- WrappingClause::PolicyName => 25,
+ WrappingClause::AlterPolicy | WrappingClause::DropPolicy => 25,
  _ => -50,
  },
 
@@ -156,6 +159,7 @@ impl CompletionScore<'_> {
  _ => -50,
  },
  CompletionRelevanceData::Function(_) => match wrapping_node {
+ WrappingNode::BinaryExpression => 15,
  WrappingNode::Relation => 10,
  _ => -50,
  },
@@ -184,7 +188,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 +352,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((
@@ -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(),

crates/pgt_treesitter/src/context/mod.rs

@@ -4,17 +4,13 @@ use std::{
 };
 mod base_parser;
 mod grant_parser;
-mod policy_parser;
 mod revoke_parser;
 
 use crate::queries::{self, QueryResult, TreeSitterQueriesExecutor};
 use pgt_text_size::{TextRange, TextSize};
 
 use crate::context::{
- base_parser::CompletionStatementParser,
- grant_parser::GrantParser,
- policy_parser::{PolicyParser, PolicyStmtKind},
- revoke_parser::RevokeParser,
+ base_parser::CompletionStatementParser, grant_parser::GrantParser, revoke_parser::RevokeParser,
 };
 
 #[derive(Debug, PartialEq, Eq, Hash, Clone)]
@@ -34,20 +30,15 @@ pub enum WrappingClause<'a> {
  DropColumn,
  AlterColumn,
  RenameColumn,
- PolicyName,
  ToRoleAssignment,
  SetStatement,
  AlterRole,
  DropRole,
 
- /// `PolicyCheck` refers to either the `WITH CHECK` or the `USING` clause
- /// in a policy statement.
- /// ```sql
- /// CREATE POLICY "my pol" ON PUBLIC.USERS
- /// FOR SELECT
- /// USING (...) -- this one!
- /// ```
- PolicyCheck,
+ CreatePolicy,
+ AlterPolicy,
+ DropPolicy,
+ CheckOrUsingClause,
 }
 
 #[derive(PartialEq, Eq, Hash, Debug, Clone)]
@@ -205,12 +196,7 @@ impl<'a> TreesitterContext<'a> {
  mentioned_columns: HashMap::new(),
  };
 
- // policy handling is important to Supabase, but they are a PostgreSQL specific extension,
- // so the pgt_treesitter_grammar language does not support it.
- // We infer the context manually.
- if PolicyParser::looks_like_matching_stmt(params.text) {
- ctx.gather_policy_context();
- } else if GrantParser::looks_like_matching_stmt(params.text) {
+ if GrantParser::looks_like_matching_stmt(params.text) {
  ctx.gather_grant_context();
  } else if RevokeParser::looks_like_matching_stmt(params.text) {
  ctx.gather_revoke_context();
@@ -278,43 +264,6 @@ impl<'a> TreesitterContext<'a> {
  };
  }
 
- fn gather_policy_context(&mut self) {
- let policy_context = PolicyParser::get_context(self.text, self.position);
-
- self.node_under_cursor = Some(NodeUnderCursor::CustomNode {
- text: policy_context.node_text,
- range: policy_context.node_range,
- kind: policy_context.node_kind.clone(),
- previous_node_kind: Some(policy_context.previous_node_kind),
- });
-
- if policy_context.node_kind == "policy_table" {
- self.schema_or_alias_name = policy_context.schema_name.clone();
- }
-
- if policy_context.table_name.is_some() {
- let mut new = HashSet::new();
- new.insert(policy_context.table_name.unwrap());
- self.mentioned_relations
- .insert(policy_context.schema_name, new);
- }
-
- self.wrapping_clause_type = match policy_context.node_kind.as_str() {
- "policy_name" if policy_context.statement_kind != PolicyStmtKind::Create => {
- Some(WrappingClause::PolicyName)
- }
- "policy_role" => Some(WrappingClause::ToRoleAssignment),
- "policy_table" => Some(WrappingClause::From),
- _ => {
- if policy_context.in_check_or_using_clause {
- Some(WrappingClause::PolicyCheck)
- } else {
- None
- }
- }
- };
- }
-
  fn gather_info_from_ts_queries(&mut self) {
  let stmt_range = self.wrapping_statement_range.as_ref();
  let sql = self.text;
@@ -498,13 +447,6 @@ impl<'a> TreesitterContext<'a> {
  }
  }
 
- "where" | "update" | "select" | "delete" | "from" | "join" | "column_definitions"
- | "alter_role" | "drop_role" | "set_statement" | "drop_table" | "alter_table"
- | "drop_column" | "alter_column" | "rename_column" => {
- self.wrapping_clause_type =
- self.get_wrapping_clause_from_current_node(current_node, &mut cursor);
- }
-
  "relation" | "binary_expression" | "assignment" => {
  self.wrapping_node_kind = current_node_kind.try_into().ok();
  }
@@ -518,7 +460,13 @@ impl<'a> TreesitterContext<'a> {
  }
  }
 
- _ => {}
+ _ => {
+ if let Some(clause_type) =
+ self.get_wrapping_clause_from_current_node(current_node, &mut cursor)
+ {
+ self.wrapping_clause_type = Some(clause_type);
+ }
+ }
  }
 
  // We have arrived at the leaf node
@@ -734,6 +682,10 @@ impl<'a> TreesitterContext<'a> {
  "alter_table" => Some(WrappingClause::AlterTable),
  "set_statement" => Some(WrappingClause::SetStatement),
  "column_definitions" => Some(WrappingClause::ColumnDefinitions),
+ "create_policy" => Some(WrappingClause::CreatePolicy),
+ "alter_policy" => Some(WrappingClause::AlterPolicy),
+ "drop_policy" => Some(WrappingClause::DropPolicy),
+ "check_or_using_clause" => Some(WrappingClause::CheckOrUsingClause),
  "insert" => Some(WrappingClause::Insert),
  "join" => {
  // sadly, we need to manually iterate over the children –