refactor(treesitter): rely more on treesitter grammar than business logic (#562)

Author: juleswritescode

crates/pgt_completions/src/providers/roles.rs

@@ -292,7 +292,7 @@ mod tests {
  QueryWithCursorPosition::cursor_marker()
  ),
  format!(
- "revoke all on table userse from owner, {}",
+ "revoke all on table users from owner, {}",
  QueryWithCursorPosition::cursor_marker()
  ),
  ];

crates/pgt_completions/src/relevance/filtering.rs

@@ -1,5 +1,5 @@
 use pgt_schema_cache::ProcKind;
-use pgt_treesitter::context::{NodeUnderCursor, TreesitterContext, WrappingClause, WrappingNode};
+use pgt_treesitter::context::{TreesitterContext, WrappingClause, WrappingNode};
 
 use super::CompletionRelevanceData;
 
@@ -17,7 +17,11 @@ impl<'a> From<CompletionRelevanceData<'a>> for CompletionFilter<'a> {
 impl CompletionFilter<'_> {
  pub fn is_relevant(&self, ctx: &TreesitterContext) -> Option<()> {
  self.completable_context(ctx)?;
- self.check_clause(ctx)?;
+
+ self.check_node_type(ctx)
+ // we want to rely on treesitter more, so checking the clause is a fallback
+ .or_else(|| self.check_clause(ctx))?;
+
  self.check_invocation(ctx)?;
  self.check_mentioned_schema_or_alias(ctx)?;
 
@@ -67,30 +71,42 @@ impl CompletionFilter<'_> {
  }
 
  // No autocompletions if there are two identifiers without a separator.
