fixie fixie

Author: juleswritescode

crates/pgt_completions/src/context.rs

@@ -181,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();
 

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/relevance/scoring.rs

@@ -1,4 +1,4 @@
-use crate::context::{ClauseType, CompletionContext, NodeText, WrappingNode};
+use crate::context::{ClauseType, CompletionContext, WrappingNode};
 
 use super::CompletionRelevanceData;
 
@@ -33,16 +33,8 @@ impl CompletionScore<'_> {
     }
 
     fn check_matches_query_input(&mut self, ctx: &CompletionContext) {
-        let node = match ctx.node_under_cursor {
-            Some(node) => node,
-            None => return,
-        };
-
-        let content = match ctx.get_ts_node_content(node) {
-            Some(c) => match c {
-                NodeText::Original(s) => s,
-                NodeText::Replaced => return,
-            },
+        let content = match ctx.get_node_under_cursor_content() {
+            Some(c) => c,
             None => return,
         };
 
@@ -53,7 +45,7 @@ impl CompletionScore<'_> {
             CompletionRelevanceData::Schema(s) => s.name.as_str(),
         };
 
-        if name.starts_with(content) {
+        if name.starts_with(content.as_str()) {
             let len: i32 = content
                 .len()
                 .try_into()
@@ -70,12 +62,13 @@ impl CompletionScore<'_> {
         };
 
         let has_mentioned_tables = !ctx.mentioned_relations.is_empty();
+        let has_mentioned_schema = ctx.schema_name.is_some();
 
         self.score += match self.data {
             CompletionRelevanceData::Table(_) => match clause_type {
                 ClauseType::From => 5,
-                ClauseType::Update => 15,
-                ClauseType::Delete => 15,
+                ClauseType::Update => 10,
+                ClauseType::Delete => 10,
                 _ => -50,
             },
             CompletionRelevanceData::Function(_) => match clause_type {
@@ -91,7 +84,9 @@ impl CompletionScore<'_> {
                 _ => -15,
             },
             CompletionRelevanceData::Schema(_) => match clause_type {
-                ClauseType::From => 10,
+                ClauseType::From if !has_mentioned_schema => 15,
+                ClauseType::Update if !has_mentioned_schema => 15,
+                ClauseType::Delete if !has_mentioned_schema => 15,
                 _ => -50,
             },
         }
@@ -103,11 +98,13 @@ impl CompletionScore<'_> {
             Some(wn) => wn,
         };
 
-        let has_schema = ctx.schema_name.is_some();
+        let has_mentioned_schema = ctx.schema_name.is_some();
+        let has_node_text = ctx.get_node_under_cursor_content().is_some();
 
         self.score += match self.data {
             CompletionRelevanceData::Table(_) => match wrapping_node {
-                WrappingNode::Relation => 15,
+                WrappingNode::Relation if has_mentioned_schema => 15,
+                WrappingNode::Relation if !has_mentioned_schema => 10,
                 WrappingNode::BinaryExpression => 5,
                 _ => -50,
             },
@@ -121,7 +118,8 @@ impl CompletionScore<'_> {
                 _ => -15,
             },
             CompletionRelevanceData::Schema(_) => match wrapping_node {
-                WrappingNode::Relation if !has_schema => 5,
+                WrappingNode::Relation if !has_mentioned_schema && !has_node_text => 15,
+                WrappingNode::Relation if !has_mentioned_schema && has_node_text => 0,
                 _ => -50,
             },
         }