ok

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};
+
 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(ctx, col.name.as_str(), alias.as_str())
+                });
+        }
+
         builder.add_item(item);
     }
 }

crates/pgt_completions/src/providers/helper.rs

@@ -2,6 +2,18 @@ use pgt_text_size::{TextRange, TextSize};
 
 use crate::{CompletionText, context::CompletionContext};
 
+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(
     ctx: &CompletionContext,
     item_name: &str,