- if ctx.node_under_cursor.as_ref().is_some_and(|n| match n {
- NodeUnderCursor::TsNode(node) => node.prev_sibling().is_some_and(|p| {
- (p.kind() == "identifier" || p.kind() == "object_reference")
- && n.kind() == "identifier"
- }),
- NodeUnderCursor::CustomNode { .. } => false,
+ if ctx.node_under_cursor.as_ref().is_some_and(|node| {
+ node.prev_sibling().is_some_and(|p| {
+ (p.kind() == "any_identifier" || p.kind() == "object_reference")
+ && node.kind() == "any_identifier"
+ })
  }) {
  return None;
  }
 
  // no completions if we're right after an asterisk:
  // `select * {}`
- if ctx.node_under_cursor.as_ref().is_some_and(|n| match n {
- NodeUnderCursor::TsNode(node) => node
- .prev_sibling()
- .is_some_and(|p| (p.kind() == "all_fields") && n.kind() == "identifier"),
- NodeUnderCursor::CustomNode { .. } => false,
+ if ctx.node_under_cursor.as_ref().is_some_and(|node| {
+ node.prev_sibling()
+ .is_some_and(|p| (p.kind() == "all_fields") && node.kind() == "any_identifier")
  }) {
  return None;
  }
 
  Some(())
  }
 
+ fn check_node_type(&self, ctx: &TreesitterContext) -> Option<()> {
+ let kind = ctx.node_under_cursor.as_ref().map(|n| n.kind())?;
+
+ let is_allowed = match kind {
+ "column_identifier" => {
+ matches!(self.data, CompletionRelevanceData::Column(_))
+ && !ctx.matches_ancestor_history(&["insert_values", "field"])
+ && !ctx.node_under_cursor_is_within_field_name("binary_expr_right")
+ }
+ _ => false,
+ };
+
+ if is_allowed { Some(()) } else { None }
+ }
+
  fn check_clause(&self, ctx: &TreesitterContext) -> Option<()> {
  ctx.wrapping_clause_type
  .as_ref()
@@ -99,6 +115,9 @@ impl CompletionFilter<'_> {
  CompletionRelevanceData::Table(_) => match clause {
  WrappingClause::From | WrappingClause::Update => true,
 
+ WrappingClause::RevokeStatement | WrappingClause::GrantStatement => ctx
+ .matches_ancestor_history(&["grantable_on_table", "object_reference"]),
+
  WrappingClause::Join { on_node: None } => true,
  WrappingClause::Join { on_node: Some(on) } => ctx
  .node_under_cursor
@@ -202,6 +221,14 @@ impl CompletionFilter<'_> {
  | WrappingClause::Update
  | WrappingClause::Delete => true,
 
+ WrappingClause::RevokeStatement | WrappingClause::GrantStatement => {
+ (ctx.matches_ancestor_history(&[
+ "grantable_on_table",
+ "object_reference",
+ ]) && ctx.schema_or_alias_name.is_none())
+ || ctx.matches_ancestor_history(&["grantable_on_all"])
+ }
+
  WrappingClause::Where => {
  ctx.before_cursor_matches_kind(&["keyword_and", "keyword_where"])
  }
@@ -238,13 +265,23 @@ impl CompletionFilter<'_> {
  }
 
  CompletionRelevanceData::Role(_) => match clause {
- WrappingClause::DropRole
- | WrappingClause::AlterRole
- | WrappingClause::ToRoleAssignment => true,
+ WrappingClause::DropRole | WrappingClause::AlterRole => true,
 
  WrappingClause::SetStatement => ctx
  .before_cursor_matches_kind(&["keyword_role", "keyword_authorization"]),
 
+ WrappingClause::RevokeStatement | WrappingClause::GrantStatement => {
+ ctx.matches_ancestor_history(&["role_specification"])
+ || ctx.node_under_cursor.as_ref().is_some_and(|k| {
+ k.kind() == "any_identifier"
+ && ctx.before_cursor_matches_kind(&[
+ "keyword_grant",
+ "keyword_revoke",
+ "keyword_for",
+ ])
+ })
+ }
+
  WrappingClause::AlterPolicy | WrappingClause::CreatePolicy => {
  ctx.before_cursor_matches_kind(&["keyword_to"])
  && ctx.matches_ancestor_history(&["policy_to_role"])

crates/pgt_hover/src/hovered_node.rs

@@ -32,7 +32,11 @@ impl HoveredNode {
  let under_cursor = ctx.node_under_cursor.as_ref()?;
 
  match under_cursor.kind() {
- "identifier" if ctx.matches_ancestor_history(&["relation", "object_reference"]) => {
+ "any_identifier"
+ if ctx.matches_ancestor_history(&["relation", "object_reference"])
+ || ctx
+ .matches_ancestor_history(&["grantable_on_table", "object_reference"]) =>
+ {
  let num_sibs = ctx.num_siblings();
  if ctx.node_under_cursor_is_nth_child(1) && num_sibs > 0 {
  return Some(HoveredNode::Schema(NodeIdentification::Name(node_content)));
@@ -48,7 +52,7 @@ impl HoveredNode {
  }
  }
 
- "identifier"
+ "any_identifier"
  if ctx.matches_ancestor_history(&["object_reference"])
  && ctx.wrapping_clause_type.as_ref().is_some_and(|clause| {
  matches!(
@@ -69,7 +73,7 @@ impl HoveredNode {
  }
  }
 
- "identifier" if ctx.matches_ancestor_history(&["field"]) => {
+ "column_identifier" => {
  if let Some(table_or_alias) = ctx.schema_or_alias_name.as_ref() {
  Some(HoveredNode::Column(NodeIdentification::SchemaAndName((
  table_or_alias.clone(),
@@ -80,7 +84,9 @@ impl HoveredNode {
  }
  }
 
- "identifier" if ctx.matches_ancestor_history(&["invocation", "object_reference"]) => {
+ "any_identifier"
+ if ctx.matches_ancestor_history(&["invocation", "object_reference"]) =>
+ {
  if let Some(schema) = ctx.schema_or_alias_name.as_ref() {
  Some(HoveredNode::Function(NodeIdentification::SchemaAndName((
  schema.clone(),
@@ -93,11 +99,20 @@ impl HoveredNode {
  }
  }
 
- "identifier" if ctx.matches_one_of_ancestors(&["alter_role", "policy_to_role"]) => {
+ "any_identifier"
+ if ctx.matches_one_of_ancestors(&[
+ "alter_role",
+ "policy_to_role",
+ "role_specification",
+ ]) || ctx.before_cursor_matches_kind(&["keyword_revoke"]) =>
+ {
+ Some(HoveredNode::Role(NodeIdentification::Name(node_content)))
+ }
+ "grant_role" | "policy_role" => {
  Some(HoveredNode::Role(NodeIdentification::Name(node_content)))
  }
 
- "identifier"
+ "any_identifier"
  if (
  // hover over custom type in `create table` or `returns`
  (ctx.matches_ancestor_history(&["type", "object_reference"])
@@ -127,16 +142,12 @@ impl HoveredNode {
  }
  }
 
- "revoke_role" | "grant_role" | "policy_role" => {
- Some(HoveredNode::Role(NodeIdentification::Name(node_content)))
- }
-
  // quoted columns
  "literal" if ctx.matches_ancestor_history(&["select_expression", "term"]) => {
  Some(HoveredNode::Column(NodeIdentification::Name(node_content)))
  }
 
- "revoke_table" | "grant_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/base_parser.rs

@@ -14,12 +14,6 @@ impl TokenNavigator {
  .is_some_and(|c| options.contains(&c.get_word_without_quotes().as_str()))
  }
 
- pub(crate) fn prev_matches(&self, options: &[&str]) -> bool {
- self.previous_token
- .as_ref()
- .is_some_and(|t| options.contains(&t.get_word_without_quotes().as_str()))
- }
-
  pub(crate) fn advance(&mut self) -> Option<WordWithIndex> {
  // we can't peek back n an iterator, so we'll have to keep track manually.
  self.previous_token = self.current_token.take();