From 195fcf575f12bd0940a5f88b3df9e9bc57262a96 Mon Sep 17 00:00:00 2001 From: Steve Dignam Date: Wed, 31 Dec 2025 00:57:16 -0500 Subject: [PATCH] ide: cte improvements, subqueries, values, union select --- crates/squawk_ide/src/goto_definition.rs | 227 ++++++++++++++++++ crates/squawk_ide/src/hover.rs | 88 ++++++- crates/squawk_ide/src/resolve.rs | 107 ++++++++- .../squawk_syntax/src/ast/generated/nodes.rs | 28 +++ crates/squawk_syntax/src/postgresql.ungram | 4 +- 5 files changed, 443 insertions(+), 11 deletions(-) diff --git a/crates/squawk_ide/src/goto_definition.rs b/crates/squawk_ide/src/goto_definition.rs index dd2f14a3..1dffc694 100644 --- a/crates/squawk_ide/src/goto_definition.rs +++ b/crates/squawk_ide/src/goto_definition.rs @@ -1423,6 +1423,44 @@ select a$0 from t; "); } + #[test] + fn goto_cte_with_partial_column_list() { + assert_snapshot!(goto(" +with t(x) as (select 1 as a, 2 as b) +select b$0 from t; +"), @r" + ╭▸ + 2 │ with t(x) as (select 1 as a, 2 as b) + │ ─ 2. destination + 3 │ select b from t; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_cte_with_partial_column_list_renamed() { + assert_snapshot!(goto(" +with t(x) as (select 1 as a, 2 as b) +select x$0 from t; +"), @r" + ╭▸ + 2 │ with t(x) as (select 1 as a, 2 as b) + │ ─ 2. destination + 3 │ select x from t; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_cte_column_list_overwrites_column() { + goto_not_found( + " +with t(x) as (select 1 as a) +select a$0 from t; +", + ); + } + #[test] fn goto_cte_shadows_table() { assert_snapshot!(goto(" @@ -1438,6 +1476,50 @@ select a from t; "); } + #[test] + fn goto_subquery_column() { + assert_snapshot!(goto(" +select a$0 from (select 1 a); +"), @r" + ╭▸ + 2 │ select a from (select 1 a); + ╰╴ ─ 1. source ─ 2. destination + "); + } + + #[test] + fn goto_subquery_column_with_as() { + assert_snapshot!(goto(" +select a$0 from (select 1 as a); +"), @r" + ╭▸ + 2 │ select a from (select 1 as a); + ╰╴ ─ 1. source ─ 2. destination + "); + } + + #[test] + fn goto_subquery_column_multiple_columns() { + assert_snapshot!(goto(" +select b$0 from (select 1 a, 2 b); +"), @r" + ╭▸ + 2 │ select b from (select 1 a, 2 b); + ╰╴ ─ 1. source ─ 2. destination + "); + } + + #[test] + fn goto_subquery_column_nested_parens() { + assert_snapshot!(goto(" +select a$0 from ((select 1 a)); +"), @r" + ╭▸ + 2 │ select a from ((select 1 a)); + ╰╴ ─ 1. source ─ 2. destination + "); + } + #[test] fn goto_insert_table() { assert_snapshot!(goto(" @@ -2046,6 +2128,94 @@ select foo.t.a$0 from t; "); } + #[test] + fn goto_cte_values_column1() { + assert_snapshot!(goto(" +with t as ( + values (1, 2), (3, 4) +) +select column1$0, column2 from t; +"), @r" + ╭▸ + 3 │ values (1, 2), (3, 4) + │ ─ 2. destination + 4 │ ) + 5 │ select column1, column2 from t; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_cte_values_column2() { + assert_snapshot!(goto(" +with t as ( + values (1, 2), (3, 4) +) +select column1, column2$0 from t; +"), @r" + ╭▸ + 3 │ values (1, 2), (3, 4) + │ ─ 2. destination + 4 │ ) + 5 │ select column1, column2 from t; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_cte_values_single_column() { + assert_snapshot!(goto(" +with t as ( + values (1), (2), (3) +) +select column1$0 from t; +"), @r" + ╭▸ + 3 │ values (1), (2), (3) + │ ─ 2. destination + 4 │ ) + 5 │ select column1 from t; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_cte_values_multiple_rows() { + assert_snapshot!(goto(" +with t as ( + values + (1, 2, 3), + (4, 5, 6), + (7, 8, 9) +) +select column3$0 from t; +"), @r" + ╭▸ + 4 │ (1, 2, 3), + │ ─ 2. destination + ‡ + 8 │ select column3 from t; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_cte_values_uppercase_column_names() { + assert_snapshot!(goto(" +with t as ( + values (1, 2), (3, 4) +) +select COLUMN1$0, COLUMN2 from t; +"), @r" + ╭▸ + 3 │ values (1, 2), (3, 4) + │ ─ 2. destination + 4 │ ) + 5 │ select COLUMN1, COLUMN2 from t; + ╰╴ ─ 1. source + "); + } + #[test] fn goto_qualified_column_with_schema_in_from_table() { assert_snapshot!(goto(" @@ -2073,4 +2243,61 @@ select t.a$0 from foo.t; ╰╴ ─ 1. source "); } + + #[test] + fn goto_cte_union_all_column() { + assert_snapshot!(goto(" +with t as ( + select 1 as a, 2 as b + union all + select 3, 4 +) +select a$0, b from t; +"), @r" + ╭▸ + 3 │ select 1 as a, 2 as b + │ ─ 2. destination + ‡ + 7 │ select a, b from t; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_cte_union_all_column_second() { + assert_snapshot!(goto(" +with t as ( + select 1 as a, 2 as b + union all + select 3, 4 +) +select a, b$0 from t; +"), @r" + ╭▸ + 3 │ select 1 as a, 2 as b + │ ─ 2. destination + ‡ + 7 │ select a, b from t; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_cte_union_column() { + assert_snapshot!(goto(" +with t as ( + select 1 as a, 2 as b + union + select 3, 4 +) +select a$0 from t; +"), @r" + ╭▸ + 3 │ select 1 as a, 2 as b + │ ─ 2. destination + ‡ + 7 │ select a from t; + ╰╴ ─ 1. source + "); + } } diff --git a/crates/squawk_ide/src/hover.rs b/crates/squawk_ide/src/hover.rs index 4efc48c6..0f597360 100644 --- a/crates/squawk_ide/src/hover.rs +++ b/crates/squawk_ide/src/hover.rs @@ -1,6 +1,6 @@ -use crate::binder; use crate::offsets::token_from_offset; use crate::resolve; +use crate::{binder, symbols::Name}; use rowan::TextSize; use squawk_syntax::ast::{self, AstNode}; @@ -69,6 +69,10 @@ pub fn hover(file: &ast::SourceFile, offset: TextSize) -> Option { return format_create_table(&create_table, &binder); } + if let Some(with_table) = name.syntax().parent().and_then(ast::WithTable::cast) { + return format_with_table(&with_table); + } + if let Some(create_index) = name.syntax().ancestors().find_map(ast::CreateIndex::cast) { return format_create_index(&create_index, &binder); } @@ -110,7 +114,14 @@ fn hover_column( if let Some(with_table) = column_name_node.ancestors().find_map(ast::WithTable::cast) { let cte_name = with_table.name()?.syntax().text().to_string(); - let column_name = column_name_node.text().to_string(); + let column_name = if column_name_node + .ancestors() + .any(|a| ast::Values::can_cast(a.kind())) + { + Name::new(name_ref.syntax().text().to_string()) + } else { + Name::new(column_name_node.text().to_string()) + }; return Some(format!("column {}.{}", cte_name, column_name)); } @@ -1595,4 +1606,77 @@ select a$0 from t; ╰╴ ─ hover "); } + + #[test] + fn hover_on_cte_definition() { + assert_snapshot!(check_hover(" +with t$0 as (select 1 a) +select a from t; +"), @r" + hover: with t as (select 1 a) + ╭▸ + 2 │ with t as (select 1 a) + ╰╴ ─ hover + "); + } + + #[test] + fn hover_on_cte_values_column1() { + assert_snapshot!(check_hover(" +with t as ( + values (1, 2), (3, 4) +) +select column1$0, column2 from t; +"), @r" + hover: column t.column1 + ╭▸ + 5 │ select column1, column2 from t; + ╰╴ ─ hover + "); + } + + #[test] + fn hover_on_cte_values_column2() { + assert_snapshot!(check_hover(" +with t as ( + values (1, 2), (3, 4) +) +select column1, column2$0 from t; +"), @r" + hover: column t.column2 + ╭▸ + 5 │ select column1, column2 from t; + ╰╴ ─ hover + "); + } + + #[test] + fn hover_on_cte_values_single_column() { + assert_snapshot!(check_hover(" +with t as ( + values (1), (2), (3) +) +select column1$0 from t; +"), @r" + hover: column t.column1 + ╭▸ + 5 │ select column1 from t; + ╰╴ ─ hover + "); + } + + #[test] + fn hover_on_cte_values_uppercase_column_names() { + assert_snapshot!(check_hover(" +with t as ( + values (1, 2), (3, 4) +) +select COLUMN1$0, COLUMN2 from t; +"), @r" + hover: column t.column1 + ╭▸ + 5 │ select COLUMN1, COLUMN2 from t; + ╰╴ ─ hover + "); + } } diff --git a/crates/squawk_ide/src/resolve.rs b/crates/squawk_ide/src/resolve.rs index 229ac152..2e8074c3 100644 --- a/crates/squawk_ide/src/resolve.rs +++ b/crates/squawk_ide/src/resolve.rs @@ -589,6 +589,14 @@ fn resolve_select_column(binder: &Binder, name_ref: &ast::NameRef) -> Option s, - ast::WithQuery::ParenSelect(ps) => match ps.select()? { + if let ast::WithQuery::Values(values) = query { + 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) + { + return Some(SyntaxNodePtr::new(expr.syntax())); + } + continue; + } + + 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()?, + _ => continue, + }; + + let cte_select = match select_variant { + ast::SelectVariant::Select(s) => s, + ast::SelectVariant::CompoundSelect(compound) => match compound.lhs()? { ast::SelectVariant::Select(s) => s, _ => continue, }, @@ -836,7 +866,12 @@ fn resolve_cte_column( let select_clause = cte_select.select_clause()?; let target_list = select_clause.target_list()?; - for target in target_list.targets() { + for (idx, target) in target_list.targets().enumerate() { + // Skip targets that are covered by the column list + if idx < column_list_len { + continue; + } + if let Some((col_name, node)) = ColumnName::from_target(target.clone()) { if let Some(col_name_str) = col_name.to_string() && Name::new(col_name_str) == *column_name @@ -866,6 +901,64 @@ fn resolve_cte_column( None } +fn resolve_subquery_column( + paren_select: &ast::ParenSelect, + column_name: &Name, +) -> Option { + let select_variant = paren_select.select()?; + let ast::SelectVariant::Select(subquery_select) = select_variant else { + return None; + }; + + let select_clause = subquery_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.clone()) { + if let Some(col_name_str) = col_name.to_string() + && Name::new(col_name_str) == *column_name + { + return Some(SyntaxNodePtr::new(&node)); + } + } + } + + None +} + +fn resolve_column_from_paren_expr( + paren_expr: &ast::ParenExpr, + 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::new(col_name_str) == *column_name + { + return Some(SyntaxNodePtr::new(&node)); + } + } + } + return None; + } + + if let Some(ast::Expr::ParenExpr(paren_expr)) = paren_expr.expr() { + return resolve_column_from_paren_expr(&paren_expr, column_name); + } + + if let Some(from_item) = paren_expr.from_item() + && let Some(paren_select) = from_item.paren_select() + { + return resolve_subquery_column(&paren_select, column_name); + } + + None +} + pub(crate) fn resolve_insert_table_columns( file: &ast::SourceFile, binder: &Binder, diff --git a/crates/squawk_syntax/src/ast/generated/nodes.rs b/crates/squawk_syntax/src/ast/generated/nodes.rs index 072eef20..92482441 100644 --- a/crates/squawk_syntax/src/ast/generated/nodes.rs +++ b/crates/squawk_syntax/src/ast/generated/nodes.rs @@ -7772,6 +7772,10 @@ impl FromItem { support::child(&self.syntax) } #[inline] + pub fn paren_expr(&self) -> Option { + support::child(&self.syntax) + } + #[inline] pub fn paren_select(&self) -> Option { support::child(&self.syntax) } @@ -10310,6 +10314,22 @@ pub struct MergePartitions { pub(crate) syntax: SyntaxNode, } impl MergePartitions { + #[inline] + pub fn path(&self) -> Option { + support::child(&self.syntax) + } + #[inline] + pub fn l_paren_token(&self) -> Option { + support::token(&self.syntax, SyntaxKind::L_PAREN) + } + #[inline] + pub fn r_paren_token(&self) -> Option { + support::token(&self.syntax, SyntaxKind::R_PAREN) + } + #[inline] + pub fn into_token(&self) -> Option { + support::token(&self.syntax, SyntaxKind::INTO_KW) + } #[inline] pub fn merge_token(&self) -> Option { support::token(&self.syntax, SyntaxKind::MERGE_KW) @@ -11720,6 +11740,10 @@ impl ParenExpr { support::child(&self.syntax) } #[inline] + pub fn from_item(&self) -> Option { + support::child(&self.syntax) + } + #[inline] pub fn select(&self) -> Option