add filtering

Author: juleswritescode

crates/pgt_completions/src/relevance/filtering.rs

@@ -18,7 +18,7 @@ impl CompletionFilter<'_> {
         self.completable_context(ctx)?;
         self.check_clause(ctx)?;
         self.check_invocation(ctx)?;
-        self.check_mentioned_schema(ctx)?;
+        self.check_mentioned_schema_or_alias(ctx)?;
 
         Some(())
     }
@@ -86,54 +86,28 @@ impl CompletionFilter<'_> {
         Some(())
     }
 
-    fn check_mentioned_schema(&self, ctx: &CompletionContext) -> Option<()> {
+    fn check_mentioned_schema_or_alias(&self, ctx: &CompletionContext) -> Option<()> {
         if ctx.schema_or_alias_name.is_none() {
             return Some(());
         }
 
-        let name = ctx.schema_or_alias_name.as_ref().unwrap();
+        let schema_or_alias = ctx.schema_or_alias_name.as_ref().unwrap();
 
-        let does_not_match = match self.data {
-            CompletionRelevanceData::Table(table) => &table.schema != name,
-            CompletionRelevanceData::Function(f) => &f.schema != name,
-            CompletionRelevanceData::Column(_) => {
-                // columns belong to tables, not schemas
-                true
-            }
-            CompletionRelevanceData::Schema(_) => {
-                // we should never allow schema suggestions if there already was one.
-                true
-            }
-        };
+        let matches = match self.data {
+            CompletionRelevanceData::Table(table) => &table.schema == schema_or_alias,
+            CompletionRelevanceData::Function(f) => &f.schema == schema_or_alias,
+            CompletionRelevanceData::Column(col) => ctx
+                .mentioned_table_aliases
+                .get(schema_or_alias)
+                .is_some_and(|t| t == &col.table_name),
 
-        if does_not_match {
-            return None;
-        }
-
-        Some(())
-    }
-
-    fn check_mentioned_alias(&self, ctx: &CompletionContext) -> Option<()> {
-        if ctx.schema_or_alias_name.is_none() {
-            return Some(());
-        }
-
-        let name = ctx.schema_or_alias_name.as_ref().unwrap();
-
-        let does_not_match = match self.data {
-            CompletionRelevanceData::Table(table) => &table.schema != name,
-            CompletionRelevanceData::Function(f) => &f.schema != name,
-            CompletionRelevanceData::Column(_) => {
-                // columns belong to tables, not schemas
-                true
-            }
             CompletionRelevanceData::Schema(_) => {
                 // we should never allow schema suggestions if there already was one.
-                true
+                false
             }
         };
 
-        if does_not_match {
+        if !matches {
             return None;
         }