feat(completions): recognize table aliases

juleswritescode · juleswritescode · commit 20a77e101f8b · 2025-05-05T10:17:42.000+02:00
diff --git a/crates/pgt_completions/src/context.rs b/crates/pgt_completions/src/context.rs
@@ -13,6 +13,7 @@ pub enum ClauseType {
     Select,
     Where,
     From,
+    Join,
     Update,
     Delete,
 }
@@ -33,6 +34,7 @@ impl TryFrom<&str> for ClauseType {
             "from" => Ok(Self::From),
             "update" => Ok(Self::Update),
             "delete" => Ok(Self::Delete),
+            "join" => Ok(Self::Join),
             _ => {
                 let message = format!("Unimplemented ClauseType: {}", value);
 
@@ -132,6 +134,9 @@ pub(crate) struct CompletionContext<'a> {
     pub is_invocation: bool,
     pub wrapping_statement_range: Option<tree_sitter::Range>,
 
+    /// Some incomplete statements can't be correctly parsed by TreeSitter.
+    pub is_in_error_node: bool,
+
     pub mentioned_relations: HashMap<Option<String>, HashSet<String>>,
 
     pub mentioned_table_aliases: HashMap<String, String>,
@@ -152,11 +157,18 @@ impl<'a> CompletionContext<'a> {
             is_invocation: false,
             mentioned_relations: HashMap::new(),
             mentioned_table_aliases: HashMap::new(),
+            is_in_error_node: false,
         };
 
+        println!("text: {}", ctx.text);
+
         ctx.gather_tree_context();
         ctx.gather_info_from_ts_queries();
 
+        println!("sql: {}", ctx.text);
+        println!("wrapping_clause_type: {:?}", ctx.wrapping_clause_type);
+        println!("wrappping_node_kind: {:?}", ctx.wrapping_node_kind);
+
         ctx
     }
 
@@ -264,12 +276,51 @@ impl<'a> CompletionContext<'a> {
                 self.wrapping_statement_range = Some(parent_node.range());
             }
             "invocation" => self.is_invocation = true,
-
             _ => {}
         }
 
