ok fine

Author: juleswritescode

crates/pgt_completions/src/context.rs

@@ -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: 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: HashSet::new(),
             is_in_error_node: false,
         };
 
@@ -144,6 +151,7 @@ impl<'a> CompletionContext<'a> {
 
         executor.add_query_results::<queries::RelationMatch>();
         executor.add_query_results::<queries::TableAliasMatch>();
+        executor.add_query_results::<queries::ColumnMatch>();
 
         for relation_match in executor.get_iter(stmt_range) {
             match relation_match {
@@ -170,8 +178,12 @@ impl<'a> CompletionContext<'a> {
                         table_alias_match.get_table(sql),
                     );
                 }
-
-                QueryResult::Column(_) => todo!(),
+                QueryResult::Column(c) => {
+                    self.mentioned_columns.insert(MentionedColumn {
+                        column: c.get_column(sql),
+                        alias: c.get_alias(sql),
+                    });
+                }
             };
         }
     }

crates/pgt_completions/src/relevance/scoring.rs

@@ -32,6 +32,7 @@ impl CompletionScore<'_> {
         self.check_matching_clause_type(ctx);
         self.check_matching_wrapping_node(ctx);
         self.check_relations_in_stmt(ctx);
+        self.check_columns_in_stmt(ctx);
     }
 
     fn check_matches_query_input(&mut self, ctx: &CompletionContext) {
@@ -235,4 +236,50 @@ impl CompletionScore<'_> {
             self.score += 2;
         }
     }
+
+    fn check_columns_in_stmt(&mut self, ctx: &CompletionContext) {
+        // we only want to consider mentioned columns in a select statement.
+        if ctx
+            .wrapping_clause_type
+            .as_ref()
+            .is_none_or(|ct| ct != &WrappingClause::Select)
+        {
+            return;
+        }
+
+        match self.data {
+            CompletionRelevanceData::Column(column) => {
+                /*
+                 * Columns can be mentioned in one of two ways:
+                 *
+                 * 1) With an alias: `select u.id`.
+                 * If the currently investigated suggestion item is "id" of the "users" table,
+                 * we want to check
+                 * a) whether the name of the column matches.
+                 * b) whether we know which table is aliased by "u" (if we don't, we ignore the alias).
+                 * c) whether the aliased table matches the currently investigated suggestion item's table.
+                 *
+                 * 2) Without an alias: `select id`.
+                 * In that case, we only check whether the mentioned column fits our currently investigated
+                 * suggestion item's name.
+                 *
+                 */
+                if ctx
+                    .mentioned_columns
+                    .iter()
+                    .any(|mentioned| match mentioned.alias.as_ref() {
+                        Some(als) => {
+                            let aliased_table = ctx.mentioned_table_aliases.get(als.as_str());
+                            column.name == mentioned.column
+                                && aliased_table.is_none_or(|t| t == &column.table_name)
+                        }
+                        None => mentioned.column == column.name,
+                    })
+                {
+                    self.score -= 10;
+                }
+            }
+            _ => {}
+        }
+    }
 }