Skip to content

Commit a8e2135

Browse files
sbdchdsbdchd
authored
ide: support ctes with update (#797)
1 parent e0166d3 commit a8e2135

File tree

4 files changed

+148
-4
lines changed

4 files changed

+148
-4
lines changed

‎crates/squawk_ide/src/goto_definition.rs‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3252,4 +3252,76 @@ update users set email = messages.email from messages$0 where users.id = message
32523252
╰╴ ─ 1. source
32533253
");
32543254
}
3255+
3256+
#[test]
3257+
fn goto_update_with_cte_table() {
3258+
assert_snapshot!(goto("
3259+
create table users(id int, email text);
3260+
with new_data as (
3261+
select 1 as id, '[email protected]' as email
3262+
)
3263+
update users set email = new_data.email from new_data$0 where users.id = new_data.id;
3264+
"), @r"
3265+
╭▸
3266+
3 │ with new_data as (
3267+
│ ──────── 2. destination
3268+
‡
3269+
6 │ update users set email = new_data.email from new_data where users.id = new_data.id;
3270+
╰╴ ─ 1. source
3271+
");
3272+
}
3273+
3274+
#[test]
3275+
fn goto_update_with_cte_column_in_set() {
3276+
assert_snapshot!(goto("
3277+
create table users(id int, email text);
3278+
with new_data as (
3279+
select 1 as id, '[email protected]' as email
3280+
)
3281+
update users set email = new_data.email$0 from new_data where users.id = new_data.id;
3282+
"), @r"
3283+
╭▸
3284+
4 │ select 1 as id, '[email protected]' as email
3285+
│ ───── 2. destination
3286+
5 │ )
3287+
6 │ update users set email = new_data.email from new_data where users.id = new_data.id;
3288+
╰╴ ─ 1. source
3289+
");
3290+
}
3291+
3292+
#[test]
3293+
fn goto_update_with_cte_column_in_where() {
3294+
assert_snapshot!(goto("
3295+
create table users(id int, email text);
3296+
with new_data as (
3297+
select 1 as id, '[email protected]' as email
3298+
)
3299+
update users set email = new_data.email from new_data where new_data.id$0 = users.id;
3300+
"), @r"
3301+
╭▸
3302+
4 │ select 1 as id, '[email protected]' as email
3303+
│ ── 2. destination
3304+
5 │ )
3305+
6 │ update users set email = new_data.email from new_data where new_data.id = users.id;
3306+
╰╴ ─ 1. source
3307+
");
3308+
}
3309+
3310+
#[test]
3311+
fn goto_update_with_cte_values() {
3312+
assert_snapshot!(goto("
3313+
create table users(id int, email text);
3314+
with new_data as (
3315+
values (1, '[email protected]')
3316+
)
3317+
update users set email = new_data.column2$0 from new_data where users.id = new_data.column1;
3318+
"), @r"
3319+
╭▸
3320+
4 │ values (1, '[email protected]')
3321+
│ ───────────────── 2. destination
3322+
5 │ )
3323+
6 │ update users set email = new_data.column2 from new_data where users.id = new_data.column1;
3324+
╰╴ ─ 1. source
3325+
");
3326+
}
32553327
}

‎crates/squawk_ide/src/hover.rs‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2380,4 +2380,52 @@ update users set email = messages.email from public.messages$0 where users.id =
23802380
╰╴ ─ hover
23812381
");
23822382
}
2383+
2384+
#[test]
2385+
fn hover_on_update_with_cte_table() {
2386+
assert_snapshot!(check_hover("
2387+
create table users(id int, email text);
2388+
with new_data as (
2389+
select 1 as id, '[email protected]' as email
2390+
)
2391+
update users set email = new_data.email from new_data$0 where users.id = new_data.id;
2392+
"), @r"
2393+
hover: with new_data as (select 1 as id, '[email protected]' as email)
2394+
╭▸
2395+
6 │ update users set email = new_data.email from new_data where users.id = new_data.id;
2396+
╰╴ ─ hover
2397+
");
2398+
}
2399+
2400+
#[test]
2401+
fn hover_on_update_with_cte_column_in_set() {
2402+
assert_snapshot!(check_hover("
2403+
create table users(id int, email text);
2404+
with new_data as (
2405+
select 1 as id, '[email protected]' as email
2406+
)
2407+
update users set email = new_data.email$0 from new_data where users.id = new_data.id;
2408+
"), @r"
2409+
hover: column new_data.email
2410+
╭▸
2411+
6 │ update users set email = new_data.email from new_data where users.id = new_data.id;
2412+
╰╴ ─ hover
2413+
");
2414+
}
2415+
2416+
#[test]
2417+
fn hover_on_update_with_cte_column_in_where() {
2418+
assert_snapshot!(check_hover("
2419+
create table users(id int, email text);
2420+
with new_data as (
2421+
select 1 as id, '[email protected]' as email
2422+
)
2423+
update users set email = new_data.email from new_data where new_data.id$0 = users.id;
2424+
"), @r"
2425+
hover: column new_data.id
2426+
╭▸
2427+
6 │ update users set email = new_data.email from new_data where new_data.id = users.id;
2428+
╰╴ ─ hover
2429+
");
2430+
}
23832431
}

‎crates/squawk_ide/src/inlay_hints.rs‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -85,9 +85,10 @@ fn inlay_hint_insert(
8585
) -> Option<()> {
8686
let values = insert.values()?;
8787
let row_list = values.row_list()?;
88+
let create_table = resolve::resolve_insert_create_table(file, binder, &insert);
8889

8990
let columns: Vec<(Name, Option<TextRange>)> = if let Some(column_list) = insert.column_list() {
90-
let create_table = resolve::resolve_insert_create_table(file, binder, &insert);
91+
// `insert into t(a, b, c) values (1, 2, 3)`
9192
column_list
9293
.columns()
9394
.filter_map(|col| {
@@ -100,8 +101,8 @@ fn inlay_hint_insert(
100101
})
101102
.collect()
102103
} else {
103-
let create_table = resolve::resolve_insert_create_table(file, binder, &insert)?;
104-
create_table
104+
// `insert into t values (1, 2, 3)`
105+
create_table?
105106
.table_arg_list()?
106107
.args()
107108
.filter_map(|arg| {

‎crates/squawk_ide/src/resolve.rs‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -925,6 +925,27 @@ fn resolve_update_where_column(binder: &Binder, name_ref: &ast::NameRef) -> Opti
925925
let column_name = Name::from_node(name_ref);
926926

927927
let update = name_ref.syntax().ancestors().find_map(ast::Update::cast)?;
928+
929+
// `update t set a = b from u`
930+
if let Some(from_clause) = update.from_clause() {
931+
for from_item in from_clause.from_items() {
932+
if let Some(result) = resolve_from_item_for_column(binder, &from_item, name_ref) {
933+
return Some(result);
934+
}
935+
}
936+
937+
for join_expr in from_clause.join_exprs() {
938+
if let Some(result) =
939+
resolve_from_join_expr(&join_expr, &|from_item: &ast::FromItem| {
940+
resolve_from_item_for_column(binder, from_item, name_ref)
941+
})
942+
{
943+
return Some(result);
944+
}
945+
}
946+
}
947+
948+
// `update t set a = b`
928949
let relation_name = update.relation_name()?;
929950
let path = relation_name.path()?;
930951

@@ -1131,8 +1152,10 @@ fn find_parent_with_clause(node: &SyntaxNode) -> Option<ast::WithClause> {
11311152
select.with_clause()
11321153
} else if let Some(delete) = ast::Delete::cast(x.clone()) {
11331154
delete.with_clause()
1134-
} else if let Some(insert) = ast::Insert::cast(x) {
1155+
} else if let Some(insert) = ast::Insert::cast(x.clone()) {
11351156
insert.with_clause()
1157+
} else if let Some(update) = ast::Update::cast(x) {
1158+
update.with_clause()
11361159
} else {
11371160
None
11381161
}

0 commit comments

Comments
 (0)