From 0df1b2a3142415a32bf0b0e9155c73f255c924bd Mon Sep 17 00:00:00 2001 From: Steve Dignam Date: Fri, 6 Mar 2026 09:33:26 -0500 Subject: [PATCH 1/2] ide: goto def subquery compound selects --- crates/squawk_ide/src/goto_definition.rs | 37 ++++++ crates/squawk_ide/src/resolve.rs | 146 +++++++++++++---------- 2 files changed, 123 insertions(+), 60 deletions(-) diff --git a/crates/squawk_ide/src/goto_definition.rs b/crates/squawk_ide/src/goto_definition.rs index 47b75a37..466e57a2 100644 --- a/crates/squawk_ide/src/goto_definition.rs +++ b/crates/squawk_ide/src/goto_definition.rs @@ -4840,6 +4840,43 @@ select a$0 from (select 1 as a); "); } + #[test] + fn goto_subquery_compound_select_column() { + assert_snapshot!(goto(" +select c$0 from (select 1 c union select 2 c); +"), @r" + ╭▸ + 2 │ select c from (select 1 c union select 2 c); + ╰╴ ─ 1. source ─ 2. destination + "); + } + + #[test] + fn goto_subquery_compound_select_column_with_nested_parens() { + assert_snapshot!(goto(" +with t as ( + select 1 as c +) +select c$0 from ((select * from t) union all (select * from t)); +"), @r" + ╭▸ + 3 │ select 1 as c + │ ─ 2. destination + 4 │ ) + 5 │ select c from ((select * from t) union all (select * from t)); + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_subquery_compound_select_nested_parens_parses() { + goto_not_found( + " +select *$0 from (((select 1)) union (select 2)); +", + ); + } + #[test] fn goto_subquery_column_multiple_columns() { assert_snapshot!(goto(" diff --git a/crates/squawk_ide/src/resolve.rs b/crates/squawk_ide/src/resolve.rs index 58be843a..fdcba651 100644 --- a/crates/squawk_ide/src/resolve.rs +++ b/crates/squawk_ide/src/resolve.rs @@ -2210,14 +2210,30 @@ fn count_columns_for_table_name( } fn resolve_select_clause(query: SelectVariant) -> Option { - match query { - ast::SelectVariant::Select(s) => Some(s), - ast::SelectVariant::ParenSelect(ps) => match ps.select()? { - ast::SelectVariant::Select(s) => Some(s), - _ => return None, - }, - _ => return None, + leftmost_select_from_variant(query) +} + +fn leftmost_select_from_variant(select_variant: ast::SelectVariant) -> Option { + let mut current = select_variant; + + for _ in 0..10_000 { + match current { + ast::SelectVariant::Select(select) => return Some(select), + ast::SelectVariant::CompoundSelect(compound) => { + current = compound.lhs()?; + } + ast::SelectVariant::ParenSelect(paren_select) => { + current = paren_select.select()?; + } + ast::SelectVariant::SelectInto(_) + | ast::SelectVariant::Table(_) + | ast::SelectVariant::Values(_) => { + return None; + } + } } + + None } fn resolve_cte_column( @@ -2496,46 +2512,57 @@ fn resolve_subquery_column_ptr( return None; } - let ast::SelectVariant::Select(subquery_select) = select_variant else { - return None; - }; + let subquery_select = leftmost_select_from_variant(select_variant)?; + resolve_column_from_select_targets( + binder, + root, + &subquery_select, + name_ref, + column_name, + column_list_len, + ) +} - let select_clause = subquery_select.select_clause()?; +fn resolve_column_from_select_targets( + binder: &Binder, + root: &SyntaxNode, + select: &ast::Select, + name_ref: &ast::NameRef, + column_name: &Name, + skip_target_count: usize, +) -> Option { + let select_clause = select.select_clause()?; let target_list = select_clause.target_list()?; + let from_clause = select.from_clause(); for (i, target) in target_list.targets().enumerate() { if let Some((col_name, node)) = ColumnName::from_target(target.clone()) { - // skip targets that have been renamed by the column list - if i < column_list_len { - if matches!(col_name, ColumnName::Star) { - // star expands to multiple columns, handled below - } else { - continue; - } + if i < skip_target_count && !matches!(col_name, ColumnName::Star) { + continue; } if let Some(col_name_str) = col_name.to_string() && Name::from_string(col_name_str) == *column_name { return Some(SyntaxNodePtr::new(&node)); } - if matches!(col_name, ColumnName::Star) { - if let Some(from_clause) = subquery_select.from_clause() { - for from_item in from_clause.from_items() { - if let Some(result) = - resolve_from_item_column_ptr(binder, root, &from_item, name_ref) - { - return Some(result); - } + if matches!(col_name, ColumnName::Star) + && let Some(from_clause) = &from_clause + { + for from_item in from_clause.from_items() { + if let Some(result) = + resolve_from_item_column_ptr(binder, root, &from_item, name_ref) + { + return Some(result); } + } - for join_expr in from_clause.join_exprs() { - if let Some(result) = - resolve_from_join_expr(&join_expr, &|from_item: &ast::FromItem| { - resolve_from_item_column_ptr(binder, root, from_item, name_ref) - }) - { - return Some(result); - } + for join_expr in from_clause.join_exprs() { + if let Some(result) = + resolve_from_join_expr(&join_expr, &|from_item: &ast::FromItem| { + resolve_from_item_column_ptr(binder, root, from_item, name_ref) + }) + { + return Some(result); } } } @@ -2544,8 +2571,8 @@ fn resolve_subquery_column_ptr( if let Some(expr) = target.expr() && let ast::Expr::FieldExpr(field_expr) = expr && let Some(table_name) = qualified_star_table_name(&field_expr) - && let Some(from_clause) = subquery_select.from_clause() - && let Some(from_item) = find_from_item_in_from_clause(&from_clause, &table_name) + && let Some(from_clause) = &from_clause + && let Some(from_item) = find_from_item_in_from_clause(from_clause, &table_name) && let Some(result) = resolve_from_item_column_ptr(binder, root, &from_item, name_ref) { return Some(result); @@ -2933,20 +2960,15 @@ pub(crate) fn find_table_source(node: &SyntaxNode) -> Option { pub(crate) fn select_from_with_query(query: ast::WithQuery) -> Option { let select_variant = match query { - ast::WithQuery::Select(s) => ast::SelectVariant::Select(s), - ast::WithQuery::ParenSelect(ps) => ps.select()?, - ast::WithQuery::CompoundSelect(compound) => compound.lhs()?, + ast::WithQuery::Select(select) => ast::SelectVariant::Select(select), + ast::WithQuery::ParenSelect(paren_select) => paren_select.select()?, + ast::WithQuery::CompoundSelect(compound_select) => { + ast::SelectVariant::CompoundSelect(compound_select) + } _ => return None, }; - match select_variant { - ast::SelectVariant::Select(s) => Some(s), - ast::SelectVariant::CompoundSelect(compound) => match compound.lhs()? { - ast::SelectVariant::Select(s) => Some(s), - _ => None, - }, - _ => None, - } + leftmost_select_from_variant(select_variant) } fn count_columns_for_target( @@ -3150,19 +3172,7 @@ fn resolve_column_from_paren_expr( column_name: &Name, ) -> Option { if let Some(select) = paren_expr.select() { - if let Some(select_clause) = select.select_clause() - && let Some(target_list) = select_clause.target_list() - { - for target in target_list.targets() { - if let Some((col_name, node)) = ColumnName::from_target(target.clone()) - && let Some(col_name_str) = col_name.to_string() - && Name::from_string(col_name_str) == *column_name - { - return Some(SyntaxNodePtr::new(&node)); - } - } - } - return None; + return resolve_column_from_select_targets(binder, root, &select, name_ref, column_name, 0); } if let Some(ast::Expr::CallExpr(call_expr)) = paren_expr.expr() @@ -3195,6 +3205,22 @@ fn resolve_column_from_paren_expr( ); } + if let Some(select_variant) = from_item + .syntax() + .children() + .find_map(ast::SelectVariant::cast) + { + let select = leftmost_select_from_variant(select_variant)?; + return resolve_column_from_select_targets( + binder, + root, + &select, + name_ref, + column_name, + 0, + ); + } + if let Some(nested_paren_expr) = from_item.paren_expr() { return resolve_column_from_paren_expr( binder, From 865d092717eb44cd7711d5eb272a72eedb7b8a96 Mon Sep 17 00:00:00 2001 From: Steve Dignam Date: Fri, 6 Mar 2026 09:36:52 -0500 Subject: [PATCH 2/2] fix --- crates/squawk_ide/src/goto_definition.rs | 9 --------- 1 file changed, 9 deletions(-) diff --git a/crates/squawk_ide/src/goto_definition.rs b/crates/squawk_ide/src/goto_definition.rs index 466e57a2..1e27d593 100644 --- a/crates/squawk_ide/src/goto_definition.rs +++ b/crates/squawk_ide/src/goto_definition.rs @@ -4868,15 +4868,6 @@ select c$0 from ((select * from t) union all (select * from t)); "); } - #[test] - fn goto_subquery_compound_select_nested_parens_parses() { - goto_not_found( - " -select *$0 from (((select 1)) union (select 2)); -", - ); - } - #[test] fn goto_subquery_column_multiple_columns() { assert_snapshot!(goto("