From 5aba8360bb4f5d4cda7d0bf791a9e3ed13dfca1c Mon Sep 17 00:00:00 2001 From: Steve Dignam Date: Mon, 28 Sep 2026 19:44:41 -0400 Subject: [PATCH] ide: fix goto def bug with * expansion --- crates/squawk_ide/src/goto_definition.rs | 46 +++++++ crates/squawk_ide/src/resolve.rs | 158 ++++++++++++++++++----- 2 files changed, 174 insertions(+), 30 deletions(-) diff --git a/crates/squawk_ide/src/goto_definition.rs b/crates/squawk_ide/src/goto_definition.rs index bcd232dd..3ae5873c 100644 --- a/crates/squawk_ide/src/goto_definition.rs +++ b/crates/squawk_ide/src/goto_definition.rs @@ -1406,6 +1406,16 @@ select * from ( "); } + #[test] + fn goto_recursive_cte_star_column_count_cycle_not_found() { + goto_not_found( + " +with recursive t as (select * from t) +select x$0 from t as t(a); +", + ); + } + #[test] fn goto_cte_forward_ref_ignored_for_qualified_star() { assert_snapshot!(goto(" @@ -9567,6 +9577,42 @@ select a$0 from (select * from t) u(x); ); } + #[test] + fn goto_cte_star_column_count_through_nested_ctes() { + assert_snapshot!(goto(" +with s as (select 1 a, 2 b), + u as (select * from s), + t as (select 3 c), + v as (select * from u, t) +select c$0 from v as x(p, q); +"), @" + ╭▸ + 4 │ t as (select 3 c), + │ ─ 2. destination + 5 │ v as (select * from u, t) + 6 │ select c from v as x(p, q); + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_view_star_column_count_through_nested_views() { + assert_snapshot!(goto(" +create table s (a int, b int); +create view u as select * from s; +create table t (c int); +create view v as select * from u, t; +select c$0 from v as x(p, q); +"), @" + ╭▸ + 4 │ create table t (c int); + │ ─ 2. destination + 5 │ create view v as select * from u, t; + 6 │ select c from v as x(p, q); + ╰╴ ─ 1. source + "); + } + #[test] fn goto_subquery_column_with_as() { assert_snapshot!(goto(" diff --git a/crates/squawk_ide/src/resolve.rs b/crates/squawk_ide/src/resolve.rs index f49e7324..31252edf 100644 --- a/crates/squawk_ide/src/resolve.rs +++ b/crates/squawk_ide/src/resolve.rs @@ -3,6 +3,7 @@ use crate::{ db::{ancestors_with_embedded, binders}, }; use rowan::TextSize; +use rustc_hash::FxHashSet; use smallvec::{SmallVec, smallvec}; use squawk_syntax::{ SyntaxKind, SyntaxNode, SyntaxNodePtr, @@ -3988,20 +3989,23 @@ fn count_columns_for_path(db: &dyn Db, path: InFile<&ast::PathRef>) -> Option>; + fn count_columns_for_cte_or_table_name( db: &dyn Db, table_name: &Name, schemas: &ResolvedSchemas, file: FileId, name_ref: &impl ast::NameLike, + visited: &mut ColumnCountVisited, ) -> Option { if schemas.unqualified() && let Some(with_table) = ast_nav::find_cte_with_table(name_ref, table_name) { - return count_columns_for_with_table(db, file, with_table); + return count_columns_for_with_table(db, file, with_table, visited); } - count_columns_for_table_name(db, table_name, schemas, file) + count_columns_for_table_name_impl(db, table_name, schemas, file, visited) } fn count_columns_for_table_name( @@ -4009,10 +4013,33 @@ fn count_columns_for_table_name( table_name: &Name, schemas: &ResolvedSchemas, file: FileId, +) -> Option { + count_columns_for_table_name_impl(db, table_name, schemas, file, &mut FxHashSet::default()) +} + +fn count_columns_for_table_name_impl( + db: &dyn Db, + table_name: &Name, + schemas: &ResolvedSchemas, + file: FileId, + visited: &mut ColumnCountVisited, ) -> Option { let (table_like_ptr, kind) = resolve_view_or_table(db, table_name, schemas, file)?; - let file = table_like_ptr.file_id; + if !visited.insert(table_like_ptr) { + return None; + } + let result = count_columns_for_table_like(db, table_like_ptr, kind, visited); + visited.remove(&table_like_ptr); + result +} +fn count_columns_for_table_like( + db: &dyn Db, + table_like_ptr: InFile, + kind: LocationKind, + visited: &mut ColumnCountVisited, +) -> Option { + let file = table_like_ptr.file_id; match kind { LocationKind::Table => { let table_like_node = table_like_ptr.to_node(db); @@ -4049,14 +4076,19 @@ fn count_columns_for_table_name( db, file, &create_table_as.query()?.select_variant()?, + visited, ); } if let Some(select_into) = table_like_node.ancestors().find_map(ast::SelectInto::cast) { - if let Some(target_list) = select_into.select_clause().and_then(|c| c.target_list()) - { - return Some(target_list.targets().count()); - } + let target_list = select_into.select_clause()?.target_list()?; + return count_columns_for_target_list( + db, + file, + &target_list, + select_into.from_clause().as_ref(), + visited, + ); } None @@ -4070,7 +4102,7 @@ fn count_columns_for_table_name( return Some(column_list.column_names().count()); } - count_columns_for_select_variant(db, file, &create_view.query()?) + count_columns_for_select_variant(db, file, &create_view.query()?, visited) } _ => None, } @@ -4080,26 +4112,43 @@ fn count_columns_for_with_table( db: &dyn Db, file: FileId, with_table: ast::WithTable, + visited: &mut ColumnCountVisited, ) -> Option { if let Some(column_list) = with_table.column_list() { return Some(column_list.column_names().count()); } - count_columns_for_with_query(db, file, with_table.query()?) + let query = with_table.query()?; + let with_table_ptr = InFile::new(file, SyntaxNodePtr::new(with_table.syntax())); + if !visited.insert(with_table_ptr) { + return None; + } + let result = count_columns_for_with_query(db, file, query, visited); + visited.remove(&with_table_ptr); + result } -fn count_columns_for_with_query(db: &dyn Db, file: FileId, query: ast::WithQuery) -> Option { +fn count_columns_for_with_query( + db: &dyn Db, + file: FileId, + query: ast::WithQuery, + visited: &mut ColumnCountVisited, +) -> Option { match query { ast::WithQuery::CompoundSelect(compound_select) => { - count_columns_for_select_variant(db, file, &compound_select.lhs()?) + count_columns_for_select_variant(db, file, &compound_select.lhs()?, visited) } ast::WithQuery::ParenSelect(paren_select) => { - count_columns_for_select_variant(db, file, &paren_select.select()?) + count_columns_for_select_variant(db, file, &paren_select.select()?, visited) } - ast::WithQuery::Select(select) => Some(count_columns_for_target_list( + ast::WithQuery::Select(select) => count_columns_for_target_list( + db, + file, &select.select_clause()?.target_list()?, - )), - ast::WithQuery::Table(table) => count_columns_for_table_query(db, file, &table), + select.from_clause().as_ref(), + visited, + ), + ast::WithQuery::Table(table) => count_columns_for_table_query(db, file, &table, visited), ast::WithQuery::Values(values) => count_columns_for_values(&values), ast::WithQuery::Delete(_) | ast::WithQuery::Insert(_) @@ -4112,36 +4161,65 @@ fn count_columns_for_select_variant( db: &dyn Db, file: FileId, select_variant: &ast::SelectVariant, + visited: &mut ColumnCountVisited, ) -> Option { match select_variant { ast::SelectVariant::CompoundSelect(compound_select) => { - count_columns_for_select_variant(db, file, &compound_select.lhs()?) + count_columns_for_select_variant(db, file, &compound_select.lhs()?, visited) } ast::SelectVariant::ParenSelect(paren_select) => { - count_columns_for_select_variant(db, file, &paren_select.select()?) + count_columns_for_select_variant(db, file, &paren_select.select()?, visited) } - ast::SelectVariant::Select(select) => Some(count_columns_for_target_list( + ast::SelectVariant::Select(select) => count_columns_for_target_list( + db, + file, &select.select_clause()?.target_list()?, - )), - ast::SelectVariant::SelectInto(select_into) => Some(count_columns_for_target_list( + select.from_clause().as_ref(), + visited, + ), + ast::SelectVariant::SelectInto(select_into) => count_columns_for_target_list( + db, + file, &select_into.select_clause()?.target_list()?, - )), - ast::SelectVariant::Table(table) => count_columns_for_table_query(db, file, table), + select_into.from_clause().as_ref(), + visited, + ), + ast::SelectVariant::Table(table) => count_columns_for_table_query(db, file, table, visited), ast::SelectVariant::Values(values) => count_columns_for_values(values), } } -fn count_columns_for_target_list(target_list: &ast::TargetList) -> usize { - target_list.targets().count() +fn count_columns_for_target_list( + db: &dyn Db, + file: FileId, + target_list: &ast::TargetList, + from_clause: Option<&ast::FromClause>, + visited: &mut ColumnCountVisited, +) -> Option { + let mut total = 0usize; + for target in target_list.targets() { + let count = if let Some(from_clause) = from_clause { + count_columns_for_target_impl(db, InFile::new(file, &target), from_clause, visited)? + } else { + 1 + }; + total = total.saturating_add(count); + } + Some(total) } -fn count_columns_for_table_query(db: &dyn Db, file: FileId, table: &ast::Table) -> Option { +fn count_columns_for_table_query( + db: &dyn Db, + file: FileId, + table: &ast::Table, + visited: &mut ColumnCountVisited, +) -> Option { let path = table.relation_name()?.relation_name_ref()?.path_ref()?; let (schema, table_name) = name::schema_and_name_path(&path)?; let table_name_ref = relation_name_ref_from_table(table)?; let position = table_name_ref.syntax().text_range().start(); let schemas = bind(db, file).resolved_schemas(position, schema.as_ref()); - count_columns_for_cte_or_table_name(db, &table_name, &schemas, file, &table_name_ref) + count_columns_for_cte_or_table_name(db, &table_name, &schemas, file, &table_name_ref, visited) } fn count_columns_for_values(values: &ast::Values) -> Option { @@ -4751,18 +4829,27 @@ fn count_columns_for_target( db: &dyn Db, target: InFile<&ast::Target>, from_clause: &ast::FromClause, +) -> Option { + count_columns_for_target_impl(db, target, from_clause, &mut FxHashSet::default()) +} + +fn count_columns_for_target_impl( + db: &dyn Db, + target: InFile<&ast::Target>, + from_clause: &ast::FromClause, + visited: &mut ColumnCountVisited, ) -> Option { let file = target.file_id; let target = target.value; if target.star_token().is_some() { - return count_columns_for_from_clause(db, InFile::new(file, from_clause)); + return count_columns_for_from_clause(db, InFile::new(file, from_clause), visited); } if let Some(ast::Expr::FieldExpr(field_expr)) = target.expr() && let Some(table_name) = qualified_star_table_name(&field_expr) && let Some(from_item) = find_from_item_in_from_clause(from_clause, &table_name) { - return count_columns_for_from_item(db, InFile::new(file, &from_item)); + return count_columns_for_from_item_impl(db, InFile::new(file, &from_item), visited); } Some(1) @@ -4771,6 +4858,7 @@ fn count_columns_for_target( fn count_columns_for_from_clause( db: &dyn Db, from_clause: InFile<&ast::FromClause>, + visited: &mut ColumnCountVisited, ) -> Option { let file = from_clause.file_id; let from_clause = from_clause.value; @@ -4778,7 +4866,9 @@ fn count_columns_for_from_clause( let mut found = false; for from_item in ast_nav::iter_from_clause(from_clause) { - if let Some(count) = count_columns_for_from_item(db, InFile::new(file, &from_item)) { + if let Some(count) = + count_columns_for_from_item_impl(db, InFile::new(file, &from_item), visited) + { total = total.saturating_add(count); found = true; } @@ -4788,13 +4878,21 @@ fn count_columns_for_from_clause( } fn count_columns_for_from_item(db: &dyn Db, from_item: InFile<&ast::FromItem>) -> Option { + count_columns_for_from_item_impl(db, from_item, &mut FxHashSet::default()) +} + +fn count_columns_for_from_item_impl( + db: &dyn Db, + from_item: InFile<&ast::FromItem>, + visited: &mut ColumnCountVisited, +) -> Option { let file = from_item.file_id; let from_item = from_item.value; let (schema, table_name) = name::schema_and_table_from_from_item(from_item)?; let scope_name_ref = relation_name_ref_from_from_item(from_item)?; let position = scope_name_ref.syntax().text_range().start(); let schemas = bind(db, file).resolved_schemas(position, schema.as_ref()); - count_columns_for_cte_or_table_name(db, &table_name, &schemas, file, &scope_name_ref) + count_columns_for_cte_or_table_name(db, &table_name, &schemas, file, &scope_name_ref, visited) } fn resolve_from_clause_column_after_index(