feat(completions): filter invalid completion items

crates/pgt_completions/src/builder.rs

@@ -1,47 +1,80 @@
-use crate::item::CompletionItem;
+use crate::{
+    CompletionItemKind,
+    context::CompletionContext,
+    item::CompletionItem,
+    relevance::{filtering::CompletionFilter, scoring::CompletionScore},
+};
 
-pub(crate) struct CompletionBuilder {
-    items: Vec<CompletionItem>,
+pub(crate) struct PossibleCompletionItem<'a> {
+    pub label: String,
+    pub description: String,
+    pub kind: CompletionItemKind,
+    pub score: CompletionScore<'a>,
+    pub filter: CompletionFilter<'a>,
 }
 
-impl CompletionBuilder {
-    pub fn new() -> Self {
-        CompletionBuilder { items: vec![] }
+pub(crate) struct CompletionBuilder<'a> {
+    items: Vec<PossibleCompletionItem<'a>>,
+    ctx: &'a CompletionContext<'a>,
+}
+
+impl<'a> CompletionBuilder<'a> {
+    pub fn new(ctx: &'a CompletionContext) -> Self {
+        CompletionBuilder { items: vec![], ctx }
     }
 
-    pub fn add_item(&mut self, item: CompletionItem) {
+    pub fn add_item(&mut self, item: PossibleCompletionItem<'a>) {
         self.items.push(item);
     }
 
-    pub fn finish(mut self) -> Vec<CompletionItem> {
-        self.items
-            .sort_by(|a, b| b.score.cmp(&a.score).then_with(|| a.label.cmp(&b.label)));
+    pub fn finish(self) -> Vec<CompletionItem> {
+        let mut items: Vec<PossibleCompletionItem> = self
+            .items
+            .into_iter()
+            .filter(|i| i.filter.is_relevant(self.ctx).is_some())
+            .collect();
+
+        for item in items.iter_mut() {
+            item.score.calc_score(self.ctx);
+        }
 
-        self.items.dedup_by(|a, b| a.label == b.label);
-        self.items.truncate(crate::LIMIT);
+        items.sort_by(|a, b| {
+            b.score
+                .get_score()
+                .cmp(&a.score.get_score())
+                .then_with(|| a.label.cmp(&b.label))
+        });
 
-        let should_preselect_first_item = self.should_preselect_first_item();
+        items.dedup_by(|a, b| a.label == b.label);
+        items.truncate(crate::LIMIT);
 
-        self.items
+        let should_preselect_first_item = should_preselect_first_item(&items);
+
+        items
             .into_iter()
             .enumerate()
-            .map(|(idx, mut item)| {
-                if idx == 0 {
-                    item.preselected = should_preselect_first_item;
+            .map(|(idx, item)| {
+                let preselected = idx == 0 && should_preselect_first_item;
+
+                CompletionItem {
+                    description: item.description,
+                    kind: item.kind,
+                    label: item.label,
+                    preselected,
+                    score: item.score.get_score(),
                 }
-                item
             })
             .collect()
     }
+}
 
-    fn should_preselect_first_item(&mut self) -> bool {
-        let mut items_iter = self.items.iter();
-        let first = items_iter.next();
-        let second = items_iter.next();
+fn should_preselect_first_item(items: &Vec<PossibleCompletionItem>) -> bool {
+    let mut items_iter = items.iter();
+    let first = items_iter.next();
+    let second = items_iter.next();
 
-        first.is_some_and(|f| match second {
-            Some(s) => (f.score - s.score) > 10,
-            None => true,
-        })
-    }
+    first.is_some_and(|f| match second {
+        Some(s) => (f.score.get_score() - s.score.get_score()) > 10,
+        None => true,
+    })
 }

crates/pgt_completions/src/complete.rs

@@ -27,7 +27,7 @@ pub fn complete(params: CompletionParams) -> Vec<CompletionItem> {
 
     let ctx = CompletionContext::new(&sanitized_params);
 
-    let mut builder = CompletionBuilder::new();
+    let mut builder = CompletionBuilder::new(&ctx);
 
     complete_tables(&ctx, &mut builder);
     complete_functions(&ctx, &mut builder);

crates/pgt_completions/src/providers/columns.rs

@@ -1,17 +1,21 @@
 use crate::{
-    CompletionItem, CompletionItemKind, builder::CompletionBuilder, context::CompletionContext,
-    relevance::CompletionRelevanceData,
+    CompletionItemKind,
+    builder::{CompletionBuilder, PossibleCompletionItem},
+    context::CompletionContext,
+    relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
 };
 
-pub fn complete_columns(ctx: &CompletionContext, builder: &mut CompletionBuilder) {
+pub fn complete_columns<'a>(ctx: &CompletionContext<'a>, builder: &mut CompletionBuilder<'a>) {
     let available_columns = &ctx.schema_cache.columns;
 
     for col in available_columns {
-        let item = CompletionItem {
+        let relevance = CompletionRelevanceData::Column(col);
+
+        let item = PossibleCompletionItem {
             label: col.name.clone(),
-            score: CompletionRelevanceData::Column(col).get_score(ctx),
+            score: CompletionScore::from(relevance.clone()),
+            filter: CompletionFilter::from(relevance),
             description: format!("Table: {}.{}", col.schema_name, col.table_name),
-            preselected: false,
             kind: CompletionItemKind::Column,
         };
 
@@ -22,7 +26,7 @@ pub fn complete_columns(ctx: &CompletionContext, builder: &mut CompletionBuilder
 #[cfg(test)]
 mod tests {
     use crate::{
-        CompletionItem, complete,
+        CompletionItem, CompletionItemKind, complete,
         test_helper::{CURSOR_POS, InputQuery, get_test_deps, get_test_params},
     };
 
@@ -224,4 +228,100 @@ mod tests {
             "`email` not present in first four completion items."
         );
     }
+
+    #[tokio::test]
+    async fn ignores_cols_in_from_clause() {
+        let setup = r#"
+        create schema private;
+
+        create table private.users (
+            id serial primary key,
+            name text,
+            address text,
+            email text
+        );
+    "#;
+
+        let test_case = TestCase {
+            message: "suggests user created tables first",
+            query: format!(r#"select * from private.{}"#, CURSOR_POS),
+            label: "",
+            description: "",
+        };
+
+        let (tree, cache) = get_test_deps(setup, test_case.get_input_query()).await;
+        let params = get_test_params(&tree, &cache, test_case.get_input_query());
+        let results = complete(params);
+
+        assert!(
+            !results
+                .into_iter()
+                .any(|item| item.kind == CompletionItemKind::Column)
+        );
+    }
+
+    #[tokio::test]
+    async fn prefers_columns_of_mentioned_tables() {
+        let setup = r#"
+        create schema private;
+
+        create table private.users (
+            id1 serial primary key,
+            name1 text,
+            address1 text,
+            email1 text
+        );
+
+        create table public.users (
+            id2 serial primary key,
+            name2 text,
+            address2 text,
+            email2 text
+        );
+    "#;
+
+        {
+            let test_case = TestCase {
+                message: "",
+                query: format!(r#"select {} from users"#, CURSOR_POS),
+                label: "suggests from table",
+                description: "",
+            };
+
+            let (tree, cache) = get_test_deps(setup, test_case.get_input_query()).await;
+            let params = get_test_params(&tree, &cache, test_case.get_input_query());
+            let results = complete(params);
+
+            assert_eq!(
+                results
+                    .into_iter()
+                    .take(4)
+                    .map(|item| item.label)
+                    .collect::<Vec<String>>(),
+                vec!["address2", "email2", "id2", "name2"]
+            );
+        }
+
+        {
+            let test_case = TestCase {
+                message: "",
+                query: format!(r#"select {} from private.users"#, CURSOR_POS),
+                label: "suggests from table",
+                description: "",
+            };
+
+            let (tree, cache) = get_test_deps(setup, test_case.get_input_query()).await;
+            let params = get_test_params(&tree, &cache, test_case.get_input_query());
+            let results = complete(params);
+
+            assert_eq!(
+                results
+                    .into_iter()
+                    .take(4)
+                    .map(|item| item.label)
+                    .collect::<Vec<String>>(),
+                vec!["address1", "email1", "id1", "name1"]
+            );
+        }
+    }
 }

crates/pgt_completions/src/providers/functions.rs

@@ -1,17 +1,21 @@
 use crate::{
-    CompletionItem, CompletionItemKind, builder::CompletionBuilder, context::CompletionContext,
-    relevance::CompletionRelevanceData,
+    CompletionItemKind,
+    builder::{CompletionBuilder, PossibleCompletionItem},
+    context::CompletionContext,
+    relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
 };
 
-pub fn complete_functions(ctx: &CompletionContext, builder: &mut CompletionBuilder) {
+pub fn complete_functions<'a>(ctx: &'a CompletionContext, builder: &mut CompletionBuilder<'a>) {
     let available_functions = &ctx.schema_cache.functions;
 
     for func in available_functions {
-        let item = CompletionItem {
+        let relevance = CompletionRelevanceData::Function(func);
+
+        let item = PossibleCompletionItem {
             label: func.name.clone(),
-            score: CompletionRelevanceData::Function(func).get_score(ctx),
+            score: CompletionScore::from(relevance.clone()),
+            filter: CompletionFilter::from(relevance),
             description: format!("Schema: {}", func.schema),
-            preselected: false,
             kind: CompletionItemKind::Function,
         };
 

crates/pgt_completions/src/providers/schemas.rs

@@ -1,20 +1,21 @@
 use crate::{
-    CompletionItem, builder::CompletionBuilder, context::CompletionContext,
-    relevance::CompletionRelevanceData,
+    builder::{CompletionBuilder, PossibleCompletionItem},
+    context::CompletionContext,
+    relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
 };
 
-pub fn complete_schemas(ctx: &CompletionContext, builder: &mut CompletionBuilder) {
+pub fn complete_schemas<'a>(ctx: &'a CompletionContext, builder: &mut CompletionBuilder<'a>) {
     let available_schemas = &ctx.schema_cache.schemas;
 
     for schema in available_schemas {
         let relevance = CompletionRelevanceData::Schema(schema);
 
-        let item = CompletionItem {
+        let item = PossibleCompletionItem {
             label: schema.name.clone(),
             description: "Schema".into(),
-            preselected: false,
             kind: crate::CompletionItemKind::Schema,
-            score: relevance.get_score(ctx),
+            score: CompletionScore::from(relevance.clone()),
+            filter: CompletionFilter::from(relevance),
         };
 
         builder.add_item(item);

crates/pgt_completions/src/providers/tables.rs

@@ -1,19 +1,21 @@
 use crate::{
-    builder::CompletionBuilder,
+    builder::{CompletionBuilder, PossibleCompletionItem},
     context::CompletionContext,
-    item::{CompletionItem, CompletionItemKind},
-    relevance::CompletionRelevanceData,
+    item::CompletionItemKind,
+    relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
 };
 
-pub fn complete_tables(ctx: &CompletionContext, builder: &mut CompletionBuilder) {
+pub fn complete_tables<'a>(ctx: &'a CompletionContext, builder: &mut CompletionBuilder<'a>) {
     let available_tables = &ctx.schema_cache.tables;
 
     for table in available_tables {
-        let item = CompletionItem {
+        let relevance = CompletionRelevanceData::Table(table);
+
+        let item = PossibleCompletionItem {
             label: table.name.clone(),
-            score: CompletionRelevanceData::Table(table).get_score(ctx),
+            score: CompletionScore::from(relevance.clone()),
+            filter: CompletionFilter::from(relevance),
             description: format!("Schema: {}", table.schema),
-            preselected: false,
             kind: CompletionItemKind::Table,
         };