feat(completions): lower priority of already mentioned columns in SELECT

crates/pgt_completions/src/context.rs

@@ -8,7 +8,7 @@ use pgt_treesitter_queries::{
 
 use crate::sanitization::SanitizedCompletionParams;
 
-#[derive(Debug, PartialEq, Eq)]
+#[derive(Debug, PartialEq, Eq, Hash)]
 pub enum WrappingClause<'a> {
     Select,
     Where,
@@ -26,6 +26,12 @@ pub(crate) enum NodeText<'a> {
     Original(&'a str),
 }
 
+#[derive(PartialEq, Eq, Hash, Debug)]
+pub(crate) struct MentionedColumn {
+    pub(crate) column: String,
+    pub(crate) alias: Option<String>,
+}
+
 /// 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
@@ -108,8 +114,8 @@ pub(crate) struct CompletionContext<'a> {
     pub is_in_error_node: bool,
 
     pub mentioned_relations: HashMap<Option<String>, HashSet<String>>,
-
     pub mentioned_table_aliases: HashMap<String, String>,
+    pub mentioned_columns: HashMap<Option<WrappingClause<'a>>, HashSet<MentionedColumn>>,
 }
 
 impl<'a> CompletionContext<'a> {
@@ -127,6 +133,7 @@ impl<'a> CompletionContext<'a> {
             is_invocation: false,
             mentioned_relations: HashMap::new(),
             mentioned_table_aliases: HashMap::new(),
+            mentioned_columns: HashMap::new(),
             is_in_error_node: false,
         };
 
@@ -144,33 +151,46 @@ impl<'a> CompletionContext<'a> {
 
         executor.add_query_results::<queries::RelationMatch>();
         executor.add_query_results::<queries::TableAliasMatch>();
+        executor.add_query_results::<queries::SelectColumnMatch>();
 
         for relation_match in executor.get_iter(stmt_range) {
             match relation_match {
                 QueryResult::Relation(r) => {
                     let schema_name = r.get_schema(sql);
                     let table_name = r.get_table(sql);
 
-                    let current = self.mentioned_relations.get_mut(&schema_name);
-
-                    match current {
-                        Some(c) => {
-                            c.insert(table_name);
-                        }
-                        None => {
-                            let mut new = HashSet::new();
-                            new.insert(table_name);
-                            self.mentioned_relations.insert(schema_name, new);
-                        }
-                    };
+                    if let Some(c) = self.mentioned_relations.get_mut(&schema_name) {
+                        c.insert(table_name);
+                    } else {
+                        let mut new = HashSet::new();
+                        new.insert(table_name);
+                        self.mentioned_relations.insert(schema_name, new);
+                    }
                 }
-
                 QueryResult::TableAliases(table_alias_match) => {
                     self.mentioned_table_aliases.insert(
                         table_alias_match.get_alias(sql),
                         table_alias_match.get_table(sql),
                     );
                 }
+                QueryResult::SelectClauseColumns(c) => {
+                    let mentioned = MentionedColumn {
+                        column: c.get_column(sql),
+                        alias: c.get_alias(sql),
+                    };
+
+                    if let Some(cols) = self
+                        .mentioned_columns
+                        .get_mut(&Some(WrappingClause::Select))
+                    {
+                        cols.insert(mentioned);
+                    } else {
+                        let mut new = HashSet::new();
+                        new.insert(mentioned);
+                        self.mentioned_columns
+                            .insert(Some(WrappingClause::Select), new);
+                    }
+                }
             };
         }
     }

crates/pgt_completions/src/providers/columns.rs

@@ -484,4 +484,93 @@ mod tests {
         )
         .await;
     }
+
+    #[tokio::test]
+    async fn prefers_not_mentioned_columns() {
+        let setup = r#"
+            create schema auth;
+
+            create table public.one (
+                id serial primary key,
+                a text,
+                b text,
+                z text
+            );
+
+            create table public.two (
+                id serial primary key,
+                c text,
+                d text,
+                e text
+            );
+        "#;
+
+        assert_complete_results(
+            format!(
+                "select {} from public.one o join public.two on o.id = t.id;",
+                CURSOR_POS
+            )
+            .as_str(),
+            vec![
+                CompletionAssertion::Label("a".to_string()),
+                CompletionAssertion::Label("b".to_string()),
+                CompletionAssertion::Label("c".to_string()),
+                CompletionAssertion::Label("d".to_string()),
+                CompletionAssertion::Label("e".to_string()),
+            ],
+            setup,
+        )
+        .await;
+
+        // "a" is already mentioned, so it jumps down
+        assert_complete_results(
+            format!(
+                "select a, {} from public.one o join public.two on o.id = t.id;",
+                CURSOR_POS
+            )
+            .as_str(),
+            vec![
+                CompletionAssertion::Label("b".to_string()),
+                CompletionAssertion::Label("c".to_string()),
+                CompletionAssertion::Label("d".to_string()),
+                CompletionAssertion::Label("e".to_string()),
+                CompletionAssertion::Label("id".to_string()),
+                CompletionAssertion::Label("z".to_string()),
+                CompletionAssertion::Label("a".to_string()),
+            ],
+            setup,
+        )
+        .await;
+
+        // "id" of table one is mentioned, but table two isn't –
+        // its priority stays up
+        assert_complete_results(
+            format!(
+                "select o.id, a, b, c, d, e, {} from public.one o join public.two on o.id = t.id;",
+                CURSOR_POS
+            )
+            .as_str(),
+            vec![
+                CompletionAssertion::LabelAndDesc(
+                    "id".to_string(),
+                    "Table: public.two".to_string(),
+                ),
+                CompletionAssertion::Label("z".to_string()),
+            ],
+            setup,
+        )
+        .await;
+
+        // "id" is ambiguous, so both "id" columns are lowered in priority
+        assert_complete_results(
+            format!(
+                "select id, a, b, c, d, e, {} from public.one o join public.two on o.id = t.id;",
+                CURSOR_POS
+            )
+            .as_str(),
+            vec![CompletionAssertion::Label("z".to_string())],
+            setup,
+        )
+        .await;
+    }
 }

