diff --git a/crates/squawk_ide/src/inlay_hints.rs b/crates/squawk_ide/src/inlay_hints.rs index fd3cf3c7..794f7385 100644 --- a/crates/squawk_ide/src/inlay_hints.rs +++ b/crates/squawk_ide/src/inlay_hints.rs @@ -88,11 +88,9 @@ fn inlay_hint_insert( binder: &Binder, insert: ast::Insert, ) -> Option<()> { - let values = insert.values()?; - let row_list = values.row_list()?; let create_table = resolve::resolve_insert_create_table(root, binder, &insert); - let columns: Vec<(Name, Option)> = if let Some(column_list) = insert.column_list() { + let columns = if let Some(column_list) = insert.column_list() { // `insert into t(a, b, c) values (1, 2, 3)` column_list .columns() @@ -124,6 +122,11 @@ fn inlay_hint_insert( .collect() }; + let Some(values) = insert.values() else { + return inlay_hint_insert_select(hints, columns, insert.stmt()?); + }; + let row_list = values.row_list()?; + for row in row_list.rows() { for ((column_name, target), expr) in columns.iter().zip(row.exprs()) { let expr_start = expr.syntax().text_range().start(); @@ -139,6 +142,53 @@ fn inlay_hint_insert( Some(()) } +fn inlay_hint_insert_select( + hints: &mut Vec, + columns: Vec<(Name, Option)>, + stmt: ast::Stmt, +) -> Option<()> { + let target_list = match stmt { + ast::Stmt::Select(select) => select.select_clause()?.target_list(), + ast::Stmt::SelectInto(select_into) => select_into.select_clause()?.target_list(), + ast::Stmt::ParenSelect(paren_select) => { + target_list_from_select_variant(paren_select.select()?) + } + _ => None, + }?; + + for ((column_name, target), target_expr) in columns.iter().zip(target_list.targets()) { + let expr = target_expr.expr()?; + let expr_start = expr.syntax().text_range().start(); + hints.push(InlayHint { + position: expr_start, + label: format!("{}: ", column_name), + kind: InlayHintKind::Parameter, + target: *target, + }); + } + + Some(()) +} + +fn target_list_from_select_variant(select: ast::SelectVariant) -> Option { + let mut current = select; + for _ in 0..100 { + match current { + ast::SelectVariant::Select(select) => { + return select.select_clause()?.target_list(); + } + ast::SelectVariant::SelectInto(select_into) => { + return select_into.select_clause()?.target_list(); + } + ast::SelectVariant::ParenSelect(paren_select) => { + current = paren_select.select()?; + } + _ => return None, + } + } + None +} + #[cfg(test)] mod test { use crate::inlay_hints::inlay_hints; @@ -348,4 +398,17 @@ insert into t values (1, 2, 3); ╰╴ ─── ─── "); } + + #[test] + fn insert_select() { + assert_snapshot!(check_inlay_hints(" +create table t (a int, b int); +insert into t select 1, 2; +"), @r" + inlay hints: + ╭▸ + 3 │ insert into t select a: 1, b: 2; + ╰╴ ─── ─── + "); + } } diff --git a/crates/squawk_ide/src/resolve.rs b/crates/squawk_ide/src/resolve.rs index bc6128ca..8648bdc9 100644 --- a/crates/squawk_ide/src/resolve.rs +++ b/crates/squawk_ide/src/resolve.rs @@ -1791,8 +1791,7 @@ fn resolve_cte_column( } if let ast::WithQuery::Table(table) = query { - let relation_name = table.relation_name()?; - let path = relation_name.path()?; + let path = table.relation_name()?.path()?; let (table_name, schema) = extract_table_schema_from_path(&path)?; if schema.is_none() @@ -1975,8 +1974,7 @@ fn resolve_subquery_column_ptr( } if let ast::SelectVariant::Table(table) = select_variant { - let relation_name = table.relation_name()?; - let path = relation_name.path()?; + let path = table.relation_name()?.path()?; let (table_name, schema) = extract_table_schema_from_path(&path)?; if schema.is_none() diff --git a/crates/squawk_syntax/src/ast/generated/nodes.rs b/crates/squawk_syntax/src/ast/generated/nodes.rs index c791d0d3..8bb078ca 100644 --- a/crates/squawk_syntax/src/ast/generated/nodes.rs +++ b/crates/squawk_syntax/src/ast/generated/nodes.rs @@ -2753,6 +2753,10 @@ pub struct CompoundSelect { pub(crate) syntax: SyntaxNode, } impl CompoundSelect { + #[inline] + pub fn all_token(&self) -> Option { + support::token(&self.syntax, SyntaxKind::ALL_KW) + } #[inline] pub fn except_token(&self) -> Option { support::token(&self.syntax, SyntaxKind::EXCEPT_KW) diff --git a/crates/squawk_syntax/src/postgresql.ungram b/crates/squawk_syntax/src/postgresql.ungram index 350c0bc8..0d61527c 100644 --- a/crates/squawk_syntax/src/postgresql.ungram +++ b/crates/squawk_syntax/src/postgresql.ungram @@ -1644,7 +1644,8 @@ SelectVariant = CompoundSelect = lhs:SelectVariant - ('union' | 'intersect' | 'except') + ('union' | 'intersect' | 'except' ) + 'all'? rhs:SelectVariant SelectInto =