Merge branch 'main' of https://github.com/supabase-community/postgres_lsp into feat/policy-completions

Author: juleswritescode

.sqlx/query-df57cc22f7d63847abce1d0d15675ba8951faa1be2ea6b2bf6714b1aa9127a6f.json

@@ -15,7 +15,7 @@ use crate::{
     sanitization::SanitizedCompletionParams,
 };
 
-#[derive(Debug, PartialEq, Eq)]
+#[derive(Debug, PartialEq, Eq, Hash)]
 pub enum WrappingClause<'a> {
     Select,
     Where,
@@ -29,6 +29,12 @@ pub enum WrappingClause<'a> {
     ToRoleAssignment,
 }
 
+#[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
@@ -150,8 +156,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> {
@@ -169,6 +175,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,
         };
 
@@ -223,33 +230,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);
+                    }
+                }
             };
         }
     }

Cargo.lock

@@ -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;
+    }
 }