feat(completions): fill in alias for columns in join clauses (#392)

Given a query

```sql
select 
  email, id
from
  auth.users u
  join auth.identities i on {}
```

if you pick `user_id` from the list, it will fill in `i.user_id` now.

Author: juleswritescode

crates/pgt_completions/src/providers/columns.rs

@@ -1,17 +1,19 @@
 use crate::{
     CompletionItemKind,
     builder::{CompletionBuilder, PossibleCompletionItem},
-    context::CompletionContext,
+    context::{CompletionContext, WrappingClause},
     relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
 };
 
+use super::helper::{find_matching_alias_for_table, get_completion_text_with_schema_or_alias};
+
 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 relevance = CompletionRelevanceData::Column(col);
 
-        let item = PossibleCompletionItem {
+        let mut item = PossibleCompletionItem {
             label: col.name.clone(),
             score: CompletionScore::from(relevance.clone()),
             filter: CompletionFilter::from(relevance),
@@ -20,6 +22,14 @@ pub fn complete_columns<'a>(ctx: &CompletionContext<'a>, builder: &mut Completio
             completion_text: None,
         };
 
+        // autocomplete with the alias in a join clause if we find one
+        if matches!(ctx.wrapping_clause_type, Some(WrappingClause::Join { .. })) {
+            item.completion_text = find_matching_alias_for_table(ctx, col.table_name.as_str())
+                .and_then(|alias| {
+                    get_completion_text_with_schema_or_alias(ctx, col.name.as_str(), alias.as_str())
+                });
+        }
+
         builder.add_item(item);
     }
 }

crates/pgt_completions/src/providers/functions.rs

@@ -5,7 +5,7 @@ use crate::{
     relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
 };
 
-use super::helper::get_completion_text_with_schema;
+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;
@@ -19,7 +19,11 @@ pub fn complete_functions<'a>(ctx: &'a CompletionContext, builder: &mut Completi
             filter: CompletionFilter::from(relevance),
             description: format!("Schema: {}", func.schema),
             kind: CompletionItemKind::Function,
-            completion_text: get_completion_text_with_schema(ctx, &func.name, &func.schema),
+            completion_text: get_completion_text_with_schema_or_alias(
+                ctx,
+                &func.name,
+                &func.schema,
+            ),
         };
 
         builder.add_item(item);

crates/pgt_completions/src/providers/helper.rs

@@ -2,12 +2,24 @@ use pgt_text_size::{TextRange, TextSize};
 
 use crate::{CompletionText, context::CompletionContext};
 
-pub(crate) fn get_completion_text_with_schema(
+pub(crate) fn find_matching_alias_for_table(
+    ctx: &CompletionContext,
+    table_name: &str,
+) -> Option<String> {
+    for (alias, table) in ctx.mentioned_table_aliases.iter() {
+        if table == table_name {
+            return Some(alias.to_string());
+        }
+    }
+    None
+}
+
+pub(crate) fn get_completion_text_with_schema_or_alias(
     ctx: &CompletionContext,
     item_name: &str,
-    item_schema_name: &str,
+    schema_or_alias_name: &str,
 ) -> Option<CompletionText> {
-    if item_schema_name == "public" || ctx.schema_or_alias_name.is_some() {
+    if schema_or_alias_name == "public" || ctx.schema_or_alias_name.is_some() {
         None
     } else {
         let node = ctx.node_under_cursor.unwrap();
@@ -18,7 +30,7 @@ pub(crate) fn get_completion_text_with_schema(
         );
 
         Some(CompletionText {
-            text: format!("{}.{}", item_schema_name, item_name),
+            text: format!("{}.{}", schema_or_alias_name, item_name),
             range,
         })
     }

crates/pgt_completions/src/providers/tables.rs

@@ -5,7 +5,7 @@ use crate::{
     relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
 };
 
-use super::helper::get_completion_text_with_schema;
+use super::helper::get_completion_text_with_schema_or_alias;
 
 pub fn complete_tables<'a>(ctx: &'a CompletionContext, builder: &mut CompletionBuilder<'a>) {
     let available_tables = &ctx.schema_cache.tables;
@@ -19,7 +19,11 @@ pub fn complete_tables<'a>(ctx: &'a CompletionContext, builder: &mut CompletionB
             filter: CompletionFilter::from(relevance),
             description: format!("Schema: {}", table.schema),
             kind: CompletionItemKind::Table,
-            completion_text: get_completion_text_with_schema(ctx, &table.name, &table.schema),
+            completion_text: get_completion_text_with_schema_or_alias(
+                ctx,
+                &table.name,
+                &table.schema,
+            ),
         };
 
         builder.add_item(item);