diff --git a/crates/squawk_ide/src/goto_definition.rs b/crates/squawk_ide/src/goto_definition.rs index b0888dc5..1ca80caa 100644 --- a/crates/squawk_ide/src/goto_definition.rs +++ b/crates/squawk_ide/src/goto_definition.rs @@ -1307,6 +1307,77 @@ select * from t group by t.b$0; "); } + #[test] + fn goto_cte_table() { + assert_snapshot!(goto(" +with x as (select 1 as a) +select a from x$0; +"), @r" + ╭▸ + 2 │ with x as (select 1 as a) + │ ─ 2. destination + 3 │ select a from x; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_cte_column() { + assert_snapshot!(goto(" +with x as (select 1 as a) +select a$0 from x; +"), @r" + ╭▸ + 2 │ with x as (select 1 as a) + │ ─ 2. destination + 3 │ select a from x; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_cte_multiple_columns() { + assert_snapshot!(goto(" +with x as (select 1 as a, 2 as b) +select b$0 from x; +"), @r" + ╭▸ + 2 │ with x as (select 1 as a, 2 as b) + │ ─ 2. destination + 3 │ select b from x; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_cte_nested() { + assert_snapshot!(goto(" +with x as (select 1 as a), + y as (select a from x) +select a$0 from y; +"), @r" + ╭▸ + 3 │ y as (select a from x) + │ ─ 2. destination + 4 │ select a from y; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_cte_unnamed_column() { + assert_snapshot!(goto(r#" +with x as (select 1) +select "?column?"$0 from x; +"#), @r#" + ╭▸ + 2 │ with x as (select 1) + │ ─ 2. destination + 3 │ select "?column?" from x; + ╰╴ ─ 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 3c585a72..084bb198 100644 --- a/crates/squawk_ide/src/resolve.rs +++ b/crates/squawk_ide/src/resolve.rs @@ -5,6 +5,7 @@ use squawk_syntax::{ }; use crate::binder::Binder; +use crate::column_name::ColumnName; pub(crate) use crate::symbols::Schema; use crate::symbols::{Name, SymbolKind}; @@ -56,6 +57,13 @@ pub(crate) fn resolve_name_ref(binder: &Binder, name_ref: &ast::NameRef) -> Opti } else { None }; + + if schema.is_none() + && let Some(cte_ptr) = resolve_cte_table(name_ref, &table_name) + { + return Some(cte_ptr); + } + let position = name_ref.syntax().text_range().start(); resolve_table(binder, &table_name, &schema, position) } @@ -593,6 +601,12 @@ fn resolve_select_column(binder: &Binder, name_ref: &ast::NameRef) -> Option Option { + let select = name_ref.syntax().ancestors().find_map(ast::Select::cast)?; + 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 + { + return Some(SyntaxNodePtr::new(name.syntax())); + } + } + + None +} + +fn resolve_cte_column( + select: &ast::Select, + cte_name: &Name, + column_name: &Name, +) -> Option { + 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 + { + let query = with_table.query()?; + + let cte_select = match query { + ast::WithQuery::Select(s) => s, + ast::WithQuery::ParenSelect(ps) => match ps.select()? { + ast::SelectVariant::Select(s) => s, + _ => continue, + }, + _ => continue, + }; + + let select_clause = cte_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) + && let Some(col_name_str) = col_name.to_string() + && Name::new(col_name_str) == *column_name + { + return Some(SyntaxNodePtr::new(&node)); + } + } + } + } + + None +} + pub(crate) fn resolve_insert_table_columns( file: &ast::SourceFile, binder: &Binder,