diff --git a/crates/squawk_ide/src/goto_definition.rs b/crates/squawk_ide/src/goto_definition.rs index 1ca80caa..dd2f14a3 100644 --- a/crates/squawk_ide/src/goto_definition.rs +++ b/crates/squawk_ide/src/goto_definition.rs @@ -1378,6 +1378,66 @@ select "?column?"$0 from x; "#); } + #[test] + fn goto_cte_star_expansion() { + assert_snapshot!(goto(" +with t as (select 1 a), + y as (select * from t) +select a$0 from y; +"), @r" + ╭▸ + 2 │ with t as (select 1 a), + │ ─ 2. destination + 3 │ y as (select * from t) + 4 │ select a from y; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_cte_reference_inside_cte() { + assert_snapshot!(goto(" +with t as (select 1 a), + y as (select a$0 from t) +select a from y; +"), @r" + ╭▸ + 2 │ with t as (select 1 a), + │ ─ 2. destination + 3 │ y as (select a from t) + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_cte_with_column_list() { + assert_snapshot!(goto(" +with t(a) as (select 1) +select a$0 from t; +"), @r" + ╭▸ + 2 │ with t(a) as (select 1) + │ ─ 2. destination + 3 │ select a from t; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_cte_shadows_table() { + assert_snapshot!(goto(" +create table t(a int); +with t as (select a$0 from t) +select a from t; +"), @r" + ╭▸ + 2 │ create table t(a int); + │ ─ 2. destination + 3 │ with t as (select a from t) + ╰╴ ─ 1. source + "); + } + #[test] fn goto_insert_table() { assert_snapshot!(goto(" diff --git a/crates/squawk_ide/src/resolve.rs b/crates/squawk_ide/src/resolve.rs index 084bb198..6765ac2a 100644 --- a/crates/squawk_ide/src/resolve.rs +++ b/crates/squawk_ide/src/resolve.rs @@ -602,7 +602,7 @@ fn resolve_select_column(binder: &Binder, name_ref: &ast::NameRef) -> Option Option Option { + let select = name_ref + .syntax() + .ancestors() + .filter_map(ast::Select::cast) + .find(|s| s.with_clause().is_some())?; let with_clause = select.with_clause()?; for with_table in with_clause.with_tables() { if let Some(name) = with_table.name() && Name::new(name.syntax().text().to_string()) == *cte_name { + // Skip if we're inside this CTE's definition (CTE doesn't shadow itself) + if with_table + .syntax() + .text_range() + .contains_range(name_ref.syntax().text_range()) + { + continue; + } + + if let Some(column_list) = with_table.column_list() { + for column in column_list.columns() { + if let Some(col_name) = column.name() + && Name::new(col_name.syntax().text().to_string()) == *column_name + { + return Some(SyntaxNodePtr::new(col_name.syntax())); + } + } + return None; + } + let query = with_table.query()?; let cte_select = match query { @@ -812,11 +837,23 @@ fn resolve_cte_column( let target_list = select_clause.target_list()?; for target in target_list.targets() { - if let Some((col_name, node)) = ColumnName::from_target(target) - && let Some(col_name_str) = col_name.to_string() - && Name::new(col_name_str) == *column_name - { - return Some(SyntaxNodePtr::new(&node)); + if let Some((col_name, node)) = ColumnName::from_target(target.clone()) { + if let Some(col_name_str) = col_name.to_string() + && Name::new(col_name_str) == *column_name + { + return Some(SyntaxNodePtr::new(&node)); + } + + if matches!(col_name, ColumnName::Star) { + if let Some(from_clause) = cte_select.from_clause() + && let Some(from_item) = from_clause.from_items().next() + && let Some(from_name_ref) = from_item.name_ref() + { + let from_table_name = + Name::new(from_name_ref.syntax().text().to_string()); + return resolve_cte_column(name_ref, &from_table_name, column_name); + } + } } } }