ide: goto def nested table/values

Author: sbdchd

crates/squawk_ide/src/goto_definition.rs

@@ -3316,6 +3316,138 @@ select b$0 from x t(a, b);
         ");
     }
 
+    #[test]
+    fn goto_values_column_alias_list() {
+        assert_snapshot!(goto("
+select c$0 from (values (1)) t(c);
+"), @r"
+          ╭▸ 
+        2 │ select c from (values (1)) t(c);
+          ╰╴       ─ 1. source           ─ 2. destination
+        ");
+    }
+
+    #[test]
+    fn goto_values_column_alias_list_qualified() {
+        assert_snapshot!(goto("
+select t.c$0 from (values (1)) t(c);
+"), @r"
+          ╭▸ 
+        2 │ select t.c from (values (1)) t(c);
+          ╰╴         ─ 1. source           ─ 2. destination
+        ");
+    }
+
+    #[test]
+    fn goto_values_column_alias_list_multiple() {
+        assert_snapshot!(goto("
+select b$0 from (values (1, 2)) t(a, b);
+"), @r"
+          ╭▸ 
+        2 │ select b from (values (1, 2)) t(a, b);
+          ╰╴       ─ 1. source                 ─ 2. destination
+        ");
+    }
+
+    #[test]
+    fn goto_table_expr_column() {
+        assert_snapshot!(goto("
+create table t(a int, b int);
+select a$0 from (table t);
+"), @r"
+          ╭▸ 
+        2 │ create table t(a int, b int);
+          │                ─ 2. destination
+        3 │ select a from (table t);
+          ╰╴       ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_table_expr_column_with_cte() {
+        assert_snapshot!(goto("
+with x as (select 1 a)
+select a$0 from (table x);
+"), @r"
+          ╭▸ 
+        2 │ with x as (select 1 a)
+          │                     ─ 2. destination
+        3 │ select a from (table x);
+          ╰╴       ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_table_expr_partial_column_alias_list() {
+        assert_snapshot!(goto("
+with t as (select 1 a, 2 b)
+select c, b$0 from (table t) u(c);
+"), @r"
+          ╭▸ 
+        2 │ with t as (select 1 a, 2 b)
+          │                          ─ 2. destination
+        3 │ select c, b from (table t) u(c);
+          ╰╴          ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_subquery_partial_column_alias_list() {
+        assert_snapshot!(goto("
+select x, b$0 from (select 1 a, 2 b) t(x);
+"), @r"
+          ╭▸ 
+        2 │ select x, b from (select 1 a, 2 b) t(x);
+          ╰╴          ─ 1. source           ─ 2. destination
+        ");
+    }
+
+    #[test]
+    fn goto_table_expr_values_cte_partial_alias() {
+        assert_snapshot!(goto("
+with t as (values (1, 2), (3, 4))
+select column2$0 from (table t) u(a);
+"), @r"
+          ╭▸ 
+        2 │ with t as (values (1, 2), (3, 4))
+          │                       ─ 2. destination
+        3 │ select column2 from (table t) u(a);
+          ╰╴             ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_cte_with_table_expr() {
+        assert_snapshot!(goto("
+create table t(a int, b int);
+with u as (table t)
+select a$0 from u;
+"), @r"
+          ╭▸ 
+        2 │ create table t(a int, b int);
+          │                ─ 2. destination
+        3 │ with u as (table t)
+        4 │ select a from u;
+          ╰╴       ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_cte_with_table_expr_nested() {
+        assert_snapshot!(goto("
+with t as (select 1 a, 2 b),
+     u as (table t)
+select b$0 from u;
+"), @r"
+          ╭▸ 
+        2 │ with t as (select 1 a, 2 b),
+          │                          ─ 2. destination
+        3 │      u as (table t)
+        4 │ select b from u;
+          ╰╴       ─ 1. source
+        ");
+    }
+
     #[test]
     fn goto_insert_table() {
         assert_snapshot!(goto("

crates/squawk_ide/src/resolve.rs

@@ -857,7 +857,7 @@ fn resolve_select_qualified_column_ptr(
             && Name::from_node(&alias_name) == column_table_name
         {
             if let Some(paren_select) = from_item.paren_select() {
-                return resolve_subquery_column(
+                return resolve_subquery_column_ptr(
                     binder,
                     root,
                     &paren_select,
@@ -1031,7 +1031,7 @@ fn resolve_from_item_column_ptr(
     let column_name = Name::from_node(column_name_ref);
     if let Some(paren_select) = from_item.paren_select() {
         let alias = from_item.alias();
-        return resolve_subquery_column(
+        return resolve_subquery_column_ptr(
             binder,
             root,
             &paren_select,
@@ -1790,6 +1790,28 @@ fn resolve_cte_column(
                 continue;
             }
 
+            if let ast::WithQuery::Table(table) = query {
+                let relation_name = table.relation_name()?;
+                let path = relation_name.path()?;
+                let (table_name, schema) = extract_table_schema_from_path(&path)?;
+
+                if schema.is_none()
+                    && let Some(nested_cte_column) =
+                        resolve_cte_column(binder, root, name_ref, &table_name, column_name)
+                {
+                    return Some(nested_cte_column);
+                }
+
+                return resolve_column_from_table_or_view(
+                    binder,
+                    root,
+                    name_ref,
+                    &table_name,
+                    &schema,
+                    column_name,
+                );
+            }
+
             if let Some(column) =
                 column_in_with_query(&query, binder, root, column_name, column_list_len)
             {
@@ -1927,7 +1949,7 @@ fn column_in_with_query(
     None
 }
 
-fn resolve_subquery_column(
+fn resolve_subquery_column_ptr(
     binder: &Binder,
     root: &SyntaxNode,
     paren_select: &ast::ParenSelect,
@@ -1936,12 +1958,6 @@ fn resolve_subquery_column(
     alias: Option<&ast::Alias>,
 ) -> Option<SyntaxNodePtr> {
     let select_variant = paren_select.select()?;
-    let ast::SelectVariant::Select(subquery_select) = select_variant else {
-        return None;
-    };
-
-    let select_clause = subquery_select.select_clause()?;
-    let target_list = select_clause.target_list()?;
 
     if let Some(alias) = alias
         && let Some(column_list) = alias.column_list()
@@ -1953,9 +1969,40 @@ fn resolve_subquery_column(
                 return Some(SyntaxNodePtr::new(col_name.syntax()));
             }
         }
-        return None;
+        if matches!(select_variant, ast::SelectVariant::Values(_)) {
+            return None;
+        }
+    }
+
+    if let ast::SelectVariant::Table(table) = select_variant {
+        let relation_name = table.relation_name()?;
+        let path = relation_name.path()?;
+        let (table_name, schema) = extract_table_schema_from_path(&path)?;
+
+        if schema.is_none()
+            && let Some(cte_column_ptr) =
+                resolve_cte_column(binder, root, name_ref, &table_name, column_name)
+        {
+            return Some(cte_column_ptr);
+        }
+
+        return resolve_column_from_table_or_view(
+            binder,
+            root,
+            name_ref,
+            &table_name,
+            &schema,
+            column_name,
+        );
     }
 
+    let ast::SelectVariant::Select(subquery_select) = select_variant else {
+        return None;
+    };
+
+    let select_clause = subquery_select.select_clause()?;
+    let target_list = select_clause.target_list()?;
+
     for target in target_list.targets() {
         if let Some((col_name, node)) = ColumnName::from_target(target.clone()) {
             if let Some(col_name_str) = col_name.to_string()
@@ -2547,7 +2594,7 @@ fn resolve_column_from_paren_expr(
         && let Some(paren_select) = from_item.paren_select()
     {
         let alias = from_item.alias();
-        return resolve_subquery_column(
+        return resolve_subquery_column_ptr(
             binder,
             root,
             &paren_select,