diff --git a/crates/squawk_ide/src/goto_definition.rs b/crates/squawk_ide/src/goto_definition.rs index 0c6c4b8e..19e28795 100644 --- a/crates/squawk_ide/src/goto_definition.rs +++ b/crates/squawk_ide/src/goto_definition.rs @@ -3002,4 +3002,79 @@ select messages.message$0 from users left join messages on users.id = messages.u ╰╴ ─ 1. source "); } + + #[test] + fn goto_insert_select_cte_column() { + assert_snapshot!(goto(" +create table users(id int, email text); +with new_data as ( + select 1 as id, 'test@example.com' as email +) +insert into users (id, email) +select id$0, email from new_data; +"), @r" + ╭▸ + 4 │ select 1 as id, 'test@example.com' as email + │ ── 2. destination + ‡ + 7 │ select id, email from new_data; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_insert_select_cte_column_second() { + assert_snapshot!(goto(" +create table users(id int, email text); +with new_data as ( + select 1 as id, 'test@example.com' as email +) +insert into users (id, email) +select id, email$0 from new_data; +"), @r" + ╭▸ + 4 │ select 1 as id, 'test@example.com' as email + │ ───── 2. destination + ‡ + 7 │ select id, email from new_data; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_insert_select_cte_table() { + assert_snapshot!(goto(" +create table users(id int, email text); +with new_data as ( + select 1 as id, 'test@example.com' as email +) +insert into users (id, email) +select id, email from new_data$0; +"), @r" + ╭▸ + 3 │ with new_data as ( + │ ──────── 2. destination + ‡ + 7 │ select id, email from new_data; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_delete_cte_column() { + assert_snapshot!(goto(" +create table users(id int, email text); +with old_data as ( + select 1 as id +) +delete from users where id in (select id$0 from old_data); +"), @r" + ╭▸ + 4 │ select 1 as id + │ ── 2. destination + 5 │ ) + 6 │ delete from users where id in (select id from old_data); + ╰╴ ─ 1. source + "); + } } diff --git a/crates/squawk_ide/src/resolve.rs b/crates/squawk_ide/src/resolve.rs index 01a906a5..7ca4f246 100644 --- a/crates/squawk_ide/src/resolve.rs +++ b/crates/squawk_ide/src/resolve.rs @@ -1,6 +1,6 @@ use rowan::{TextRange, TextSize}; use squawk_syntax::{ - SyntaxNodePtr, + SyntaxNode, SyntaxNodePtr, ast::{self, AstNode}, }; @@ -1120,9 +1120,7 @@ pub(crate) fn find_column_in_table( } fn resolve_cte_table(name_ref: &ast::NameRef, cte_name: &Name) -> Option { - let select = name_ref.syntax().ancestors().find_map(ast::Select::cast)?; - let with_clause = select.with_clause()?; - + let with_clause = find_parent_with_clause(name_ref.syntax())?; for with_table in with_clause.with_tables() { if let Some(name) = with_table.name() && Name::from_node(&name) == *cte_name @@ -1134,17 +1132,26 @@ fn resolve_cte_table(name_ref: &ast::NameRef, cte_name: &Name) -> Option Option { + node.ancestors().find_map(|x| { + if let Some(select) = ast::Select::cast(x.clone()) { + select.with_clause() + } else if let Some(delete) = ast::Delete::cast(x.clone()) { + delete.with_clause() + } else if let Some(insert) = ast::Insert::cast(x) { + insert.with_clause() + } else { + None + } + }) +} + fn resolve_cte_column( name_ref: &ast::NameRef, cte_name: &Name, column_name: &Name, ) -> Option { - let select = name_ref - .syntax() - .ancestors() - .filter_map(ast::Select::cast) - .find(|s| s.with_clause().is_some())?; - let with_clause = select.with_clause()?; + let with_clause = find_parent_with_clause(name_ref.syntax())?; for with_table in with_clause.with_tables() { if let Some(name) = with_table.name() diff --git a/crates/squawk_syntax/src/ast/generated/nodes.rs b/crates/squawk_syntax/src/ast/generated/nodes.rs index 92482441..a1df1a70 100644 --- a/crates/squawk_syntax/src/ast/generated/nodes.rs +++ b/crates/squawk_syntax/src/ast/generated/nodes.rs @@ -5247,6 +5247,10 @@ impl Delete { support::child(&self.syntax) } #[inline] + pub fn with_clause(&self) -> Option { + support::child(&self.syntax) + } + #[inline] pub fn delete_token(&self) -> Option { support::token(&self.syntax, SyntaxKind::DELETE_KW) } @@ -8433,6 +8437,10 @@ impl Insert { support::child(&self.syntax) } #[inline] + pub fn with_clause(&self) -> Option { + support::child(&self.syntax) + } + #[inline] pub fn default_token(&self) -> Option { support::token(&self.syntax, SyntaxKind::DEFAULT_KW) } @@ -15394,6 +15402,10 @@ impl Update { support::child(&self.syntax) } #[inline] + pub fn with_clause(&self) -> Option { + support::child(&self.syntax) + } + #[inline] pub fn update_token(&self) -> Option { support::token(&self.syntax, SyntaxKind::UPDATE_KW) } diff --git a/crates/squawk_syntax/src/postgresql.ungram b/crates/squawk_syntax/src/postgresql.ungram index 5203ee91..93bb7b04 100644 --- a/crates/squawk_syntax/src/postgresql.ungram +++ b/crates/squawk_syntax/src/postgresql.ungram @@ -1263,6 +1263,7 @@ Table = 'table' RelationName Insert = + WithClause? 'insert' 'into' Path Alias? ColumnList? ('overriding' ('system' | 'user') 'value')? ('default' 'values' | Values | Stmt) @@ -1326,6 +1327,7 @@ SetExpr = | 'default' Update = + WithClause? 'update' RelationName Alias? @@ -1348,6 +1350,7 @@ ReturningOption = Name Delete = + WithClause? 'delete' 'from' RelationName Alias? UsingClause? (WhereClause | WhereCurrentOf)?