fix(completions): improved completion in delete/update clauses

crates/pgt_cli/src/execute/mod.rs

@@ -76,12 +76,11 @@ pub enum TraversalMode {
  Dummy,
  /// This mode is enabled when running the command `check`
  Check {
- /// The type of fixes that should be applied when analyzing a file.
- ///
- /// It's [None] if the `check` command is called without `--apply` or `--apply-suggested`
- /// arguments.
+ // The type of fixes that should be applied when analyzing a file.
+ //
+ // It's [None] if the `check` command is called without `--apply` or `--apply-suggested`
+ // arguments.
  // fix_file_mode: Option<FixFileMode>,
-
  /// An optional tuple.
  /// 1. The virtual path to the file
  /// 2. The content of the file

crates/pgt_completions/src/context.rs

@@ -30,7 +30,7 @@ impl TryFrom<&str> for ClauseType {
  match value {
  "select" => Ok(Self::Select),
  "where" => Ok(Self::Where),
- "from" | "keyword_from" => Ok(Self::From),
+ "from" => Ok(Self::From),
  "update" => Ok(Self::Update),
  "delete" => Ok(Self::Delete),
  _ => {
@@ -49,8 +49,52 @@ impl TryFrom<&str> for ClauseType {
 
 impl TryFrom<String> for ClauseType {
  type Error = String;
- fn try_from(value: String) -> Result<ClauseType, Self::Error> {
- ClauseType::try_from(value.as_str())
+ fn try_from(value: String) -> Result<Self, Self::Error> {
+ Self::try_from(value.as_str())
+ }
+}
+
+/// We can map a few nodes, such as the "update" node, to actual SQL clauses.
+/// That gives us a lot of insight for completions.
+/// Other nodes, such as the "relation" node, gives us less but still
+/// relevant information.
+/// `WrappingNode` maps to such nodes.
+///
+/// Note: This is not the direct parent of the `node_under_cursor`, but the closest
+/// *relevant* parent.
+#[derive(Debug, PartialEq, Eq)]
+pub enum WrappingNode {
+ Relation,
+ BinaryExpression,
+ Assignment,
+}
+
+impl TryFrom<&str> for WrappingNode {
+ type Error = String;
+
+ fn try_from(value: &str) -> Result<Self, Self::Error> {
+ match value {
+ "relation" => Ok(Self::Relation),
+ "assignment" => Ok(Self::Assignment),
+ "binary_expression" => Ok(Self::BinaryExpression),
+ _ => {
+ let message = format!("Unimplemented Relation: {}", value);
+
+ // Err on tests, so we notice that we're lacking an implementation immediately.
+ if cfg!(test) {
+ panic!("{}", message);
+ }
+
+ Err(message)
+ }
+ }
+ }
+}
+
+impl TryFrom<String> for WrappingNode {
+ type Error = String;
+ fn try_from(value: String) -> Result<Self, Self::Error> {
+ Self::try_from(value.as_str())
  }
 }
 
@@ -64,6 +108,9 @@ pub(crate) struct CompletionContext<'a> {
 
  pub schema_name: Option<String>,
  pub wrapping_clause_type: Option<ClauseType>,
+
+ pub wrapping_node_kind: Option<WrappingNode>,
+
  pub is_invocation: bool,
  pub wrapping_statement_range: Option<tree_sitter::Range>,
 
@@ -80,6 +127,7 @@ impl<'a> CompletionContext<'a> {
  node_under_cursor: None,
  schema_name: None,
  wrapping_clause_type: None,
+ wrapping_node_kind: None,
  wrapping_statement_range: None,
  is_invocation: false,
  mentioned_relations: HashMap::new(),
@@ -133,6 +181,15 @@ impl<'a> CompletionContext<'a> {
  })
  }
 
+ pub fn get_node_under_cursor_content(&self) -> Option<String> {
+ self.node_under_cursor
+ .and_then(|n| self.get_ts_node_content(n))
+ .and_then(|txt| match txt {
+ NodeText::Replaced => None,
+ NodeText::Original(c) => Some(c.to_string()),
+ })
+ }
+
  fn gather_tree_context(&mut self) {
  let mut cursor = self.tree.root_node().walk();
 
@@ -163,23 +220,26 @@ impl<'a> CompletionContext<'a> {
  ) {
  let current_node = cursor.node();
 
+ let parent_node_kind = parent_node.kind();
+ let current_node_kind = current_node.kind();
+
  // prevent infinite recursion – this can happen if we only have a PROGRAM node
- if current_node.kind() == parent_node.kind() {
+ if current_node_kind == parent_node_kind {
  self.node_under_cursor = Some(current_node);
  return;
  }
 
- match parent_node.kind() {
+ match parent_node_kind {
  "statement" | "subquery" => {
- self.wrapping_clause_type = current_node.kind().try_into().ok();
+ self.wrapping_clause_type = current_node_kind.try_into().ok();
  self.wrapping_statement_range = Some(parent_node.range());
  }
  "invocation" => self.is_invocation = true,
 
  _ => {}
  }
 
- match current_node.kind() {
+ match current_node_kind {
  "object_reference" => {
  let content = self.get_ts_node_content(current_node);
  if let Some(node_txt) = content {
@@ -195,13 +255,12 @@ impl<'a> CompletionContext<'a> {
  }
  }
 
- // in Treesitter, the Where clause is nested inside other clauses
- "where" => {
- self.wrapping_clause_type = "where".try_into().ok();
+ "where" | "update" | "select" | "delete" | "from" => {
+ self.wrapping_clause_type = current_node_kind.try_into().ok();
  }
 
- "keyword_from" => {
- self.wrapping_clause_type = "keyword_from".try_into().ok();
+ "relation" | "binary_expression" | "assignment" => {
+ self.wrapping_node_kind = current_node_kind.try_into().ok();
  }
 
  _ => {}
@@ -406,10 +465,6 @@ mod tests {
  ctx.get_ts_node_content(node),
  Some(NodeText::Original("from"))
  );
- assert_eq!(
- ctx.wrapping_clause_type,
- Some(crate::context::ClauseType::From)
- );
  }
 
  #[test]

crates/pgt_completions/src/providers/schemas.rs

@@ -27,8 +27,8 @@ pub fn complete_schemas<'a>(ctx: &'a CompletionContext, builder: &mut Completion
 mod tests {
 
  use crate::{
- CompletionItemKind, complete,
- test_helper::{CURSOR_POS, get_test_deps, get_test_params},
+ CompletionItemKind,
+ test_helper::{CURSOR_POS, CompletionAssertion, assert_complete_results},
  };
 
  #[tokio::test]
@@ -46,27 +46,59 @@ mod tests {
  );
  "#;
 
- let query = format!("select * from {}", CURSOR_POS);
+ assert_complete_results(
+ format!("select * from {}", CURSOR_POS).as_str(),
+ vec![
+ CompletionAssertion::LabelAndKind("public".to_string(), CompletionItemKind::Schema),
+ CompletionAssertion::LabelAndKind("auth".to_string(), CompletionItemKind::Schema),
+ CompletionAssertion::LabelAndKind(
+ "internal".to_string(),
+ CompletionItemKind::Schema,
+ ),
+ CompletionAssertion::LabelAndKind(
+ "private".to_string(),
+ CompletionItemKind::Schema,
+ ),
+ CompletionAssertion::LabelAndKind(
+ "information_schema".to_string(),
+ CompletionItemKind::Schema,
+ ),
+ CompletionAssertion::LabelAndKind(
+ "pg_catalog".to_string(),
+ CompletionItemKind::Schema,
+ ),
+ CompletionAssertion::LabelAndKind(
+ "pg_toast".to_string(),
+ CompletionItemKind::Schema,
+ ),
+ CompletionAssertion::LabelAndKind("users".to_string(), CompletionItemKind::Table),
+ ],
+ setup,
+ )
+ .await;
+ }
 
- let (tree, cache) = get_test_deps(setup, query.as_str().into()).await;
- let params = get_test_params(&tree, &cache, query.as_str().into());
- let items = complete(params);
+ #[tokio::test]
+ async fn suggests_tables_and_schemas_with_matching_keys() {
+ let setup = r#"
+ create schema ultimate;
 
- assert!(!items.is_empty());
+ -- add a table to compete against schemas
+ create table users (
+ id serial primary key,
+ name text,
+ password text
+ );
+ "#;
 
- assert_eq!(
- items
- .into_iter()
- .take(5)
- .map(|i| (i.label, i.kind))
- .collect::<Vec<(String, CompletionItemKind)>>(),
+ assert_complete_results(
+ format!("select * from u{}", CURSOR_POS).as_str(),
  vec![
- ("public".to_string(), CompletionItemKind::Schema),
- ("auth".to_string(), CompletionItemKind::Schema),
- ("internal".to_string(), CompletionItemKind::Schema),
- ("private".to_string(), CompletionItemKind::Schema),
- ("users".to_string(), CompletionItemKind::Table),
- ]
- );
+ CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
+ CompletionAssertion::LabelAndKind("ultimate".into(), CompletionItemKind::Schema),
+ ],
+ setup,
+ )
+ .await;
  }
 }

crates/pgt_completions/src/providers/tables.rs

@@ -31,7 +31,10 @@ mod tests {
 
  use crate::{
  CompletionItem, CompletionItemKind, complete,
- test_helper::{CURSOR_POS, get_test_deps, get_test_params},
+ test_helper::{
+ CURSOR_POS, CompletionAssertion, assert_complete_results, assert_no_complete_results,
+ get_test_deps, get_test_params,
+ },
  };
 
  #[tokio::test]
@@ -178,4 +181,96 @@ mod tests {
  assert_eq!(label, "coos");
  assert_eq!(kind, CompletionItemKind::Table);
  }
+
+ #[tokio::test]
+ async fn suggests_tables_in_update() {
+ let setup = r#"
+ create table coos (
+ id serial primary key,
+ name text
+ );
+ "#;
+
+ assert_complete_results(
+ format!("update {}", CURSOR_POS).as_str(),
+ vec![CompletionAssertion::LabelAndKind(
+ "public".into(),
+ CompletionItemKind::Schema,
+ )],
+ setup,
+ )
+ .await;
+
+ assert_complete_results(
+ format!("update public.{}", CURSOR_POS).as_str(),
+ vec![CompletionAssertion::LabelAndKind(
+ "coos".into(),
+ CompletionItemKind::Table,
+ )],
+ setup,
+ )
+ .await;
+
+ assert_no_complete_results(format!("update public.coos {}", CURSOR_POS).as_str(), setup)
+ .await;
+
+ assert_complete_results(
+ format!("update coos set {}", CURSOR_POS).as_str(),
+ vec![
+ CompletionAssertion::Label("id".into()),
+ CompletionAssertion::Label("name".into()),
+ ],
+ setup,
+ )
+ .await;
+
+ assert_complete_results(
+ format!("update coos set name = 'cool' where {}", CURSOR_POS).as_str(),
+ vec![
+ CompletionAssertion::Label("id".into()),
+ CompletionAssertion::Label("name".into()),
+ ],
+ setup,
+ )
+ .await;
+ }
+
+ #[tokio::test]
+ async fn suggests_tables_in_delete() {
+ let setup = r#"
+ create table coos (
+ id serial primary key,
+ name text
+ );
+ "#;
+
+ assert_no_complete_results(format!("delete {}", CURSOR_POS).as_str(), setup).await;
+
+ assert_complete_results(
+ format!("delete from {}", CURSOR_POS).as_str(),
+ vec![
+ CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
+ CompletionAssertion::LabelAndKind("coos".into(), CompletionItemKind::Table),
+ ],
+ setup,
+ )
+ .await;
+
+ assert_complete_results(
+ format!("delete from public.{}", CURSOR_POS).as_str(),
+ vec![CompletionAssertion::Label("coos".into())],
+ setup,
+ )
+ .await;
+
+ assert_complete_results(
+ format!("delete from public.coos where {}", CURSOR_POS).as_str(),
+ vec![
+ CompletionAssertion::Label("id".into()),
+ CompletionAssertion::Label("name".into()),
+ ],
+ setup,
+ )
+ .await;
+ }
 }

crates/pgt_completions/src/relevance/filtering.rs

@@ -35,6 +35,16 @@ impl CompletionFilter<'_> {
  return None;
  }
 
+ // No autocompletions if there are two identifiers without a separator.
+ if ctx.node_under_cursor.is_some_and(|n| {
+ n.prev_sibling().is_some_and(|p| {
+ (p.kind() == "identifier" || p.kind() == "object_reference")
+ && n.kind() == "identifier"
+ })
+ }) {
+ return None;
+ }
+
  Some(())
  }