From 84d80996fa2d81422852f143176b821240fc229b Mon Sep 17 00:00:00 2001 From: Steve Dignam Date: Wed, 23 Sep 2026 21:12:00 -0400 Subject: [PATCH] ide: goto def for begin atomic function defs --- crates/squawk_ide/src/ast_nav.rs | 11 ++ crates/squawk_ide/src/classify.rs | 3 + crates/squawk_ide/src/goto_definition.rs | 156 +++++++++++++++++++++++ crates/squawk_ide/src/resolve.rs | 17 ++- 4 files changed, 183 insertions(+), 4 deletions(-) diff --git a/crates/squawk_ide/src/ast_nav.rs b/crates/squawk_ide/src/ast_nav.rs index da4900b9..f37f75a7 100644 --- a/crates/squawk_ide/src/ast_nav.rs +++ b/crates/squawk_ide/src/ast_nav.rs @@ -40,6 +40,17 @@ pub(crate) fn find_cte_with_table( None } +pub(crate) fn ancestors_outside_own_with_clause( + node: &SyntaxNode, +) -> impl Iterator { + let mut prev_was_with_clause = false; + node.ancestors().filter(move |ancestor| { + let skip = prev_was_with_clause; + prev_was_with_clause = ast::WithClause::can_cast(ancestor.kind()); + !skip + }) +} + pub(crate) fn iter_values_columns(values: &ast::Values) -> impl Iterator { values .row_list() diff --git a/crates/squawk_ide/src/classify.rs b/crates/squawk_ide/src/classify.rs index 92dbda27..100d79fc 100644 --- a/crates/squawk_ide/src/classify.rs +++ b/crates/squawk_ide/src/classify.rs @@ -1004,6 +1004,9 @@ pub(crate) fn classify_name_ref(node: &SyntaxNode) -> Option { // (anything in SELECT except FROM clause) return Some(NameRefClass::SelectColumn); } + if ast::Values::can_cast(ancestor.kind()) { + return Some(NameRefClass::SelectColumn); + } if ast::CompoundSelect::can_cast(ancestor.kind()) && in_order_by_clause && let Some(parent) = node.parent() diff --git a/crates/squawk_ide/src/goto_definition.rs b/crates/squawk_ide/src/goto_definition.rs index 78fba113..1e24c452 100644 --- a/crates/squawk_ide/src/goto_definition.rs +++ b/crates/squawk_ide/src/goto_definition.rs @@ -1975,6 +1975,162 @@ end; "); } + #[test] + fn goto_function_param_in_begin_atomic_insert_values() { + assert_snapshot!(goto(" +create table t (a int); +create function f(x int) returns void + begin atomic + insert into t values (x$0); + end; +"), @" + ╭▸ + 3 │ create function f(x int) returns void + │ ─ 2. destination + 4 │ begin atomic + 5 │ insert into t values (x); + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_function_param_in_begin_atomic_update() { + assert_snapshot!(goto(" +create table t (a int); +create function f(x int) returns void + begin atomic + update t set a = x$0; + end; +"), @" + ╭▸ + 3 │ create function f(x int) returns void + │ ─ 2. destination + 4 │ begin atomic + 5 │ update t set a = x; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_function_param_in_begin_atomic_delete_where() { + assert_snapshot!(goto(" +create table t (a int); +create function f(x int) returns void + begin atomic + delete from t where a = x$0; + end; +"), @" + ╭▸ + 3 │ create function f(x int) returns void + │ ─ 2. destination + 4 │ begin atomic + 5 │ delete from t where a = x; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_function_param_in_begin_atomic_merge() { + assert_snapshot!(goto(" +create table t (a int); +create function f(x int) returns void + begin atomic + merge into t + using t u on t.a = x$0 + when matched then + do nothing; + end; +"), @" + ╭▸ + 3 │ create function f(x int) returns void + │ ─ 2. destination + ‡ + 6 │ using t u on t.a = x + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_function_param_in_begin_atomic_values() { + assert_snapshot!(goto(" +create function f(x int) returns int + begin atomic + values (x$0); + end; +"), @" + ╭▸ + 2 │ create function f(x int) returns int + │ ─ 2. destination + 3 │ begin atomic + 4 │ values (x); + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_begin_atomic_column_shadows_function_param() { + assert_snapshot!(goto(" +create table t (x int); +create function f(x int) returns void + begin atomic + delete from t where x$0 = 1; + end; +"), @" + ╭▸ + 2 │ create table t (x int); + │ ─ 2. destination + ‡ + 5 │ delete from t where x = 1; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_function_param_in_cte_insert_values_not_outer_select_column() { + assert_snapshot!(goto(" +create table t (a int); +create table u (x int); +create function f(x int) returns int + begin atomic + with c as ( + insert into t values (x$0) returning 1 + ) + select 1 + from u; + end; +"), @" + ╭▸ + 4 │ create function f(x int) returns int + │ ─ 2. destination + ‡ + 7 │ insert into t values (x) returning 1 + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_function_param_in_cte_insert_values_not_outer_update_column() { + assert_snapshot!(goto(" +create table t (a int); +create table u (x int); +create function f(x int) returns void + begin atomic + with c as ( + insert into t values (x$0) + ) + update u + set x = 1; + end; +"), @" + ╭▸ + 4 │ create function f(x int) returns void + │ ─ 2. destination + ‡ + 7 │ insert into t values (x) + ╰╴ ─ 1. source + "); + } + #[test] fn goto_positional_param_unnamed_param() { assert_snapshot!(goto(" diff --git a/crates/squawk_ide/src/resolve.rs b/crates/squawk_ide/src/resolve.rs index c73d5ece..3169886e 100644 --- a/crates/squawk_ide/src/resolve.rs +++ b/crates/squawk_ide/src/resolve.rs @@ -2813,7 +2813,7 @@ fn resolve_select_column_ptr( // Walk up through enclosing selects so a column in a correlated subquery // resolves against an outer query's from clause: // `select (select (select a)) from foo` - for ancestor in column_name_ref.syntax().ancestors() { + for ancestor in ast_nav::ancestors_outside_own_with_clause(column_name_ref.syntax()) { let Some(from_clause) = select_like_from_clause(&ancestor) else { continue; }; @@ -2835,7 +2835,7 @@ fn resolve_select_column_ptr( // A correlated subquery can reference the target relation of an enclosing // DML statement, e.g. `update foo set a = (select b)` where `b` is `foo.b` let in_file = InFile::new(file, column_name_ref); - for ancestor in column_name_ref.syntax().ancestors() { + for ancestor in ast_nav::ancestors_outside_own_with_clause(column_name_ref.syntax()) { match ancestor.kind() { SyntaxKind::UPDATE => return resolve_update_column_ptr(db, in_file), SyntaxKind::DELETE => return resolve_delete_column_ptr(db, in_file), @@ -5497,7 +5497,11 @@ fn resolve_update_column_ptr( )?; } - resolve_column_for_path(db, InFile::new(file, &path), column_name) + resolve_column_for_path(db, InFile::new(file, &path), column_name).or_else(|| { + (!is_set_target) + .then(|| resolve_enclosing_function_param(InFile::new(file, column_name_ref))) + .flatten() + }) } fn resolve_delete_column_ptr( @@ -5535,6 +5539,7 @@ fn resolve_delete_column_ptr( } resolve_column_for_path(db, InFile::new(file, &path), column_name) + .or_else(|| resolve_enclosing_function_param(InFile::new(file, column_name_ref))) } fn resolve_delete_table_name_ptr( @@ -5622,7 +5627,11 @@ fn resolve_merge_column_ptr( } let path = merge.table_relation_name()?.table_name_ref()?.path_ref()?; - resolve_column_for_path(db, InFile::new(file, &path), column_name) + resolve_column_for_path(db, InFile::new(file, &path), column_name).or_else(|| { + (!is_set_target && !in_insert_column_list) + .then(|| resolve_enclosing_function_param(InFile::new(file, column_name_ref))) + .flatten() + }) } // TODO: I think we could use trait(s) here to simplify this and have the