diff --git a/crates/squawk_ide/src/goto_definition.rs b/crates/squawk_ide/src/goto_definition.rs index ba7ba1f2..21b2bfa3 100644 --- a/crates/squawk_ide/src/goto_definition.rs +++ b/crates/squawk_ide/src/goto_definition.rs @@ -3252,4 +3252,76 @@ update users set email = messages.email from messages$0 where users.id = message ╰╴ ─ 1. source "); } + + #[test] + fn goto_update_with_cte_table() { + assert_snapshot!(goto(" +create table users(id int, email text); +with new_data as ( + select 1 as id, 'new@example.com' as email +) +update users set email = new_data.email from new_data$0 where users.id = new_data.id; +"), @r" + ╭▸ + 3 │ with new_data as ( + │ ──────── 2. destination + ‡ + 6 │ update users set email = new_data.email from new_data where users.id = new_data.id; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_update_with_cte_column_in_set() { + assert_snapshot!(goto(" +create table users(id int, email text); +with new_data as ( + select 1 as id, 'new@example.com' as email +) +update users set email = new_data.email$0 from new_data where users.id = new_data.id; +"), @r" + ╭▸ + 4 │ select 1 as id, 'new@example.com' as email + │ ───── 2. destination + 5 │ ) + 6 │ update users set email = new_data.email from new_data where users.id = new_data.id; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_update_with_cte_column_in_where() { + assert_snapshot!(goto(" +create table users(id int, email text); +with new_data as ( + select 1 as id, 'new@example.com' as email +) +update users set email = new_data.email from new_data where new_data.id$0 = users.id; +"), @r" + ╭▸ + 4 │ select 1 as id, 'new@example.com' as email + │ ── 2. destination + 5 │ ) + 6 │ update users set email = new_data.email from new_data where new_data.id = users.id; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_update_with_cte_values() { + assert_snapshot!(goto(" +create table users(id int, email text); +with new_data as ( + values (1, 'new@example.com') +) +update users set email = new_data.column2$0 from new_data where users.id = new_data.column1; +"), @r" + ╭▸ + 4 │ values (1, 'new@example.com') + │ ───────────────── 2. destination + 5 │ ) + 6 │ update users set email = new_data.column2 from new_data where users.id = new_data.column1; + ╰╴ ─ 1. source + "); + } } diff --git a/crates/squawk_ide/src/hover.rs b/crates/squawk_ide/src/hover.rs index 18a550cf..16da7001 100644 --- a/crates/squawk_ide/src/hover.rs +++ b/crates/squawk_ide/src/hover.rs @@ -2380,4 +2380,52 @@ update users set email = messages.email from public.messages$0 where users.id = ╰╴ ─ hover "); } + + #[test] + fn hover_on_update_with_cte_table() { + assert_snapshot!(check_hover(" +create table users(id int, email text); +with new_data as ( + select 1 as id, 'new@example.com' as email +) +update users set email = new_data.email from new_data$0 where users.id = new_data.id; +"), @r" + hover: with new_data as (select 1 as id, 'new@example.com' as email) + ╭▸ + 6 │ update users set email = new_data.email from new_data where users.id = new_data.id; + ╰╴ ─ hover + "); + } + + #[test] + fn hover_on_update_with_cte_column_in_set() { + assert_snapshot!(check_hover(" +create table users(id int, email text); +with new_data as ( + select 1 as id, 'new@example.com' as email +) +update users set email = new_data.email$0 from new_data where users.id = new_data.id; +"), @r" + hover: column new_data.email + ╭▸ + 6 │ update users set email = new_data.email from new_data where users.id = new_data.id; + ╰╴ ─ hover + "); + } + + #[test] + fn hover_on_update_with_cte_column_in_where() { + assert_snapshot!(check_hover(" +create table users(id int, email text); +with new_data as ( + select 1 as id, 'new@example.com' as email +) +update users set email = new_data.email from new_data where new_data.id$0 = users.id; +"), @r" + hover: column new_data.id + ╭▸ + 6 │ update users set email = new_data.email from new_data where new_data.id = users.id; + ╰╴ ─ hover + "); + } } diff --git a/crates/squawk_ide/src/inlay_hints.rs b/crates/squawk_ide/src/inlay_hints.rs index 78e826c5..94eecf58 100644 --- a/crates/squawk_ide/src/inlay_hints.rs +++ b/crates/squawk_ide/src/inlay_hints.rs @@ -85,9 +85,10 @@ fn inlay_hint_insert( ) -> Option<()> { let values = insert.values()?; let row_list = values.row_list()?; + let create_table = resolve::resolve_insert_create_table(file, binder, &insert); let columns: Vec<(Name, Option)> = if let Some(column_list) = insert.column_list() { - let create_table = resolve::resolve_insert_create_table(file, binder, &insert); + // `insert into t(a, b, c) values (1, 2, 3)` column_list .columns() .filter_map(|col| { @@ -100,8 +101,8 @@ fn inlay_hint_insert( }) .collect() } else { - let create_table = resolve::resolve_insert_create_table(file, binder, &insert)?; - create_table + // `insert into t values (1, 2, 3)` + create_table? .table_arg_list()? .args() .filter_map(|arg| { diff --git a/crates/squawk_ide/src/resolve.rs b/crates/squawk_ide/src/resolve.rs index 18417574..82cf1140 100644 --- a/crates/squawk_ide/src/resolve.rs +++ b/crates/squawk_ide/src/resolve.rs @@ -925,6 +925,27 @@ fn resolve_update_where_column(binder: &Binder, name_ref: &ast::NameRef) -> Opti let column_name = Name::from_node(name_ref); let update = name_ref.syntax().ancestors().find_map(ast::Update::cast)?; + + // `update t set a = b from u` + if let Some(from_clause) = update.from_clause() { + for from_item in from_clause.from_items() { + if let Some(result) = resolve_from_item_for_column(binder, &from_item, name_ref) { + return Some(result); + } + } + + for join_expr in from_clause.join_exprs() { + if let Some(result) = + resolve_from_join_expr(&join_expr, &|from_item: &ast::FromItem| { + resolve_from_item_for_column(binder, from_item, name_ref) + }) + { + return Some(result); + } + } + } + + // `update t set a = b` let relation_name = update.relation_name()?; let path = relation_name.path()?; @@ -1131,8 +1152,10 @@ fn find_parent_with_clause(node: &SyntaxNode) -> Option { 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) { + } else if let Some(insert) = ast::Insert::cast(x.clone()) { insert.with_clause() + } else if let Some(update) = ast::Update::cast(x) { + update.with_clause() } else { None }