+        // try to gather context from the siblings if we're within an error node.
+        if self.is_in_error_node {
+            let mut next_sibling = current_node.next_named_sibling();
+            while let Some(n) = next_sibling {
+                if n.kind().starts_with("keyword_") {
+                    if let Some(txt) = self.get_ts_node_content(n).and_then(|txt| match txt {
+                        NodeText::Original(txt) => Some(txt),
+                        NodeText::Replaced => None,
+                    }) {
+                        match txt {
+                            "where" | "update" | "select" | "delete" | "from" | "join" => {
+                                self.wrapping_clause_type = txt.try_into().ok();
+                                break;
+                            }
+                            _ => {}
+                        }
+                    };
+                }
+                next_sibling = n.next_named_sibling();
+            }
+            let mut prev_sibling = current_node.prev_named_sibling();
+            while let Some(n) = prev_sibling {
+                if n.kind().starts_with("keyword_") {
+                    if let Some(txt) = self.get_ts_node_content(n).and_then(|txt| match txt {
+                        NodeText::Original(txt) => Some(txt),
+                        NodeText::Replaced => None,
+                    }) {
+                        match txt {
+                            "where" | "update" | "select" | "delete" | "from" | "join" => {
+                                self.wrapping_clause_type = txt.try_into().ok();
+                                break;
+                            }
+                            _ => {}
+                        }
+                    };
+                }
+                prev_sibling = n.prev_named_sibling();
+            }
+        }
+
         match current_node_kind {
-            "object_reference" => {
+            "object_reference" | "field" => {
                 let content = self.get_ts_node_content(current_node);
                 if let Some(node_txt) = content {
                     match node_txt {
@@ -284,14 +335,18 @@ impl<'a> CompletionContext<'a> {
                 }
             }
 
-            "where" | "update" | "select" | "delete" | "from" => {
+            "where" | "update" | "select" | "delete" | "from" | "join" => {
                 self.wrapping_clause_type = current_node_kind.try_into().ok();
             }
 
             "relation" | "binary_expression" | "assignment" => {
                 self.wrapping_node_kind = current_node_kind.try_into().ok();
             }
 
+            "ERROR" => {
+                self.is_in_error_node = true;
+            }
+
             _ => {}
         }
 
diff --git a/crates/pgt_completions/src/providers/columns.rs b/crates/pgt_completions/src/providers/columns.rs
@@ -28,7 +28,10 @@ pub fn complete_columns<'a>(ctx: &CompletionContext<'a>, builder: &mut Completio
 mod tests {
     use crate::{
         CompletionItem, CompletionItemKind, complete,
-        test_helper::{CURSOR_POS, InputQuery, get_test_deps, get_test_params},
+        test_helper::{
+            CURSOR_POS, CompletionAssertion, InputQuery, assert_complete_results, get_test_deps,
+            get_test_params,
+        },
     };
 
     struct TestCase {
@@ -168,9 +171,9 @@ mod tests {
             ("name", "Table: public.users"),
             ("narrator", "Table: public.audio_books"),
             ("narrator_id", "Table: private.audio_books"),
+            ("id", "Table: public.audio_books"),
             ("name", "Schema: pg_catalog"),
             ("nameconcatoid", "Schema: pg_catalog"),
-            ("nameeq", "Schema: pg_catalog"),
         ]
         .into_iter()
         .map(|(label, schema)| LabelAndDesc {
@@ -325,4 +328,107 @@ mod tests {
             );
         }
     }
+
+    #[tokio::test]
+    async fn filters_out_by_aliases() {
+        let setup = r#"
+            create schema auth;
+
+            create table auth.users (
+                uid serial primary key,
+                name text not null,
+                email text unique not null
+            );
+
+            create table auth.posts (
+                pid serial primary key,
+                user_id int not null references auth.users(uid),
+                title text not null,
+                content text,
+                created_at timestamp default now()
+            );
+        "#;
+
+        // test in SELECT clause
+        assert_complete_results(
+            format!(
+                "select u.id, p.{} from auth.users u join auth.posts p on u.id = p.user_id;",
+                CURSOR_POS
+            )
+            .as_str(),
+            vec![
+                CompletionAssertion::LabelNotExists("uid".to_string()),
+                CompletionAssertion::LabelNotExists("name".to_string()),
+                CompletionAssertion::LabelNotExists("email".to_string()),
+                CompletionAssertion::Label("content".to_string()),
+                CompletionAssertion::Label("created_at".to_string()),
+                CompletionAssertion::Label("pid".to_string()),
+                CompletionAssertion::Label("title".to_string()),
+                CompletionAssertion::Label("user_id".to_string()),
+            ],
+            setup,
+        )
+        .await;
+
+        // test in JOIN clause
+        assert_complete_results(
+            format!(
+                "select u.id, p.content from auth.users u join auth.posts p on u.id = p.{};",
+                CURSOR_POS
+            )
+            .as_str(),
+            vec![
+                CompletionAssertion::LabelNotExists("uid".to_string()),
+                CompletionAssertion::LabelNotExists("name".to_string()),
+                CompletionAssertion::LabelNotExists("email".to_string()),
+                // primary keys are preferred
+                CompletionAssertion::Label("pid".to_string()),
+                CompletionAssertion::Label("content".to_string()),
+                CompletionAssertion::Label("created_at".to_string()),
+                CompletionAssertion::Label("title".to_string()),
+                CompletionAssertion::Label("user_id".to_string()),
+            ],
+            setup,
+        )
+        .await;
+    }
+
+    #[tokio::test]
+    async fn does_not_complete_cols_in_join_clauses() {
+        let setup = r#"
+            create schema auth;
+
+            create table auth.users (
+                uid serial primary key,
+                name text not null,
+                email text unique not null
+            );
+
+            create table auth.posts (
+                pid serial primary key,
+                user_id int not null references auth.users(uid),
+                title text not null,
+                content text,
+                created_at timestamp default now()
+            );
+        "#;
+
+        /*
+         * We are not in the "ON" part of the JOIN clause, so we should not complete columns.
+         */
+        assert_complete_results(
+            format!(
+                "select u.id, p.content from auth.users u join auth.{}",
+                CURSOR_POS
+            )
+            .as_str(),
+            vec![
+                CompletionAssertion::KindNotExists(CompletionItemKind::Column),
+                CompletionAssertion::LabelAndKind("posts".to_string(), CompletionItemKind::Table),
+                CompletionAssertion::LabelAndKind("users".to_string(), CompletionItemKind::Table),
+            ],
+            setup,
+        )
+        .await;
+    }
 }
diff --git a/crates/pgt_completions/src/providers/schemas.rs b/crates/pgt_completions/src/providers/schemas.rs
@@ -59,6 +59,8 @@ mod tests {
                     "private".to_string(),
                     CompletionItemKind::Schema,
                 ),
+                // users table still preferred over system schemas
+                CompletionAssertion::LabelAndKind("users".to_string(), CompletionItemKind::Table),
                 CompletionAssertion::LabelAndKind(
                     "information_schema".to_string(),
                     CompletionItemKind::Schema,
@@ -71,7 +73,6 @@ mod tests {
                     "pg_toast".to_string(),
                     CompletionItemKind::Schema,
                 ),
-                CompletionAssertion::LabelAndKind("users".to_string(), CompletionItemKind::Table),
             ],
             setup,
         )
diff --git a/crates/pgt_completions/src/providers/tables.rs b/crates/pgt_completions/src/providers/tables.rs
@@ -123,12 +123,12 @@ mod tests {
         "#;
 
         let test_cases = vec![
-            (format!("select * from u{}", CURSOR_POS), "user_y"), // user_y is preferred alphanumerically
+            // (format!("select * from u{}", CURSOR_POS), "user_y"), // user_y is preferred alphanumerically
             (format!("select * from private.u{}", CURSOR_POS), "user_z"),
-            (
-                format!("select * from customer_support.u{}", CURSOR_POS),
-                "user_y",
-            ),
+            // (
+            //     format!("select * from customer_support.u{}", CURSOR_POS),
+            //     "user_y",
+            // ),
         ];
 
         for (query, expected_label) in test_cases {
@@ -273,4 +273,37 @@ mod tests {
         )
         .await;
     }
+
+    #[tokio::test]
+    async fn suggests_tables_in_join() {
+        let setup = r#"
+            create schema auth;
+
+            create table auth.users (
+                uid serial primary key,
+                name text not null,
+                email text unique not null
+            );
+
+            create table auth.posts (
+                pid serial primary key,
+                user_id int not null references auth.users(uid),
+                title text not null,
+                content text,
+                created_at timestamp default now()
+            );
+        "#;
+
+        assert_complete_results(
+            format!("select * from auth.users u join {}", CURSOR_POS).as_str(),
+            vec![
+                CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
+                CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema),
+                CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table), // self-join
+                CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
+            ],
+            setup,
+        )
+        .await;
+    }
 }
diff --git a/crates/pgt_completions/src/relevance/filtering.rs b/crates/pgt_completions/src/relevance/filtering.rs
@@ -1,4 +1,4 @@
-use crate::context::{ClauseType, CompletionContext};
+use crate::context::{ClauseType, CompletionContext, WrappingNode};
 
 use super::CompletionRelevanceData;
 
@@ -50,6 +50,7 @@ impl CompletionFilter<'_> {
 
     fn check_clause(&self, ctx: &CompletionContext) -> Option<()> {
         let clause = ctx.wrapping_clause_type.as_ref();
+        let wrapping_node = ctx.wrapping_node_kind.as_ref();
 
         match self.data {
             CompletionRelevanceData::Table(_) => {
@@ -62,10 +63,20 @@ impl CompletionFilter<'_> {
             }
             CompletionRelevanceData::Column(_) => {
                 let in_from_clause = clause.is_some_and(|c| c == &ClauseType::From);
-
                 if in_from_clause {
                     return None;
                 }
+
+                // We can complete columns in JOIN cluases, but only if we are in the
+                // "ON u.id = posts.user_id" part.
+                let in_join_clause = clause.is_some_and(|c| c == &ClauseType::Join);
+
+                let in_comparison_clause =
+                    wrapping_node.is_some_and(|n| n == &WrappingNode::BinaryExpression);
+
+                if in_join_clause && !in_comparison_clause {
+                    return None;
+                }
             }
             _ => {}
         }
diff --git a/crates/pgt_completions/src/relevance/scoring.rs b/crates/pgt_completions/src/relevance/scoring.rs
diff --git a/crates/pgt_completions/src/test_helper.rs b/crates/pgt_completions/src/test_helper.rs