crates/pgt_completions/src/providers/triggers.rs

@@ -0,0 +1,169 @@
+use crate::{
+    CompletionItemKind,
+    builder::{CompletionBuilder, PossibleCompletionItem},
+    context::CompletionContext,
+    relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
+};
+
+use super::helper::get_completion_text_with_schema_or_alias;
+
+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 relevance = CompletionRelevanceData::Function(func);
+
+        let item = PossibleCompletionItem {
+            label: func.name.clone(),
+            score: CompletionScore::from(relevance.clone()),
+            filter: CompletionFilter::from(relevance),
+            description: format!("Schema: {}", func.schema),
+            kind: CompletionItemKind::Function,
+            completion_text: get_completion_text_with_schema_or_alias(
+                ctx,
+                &func.name,
+                &func.schema,
+            ),
+        };
+
+        builder.add_item(item);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{
+        CompletionItem, CompletionItemKind, complete,
+        test_helper::{CURSOR_POS, get_test_deps, get_test_params},
+    };
+
+    #[tokio::test]
+    async fn completes_fn() {
+        let setup = r#"
+          create or replace function cool()
+          returns trigger
+          language plpgsql
+          security invoker
+          as $$
+          begin
+            raise exception 'dont matter';
+          end;
+          $$;
+        "#;
+
+        let query = format!("select coo{}", CURSOR_POS);
+
+        let (tree, cache) = get_test_deps(setup, query.as_str().into()).await;
+        let params = get_test_params(&tree, &cache, query.as_str().into());
+        let results = complete(params);
+
+        let CompletionItem { label, .. } = results
+            .into_iter()
+            .next()
+            .expect("Should return at least one completion item");
+
+        assert_eq!(label, "cool");
+    }
+
+    #[tokio::test]
+    async fn prefers_fn_if_invocation() {
+        let setup = r#"
+          create table coos (
+            id serial primary key,
+            name text
+          );
+
+          create or replace function cool()
+          returns trigger
+          language plpgsql
+          security invoker
+          as $$
+          begin
+            raise exception 'dont matter';
+          end;
+          $$;
+        "#;
+
+        let query = format!(r#"select * from coo{}()"#, CURSOR_POS);
+
+        let (tree, cache) = get_test_deps(setup, query.as_str().into()).await;
+        let params = get_test_params(&tree, &cache, query.as_str().into());
+        let results = complete(params);
+
+        let CompletionItem { label, kind, .. } = results
+            .into_iter()
+            .next()
+            .expect("Should return at least one completion item");
+
+        assert_eq!(label, "cool");
+        assert_eq!(kind, CompletionItemKind::Function);
+    }
+
+    #[tokio::test]
+    async fn prefers_fn_in_select_clause() {
+        let setup = r#"
+          create table coos (
+            id serial primary key,
+            name text
+          );
+
+          create or replace function cool()
+          returns trigger
+          language plpgsql
+          security invoker
+          as $$
+          begin
+            raise exception 'dont matter';
+          end;
+          $$;
+        "#;
+
+        let query = format!(r#"select coo{}"#, CURSOR_POS);
+
+        let (tree, cache) = get_test_deps(setup, query.as_str().into()).await;
+        let params = get_test_params(&tree, &cache, query.as_str().into());
+        let results = complete(params);
+
+        let CompletionItem { label, kind, .. } = results
+            .into_iter()
+            .next()
+            .expect("Should return at least one completion item");
+
+        assert_eq!(label, "cool");
+        assert_eq!(kind, CompletionItemKind::Function);
+    }
+
+    #[tokio::test]
+    async fn prefers_function_in_from_clause_if_invocation() {
+        let setup = r#"
+          create table coos (
+            id serial primary key,
+            name text
+          );
+
+          create or replace function cool()
+          returns trigger
+          language plpgsql
+          security invoker
+          as $$
+          begin
+            raise exception 'dont matter';
+          end;
+          $$;
+        "#;
+
+        let query = format!(r#"select * from coo{}()"#, CURSOR_POS);
+
+        let (tree, cache) = get_test_deps(setup, query.as_str().into()).await;
+        let params = get_test_params(&tree, &cache, query.as_str().into());
+        let results = complete(params);
+
+        let CompletionItem { label, kind, .. } = results
+            .into_iter()
+            .next()
+            .expect("Should return at least one completion item");
+
+        assert_eq!(label, "cool");
+        assert_eq!(kind, CompletionItemKind::Function);
+    }
+}