ok

Author: juleswritescode

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(),
@@ -163,23 +211,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 +246,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 +456,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/helper.rs

@@ -7,9 +7,7 @@ pub(crate) fn get_completion_text_with_schema(
     item_name: &str,
     item_schema_name: &str,
 ) -> Option<CompletionText> {
-    if item_schema_name == "public" {
-        None
-    } else if ctx.schema_name.is_some() {
+    if item_schema_name == "public" || ctx.schema_name.is_some() {
         None
     } else {
         let node = ctx.node_under_cursor.unwrap();

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

crates/pgt_completions/src/relevance/scoring.rs

@@ -1,4 +1,4 @@
-use crate::context::{ClauseType, CompletionContext, NodeText};
+use crate::context::{ClauseType, CompletionContext, NodeText, WrappingNode};
 
 use super::CompletionRelevanceData;
 
@@ -28,6 +28,7 @@ impl CompletionScore<'_> {
         self.check_matches_query_input(ctx);
         self.check_is_invocation(ctx);
         self.check_matching_clause_type(ctx);
+        self.check_matching_wrapping_node(ctx);
         self.check_relations_in_stmt(ctx);
     }
 
@@ -96,6 +97,36 @@ impl CompletionScore<'_> {
         }
     }
 
+    fn check_matching_wrapping_node(&mut self, ctx: &CompletionContext) {
+        let wrapping_node = match ctx.wrapping_node_kind.as_ref() {
+            None => return,
+            Some(wn) => wn,
+        };
+
+        let has_schema = ctx.schema_name.is_some();
+
+        self.score += match self.data {
+            CompletionRelevanceData::Table(_) => match wrapping_node {
+                WrappingNode::Relation => 15,
+                WrappingNode::BinaryExpression => 5,
+                _ => -50,
+            },
+            CompletionRelevanceData::Function(_) => match wrapping_node {
+                WrappingNode::Relation => 10,
+                _ => -50,
+            },
+            CompletionRelevanceData::Column(_) => match wrapping_node {
+                WrappingNode::BinaryExpression => 15,
+                WrappingNode::Assignment => 15,
+                _ => -15,
+            },
+            CompletionRelevanceData::Schema(_) => match wrapping_node {
+                WrappingNode::Relation if !has_schema => 5,
+                _ => -50,
+            },
+        }
+    }
+
     fn check_is_invocation(&mut self, ctx: &CompletionContext) {
         self.score += match self.data {
             CompletionRelevanceData::Function(_) if ctx.is_invocation => 30,

crates/pgt_completions/src/test_helper.rs

@@ -2,7 +2,7 @@ use pgt_schema_cache::SchemaCache;
 use pgt_test_utils::test_database::get_new_test_db;
 use sqlx::Executor;
 
-use crate::CompletionParams;
+use crate::{CompletionItem, CompletionItemKind, CompletionParams, complete};
 
 pub static CURSOR_POS: char = '€';
 
@@ -139,3 +139,48 @@ mod tests {
         }
     }
 }
+
+#[derive(Debug, PartialEq, Eq)]
+pub(crate) enum CompletionAssertion {
+    Label(String),
+    LabelAndKind(String, CompletionItemKind),
+}
+
+impl CompletionAssertion {
+    fn assert_eq(self, item: CompletionItem) {
+        match self {
+            CompletionAssertion::Label(label) => {
+                assert_eq!(item.label, label);
+            }
+            CompletionAssertion::LabelAndKind(label, kind) => {
+                assert_eq!(item.label, label);
+                assert_eq!(item.kind, kind);
+            }
+        }
+    }
+}
+
+pub(crate) async fn assert_complete_results(
+    query: &str,
+    assertions: Vec<CompletionAssertion>,
+    setup: &str,
+) {
+    let (tree, cache) = get_test_deps(setup, query.into()).await;
+    let params = get_test_params(&tree, &cache, query.into());
+    let items = complete(params);
+
+    assertions
+        .into_iter()
+        .zip(items.into_iter())
+        .for_each(|(assertion, result)| {
+            assertion.assert_eq(result);
+        });
+}
+
+pub(crate) async fn assert_no_complete_results(query: &str, setup: &str) {
+    let (tree, cache) = get_test_deps(setup, query.into()).await;
+    let params = get_test_params(&tree, &cache, query.into());
+    let items = complete(params);
+
+    assert_eq!(items.len(), 0)
+}

crates/pgt_workspace/src/features/code_actions.rs

@@ -46,7 +46,6 @@ pub struct CommandAction {
 
 #[derive(Debug, serde::Serialize, serde::Deserialize, strum::EnumIter)]
 #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
-
 pub enum CommandActionCategory {
     ExecuteStatement(StatementId),
 }