not sure about this

juleswritescode · juleswritescode · commit d9f1c8d3d662 · 2025-05-19T09:58:23.000+02:00
diff --git a/crates/pgt_completions/src/context.rs b/crates/pgt_completions/src/context.rs
@@ -1,4 +1,7 @@
-use std::collections::{HashMap, HashSet};
+use std::{
+    cmp,
+    collections::{HashMap, HashSet},
+};
 
 use pgt_schema_cache::SchemaCache;
 use pgt_treesitter_queries::{
@@ -8,7 +11,7 @@ use pgt_treesitter_queries::{
 
 use crate::sanitization::SanitizedCompletionParams;
 
-#[derive(Debug, PartialEq, Eq, Hash)]
+#[derive(Debug, PartialEq, Eq, Hash, Clone)]
 pub enum WrappingClause<'a> {
     Select,
     Where,
@@ -19,6 +22,7 @@ pub enum WrappingClause<'a> {
     Update,
     Delete,
     ColumnDefinitions,
+    Insert,
 }
 
 #[derive(PartialEq, Eq, Debug)]
@@ -46,6 +50,7 @@ pub enum WrappingNode {
     Relation,
     BinaryExpression,
     Assignment,
+    List,
 }
 
 impl TryFrom<&str> for WrappingNode {
@@ -56,6 +61,7 @@ impl TryFrom<&str> for WrappingNode {
             "relation" => Ok(Self::Relation),
             "assignment" => Ok(Self::Assignment),
             "binary_expression" => Ok(Self::BinaryExpression),
+            "list" => Ok(Self::List),
             _ => {
                 let message = format!("Unimplemented Relation: {}", value);
 
@@ -142,13 +148,6 @@ impl<'a> CompletionContext<'a> {
         ctx.gather_tree_context();
         ctx.gather_info_from_ts_queries();
 
-        if cfg!(test) {
-            println!("{:#?}", ctx.wrapping_clause_type);
-            println!("{:#?}", ctx.wrapping_node_kind);
-            println!("{:#?}", ctx.is_in_error_node);
-            println!("{:#?}", ctx.text);
-        }
-
         ctx
     }
 
@@ -240,10 +239,20 @@ impl<'a> CompletionContext<'a> {
          * `select * from use           {}` becomes `select * from use{}`.
          */
         let current_node = cursor.node();
-        while cursor.goto_first_child_for_byte(self.position).is_none() && self.position > 0 {
-            self.position -= 1;
+
+        let mut chars = self.text.chars();
+
+        if chars
+            .nth(self.position)
+            .is_some_and(|c| !c.is_ascii_whitespace())
+        {
+            self.position = cmp::min(self.position + 1, self.text.len());
+        } else {
+            self.position = cmp::min(self.position, self.text.len());
         }
 
+        cursor.goto_first_child_for_byte(self.position);
+
         self.gather_context_from_node(cursor, current_node);
     }
 
@@ -276,23 +285,11 @@ impl<'a> CompletionContext<'a> {
 
         // try to gather context from the siblings if we're within an error node.
         if self.is_in_error_node {
-            let mut next_sibling = current_node.next_named_sibling();
-            while let Some(n) = next_sibling {
-                if let Some(clause_type) = self.get_wrapping_clause_from_keyword_node(n) {
-                    self.wrapping_clause_type = Some(clause_type);
-                    break;
-                } else {
-                    next_sibling = n.next_named_sibling();
-                }
+            if let Some(clause_type) = self.get_wrapping_clause_from_siblings(current_node) {
+                self.wrapping_clause_type = Some(clause_type);
             }
-            let mut prev_sibling = current_node.prev_named_sibling();
-            while let Some(n) = prev_sibling {
-                if let Some(clause_type) = self.get_wrapping_clause_from_keyword_node(n) {
-                    self.wrapping_clause_type = Some(clause_type);
-                    break;
-                } else {
-                    prev_sibling = n.prev_named_sibling();
-                }
+            if let Some(wrapping_node) = self.get_wrapping_node_from_siblings(current_node) {
+                self.wrapping_node_kind = Some(wrapping_node)
             }
         }
 
@@ -317,7 +314,7 @@ impl<'a> CompletionContext<'a> {
                     self.get_wrapping_clause_from_current_node(current_node, &mut cursor);
             }
 
-            "relation" | "binary_expression" | "assignment" => {
+            "relation" | "binary_expression" | "assignment" | "list" => {
                 self.wrapping_node_kind = current_node_kind.try_into().ok();
             }
 
@@ -338,31 +335,89 @@ impl<'a> CompletionContext<'a> {
         self.gather_context_from_node(cursor, current_node);
     }
 
-    fn get_wrapping_clause_from_keyword_node(
+    fn get_first_sibling(&self, node: tree_sitter::Node<'a>) -> tree_sitter::Node<'a> {
+        let mut first_sibling = node;
+        while let Some(n) = first_sibling.prev_sibling() {
+            first_sibling = n;
+        }
+        first_sibling
+    }
+
+    fn get_wrapping_node_from_siblings(&self, node: tree_sitter::Node<'a>) -> Option<WrappingNode> {
+        self.wrapping_clause_type
+            .as_ref()
+            .and_then(|clause| match clause {
+                WrappingClause::Insert => {
+                    if node.prev_sibling().is_some_and(|n| n.kind() == "(")
+                        || node.next_sibling().is_some_and(|n| n.kind() == ")")
+                    {
+                        Some(WrappingNode::List)
+                    } else {
+                        None
+                    }
+                }
+                _ => None,
+            })
+    }
+
+    fn get_wrapping_clause_from_siblings(
         &self,
         node: tree_sitter::Node<'a>,
     ) -> Option<WrappingClause<'a>> {
-        if node.kind().starts_with("keyword_") {
-            if let Some(txt) = self.get_ts_node_content(node).and_then(|txt| match txt {
-                NodeText::Original(txt) => Some(txt),
-                NodeText::Replaced => None,
-            }) {
-                match txt {
-                    "where" => return Some(WrappingClause::Where),
-                    "update" => return Some(WrappingClause::Update),
-                    "select" => return Some(WrappingClause::Select),
-                    "delete" => return Some(WrappingClause::Delete),
-                    "from" => return Some(WrappingClause::From),
-                    "join" => {
-                        // TODO: not sure if we can infer it here.
-                        return Some(WrappingClause::Join { on_node: None });
+        let clause_combinations: Vec<(WrappingClause, &[&'static str])> = vec![
+            (WrappingClause::Where, &["where"]),
+            (WrappingClause::Update, &["update"]),
+            (WrappingClause::Select, &["select"]),
+            (WrappingClause::Delete, &["delete"]),
+            (WrappingClause::Insert, &["insert", "into"]),
+            (WrappingClause::From, &["from"]),
+            (WrappingClause::Join { on_node: None }, &["join"]),
+        ];
+
+        let first_sibling = self.get_first_sibling(node);
+
+        /*
+         * For each clause, we'll iterate from first_sibling to the next ones,
+         * either until the end or until we land on the node under the cursor.
+         * We'll score the `WrappingClause` by how many tokens it matches in order.
+         */
+        let mut clauses_with_score: Vec<(WrappingClause, usize)> = clause_combinations
+            .into_iter()
+            .map(|(clause, tokens)| {
+                let mut idx = 0;
+
+                let mut sibling = Some(first_sibling);
+                while let Some(sib) = sibling {
+                    if sib.end_byte() >= node.end_byte() || idx >= tokens.len() {
+                        break;
                     }
-                    _ => {}
+
+                    if let Some(sibling_content) =
+                        self.get_ts_node_content(sib).and_then(|txt| match txt {
+                            NodeText::Original(txt) => Some(txt),
+                            NodeText::Replaced => None,
+                        })
+                    {
+                        if sibling_content == tokens[idx] {
+                            idx += 1;
+                        }
+                    } else {
+                        break;
+                    }
+
+                    sibling = sib.next_sibling();
                 }
-            };
-        }
 
-        None
+                (clause, idx)
+            })
+            .collect();
+
+        clauses_with_score.sort_by(|(_, score_a), (_, score_b)| score_b.cmp(score_a));
+        clauses_with_score
+            .iter()
+            .filter(|(_, score)| *score > 0)
+            .next()
+            .map(|c| c.0.clone())
     }
 
     fn get_wrapping_clause_from_current_node(
@@ -377,6 +432,7 @@ impl<'a> CompletionContext<'a> {
             "delete" => Some(WrappingClause::Delete),
             "from" => Some(WrappingClause::From),
             "column_definitions" => Some(WrappingClause::ColumnDefinitions),
+            "insert" => Some(WrappingClause::Insert),
             "join" => {
                 // sadly, we need to manually iterate over the children –
                 // `node.child_by_field_id(..)` does not work as expected
diff --git a/crates/pgt_completions/src/providers/columns.rs b/crates/pgt_completions/src/providers/columns.rs
@@ -573,4 +573,24 @@ mod tests {
         )
         .await;
     }
+
+    #[tokio::test]
+    async fn suggests_columns_in_insert_clause() {
+        let setup = r#"
+            create table instruments (
+                id bigint primary key generated always as identity,
+                name text not null
+            );
+        "#;
+
+        assert_complete_results(
+            format!("insert into instruments ({})", CURSOR_POS).as_str(),
+            vec![
+                CompletionAssertion::Label("id".to_string()),
+                CompletionAssertion::Label("name".to_string()),
+            ],
+            setup,
+        )
+        .await;
+    }
 }
diff --git a/crates/pgt_completions/src/relevance/filtering.rs b/crates/pgt_completions/src/relevance/filtering.rs
@@ -1,4 +1,4 @@
-use crate::context::{CompletionContext, WrappingClause};
+use crate::context::{CompletionContext, WrappingClause, WrappingNode};
 
 use super::CompletionRelevanceData;
 
@@ -73,6 +73,17 @@ impl CompletionFilter<'_> {
                         WrappingClause::Select
                         | WrappingClause::Where
                         | WrappingClause::ColumnDefinitions => false,
+
+                        WrappingClause::Insert => {
+                            ctx.wrapping_node_kind
+                                .as_ref()
+                                .is_some_and(|n| n != &WrappingNode::List)
+                                && ctx.node_under_cursor.is_some_and(|n| {
+                                    n.prev_sibling()
+                                        .is_some_and(|sib| sib.kind() == "keyword_into")
+                                })
+                        }
+
                         _ => true,
                     },
                     CompletionRelevanceData::Column(_) => {
@@ -88,6 +99,11 @@ impl CompletionFilter<'_> {
                             // we are in a JOIN, but definitely not after an ON
                             WrappingClause::Join { on_node: None } => false,
 
+                            WrappingClause::Insert => ctx
+                                .wrapping_node_kind
+                                .as_ref()
+                                .is_some_and(|n| n == &WrappingNode::List),
+
                             _ => true,
                         }
                     }
@@ -107,6 +123,16 @@ impl CompletionFilter<'_> {
                         | WrappingClause::Update
                         | WrappingClause::Delete => true,
 
+                        WrappingClause::Insert => {
+                            ctx.wrapping_node_kind
+                                .as_ref()
+                                .is_some_and(|n| n != &WrappingNode::List)
+                                && ctx.node_under_cursor.is_some_and(|n| {
+                                    n.prev_sibling()
+                                        .is_some_and(|sib| sib.kind() == "keyword_into")
+                                })
+                        }
+
                         WrappingClause::ColumnDefinitions => false,
                     },
                 }
diff --git a/crates/pgt_completions/src/sanitization.rs b/crates/pgt_completions/src/sanitization.rs
@@ -25,6 +25,7 @@ where
             || cursor_prepared_to_write_token_after_last_node(params.tree, params.position)
             || cursor_before_semicolon(params.tree, params.position)
             || cursor_on_a_dot(&params.text, params.position)
+            || cursor_between_parentheses(&params.text, params.position)
         {
             SanitizedCompletionParams::with_adjusted_sql(params)
         } else {
@@ -200,13 +201,19 @@ fn cursor_before_semicolon(tree: &tree_sitter::Tree, position: TextSize) -> bool
         .unwrap_or(false)
 }
 
+fn cursor_between_parentheses(sql: &str, position: TextSize) -> bool {
+    let position: usize = position.into();
+    let mut chars = sql.chars();
+    chars.nth(position - 1).is_some_and(|c| c == '(') && chars.next().is_some_and(|c| c == ')')
+}
+
 #[cfg(test)]
 mod tests {
     use pgt_text_size::TextSize;
 
     use crate::sanitization::{
-        cursor_before_semicolon, cursor_inbetween_nodes, cursor_on_a_dot,
-        cursor_prepared_to_write_token_after_last_node,
+        cursor_before_semicolon, cursor_between_parentheses, cursor_inbetween_nodes,
+        cursor_on_a_dot, cursor_prepared_to_write_token_after_last_node,
     };
 
     #[test]
@@ -317,4 +324,18 @@ mod tests {
         assert!(cursor_before_semicolon(&tree, TextSize::new(16)));
         assert!(cursor_before_semicolon(&tree, TextSize::new(17)));
     }
+
+    #[test]
+    fn between_parentheses() {
+        let input = "insert into instruments ()";
+
+        // insert into (|) <- right in the parentheses
+        assert!(cursor_between_parentheses(input, TextSize::new(25)));
+
+        // insert into ()| <- too late
+        assert!(!cursor_between_parentheses(input, TextSize::new(26)));
+
+        // insert into |() <- too early
+        assert!(!cursor_between_parentheses(input, TextSize::new(24)));
+    }
 }
