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(())
     }