ide: table alias goto def support (#791)

Author: sbdchd

crates/squawk_ide/src/goto_definition.rs

@@ -2452,4 +2452,56 @@ drop aggregate sum$0(bigint);
           ╰╴                 ─ 1. source
         ");
     }
+
+    #[test]
+    fn goto_table_alias_in_qualified_column() {
+        assert_snapshot!(goto("
+create table t(a int8, b text);
+select f$0.a from t as f;
+"), @r"
+          ╭▸ 
+        3 │ select f.a from t as f;
+          ╰╴       ─ 1. source   ─ 2. destination
+        ");
+    }
+
+    #[test]
+    fn goto_column_through_table_alias() {
+        assert_snapshot!(goto("
+create table t(a int8, b text);
+select f.a$0 from t as f;
+"), @r"
+          ╭▸ 
+        2 │ create table t(a int8, b text);
+          │                ─ 2. destination
+        3 │ select f.a from t as f;
+          ╰╴         ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_cte_alias_renamed_column() {
+        assert_snapshot!(goto("
+with t as (select 1 a, 2 b)
+select f.x$0 from t as f(x);
+"), @r"
+          ╭▸ 
+        3 │ select f.x from t as f(x);
+          ╰╴         ─ 1. source   ─ 2. destination
+        ");
+    }
+
+    #[test]
+    fn goto_cte_alias_unrenamed_column() {
+        assert_snapshot!(goto("
+with t as (select 1 a, 2 b)
+select f.b$0 from t as f(x);
+"), @r"
+          ╭▸ 
+        2 │ with t as (select 1 a, 2 b)
+          │                          ─ 2. destination
+        3 │ select f.b from t as f(x);
+          ╰╴         ─ 1. source
+        ");
+    }
 }

crates/squawk_ide/src/resolve.rs

@@ -552,6 +552,12 @@ fn resolve_select_qualified_column_table(
     let from_clause = select.from_clause()?;
     let from_item = from_clause.from_items().next()?;
 
+    if let Some(alias_name) = from_item.alias().and_then(|a| a.name())
+        && Name::from_node(&alias_name) == table_name
+    {
+        return Some(SyntaxNodePtr::new(alias_name.syntax()));
+    }
+
     let (table_name, schema) = if let Some(name_ref_node) = from_item.name_ref() {
         // `from foo`
         let from_table_name = Name::from_node(&name_ref_node);
@@ -618,7 +624,48 @@ fn resolve_select_qualified_column(
         let from_clause = select.from_clause()?;
         let from_item = from_clause.from_items().next()?;
 
-        if let Some(name_ref_node) = from_item.name_ref() {
+        // `from t as u`
+        // `from t as u(a, b, c)`
+        if let Some(alias) = from_item.alias()
+            && let Some(alias_name) = alias.name()
+            && Name::from_node(&alias_name) == column_table_name
+        {
+            // `from t as u(a, b, c)`
+            if let Some(column_list) = alias.column_list() {
+                for column in column_list.columns() {
+                    if let Some(col_name) = column.name()
+                        && Name::from_node(&col_name) == column_name
+                    {
+                        return Some(SyntaxNodePtr::new(col_name.syntax()));
+                    }
+                }
+
+                // ```sql
+                // create table t(a int, b int);
+                // select b from t as u(x);
+                //        ^
+                // ```
+                if let Some(name_ref_node) = from_item.name_ref() {
+                    let cte_name = Name::from_node(&name_ref_node);
+                    return resolve_cte_column(name_ref, &cte_name, &column_name);
+                }
+            }
+
+            // `from t as u`
+            if let Some(name_ref_node) = from_item.name_ref() {
+                (Name::from_node(&name_ref_node), None)
+            // `from foo.t as u`
+            } else if let Some(from_field_expr) = from_item.field_expr() {
+                let table_name = Name::from_node(&from_field_expr.field()?);
+                let ast::Expr::NameRef(schema_name_ref) = from_field_expr.base()? else {
+                    return None;
+                };
+                let schema = Schema(Name::from_node(&schema_name_ref));
+                (table_name, Some(schema))
+            } else {
+                return None;
+            }
+        } else if let Some(name_ref_node) = from_item.name_ref() {
             // `from bar`
             let from_table_name = Name::from_node(&name_ref_node);
             if from_table_name == column_table_name {