From 5400886287f7938267beb23ab4497c520c83e4cb Mon Sep 17 00:00:00 2001 From: Steve Dignam Date: Mon, 27 Apr 2026 20:02:42 -0400 Subject: [PATCH] ide: refactor resolve w/ ast_nav for create table + fix some hovers --- crates/squawk_ide/src/ast_nav.rs | 31 ++++++ crates/squawk_ide/src/collect.rs | 115 +++++++++------------- crates/squawk_ide/src/document_symbols.rs | 46 +++------ crates/squawk_ide/src/goto_definition.rs | 12 +++ crates/squawk_ide/src/hover.rs | 77 +++++++++++++-- crates/squawk_ide/src/resolve.rs | 106 ++++++++------------ 6 files changed, 220 insertions(+), 167 deletions(-) diff --git a/crates/squawk_ide/src/ast_nav.rs b/crates/squawk_ide/src/ast_nav.rs index 208e9cfc..cd7e3853 100644 --- a/crates/squawk_ide/src/ast_nav.rs +++ b/crates/squawk_ide/src/ast_nav.rs @@ -230,6 +230,37 @@ pub(crate) fn parent_source(node: &SyntaxNode) -> Option { None } +pub(crate) enum CreateTableArg { + Column(ast::Column), + Inherits(ast::Path), + LikeClause(ast::LikeClause), + TableConstraint(#[expect(unused)] ast::TableConstraint), +} + +pub(crate) fn create_table_args( + create_table: &impl ast::HasCreateTable, +) -> impl Iterator { + let inherits_iter = create_table + .inherits() + .into_iter() + .flat_map(|inherits| inherits.paths()) + .map(CreateTableArg::Inherits); + + let args_iter = create_table + .table_arg_list() + .into_iter() + .flat_map(|arg_list| arg_list.args()) + .map(|arg| match arg { + ast::TableArg::Column(column) => CreateTableArg::Column(column), + ast::TableArg::LikeClause(like_clause) => CreateTableArg::LikeClause(like_clause), + ast::TableArg::TableConstraint(constraint) => { + CreateTableArg::TableConstraint(constraint) + } + }); + + inherits_iter.chain(args_iter) +} + struct UnwrapParenExpr { current: Option, } diff --git a/crates/squawk_ide/src/collect.rs b/crates/squawk_ide/src/collect.rs index d37b5f3b..067f2b5f 100644 --- a/crates/squawk_ide/src/collect.rs +++ b/crates/squawk_ide/src/collect.rs @@ -37,48 +37,37 @@ fn columns_from_create_table_impl( return; } - if let Some(inherits) = create_table.inherits() { - for path in inherits.paths() { - if let Some((schema, table_name)) = name::schema_and_name_path(&path) { - let position = path.syntax().text_range().start(); - if let Some(ResolvedTableName::Table(parent_table)) = - resolve_table_name(db, file, &table_name, &schema, position) - { - columns_from_create_table_impl(db, file, &parent_table, columns, depth + 1); + for arg in ast_nav::create_table_args(create_table) { + match arg { + ast_nav::CreateTableArg::Inherits(path) => { + if let Some((schema, table_name)) = name::schema_and_name_path(&path) { + let position = path.syntax().text_range().start(); + if let Some(ResolvedTableName::Table(parent_table)) = + resolve_table_name(db, file, &table_name, &schema, position) + { + columns_from_create_table_impl(db, file, &parent_table, columns, depth + 1); + } } } - } - } - - if let Some(arg_list) = create_table.table_arg_list() { - for arg in arg_list.args() { - match &arg { - ast::TableArg::Column(column) => { - if let Some(name) = column.name() { - let col_name = Name::from_node(&name); - columns.push((col_name, Some(SyntaxNodePtr::new(name.syntax())))); - } + ast_nav::CreateTableArg::Column(column) => { + if let Some(name) = column.name() { + let col_name = Name::from_node(&name); + columns.push((col_name, Some(SyntaxNodePtr::new(name.syntax())))); } - ast::TableArg::LikeClause(like_clause) => { - if let Some(path) = like_clause.path() - && let Some((schema, table_name)) = name::schema_and_name_path(&path) + } + ast_nav::CreateTableArg::LikeClause(like_clause) => { + if let Some(path) = like_clause.path() + && let Some((schema, table_name)) = name::schema_and_name_path(&path) + { + let position = path.syntax().text_range().start(); + if let Some(ResolvedTableName::Table(source_table)) = + resolve_table_name(db, file, &table_name, &schema, position) { - let position = path.syntax().text_range().start(); - if let Some(ResolvedTableName::Table(source_table)) = - resolve_table_name(db, file, &table_name, &schema, position) - { - columns_from_create_table_impl( - db, - file, - &source_table, - columns, - depth + 1, - ); - } + columns_from_create_table_impl(db, file, &source_table, columns, depth + 1); } } - ast::TableArg::TableConstraint(_) => (), } + ast_nav::CreateTableArg::TableConstraint(_) => (), } } } @@ -105,45 +94,37 @@ fn table_columns_impl( let mut columns = vec![]; - if let Some(inherits) = create_table.inherits() { - for path in inherits.paths() { - if let Some((schema, table_name)) = name::schema_and_name_path(&path) { - let position = path.syntax().text_range().start(); - if let Some(ResolvedTableName::Table(parent_table)) = - resolve_table_name(db, file, &table_name, &schema, position) - { - let inherited_columns = table_columns_impl(db, file, &parent_table, depth + 1); - columns.extend(inherited_columns); + for arg in ast_nav::create_table_args(create_table) { + match arg { + ast_nav::CreateTableArg::Inherits(path) => { + if let Some((schema, table_name)) = name::schema_and_name_path(&path) { + let position = path.syntax().text_range().start(); + if let Some(ResolvedTableName::Table(parent_table)) = + resolve_table_name(db, file, &table_name, &schema, position) + { + columns.extend(table_columns_impl(db, file, &parent_table, depth + 1)); + } } } - } - } - - if let Some(arg_list) = create_table.table_arg_list() { - for arg in arg_list.args() { - match &arg { - ast::TableArg::Column(column) => { - if let Some(name) = column.name() { - let ty = column.ty().and_then(|ty| infer_type_from_ty(&ty)); - columns.push((Name::from_node(&name), ty)); - } + ast_nav::CreateTableArg::Column(column) => { + if let Some(name) = column.name() { + let ty = column.ty().and_then(|ty| infer_type_from_ty(&ty)); + columns.push((Name::from_node(&name), ty)); } - ast::TableArg::LikeClause(like_clause) => { - if let Some(path) = like_clause.path() - && let Some((schema, table_name)) = name::schema_and_name_path(&path) + } + ast_nav::CreateTableArg::LikeClause(like_clause) => { + if let Some(path) = like_clause.path() + && let Some((schema, table_name)) = name::schema_and_name_path(&path) + { + let position = path.syntax().text_range().start(); + if let Some(ResolvedTableName::Table(source_table)) = + resolve_table_name(db, file, &table_name, &schema, position) { - let position = path.syntax().text_range().start(); - if let Some(ResolvedTableName::Table(source_table)) = - resolve_table_name(db, file, &table_name, &schema, position) - { - let like_columns = - table_columns_impl(db, file, &source_table, depth + 1); - columns.extend(like_columns); - } + columns.extend(table_columns_impl(db, file, &source_table, depth + 1)); } } - ast::TableArg::TableConstraint(_) => (), } + ast_nav::CreateTableArg::TableConstraint(_) => (), } } diff --git a/crates/squawk_ide/src/document_symbols.rs b/crates/squawk_ide/src/document_symbols.rs index 116f7907..06d31c84 100644 --- a/crates/squawk_ide/src/document_symbols.rs +++ b/crates/squawk_ide/src/document_symbols.rs @@ -273,8 +273,7 @@ fn create_table_symbol( create_table: impl ast::HasCreateTable, ) -> Option { let path = create_table.path()?; - let segment = path.segment()?; - let name_node = segment.name()?; + let name_node = path.segment()?.name()?; let (schema, table_name) = resolve_table_info(db, file, &path)?; let name = format!("{}.{}", schema.0, table_name); @@ -283,13 +282,11 @@ fn create_table_symbol( let focus_range = name_node.syntax().text_range(); let mut children = vec![]; - if let Some(table_arg_list) = create_table.table_arg_list() { - for arg in table_arg_list.args() { - if let ast::TableArg::Column(column) = arg - && let Some(column_symbol) = create_column_symbol(column) - { - children.push(column_symbol); - } + for arg in create_table.table_arg_list()?.args() { + if let ast::TableArg::Column(column) = arg + && let Some(column_symbol) = create_column_symbol(column) + { + children.push(column_symbol); } } @@ -309,12 +306,7 @@ fn create_table_as_symbol( create_table_as: ast::CreateTableAs, ) -> Option { let path = create_table_as.path()?; - let segment = path.segment()?; - let name_node = if let Some(name) = segment.name() { - name.syntax().clone() - } else { - return None; - }; + let name_node = path.segment()?.name()?.syntax().clone(); let (schema, table_name) = resolve_table_info(db, file, &path)?; let name = format!("{}.{}", schema.0, table_name); @@ -340,8 +332,7 @@ fn create_view_symbol( create_view: ast::CreateView, ) -> Option { let path = create_view.path()?; - let segment = path.segment()?; - let name_node = segment.name()?; + let name_node = path.segment()?.name()?; let (schema, view_name) = resolve_view_info(db, file, &path)?; let name = format!("{}.{}", schema.0, view_name); @@ -391,8 +382,7 @@ fn create_materialized_view_symbol( create_view: ast::CreateMaterializedView, ) -> Option { let path = create_view.path()?; - let segment = path.segment()?; - let name_node = segment.name()?; + let name_node = path.segment()?.name()?; let (schema, view_name) = resolve_view_info(db, file, &path)?; let name = format!("{}.{}", schema.0, view_name); @@ -415,8 +405,7 @@ fn create_function_symbol( create_function: ast::CreateFunction, ) -> Option { let path = create_function.path()?; - let segment = path.segment()?; - let name_node = segment.name()?; + let name_node = path.segment()?.name()?; let (schema, function_name) = resolve_function_info(db, file, &path)?; let name = format!("{}.{}", schema.0, function_name); @@ -440,8 +429,7 @@ fn create_aggregate_symbol( create_aggregate: ast::CreateAggregate, ) -> Option { let path = create_aggregate.path()?; - let segment = path.segment()?; - let name_node = segment.name()?; + let name_node = path.segment()?.name()?; let (schema, aggregate_name) = resolve_aggregate_info(db, file, &path)?; let name = format!("{}.{}", schema.0, aggregate_name); @@ -465,8 +453,7 @@ fn create_procedure_symbol( create_procedure: ast::CreateProcedure, ) -> Option { let path = create_procedure.path()?; - let segment = path.segment()?; - let name_node = segment.name()?; + let name_node = path.segment()?.name()?; let (schema, procedure_name) = resolve_procedure_info(db, file, &path)?; let name = format!("{}.{}", schema.0, procedure_name); @@ -507,8 +494,7 @@ fn create_domain_symbol( create_domain: ast::CreateDomain, ) -> Option { let path = create_domain.path()?; - let segment = path.segment()?; - let name_node = segment.name()?; + let name_node = path.segment()?.name()?; let (schema, domain_name) = resolve_type_info(db, file, &path)?; let name = format!("{}.{}", schema.0, domain_name); @@ -532,8 +518,7 @@ fn create_sequence_symbol( create_sequence: ast::CreateSequence, ) -> Option { let path = create_sequence.path()?; - let segment = path.segment()?; - let name_node = segment.name()?; + let name_node = path.segment()?.name()?; let (schema, sequence_name) = resolve_sequence_info(db, file, &path)?; let name = format!("{}.{}", schema.0, sequence_name); @@ -716,8 +701,7 @@ fn create_type_symbol( create_type: ast::CreateType, ) -> Option { let path = create_type.path()?; - let segment = path.segment()?; - let name_node = segment.name()?; + let name_node = path.segment()?.name()?; let (schema, type_name) = resolve_type_info(db, file, &path)?; let name = format!("{}.{}", schema.0, type_name); diff --git a/crates/squawk_ide/src/goto_definition.rs b/crates/squawk_ide/src/goto_definition.rs index b598f6df..ef225cda 100644 --- a/crates/squawk_ide/src/goto_definition.rs +++ b/crates/squawk_ide/src/goto_definition.rs @@ -6177,6 +6177,18 @@ select COLUMN1$0, COLUMN2 from t; "); } + #[test] + fn goto_cte_values_with_explicit_column_list_column1_not_found() { + goto_not_found( + " +with t(a, b) as ( + values (1, 2), (3, 4) +) +select column1$0 from t; +", + ); + } + #[test] fn goto_qualified_column_with_schema_in_from_table() { assert_snapshot!(goto(" diff --git a/crates/squawk_ide/src/hover.rs b/crates/squawk_ide/src/hover.rs index dcdbeeff..a03faa80 100644 --- a/crates/squawk_ide/src/hover.rs +++ b/crates/squawk_ide/src/hover.rs @@ -278,8 +278,12 @@ fn column_name_from_node(column_name_node: &SyntaxNode) -> Option { .exprs() .position(|expr| expr.syntax() == column_name_node)?; Some(Name::from_string(format!("column{}", idx + 1))) + } else if let Some(name) = ast::Name::cast(column_name_node.clone()) { + Some(Name::from_node(&name)) } else { - Some(Name::from_string(column_name_node.text().to_string())) + let target = column_name_node.ancestors().find_map(ast::Target::cast)?; + let (col_name, _) = ColumnName::from_target(target)?; + Some(Name::from_string(col_name.to_string()?)) } } @@ -298,7 +302,7 @@ fn format_hover_for_column_ptr(db: &dyn Db, def: Location) -> Option { ast_nav::ParentSouce::ParenSelect(paren_select) => { // Qualified access like `t.a` if let Some(table_name) = subquery_alias_name(&paren_select) { - let column_name = Name::from_string(def_node.text().to_string()); + let column_name = column_name_from_node(def_node)?; return Some(Hover::snippet(ColumnHover::table_column( &table_name.to_string(), &column_name.to_string(), @@ -321,8 +325,7 @@ fn format_hover_for_column_ptr(db: &dyn Db, def: Location) -> Option { // select a from v; // ^ ast_nav::ParentSouce::CreateView(create_view) => { - let column_name = - ast::Name::cast(def_node.clone()).map(|name| Name::from_node(&name))?; + let column_name = column_name_from_node(def_node)?; let path = create_view.path()?; let (schema, view_name) = resolve::resolve_view_info(db, def.file, &path)?; return Some(Hover::snippet(ColumnHover::schema_table_column( @@ -342,8 +345,7 @@ fn format_hover_for_column_ptr(db: &dyn Db, def: Location) -> Option { ))); } ast_nav::ParentSouce::CreateTableAs(create_table_as) => { - let name = ast::Name::cast(def_node.clone())?; - let column_name = Name::from_node(&name); + let column_name = column_name_from_node(def_node)?; let path = create_table_as.path()?; let (schema, table_name) = resolve::resolve_table_info(db, def.file, &path)?; return Some(Hover::snippet(ColumnHover::schema_table_column( @@ -3407,6 +3409,69 @@ select *$0 from (select 1) as sub; "); } + #[test] + fn hover_on_view_inferred_column_anme() { + assert_snapshot!(check_hover(r#" +create view v as select 1; +select "?column?"$0 from v; +"#), @r#" + hover: column public.v.?column? + ╭▸ + 3 │ select "?column?" from v; + ╰╴ ─ hover + "#); + } + + #[test] + fn hover_on_cte_inferred_column_name() { + assert_snapshot!(check_hover(r#" +with x as (select 1) +select "?column?"$0 from x; +"#), @r#" + hover: column x.?column? + ╭▸ + 3 │ select "?column?" from x; + ╰╴ ─ hover + "#); + } + + #[test] + fn hover_on_create_table_as_inferred_column_name() { + assert_snapshot!(check_hover(r#" +create table t as select 1; +select "?column?"$0 from t; +"#), @r#" + hover: column public.t.?column? + ╭▸ + 3 │ select "?column?" from t; + ╰╴ ─ hover + "#); + } + + #[test] + fn hover_on_paren_select_inferred_column_name() { + assert_snapshot!(check_hover(r#" +select "?column?"$0 from (select 1); +"#), @r#" + hover: column ?column? integer + ╭▸ + 2 │ select "?column?" from (select 1); + ╰╴ ─ hover + "#); + } + + #[test] + fn hover_on_paren_select_aliased_inferred_column_name() { + assert_snapshot!(check_hover(r#" +select sub."?column?"$0 from (select 1) sub; +"#), @r#" + hover: column sub.?column? + ╭▸ + 2 │ select sub."?column?" from (select 1) sub; + ╰╴ ─ hover + "#); + } + #[test] fn hover_on_view_qualified_star() { assert_snapshot!(check_hover(" diff --git a/crates/squawk_ide/src/resolve.rs b/crates/squawk_ide/src/resolve.rs index 330010a9..851cf959 100644 --- a/crates/squawk_ide/src/resolve.rs +++ b/crates/squawk_ide/src/resolve.rs @@ -1777,37 +1777,33 @@ fn find_column_in_create_table_impl( return None; } - if let Some(inherits) = create_table.inherits() { - for path in inherits.paths() { - let (schema, table_name) = name::schema_and_name_path(&path)?; - let position = path.syntax().text_range().start(); - - if let Some(resolved) = resolve_table_name(db, file, &table_name, &schema, position) { - if let Some(ptr) = match resolved { - ResolvedTableName::Table(parent_table) => find_column_in_create_table_impl( - db, - file, - &parent_table, - column_name, - depth + 1, - ), - ResolvedTableName::TableAs(create_table_as) => { - find_column_in_create_table_as(file, &create_table_as, column_name) - } - ResolvedTableName::SelectInto(select_into) => { - find_column_in_select_into(file, &select_into, column_name) + for arg in ast_nav::create_table_args(create_table) { + match arg { + ast_nav::CreateTableArg::Inherits(path) => { + let (schema, table_name) = name::schema_and_name_path(&path)?; + let position = path.syntax().text_range().start(); + if let Some(resolved) = resolve_table_name(db, file, &table_name, &schema, position) + && let Some(ptr) = match resolved { + ResolvedTableName::Table(parent_table) => find_column_in_create_table_impl( + db, + file, + &parent_table, + column_name, + depth + 1, + ), + ResolvedTableName::TableAs(create_table_as) => { + find_column_in_create_table_as(file, &create_table_as, column_name) + } + ResolvedTableName::SelectInto(select_into) => { + find_column_in_select_into(file, &select_into, column_name) + } + ResolvedTableName::View(_) => None, } - ResolvedTableName::View(_) => None, - } { + { return Some(ptr); } } - } - } - - for arg in create_table.table_arg_list()?.args() { - match &arg { - ast::TableArg::Column(column) => { + ast_nav::CreateTableArg::Column(column) => { if let Some(name) = column.name() && Name::from_node(&name) == *column_name { @@ -1818,14 +1814,12 @@ fn find_column_in_create_table_impl( )]); } } - ast::TableArg::LikeClause(like_clause) => { + ast_nav::CreateTableArg::LikeClause(like_clause) => { let path = like_clause.path()?; let (schema, table_name) = name::schema_and_name_path(&path)?; let position = path.syntax().text_range().start(); - if let Some(resolved) = resolve_table_name(db, file, &table_name, &schema, position) - { - if let Some(ptr) = match resolved { + && let Some(ptr) = match resolved { ResolvedTableName::Table(source_table) => find_column_in_create_table_impl( db, file, @@ -1840,12 +1834,12 @@ fn find_column_in_create_table_impl( find_column_in_select_into(file, &select_into, column_name) } ResolvedTableName::View(_) => None, - } { - return Some(ptr); } + { + return Some(ptr); } } - ast::TableArg::TableConstraint(_) => (), + ast_nav::CreateTableArg::TableConstraint(_) => (), } } @@ -1877,10 +1871,7 @@ fn find_column_in_create_view_like( let select = ast_nav::select_from_variant(create_view.query()?)?; - let select_clause = select.select_clause()?; - let target_list = select_clause.target_list()?; - - for (idx, target) in target_list.targets().enumerate() { + for (idx, target) in select.select_clause()?.target_list()?.targets().enumerate() { if idx < column_list_len { continue; } @@ -1954,7 +1945,6 @@ fn resolve_cte_table(name_ref: &ast::NameRef, cte_name: &Name) -> Option Option { let (schema, table_name) = name::schema_and_name_path(path)?; let position = path.syntax().text_range().start(); - count_columns_for_table_name(db, file, &table_name, &schema, position, None) } @@ -2126,8 +2116,7 @@ fn resolve_cte_column( let cte_select = ast_nav::select_from_with_query(query)?; - let select_clause = cte_select.select_clause()?; - let target_list = select_clause.target_list()?; + let target_list = cte_select.select_clause()?.target_list()?; let from_clause = cte_select.from_clause(); let mut column_index: usize = 0; @@ -2209,19 +2198,14 @@ fn column_in_with_query( column_list_len: usize, ) -> Option> { let (returning_clause, path) = match query { - ast::WithQuery::Delete(delete) => ( - delete.returning_clause(), - delete.relation_name().and_then(|r| r.path()), - ), - ast::WithQuery::Insert(insert) => (insert.returning_clause(), insert.path()), - ast::WithQuery::Merge(merge) => ( - merge.returning_clause(), - merge.relation_name().and_then(|r| r.path()), - ), - ast::WithQuery::Update(update) => ( - update.returning_clause(), - update.relation_name().and_then(|r| r.path()), - ), + ast::WithQuery::Delete(delete) => { + (delete.returning_clause(), delete.relation_name()?.path()?) + } + ast::WithQuery::Insert(insert) => (insert.returning_clause(), insert.path()?), + ast::WithQuery::Merge(merge) => (merge.returning_clause(), merge.relation_name()?.path()?), + ast::WithQuery::Update(update) => { + (update.returning_clause(), update.relation_name()?.path()?) + } ast::WithQuery::Select(_) | ast::WithQuery::CompoundSelect(_) | ast::WithQuery::Table(_) @@ -2229,10 +2213,8 @@ fn column_in_with_query( | ast::WithQuery::ParenSelect(_) => return None, }; - let target_list = returning_clause?.target_list()?; - let path = path?; let mut column_index: usize = 0; - for target in target_list.targets() { + for target in returning_clause?.target_list()?.targets() { let target_column_count = if target.star_token().is_some() { count_columns_for_path(db, file, &path).unwrap_or(1) } else { @@ -2316,9 +2298,9 @@ fn resolve_subquery_column_ptr( if let Some(num_str) = column_name.0.strip_prefix("column") && let Ok(col_num) = num_str.parse::() && col_num > 0 - && let Some(row_list) = values.row_list() - && let Some(first_row) = row_list.rows().next() - && let Some(expr) = first_row.exprs().nth(col_num - 1) + && let Some(expr) = values + .row_list() + .and_then(|x| x.rows().next()?.exprs().nth(col_num - 1)) { return Some(smallvec![Location::new( file, @@ -2348,11 +2330,9 @@ fn resolve_column_from_select_targets( 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() { + for (i, target) in select.select_clause()?.target_list()?.targets().enumerate() { if let Some((col_name, node)) = ColumnName::from_target(target.clone()) { if i < skip_target_count && !matches!(col_name, ColumnName::Star) { continue;