From 937b8ea7ef15eba5755f4a24d6062fb621aa905e Mon Sep 17 00:00:00 2001 From: Steve Dignam Date: Thu, 26 Feb 2026 19:12:39 -0500 Subject: [PATCH] ide: lateral joins, subqueries with table, hover aliases --- crates/squawk_ide/src/completion.rs | 121 +++++- crates/squawk_ide/src/goto_definition.rs | 106 +++++ crates/squawk_ide/src/hover.rs | 400 ++++++++++++++++--- crates/squawk_ide/src/resolve.rs | 479 +++++++++++++++++++---- 4 files changed, 957 insertions(+), 149 deletions(-) diff --git a/crates/squawk_ide/src/completion.rs b/crates/squawk_ide/src/completion.rs index 2b5273c9..078eb8de 100644 --- a/crates/squawk_ide/src/completion.rs +++ b/crates/squawk_ide/src/completion.rs @@ -356,7 +356,11 @@ fn column_completions_from_clause( })); } Some(resolve::TableSource::WithTable(with_table)) => { - let columns = resolve::collect_with_table_columns_with_types(&with_table); + let columns = resolve::collect_with_table_columns_with_types( + binder, + file.syntax(), + &with_table, + ); completions.extend(columns.into_iter().map(|(name, ty)| CompletionItem { label: name.to_string(), kind: CompletionItemKind::Column, @@ -394,20 +398,42 @@ fn column_completions_from_clause( })); } Some(resolve::TableSource::Alias(alias)) => { - if let Some(column_list) = alias.column_list() { - completions.extend(column_list.columns().filter_map(|column| { - let name = column.name()?; - Some(CompletionItem { - label: Name::from_node(&name).to_string(), + let alias_columns: Vec = alias + .column_list() + .into_iter() + .flat_map(|column_list| column_list.columns()) + .filter_map(|column| column.name().map(|name| Name::from_node(&name))) + .collect(); + + let base_columns = alias_base_columns_with_types(binder, file, &alias); + + for (idx, alias_column) in alias_columns.iter().enumerate() { + completions.push(CompletionItem { + label: alias_column.to_string(), + kind: CompletionItemKind::Column, + detail: base_columns.get(idx).and_then(|(_, ty)| ty.clone()), + insert_text: None, + insert_text_format: None, + trigger_completion_after_insert: false, + sort_text: Some(format!("{idx:04}")), + }); + } + + completions.extend( + base_columns + .into_iter() + .skip(alias_columns.len()) + .enumerate() + .map(|(idx, (name, ty))| CompletionItem { + label: name.to_string(), kind: CompletionItemKind::Column, - detail: None, + detail: ty, insert_text: None, insert_text_format: None, trigger_completion_after_insert: false, - sort_text: None, - }) - })); - } + sort_text: Some(format!("{:04}", idx + alias_columns.len())), + }), + ); } Some(resolve::TableSource::ParenSelect(paren_select)) => { let columns = resolve::collect_paren_select_columns_with_types( @@ -431,6 +457,59 @@ fn column_completions_from_clause( completions } +fn alias_base_columns_with_types( + binder: &binder::Binder, + file: &ast::SourceFile, + alias: &ast::Alias, +) -> Vec<(Name, Option)> { + let Some(from_item) = alias.syntax().ancestors().find_map(ast::FromItem::cast) else { + return vec![]; + }; + let Some(table_ptr) = resolve::table_ptr_from_from_item(binder, &from_item) else { + return vec![]; + }; + + let table_node = table_ptr.to_node(file.syntax()); + + match resolve::find_table_source(&table_node) { + Some(resolve::TableSource::CreateTable(create_table)) => { + resolve::collect_table_columns(binder, file.syntax(), &create_table) + .into_iter() + .filter_map(|column| { + let name = column.name()?; + let detail = column.ty().map(|t| t.syntax().text().to_string()); + Some((Name::from_node(&name), detail)) + }) + .collect() + } + Some(resolve::TableSource::WithTable(with_table)) => { + resolve::collect_with_table_columns_with_types(binder, file.syntax(), &with_table) + .into_iter() + .map(|(name, ty)| (name, ty.map(|t| t.to_string()))) + .collect() + } + Some(resolve::TableSource::CreateView(create_view)) => { + resolve::collect_view_columns_with_types(&create_view) + .into_iter() + .map(|(name, ty)| (name, ty.map(|t| t.to_string()))) + .collect() + } + Some(resolve::TableSource::CreateMaterializedView(create_materialized_view)) => { + resolve::collect_materialized_view_columns_with_types(&create_materialized_view) + .into_iter() + .map(|(name, ty)| (name, ty.map(|t| t.to_string()))) + .collect() + } + Some(resolve::TableSource::ParenSelect(paren_select)) => { + resolve::collect_paren_select_columns_with_types(binder, file.syntax(), &paren_select) + .into_iter() + .map(|(name, ty)| (name, ty.map(|t| t.to_string()))) + .collect() + } + Some(resolve::TableSource::Alias(_)) | None => vec![], + } +} + fn schema_completions(binder: &binder::Binder) -> Vec { let builtin_schemas = [ "public", @@ -1098,6 +1177,26 @@ select $0 from t; "); } + #[test] + fn completion_after_select_with_cte_alias_column_list() { + assert_snapshot!(completions(" +with t as (select 1 a, 2 b, 3 c) +select $0 from t as u(x, y); +"), @r" + label | kind | detail + --------------------+----------+--------- + x | Column | integer + y | Column | integer + c | Column | integer + * | Operator | + public | Schema | + pg_catalog | Schema | + pg_temp | Schema | + pg_toast | Schema | + information_schema | Schema | + "); + } + #[test] fn completion_values_cte() { assert_snapshot!(completions(" diff --git a/crates/squawk_ide/src/goto_definition.rs b/crates/squawk_ide/src/goto_definition.rs index d7c9ef28..8389a14d 100644 --- a/crates/squawk_ide/src/goto_definition.rs +++ b/crates/squawk_ide/src/goto_definition.rs @@ -596,6 +596,84 @@ select * from t, json_to_recordset(t.x$0) as r(a int, b int); "#); } + #[test] + fn goto_lateral_values_alias_in_subquery() { + assert_snapshot!(goto(" +select u.n, x.val +from (values (1), (2)) u(n) +cross join lateral (select u$0.n * 10 as val) x; +"), @r" + ╭▸ + 3 │ from (values (1), (2)) u(n) + │ ─ 2. destination + 4 │ cross join lateral (select u.n * 10 as val) x; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_lateral_missing_not_found() { + // Query 1 ERROR at Line 3: : ERROR: invalid reference to FROM-clause entry for table "u" + // LINE 3: cross join (select u.n * 10 as val) x; + // ^ + // DETAIL: There is an entry for table "u", but it cannot be referenced from this part of the query. + // HINT: To reference that table, you must mark this subquery with LATERAL. + goto_not_found( + " +select u.n, x.val +from (values (1), (2)) u(n) +cross join (select u$0.n * 10 as val) x; +", + ); + } + + #[test] + fn goto_lateral_deeply_nested_paren_expr_values_alias_in_subquery() { + assert_snapshot!(goto(" +select u.n, x.val +from (values (1), (2)) u(n) +cross join lateral ((((select u$0.n * 10 as val)))) x; +"), @r" + ╭▸ + 3 │ from (values (1), (2)) u(n) + │ ─ 2. destination + 4 │ cross join lateral ((((select u.n * 10 as val)))) x; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_lateral_deeply_nested_paren_expr_values_alias_column() { + assert_snapshot!(goto(" +select u.n, x.val$0 +from (values (1), (2)) u(n) +cross join lateral ((((select u.n * 10 as val)))) x; +"), @r" + ╭▸ + 2 │ select u.n, x.val + │ ─ 1. source + 3 │ from (values (1), (2)) u(n) + 4 │ cross join lateral ((((select u.n * 10 as val)))) x; + ╰╴ ─── 2. destination + "); + } + + #[test] + fn goto_lateral_deeply_nested_paren_expr_missing_not_found() { + // Query 1 ERROR at Line 3: : ERROR: invalid reference to FROM-clause entry for table "u" + // LINE 3: cross join ((((select u.n * 10 as val)))) x; + // ^ + // DETAIL: There is an entry for table "u", but it cannot be referenced from this part of the query. + // HINT: To reference that table, you must mark this subquery with LATERAL. + goto_not_found( + " +select u.n, x.val +from (values (1), (2)) u(n) +cross join ((((select u$0.n * 10 as val)))) x; +", + ); + } + #[test] fn goto_drop_sequence() { assert_snapshot!(goto(" @@ -4822,6 +4900,20 @@ select b$0 from (values (1, 2)) t(a, b); "); } + #[test] + fn goto_values_column_alias_list_nested_parens() { + assert_snapshot!(goto(" +select n$0 +from ((values (1), (2))) u(n); +"), @r" + ╭▸ + 2 │ select n + │ ─ 1. source + 3 │ from ((values (1), (2))) u(n); + ╰╴ ─ 2. destination + "); + } + #[test] fn goto_table_expr_column() { assert_snapshot!(goto(" @@ -4850,6 +4942,20 @@ select a$0 from (table x); "); } + #[test] + fn goto_table_expr_cte_table() { + assert_snapshot!(goto(" +with t as (select 1 a, 2 b) +select * from (table t$0); +"), @r" + ╭▸ + 2 │ with t as (select 1 a, 2 b) + │ ─ 2. destination + 3 │ select * from (table t); + ╰╴ ─ 1. source + "); + } + #[test] fn goto_table_expr_partial_column_alias_list() { assert_snapshot!(goto(" diff --git a/crates/squawk_ide/src/hover.rs b/crates/squawk_ide/src/hover.rs index a4119800..09cdd653 100644 --- a/crates/squawk_ide/src/hover.rs +++ b/crates/squawk_ide/src/hover.rs @@ -426,9 +426,13 @@ fn hover_column_definition( // TODO: we should pass in the file id along with the def_node and then use // salsa to lookup the the correct binder. -fn format_table_source(source: resolve::TableSource, binder: &binder::Binder) -> Option { +fn format_table_source( + root: &SyntaxNode, + source: resolve::TableSource, + binder: &binder::Binder, +) -> Option { match source { - resolve::TableSource::Alias(alias) => format_alias_with_column_list(&alias), + resolve::TableSource::Alias(alias) => format_alias_with_column_list(root, &alias, binder), resolve::TableSource::WithTable(with_table) => format_with_table(&with_table), resolve::TableSource::CreateView(create_view) => format_create_view(&create_view, binder), resolve::TableSource::CreateMaterializedView(create_materialized_view) => { @@ -446,19 +450,48 @@ fn hover_table(binder: &binder::Binder, def_node: &SyntaxNode) -> Option return Some(result); } - if let Some(source) = resolve::find_table_source(def_node) { - return format_table_source(source, binder); + if let Some(source) = resolve::find_table_source(def_node) + && let Some(root) = def_node.ancestors().last() + { + return format_table_source(&root, source, binder); } None } -fn format_alias_with_column_list(alias: &ast::Alias) -> Option { +fn format_alias_with_column_list( + root: &SyntaxNode, + alias: &ast::Alias, + binder: &binder::Binder, +) -> Option { let alias_name = alias.name()?; - let column_list = alias.column_list()?; - let name = Name::from_node(&alias_name).to_string(); - let columns = column_list.syntax().text().to_string(); - Some(format!("table {}{}", name, columns)) + let name = Name::from_node(&alias_name); + + let mut columns: Vec = alias + .column_list()? + .columns() + .filter_map(|column| { + column + .name() + .map(|column_name| Name::from_node(&column_name)) + }) + .collect(); + + if let Some(from_item) = alias.syntax().ancestors().find_map(ast::FromItem::cast) + && let Some(table_ptr) = resolve::table_ptr_from_from_item(binder, &from_item) + { + let base_columns = collect_star_column_names(root, &table_ptr, binder); + for column in base_columns.iter().skip(columns.len()) { + columns.push(column.clone()); + } + } + + let columns = columns + .iter() + .map(|column| column.to_string()) + .collect::>() + .join(", "); + Some(format!("table {}({})", name, columns)) } fn hover_qualified_star( @@ -475,12 +508,13 @@ fn hover_unqualified_star( target: &ast::Target, binder: &binder::Binder, ) -> Option { - let table_ptrs = resolve::resolve_unqualified_star_table_ptrs(binder, target)?; - let mut results = vec![]; - for table_ptr in table_ptrs { - if let Some(columns) = hover_qualified_star_columns(root, &table_ptr, binder) { - results.push(columns); - } + let mut results = hover_unqualified_star_with_binder(root, target, binder); + + if results.is_empty() && target_has_schema_qualified_from_item(target) { + let builtins_tree = ast::SourceFile::parse(BUILTINS_SQL).tree(); + let builtins_binder = binder::bind(&builtins_tree); + results = + hover_unqualified_star_with_binder(builtins_tree.syntax(), target, &builtins_binder); } if results.is_empty() { @@ -490,6 +524,41 @@ fn hover_unqualified_star( Some(results.join("\n")) } +fn hover_unqualified_star_with_binder( + root: &SyntaxNode, + target: &ast::Target, + binder: &binder::Binder, +) -> Vec { + let mut results = vec![]; + + if let Some(table_ptrs) = resolve::resolve_unqualified_star_table_ptrs(binder, target) { + for table_ptr in table_ptrs { + if let Some(columns) = hover_qualified_star_columns(root, &table_ptr, binder) { + results.push(columns); + } + } + } + + results +} + +fn target_has_schema_qualified_from_item(target: &ast::Target) -> bool { + let Some(select) = target.syntax().ancestors().find_map(ast::Select::cast) else { + return false; + }; + let Some(from_clause) = select.from_clause() else { + return false; + }; + + for from_item in from_clause.from_items() { + if from_item.field_expr().is_some() { + return true; + } + } + + false +} + fn hover_unqualified_star_in_arg_list( root: &SyntaxNode, arg_list: &ast::ArgList, @@ -536,9 +605,11 @@ fn hover_qualified_star_columns( } match resolve::find_table_source(&table_name_node)? { - resolve::TableSource::Alias(alias) => hover_qualified_star_columns_from_alias(&alias), + resolve::TableSource::Alias(alias) => { + hover_qualified_star_columns_from_alias(root, &alias, binder) + } resolve::TableSource::WithTable(with_table) => { - hover_qualified_star_columns_from_cte(&with_table) + hover_qualified_star_columns_from_cte(root, &with_table, binder) } resolve::TableSource::CreateTable(create_table) => { hover_qualified_star_columns_from_table(root, &create_table, binder) @@ -555,25 +626,100 @@ fn hover_qualified_star_columns( } } -fn hover_qualified_star_columns_from_alias(alias: &ast::Alias) -> Option { +fn hover_qualified_star_columns_from_alias( + root: &SyntaxNode, + alias: &ast::Alias, + binder: &binder::Binder, +) -> Option { let alias_name = Name::from_node(&alias.name()?); - let columns = alias.column_list()?.columns(); - let results: Vec = columns - .filter_map(|column| { - let column_name = Name::from_node(&column.name()?); - Some(ColumnHover::table_column( - &alias_name.to_string(), - &column_name.to_string(), - )) - }) + let alias_columns: Vec = alias + .column_list()? + .columns() + .filter_map(|column| column.name().map(|name| Name::from_node(&name))) .collect(); - if results.is_empty() { + if alias_columns.is_empty() { return None; } + let mut results: Vec = alias_columns + .iter() + .map(|column_name| { + ColumnHover::table_column(&alias_name.to_string(), &column_name.to_string()) + }) + .collect(); + + let from_item = alias.syntax().ancestors().find_map(ast::FromItem::cast)?; + let table_ptr = resolve::table_ptr_from_from_item(binder, &from_item)?; + let base_column_names = collect_star_column_names(root, &table_ptr, binder); + + for column_name in base_column_names.iter().skip(alias_columns.len()) { + results.push(ColumnHover::table_column( + &alias_name.to_string(), + &column_name.to_string(), + )); + } + Some(results.join("\n")) } + +fn collect_star_column_names( + root: &SyntaxNode, + table_ptr: &squawk_syntax::SyntaxNodePtr, + binder: &binder::Binder, +) -> Vec { + let table_name_node = table_ptr.to_node(root); + + if let Some(paren_select) = ast::ParenSelect::cast(table_name_node.clone()) { + return resolve::collect_paren_select_columns_with_types(binder, root, &paren_select) + .into_iter() + .map(|(name, _ty)| name) + .collect(); + } + + match resolve::find_table_source(&table_name_node) { + Some(resolve::TableSource::Alias(alias)) => alias + .column_list() + .into_iter() + .flat_map(|column_list| column_list.columns()) + .filter_map(|column| column.name().map(|name| Name::from_node(&name))) + .collect(), + Some(resolve::TableSource::WithTable(with_table)) => { + let columns = resolve::collect_with_table_column_names(binder, root, &with_table); + if !columns.is_empty() { + return columns; + } + + let builtins_tree = ast::SourceFile::parse(BUILTINS_SQL).tree(); + let builtins_binder = binder::bind(&builtins_tree); + resolve::collect_with_table_column_names( + &builtins_binder, + builtins_tree.syntax(), + &with_table, + ) + } + Some(resolve::TableSource::CreateTable(create_table)) => { + resolve::collect_table_columns(binder, root, &create_table) + .into_iter() + .filter_map(|column| column.name().map(|name| Name::from_node(&name))) + .collect() + } + Some(resolve::TableSource::CreateView(create_view)) => { + resolve::collect_view_column_names(&create_view) + } + Some(resolve::TableSource::CreateMaterializedView(create_materialized_view)) => { + resolve::collect_materialized_view_column_names(&create_materialized_view) + } + Some(resolve::TableSource::ParenSelect(paren_select)) => { + resolve::collect_paren_select_columns_with_types(binder, root, &paren_select) + .into_iter() + .map(|(name, _ty)| name) + .collect() + } + None => vec![], + } +} + fn hover_qualified_star_columns_from_table( root: &SyntaxNode, create_table: &impl ast::HasCreateTable, @@ -604,9 +750,13 @@ fn hover_qualified_star_columns_from_table( Some(results.join("\n")) } -fn hover_qualified_star_columns_from_cte(with_table: &ast::WithTable) -> Option { +fn hover_qualified_star_columns_from_cte( + root: &SyntaxNode, + with_table: &ast::WithTable, + binder: &binder::Binder, +) -> Option { let cte_name = Name::from_node(&with_table.name()?); - let column_names = resolve::collect_with_table_column_names(with_table); + let column_names = resolve::collect_with_table_column_names(binder, root, with_table); let results: Vec = column_names .iter() .map(|column_name| { @@ -672,34 +822,58 @@ fn hover_qualified_star_columns_from_subquery( paren_select: &ast::ParenSelect, binder: &binder::Binder, ) -> Option { - let ast::SelectVariant::Select(select) = paren_select.select()? else { - return None; - }; - - let select_clause = select.select_clause()?; - let target_list = select_clause.target_list()?; - - let mut results = vec![]; - let subquery_alias = subquery_alias_name(paren_select); - - for target in target_list.targets() { - if target.star_token().is_some() { - let table_ptrs = resolve::resolve_unqualified_star_table_ptrs(binder, &target)?; - for table_ptr in table_ptrs { - if let Some(columns) = hover_qualified_star_columns(root, &table_ptr, binder) { - results.push(columns) + let select_variant = paren_select.select()?; + + if let ast::SelectVariant::Select(select) = select_variant { + let select_clause = select.select_clause()?; + let target_list = select_clause.target_list()?; + + let mut results = vec![]; + let subquery_alias = subquery_alias_name(paren_select); + + for target in target_list.targets() { + if target.star_token().is_some() { + let table_ptrs = resolve::resolve_unqualified_star_table_ptrs(binder, &target)?; + for table_ptr in table_ptrs { + if let Some(columns) = hover_qualified_star_columns(root, &table_ptr, binder) { + results.push(columns) + } } + continue; + } + + if let Some(result) = + hover_subquery_target_column(root, &target, subquery_alias.as_ref(), binder) + { + results.push(result); } - continue; } - if let Some(result) = - hover_subquery_target_column(root, &target, subquery_alias.as_ref(), binder) - { - results.push(result); + if results.is_empty() { + return None; } + + return Some(results.join("\n")); } + let subquery_alias = subquery_alias_name(paren_select); + let results: Vec = + resolve::collect_paren_select_columns_with_types(binder, root, paren_select) + .into_iter() + .map(|(column_name, ty)| { + if let Some(alias) = &subquery_alias { + return ColumnHover::table_column(&alias.to_string(), &column_name.to_string()); + } + if let Some(ty) = ty { + return ColumnHover::anon_column_type( + &column_name.to_string(), + &ty.to_string(), + ); + } + ColumnHover::anon_column(&column_name.to_string()) + }) + .collect(); + if results.is_empty() { return None; } @@ -2843,6 +3017,19 @@ select a from t$0; "); } + #[test] + fn hover_on_select_cte_table_as_column() { + assert_snapshot!(check_hover(" +with t as (select 1 a, 2 b, 3 c) +select t$0 from t; +"), @r" + hover: with t as (select 1 a, 2 b, 3 c) + ╭▸ + 3 │ select t from t; + ╰╴ ─ hover + "); + } + #[test] fn hover_on_cte_column() { assert_snapshot!(check_hover(" @@ -3038,16 +3225,91 @@ select t.*$0 from t; #[test] fn hover_on_cte_table_alias_with_column_list() { assert_snapshot!(check_hover(" -with t as (select 1 a, 2 b) +with t as (select 1 a, 2 b, 3 c) select u$0.x, u.y from t as u(x, y); "), @" - hover: table u(x, y) + hover: table u(x, y, c) ╭▸ 3 │ select u.x, u.y from t as u(x, y); ╰╴ ─ hover "); } + #[test] + fn hover_on_cte_table_alias_with_column_list_table_ref() { + assert_snapshot!(check_hover(" +with t as (select 1 a, 2 b, 3 c) +select u$0 from t as u(x, y); +"), @" + hover: table u(x, y, c) + ╭▸ + 3 │ select u from t as u(x, y); + ╰╴ ─ hover + "); + } + + #[test] + fn hover_on_cte_table_alias_with_partial_column_list_star() { + assert_snapshot!(check_hover(" +with t as (select 1 a, 2 b, 3 c) +select *$0 from t u(x, y); +"), @" + hover: column u.x + column u.y + column u.c + ╭▸ + 3 │ select * from t u(x, y); + ╰╴ ─ hover + "); + } + + #[test] + fn hover_on_cte_table_alias_with_partial_column_list_star_from_information_schema() { + assert_snapshot!(check_hover(" +with t as (select * from information_schema.sql_features) +select *$0 from t u(x); +"), @" + hover: column u.x + column u.feature_name + column u.sub_feature_id + column u.sub_feature_name + column u.is_supported + column u.is_verified_by + column u.comments + ╭▸ + 3 │ select * from t u(x); + ╰╴ ─ hover + "); + } + + #[test] + fn hover_on_cte_table_alias_with_partial_column_list_qualified_star() { + assert_snapshot!(check_hover(" +with t as (select 1 a, 2 b, 3 c) +select u.*$0 from t u(x, y); +"), @" + hover: column u.x + column u.y + column u.c + ╭▸ + 3 │ select u.* from t u(x, y); + ╰╴ ─ hover + "); + } + + #[test] + fn hover_on_star_from_cte_empty_select() { + assert!( + check_hover_( + " +with t as (select) +select *$0 from t; +", + ) + .is_none() + ); + } + #[test] fn hover_on_star_with_subquery_from_cte() { assert_snapshot!(check_hover(" @@ -3079,6 +3341,38 @@ select *$0 from (select a from t); "); } + #[test] + fn hover_on_star_with_subquery_from_table_statement() { + assert_snapshot!(check_hover(" +with t as (select 1 a, 2 b) +select *$0 from (table t); +"), @r" + hover: column a + column b + ╭▸ + 3 │ select * from (table t); + ╰╴ ─ hover + "); + } + + #[test] + fn hover_on_star_from_information_schema_table() { + assert_snapshot!(check_hover(" +select *$0 from information_schema.sql_features; +"), @" + hover: column information_schema.sql_features.feature_id information_schema.character_data + column information_schema.sql_features.feature_name information_schema.character_data + column information_schema.sql_features.sub_feature_id information_schema.character_data + column information_schema.sql_features.sub_feature_name information_schema.character_data + column information_schema.sql_features.is_supported information_schema.yes_or_no + column information_schema.sql_features.is_verified_by information_schema.character_data + column information_schema.sql_features.comments information_schema.character_data + ╭▸ + 2 │ select * from information_schema.sql_features; + ╰╴ ─ hover + "); + } + #[test] fn hover_on_star_with_subquery_literal() { assert_snapshot!(check_hover(" diff --git a/crates/squawk_ide/src/resolve.rs b/crates/squawk_ide/src/resolve.rs index 29e95214..06ae4ca9 100644 --- a/crates/squawk_ide/src/resolve.rs +++ b/crates/squawk_ide/src/resolve.rs @@ -33,6 +33,13 @@ pub(crate) fn resolve_name_ref_ptrs( NameRefClass::Table => { let (table_name, schema) = extract_table_schema_from_name_ref(name_ref)?; let position = name_ref.syntax().text_range().start(); + + if schema.is_none() + && let Some(cte_ptr) = resolve_cte_table(name_ref, &table_name) + { + return Some(smallvec![cte_ptr]); + } + resolve_table_name_ptr(binder, &table_name, &schema, position).map(|ptr| smallvec![ptr]) } NameRefClass::NamedArgParameter => { @@ -849,12 +856,7 @@ fn resolve_select_qualified_column_table_name_ptr( None }; - let select = table_name_ref - .syntax() - .ancestors() - .find_map(ast::Select::cast)?; - let from_clause = select.from_clause()?; - let from_item = find_from_item_in_from_clause(&from_clause, &table_name)?; + let from_item = find_from_item_for_select_qualified_name_ref(table_name_ref, &table_name)?; if let Some(alias_name) = from_item.alias().and_then(|a| a.name()) && Name::from_node(&alias_name) == table_name @@ -1053,12 +1055,8 @@ fn resolve_select_qualified_column_ptr( let path = delete.relation_name()?.path()?; extract_table_schema_from_path(&path)? } else { - let select = column_name_ref - .syntax() - .ancestors() - .find_map(ast::Select::cast)?; - let from_clause = select.from_clause()?; - let from_item = find_from_item_in_from_clause(&from_clause, &column_table_name)?; + let from_item = + find_from_item_for_select_qualified_name_ref(column_name_ref, &column_table_name)?; if let Some(call_expr) = from_item.call_expr() && let Some(ptr) = resolve_column_from_call_expr_return_table( @@ -1090,6 +1088,16 @@ fn resolve_select_qualified_column_ptr( ); } + if let Some(paren_expr) = from_item.paren_expr() { + return resolve_column_from_paren_expr( + binder, + root, + &paren_expr, + column_name_ref, + &column_name, + ); + } + // `from t as u(a, b, c)` if let Some(column_list) = alias.column_list() { for column in column_list.columns() { @@ -1288,6 +1296,18 @@ fn resolve_from_item_column_ptr( } if let Some(paren_expr) = from_item.paren_expr() { + if let Some(alias) = from_item.alias() + && let Some(column_list) = alias.column_list() + { + for col in column_list.columns() { + if let Some(col_name) = col.name() + && Name::from_node(&col_name) == column_name + { + return Some(SyntaxNodePtr::new(col_name.syntax())); + } + } + } + return resolve_column_from_paren_expr( binder, root, @@ -1327,23 +1347,45 @@ fn resolve_from_item_column_ptr( let (table_name, schema) = table_and_schema_from_from_item(from_item)?; - if schema.is_none() && resolve_cte_table(column_name_ref, &table_name).is_some() { + if schema.is_none() + && let Some(cte_ptr) = resolve_cte_table(column_name_ref, &table_name) + { if let Some(cte_column_ptr) = resolve_cte_column(binder, root, column_name_ref, &table_name, &column_name) { return Some(cte_column_ptr); } + if column_name == table_name { + return Some(cte_ptr); + } + if let Some(alias) = from_item.alias() + && let Some(alias_name) = alias.name() + && Name::from_node(&alias_name) == column_name + { + return Some(SyntaxNodePtr::new(alias_name.syntax())); + } return None; } - resolve_column_from_table_or_view( + if let Some(ptr) = resolve_column_from_table_or_view( binder, root, column_name_ref, &table_name, &schema, &column_name, - ) + ) { + return Some(ptr); + } + + if let Some(alias) = from_item.alias() + && let Some(alias_name) = alias.name() + && Name::from_node(&alias_name) == column_name + { + return Some(SyntaxNodePtr::new(alias_name.syntax())); + } + + None } fn resolve_column_from_table_or_view( @@ -1780,6 +1822,34 @@ pub(crate) fn find_from_item_in_from_clause( None } +fn find_from_item_for_select_qualified_name_ref( + name_ref: &ast::NameRef, + table_name: &Name, +) -> Option { + let select = name_ref.syntax().ancestors().find_map(ast::Select::cast)?; + + if let Some(from_clause) = select.from_clause() + && let Some(from_item) = find_from_item_in_from_clause(&from_clause, table_name) + { + return Some(from_item); + } + + let lateral_from_item = name_ref.syntax().ancestors().find_map(|ancestor| { + ast::FromItem::cast(ancestor).filter(|from_item| from_item.lateral_token().is_some()) + })?; + + for ancestor in lateral_from_item.syntax().ancestors() { + if let Some(outer_select) = ast::Select::cast(ancestor) + && let Some(from_clause) = outer_select.from_clause() + && let Some(outer_from_item) = find_from_item_in_from_clause(&from_clause, table_name) + { + return Some(outer_from_item); + } + } + + None +} + pub(crate) fn extract_table_name(path: &ast::Path) -> Option { let segment = path.segment()?; let name_ref = segment.name_ref()?; @@ -2480,6 +2550,13 @@ pub(crate) fn resolve_qualified_star_table_ptr( if let Some(select) = ast::Select::cast(ancestor.clone()) { let from_clause = select.from_clause()?; let from_item = find_from_item_in_from_clause(&from_clause, &table_name)?; + + if let Some(alias) = from_item.alias() + && alias.column_list().is_some() + { + return Some(SyntaxNodePtr::new(alias.syntax())); + } + let (table_name, schema) = table_and_schema_from_from_item(&from_item)?; if let Some(table_name_ptr) = @@ -2674,6 +2751,35 @@ pub(crate) fn table_ptrs_from_clause( results } +pub(crate) fn table_ptr_from_from_item( + binder: &Binder, + from_item: &ast::FromItem, +) -> Option { + if let Some(paren_select) = from_item.paren_select() { + return Some(SyntaxNodePtr::new(paren_select.syntax())); + } + + let (table_name, schema) = table_and_schema_from_from_item(from_item)?; + let position = from_item.syntax().text_range().start(); + + if let Some(table_ptr) = resolve_table_name_ptr(binder, &table_name, &schema, position) { + return Some(table_ptr); + } + + if let Some(view_name_ptr) = resolve_view_name_ptr(binder, &table_name, &schema, position) { + return Some(view_name_ptr); + } + + if schema.is_none() + && let Some(name_ref) = from_item.name_ref() + && let Some(cte_ptr) = resolve_cte_table(&name_ref, &table_name) + { + return Some(cte_ptr); + } + + None +} + fn collect_table_ptrs_from_join_expr( binder: &Binder, join_expr: &ast::JoinExpr, @@ -2699,6 +2805,13 @@ fn collect_tables_from_item( from_item: &ast::FromItem, results: &mut Vec, ) { + if let Some(alias) = from_item.alias() + && alias.column_list().is_some() + { + results.push(SyntaxNodePtr::new(alias.syntax())); + return; + } + if let Some(paren_select) = from_item.paren_select() { results.push(SyntaxNodePtr::new(paren_select.syntax())); return; @@ -3023,18 +3136,28 @@ fn resolve_column_from_paren_expr( return resolve_column_from_paren_expr(binder, root, &paren_expr, name_ref, column_name); } - if let Some(from_item) = paren_expr.from_item() - && let Some(paren_select) = from_item.paren_select() - { - let alias = from_item.alias(); - return resolve_subquery_column_ptr( - binder, - root, - &paren_select, - name_ref, - column_name, - alias.as_ref(), - ); + if let Some(from_item) = paren_expr.from_item() { + if let Some(paren_select) = from_item.paren_select() { + let alias = from_item.alias(); + return resolve_subquery_column_ptr( + binder, + root, + &paren_select, + name_ref, + column_name, + alias.as_ref(), + ); + } + + if let Some(nested_paren_expr) = from_item.paren_expr() { + return resolve_column_from_paren_expr( + binder, + root, + &nested_paren_expr, + name_ref, + column_name, + ); + } } None @@ -3264,14 +3387,87 @@ fn select_from_materialized_view_query( } } -pub(crate) fn collect_with_table_column_names(with_table: &ast::WithTable) -> Vec { - if let Some(column_list) = with_table.column_list() { - let columns = collect_column_names_from_column_list(&column_list); - if !columns.is_empty() { - return columns; +pub(crate) fn collect_with_table_column_names( + binder: &Binder, + root: &SyntaxNode, + with_table: &ast::WithTable, +) -> Vec { + collect_with_table_columns_with_types(binder, root, with_table) + .into_iter() + .map(|(name, _)| name) + .collect() +} + +fn resolve_symbol_info( + binder: &Binder, + path: &ast::Path, + kind: SymbolKind, +) -> Option<(Schema, String)> { + let name_str = extract_table_name_from_path(path)?; + let schema = extract_schema_from_path(path); + let position = path.syntax().text_range().start(); + binder.lookup_info(name_str, &schema, kind, position) +} + +fn collect_column_names_from_column_list(column_list: &ast::ColumnList) -> Vec { + let mut columns = vec![]; + for column in column_list.columns() { + if let Some(name) = column.name() { + columns.push(Name::from_node(&name)); } } + columns +} + +fn collect_target_list_column_names(target_list: &ast::TargetList) -> Vec { + let mut columns = vec![]; + 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() + { + columns.push(Name::from_string(col_name_str)); + } + } + columns +} + +pub(crate) fn collect_with_table_columns_with_types( + binder: &Binder, + root: &SyntaxNode, + with_table: &ast::WithTable, +) -> Vec<(Name, Option)> { + let base_columns = collect_with_table_query_columns_with_types(binder, root, with_table); + + let alias_columns: Vec = with_table + .column_list() + .into_iter() + .flat_map(|column_list| column_list.columns()) + .filter_map(|column| column.name().map(|name| Name::from_node(&name))) + .collect(); + + if alias_columns.is_empty() { + return base_columns; + } + let mut results = vec![]; + + for (idx, alias_name) in alias_columns.iter().enumerate() { + results.push(( + alias_name.clone(), + base_columns.get(idx).and_then(|(_, ty)| ty.clone()), + )); + } + + results.extend(base_columns.into_iter().skip(alias_columns.len())); + + results +} + +fn collect_with_table_query_columns_with_types( + binder: &Binder, + root: &SyntaxNode, + with_table: &ast::WithTable, +) -> Vec<(Name, Option)> { let Some(query) = with_table.query() else { return vec![]; }; @@ -3281,14 +3477,16 @@ pub(crate) fn collect_with_table_column_names(with_table: &ast::WithTable) -> Ve if let Some(row_list) = values.row_list() && let Some(first_row) = row_list.rows().next() { - for (idx, _expr) in first_row.exprs().enumerate() { - results.push(Name::from_string(format!("column{}", idx + 1))); + for (idx, expr) in first_row.exprs().enumerate() { + let name = Name::from_string(format!("column{}", idx + 1)); + let ty = infer_type_from_expr(&expr); + results.push((name, ty)); } } return results; } - if let Some(columns) = columns_from_returning_clause(&query) { + if let Some(columns) = columns_from_returning_clause_with_types(&query) { return columns; } @@ -3302,10 +3500,47 @@ pub(crate) fn collect_with_table_column_names(with_table: &ast::WithTable) -> Ve return vec![]; }; - collect_target_list_column_names(&target_list) + let from_clause = cte_select.from_clause(); + let mut columns = vec![]; + + for target in target_list.targets() { + if let Some((col_name, _node)) = ColumnName::from_target(target.clone()) { + if let Some(col_name_str) = col_name.to_string() { + let ty = target.expr().and_then(|e| infer_type_from_expr(&e)); + columns.push((Name::from_string(col_name_str), ty)); + continue; + } + + if target.star_token().is_some() + && let Some(from_clause) = &from_clause + { + columns.extend(collect_columns_for_star_from_clause( + binder, + root, + from_clause, + )); + continue; + } + } + + 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) = &from_clause + && let Some(from_item) = find_from_item_in_from_clause(from_clause, &table_name) + { + columns.extend(collect_columns_for_star_from_from_item( + binder, root, &from_item, + )); + } + } + + columns } -fn columns_from_returning_clause(query: &ast::WithQuery) -> Option> { +fn columns_from_returning_clause_with_types( + query: &ast::WithQuery, +) -> Option)>> { let returning_clause = match query { ast::WithQuery::Delete(delete) => delete.returning_clause(), ast::WithQuery::Insert(insert) => insert.returning_clause(), @@ -3317,80 +3552,139 @@ fn columns_from_returning_clause(query: &ast::WithQuery) -> Option> { | ast::WithQuery::Values(_) | ast::WithQuery::ParenSelect(_) => None, }; + if let Some(returning_clause) = returning_clause { if let Some(target_list) = returning_clause.target_list() { - return Some(collect_target_list_column_names(&target_list)); + return Some(collect_target_list_columns_with_types(&target_list)); } return Some(vec![]); } + None } -fn resolve_symbol_info( +fn collect_columns_for_star_from_clause( binder: &Binder, - path: &ast::Path, - kind: SymbolKind, -) -> Option<(Schema, String)> { - let name_str = extract_table_name_from_path(path)?; - let schema = extract_schema_from_path(path); - let position = path.syntax().text_range().start(); - binder.lookup_info(name_str, &schema, kind, position) -} - -fn collect_column_names_from_column_list(column_list: &ast::ColumnList) -> Vec { + root: &SyntaxNode, + from_clause: &ast::FromClause, +) -> Vec<(Name, Option)> { let mut columns = vec![]; - for column in column_list.columns() { - if let Some(name) = column.name() { - columns.push(Name::from_node(&name)); - } + + for from_item in from_clause.from_items() { + columns.extend(collect_columns_for_star_from_from_item( + binder, root, &from_item, + )); } - columns -} -fn collect_target_list_column_names(target_list: &ast::TargetList) -> Vec { - let mut columns = vec![]; - 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() + for join_expr in from_clause.join_exprs() { + if let Some(from_item) = join_expr.from_item() { + columns.extend(collect_columns_for_star_from_from_item( + binder, root, &from_item, + )); + } + + if let Some(join) = join_expr.join() + && let Some(from_item) = join.from_item() { - columns.push(Name::from_string(col_name_str)); + columns.extend(collect_columns_for_star_from_from_item( + binder, root, &from_item, + )); } } + columns } -pub(crate) fn collect_with_table_columns_with_types( - with_table: &ast::WithTable, +fn collect_columns_for_star_from_from_item( + binder: &Binder, + root: &SyntaxNode, + from_item: &ast::FromItem, ) -> Vec<(Name, Option)> { - let Some(query) = with_table.query() else { - return vec![]; - }; - - if let ast::WithQuery::Values(values) = query { - let mut results = vec![]; - if let Some(row_list) = values.row_list() - && let Some(first_row) = row_list.rows().next() - { - for (idx, expr) in first_row.exprs().enumerate() { - let name = Name::from_string(format!("column{}", idx + 1)); - let ty = infer_type_from_expr(&expr); - results.push((name, ty)); - } - } - return results; + if let Some(alias) = from_item.alias() + && alias.column_list().is_some() + { + return collect_columns_for_star_from_alias(binder, root, from_item, &alias); } - let Some(cte_select) = select_from_with_query(query) else { + let Some(table_ptr) = table_ptr_from_from_item(binder, from_item) else { return vec![]; }; - let Some(select_clause) = cte_select.select_clause() else { - return vec![]; - }; - let Some(target_list) = select_clause.target_list() else { + + collect_columns_for_star_from_table_ptr(binder, root, &table_ptr) +} + +fn collect_columns_for_star_from_alias( + binder: &Binder, + root: &SyntaxNode, + from_item: &ast::FromItem, + alias: &ast::Alias, +) -> Vec<(Name, Option)> { + let alias_columns: Vec = alias + .column_list() + .into_iter() + .flat_map(|column_list| column_list.columns()) + .filter_map(|column| column.name().map(|name| Name::from_node(&name))) + .collect(); + + let Some(table_ptr) = table_ptr_from_from_item(binder, from_item) else { return vec![]; }; - collect_target_list_columns_with_types(&target_list) + let base_columns = collect_columns_for_star_from_table_ptr(binder, root, &table_ptr); + let mut results = vec![]; + + for (idx, alias_name) in alias_columns.iter().enumerate() { + results.push(( + alias_name.clone(), + base_columns.get(idx).and_then(|(_, ty)| ty.clone()), + )); + } + + results.extend(base_columns.into_iter().skip(alias_columns.len())); + + results +} + +fn collect_columns_for_star_from_table_ptr( + binder: &Binder, + root: &SyntaxNode, + table_ptr: &SyntaxNodePtr, +) -> Vec<(Name, Option)> { + let table_node = table_ptr.to_node(root); + + if let Some(paren_select) = ast::ParenSelect::cast(table_node.clone()) { + return collect_paren_select_columns_with_types(binder, root, &paren_select); + } + + match find_table_source(&table_node) { + Some(TableSource::Alias(alias)) => { + let Some(from_item) = alias.syntax().ancestors().find_map(ast::FromItem::cast) else { + return vec![]; + }; + collect_columns_for_star_from_alias(binder, root, &from_item, &alias) + } + Some(TableSource::WithTable(with_table)) => { + collect_with_table_columns_with_types(binder, root, &with_table) + } + Some(TableSource::CreateTable(create_table)) => { + collect_table_columns(binder, root, &create_table) + .into_iter() + .filter_map(|column| { + let name = Name::from_node(&column.name()?); + let ty = column.ty().and_then(|ty| infer_type_from_ty(&ty)); + Some((name, ty)) + }) + .collect() + } + Some(TableSource::CreateView(create_view)) => collect_view_columns_with_types(&create_view), + Some(TableSource::CreateMaterializedView(create_materialized_view)) => { + collect_materialized_view_columns_with_types(&create_materialized_view) + } + Some(TableSource::ParenSelect(paren_select)) => { + collect_paren_select_columns_with_types(binder, root, &paren_select) + } + None => vec![], + } } fn collect_target_list_columns_with_types( @@ -3472,6 +3766,21 @@ fn collect_select_variant_columns_with_types( let Some((table_name, schema)) = extract_table_schema_from_path(&path) else { return vec![]; }; + + if schema.is_none() + && let Some(name_ref) = path.segment().and_then(|segment| segment.name_ref()) + && let Some(with_table_ptr) = resolve_cte_table(&name_ref, &table_name) + && let Some(with_table) = with_table_ptr + .to_node(root) + .ancestors() + .find_map(ast::WithTable::cast) + { + return collect_with_table_column_names(binder, root, &with_table) + .into_iter() + .map(|name| (name, None)) + .collect(); + } + let position = table.syntax().text_range().start(); let Some(table_ptr) = binder.lookup_with(&table_name, SymbolKind::Table, position, &schema)