Skip to content

Commit 3c65a4e

Browse files
juleswritescodejuleswritescode
authored
feat(grammar): partial update statements (#697)
1 parent 5291a1a commit 3c65a4e

9 files changed

Lines changed: 315 additions & 53 deletions

File tree

‎crates/pgls_completions/src/snapshots/pgls_completions__test_helper__suggests_tables_in_update.snap‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -76,9 +76,9 @@ update public.coos set |
7676
Results:
7777
id - public.coos.id (Column)
7878
name - public.coos.name (Column)
79+
public - public (Schema)
7980
information_schema - information_schema (Schema)
8081
pg_catalog - pg_catalog (Schema)
81-
pg_toast - pg_toast (Schema)
8282

8383
--------------
8484

@@ -107,7 +107,20 @@ information_schema - information_schema (Schema)
107107

108108
update public.coos set name = '|
109109
update public.coos set name = 'cool' |
110+
111+
Results:
112+
from - from (Keyword)
113+
where - where (Keyword)
114+
115+
--------------
116+
110117
update public.coos set name = 'cool' w|
118+
119+
Results:
120+
where - where (Keyword)
121+
122+
--------------
123+
111124
update public.coos set name = 'cool' where |
112125

113126
Results:

‎crates/pgls_tokenizer/src/lib.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -690,12 +690,12 @@ mod tests {
690690
#[test]
691691
fn debug_simple_cast() {
692692
let result = lex("::test");
693-
assert_debug_snapshot!(result, @r###"
693+
assert_debug_snapshot!(result, @r#"
694694
[
695695
"::" @ DoubleColon,
696696
"test" @ Ident,
697697
]
698-
"###);
698+
"#);
699699
}
700700

701701
#[test]

‎crates/pgls_treesitter_grammar/grammar.js‎

Lines changed: 103 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -2357,7 +2357,7 @@ module.exports = grammar({
23572357

23582358
_column_indirection: ($) =>
23592359
seq(
2360-
field("end", $.column_identifier),
2360+
$.column_identifier,
23612361
repeat(
23622362
choice(
23632363
prec.right($.column_indirection_array_access),
@@ -2389,8 +2389,7 @@ module.exports = grammar({
23892389
),
23902390
),
23912391

2392-
_set_values: ($) =>
2393-
partialSeq($.keyword_set, comma_list($.assignment, true)),
2392+
_set_values: ($) => $.update_set_values,
23942393

23952394
_column_list: ($) => paren_list(alias($._column, $.column), false),
23962395
_column: ($) =>
@@ -2524,10 +2523,54 @@ module.exports = grammar({
25242523
update: ($) =>
25252524
partialSeq(
25262525
$.keyword_update,
2527-
optional($.keyword_only),
2528-
$.relation,
2529-
$._set_values,
2530-
optional($.where),
2526+
$.update_target,
2527+
$.update_set_values,
2528+
optional($.update_from),
2529+
optional(choice($.where, $.where_current_of)),
2530+
),
2531+
2532+
update_target: ($) =>
2533+
choice(
2534+
partialSeq(
2535+
$.keyword_only,
2536+
field("end", $.table_reference),
2537+
optional("*"),
2538+
optional($.alias),
2539+
),
2540+
seq(field("end", $.table_reference), optional("*"), optional($.alias)),
2541+
),
2542+
2543+
update_set_values: ($) =>
2544+
partialSeq($.keyword_set, field("end", comma_list($.assignment, true))),
2545+
2546+
update_from: ($) =>
2547+
partialSeq(
2548+
$.keyword_from,
2549+
field(
2550+
"end",
2551+
seq(
2552+
comma_list($.relation, true),
2553+
repeat(
2554+
choice(
2555+
$.join,
2556+
$.cross_join,
2557+
$.lateral_join,
2558+
$.lateral_cross_join,
2559+
),
2560+
),
2561+
),
2562+
),
2563+
),
2564+
2565+
where_current_of: ($) =>
2566+
prec(
2567+
1,
2568+
partialSeq(
2569+
$.keyword_where,
2570+
$.keyword_current,
2571+
$.keyword_of,
2572+
field("end", $.any_identifier),
2573+
),
25312574
),
25322575

25332576
table_partition: ($) =>
@@ -2566,10 +2609,40 @@ module.exports = grammar({
25662609
),
25672610

25682611
assignment: ($) =>
2569-
partialSeq(
2570-
field("left", $.column_reference),
2571-
"=",
2572-
field("right", $._expression),
2612+
choice(
2613+
partialSeq(
2614+
field("left", $._column_indirection),
2615+
"=",
2616+
field("end", choice($._expression, $.all_fields, $.keyword_default)),
2617+
),
2618+
partialSeq(
2619+
field("left", $.lhs_column_list),
2620+
"=",
2621+
field("end", $.rhs_column_list),
2622+
),
2623+
),
2624+
2625+
lhs_column_list: ($) =>
2626+
partialSeq("(", comma_list($._column_indirection, false), ")"),
2627+
2628+
rhs_column_list: ($) =>
2629+
choice(
2630+
wrapped_in_parenthesis(
2631+
comma_list(
2632+
choice($._expression, $.all_fields, $.keyword_default),
2633+
true,
2634+
),
2635+
),
2636+
partialSeq(
2637+
$.keyword_row,
2638+
wrapped_in_parenthesis(
2639+
comma_list(
2640+
choice($._expression, $.all_fields, $.keyword_default),
2641+
false,
2642+
),
2643+
),
2644+
),
2645+
$.subquery,
25732646
),
25742647

25752648
table_option: ($) =>
@@ -3082,7 +3155,25 @@ module.exports = grammar({
30823155

30833156
offset: ($) => partialSeq($.keyword_offset, field("end", $.literal)),
30843157

3085-
returning: ($) => partialSeq($.keyword_returning, $.select_expression),
3158+
returning: ($) =>
3159+
partialSeq(
3160+
$.keyword_returning,
3161+
optional($.returning_with),
3162+
field("end", $.select_expression),
3163+
),
3164+
3165+
returning_with: ($) =>
3166+
partialSeq(
3167+
$.keyword_with,
3168+
wrapped_in_parenthesis(comma_list($.returning_with_item, true)),
3169+
),
3170+
3171+
returning_with_item: ($) =>
3172+
partialSeq(
3173+
choice($.keyword_old, $.keyword_new),
3174+
$.keyword_as,
3175+
field("end", $.any_identifier),
3176+
),
30863177

30873178
grant_statement: ($) =>
30883179
prec.left(

‎crates/pgls_treesitter_grammar/tests/grammar_tests.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,8 +36,8 @@ fn test_1() {
3636
#[test]
3737
fn test_2() {
3838
let sql1 = "update auth.users set email = 'my@mail.com';";
39-
let sql2 = "update auth.users set users.email = 'my@mail.com';";
40-
let sql3 = "update auth.users set auth.users.email = 'my@mail.com';";
39+
let sql2 = "update auth.users as u set email = 'my@mail.com' where u.id = 1;";
40+
let sql3 = "update auth.users set (email, id) = ('my@mail.com', 1);";
4141

4242
file_snapshot("test_2_sql1", sql1);
4343
file_snapshot("test_2_sql2", sql2);

‎crates/pgls_treesitter_grammar/tests/partial-sqls/update.sql‎

Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,141 @@
1+
UPDATE update_test SET a = DEFAULT, b = DEFAULT;
2+
3+
UPDATE update_test AS t SET b = 10 WHERE t.a = 10;
4+
5+
UPDATE update_test t SET b = t.b + 10 WHERE t.a = 10;
6+
7+
UPDATE update_test t SET t.b = t.b + 10 WHERE t.a = 10;
8+
9+
UPDATE update_test SET a=v.i FROM (VALUES(100, 20)) AS v(i, j)
10+
WHERE update_test.b = v.j;
11+
12+
UPDATE update_test SET a = v.* FROM (VALUES(100, 20)) AS v(i, j)
13+
WHERE update_test.b = v.j;
14+
15+
UPDATE update_test SET (c,b,a) = ('bugle', b+11, DEFAULT) WHERE c = 'foo';
16+
17+
UPDATE update_test SET (c,b) = ('car', a+b), a = a + 1 WHERE a = 10;
18+
19+
UPDATE update_test SET (c,b) = ('car', a+b), b = a + 1 WHERE a = 10;
20+
21+
UPDATE update_test
22+
SET (b,a) = (select a,b from update_test where b = 41 and c = 'car')
23+
WHERE a = 100 AND b = 20;
24+
25+
UPDATE update_test o
26+
SET (b,a) = (select a+1,b from update_test i
27+
where i.a=o.a and i.b=o.b and i.c is not distinct from o.c);
28+
29+
UPDATE update_test SET (b,a) = (select a+1,b from update_test);
30+
31+
UPDATE update_test SET (b,a) = (select a+1,b from update_test where a = 1000)
32+
WHERE a = 11;
33+
34+
UPDATE update_test SET (a,b) = ROW(v.*) FROM (VALUES(21, 100)) AS v(i, j)
35+
WHERE update_test.a = v.i;
36+
37+
UPDATE update_test SET (a,b) = (v.*) FROM (VALUES(21, 101)) AS v(i, j)
38+
WHERE update_test.a = v.i;
39+
40+
UPDATE update_test AS t SET b = update_test.b + 10 WHERE t.a = 10;
41+
42+
UPDATE update_test SET c = repeat('x', 10000) WHERE c = 'car';
43+
44+
UPDATE update_test t
45+
SET (a, b) = (SELECT b, a FROM update_test s WHERE s.a = t.a)
46+
WHERE CURRENT_USER = SESSION_USER;
47+
48+
UPDATE part_b_10_b_20 set b = b - 6;
49+
50+
UPDATE part_c_100_200 set c = c - 20, d = c WHERE c = 105;
51+
52+
UPDATE part_b_10_b_20 set a = 'a';
53+
54+
UPDATE range_parted set d = d - 10 WHERE d > 10;
55+
56+
UPDATE range_parted set e = d;
57+
58+
UPDATE part_c_1_100 set c = c + 20 WHERE c = 98;
59+
60+
UPDATE part_b_10_b_20 set c = c + 20 returning c, b, a;
61+
62+
UPDATE part_b_10_b_20 set b = b - 6 WHERE c > 116 returning *;
63+
64+
UPDATE range_parted set b = b - 6 WHERE c > 116 returning a, b + c;
65+
66+
UPDATE upview set c = 199 WHERE b = 4;
67+
68+
UPDATE upview set c = 120 WHERE b = 4;
69+
70+
UPDATE upview set a = 'b', b = 15, c = 120 WHERE b = 4;
71+
72+
UPDATE upview set a = 'b', b = 15 WHERE b = 4;
73+
74+
UPDATE range_parted set c = 95 WHERE a = 'b' and b > 10 and c > 100 returning (range_parted), *;
75+
76+
UPDATE range_parted set c = c + 50 WHERE a = 'b' and b > 10 and c >= 96;
77+
78+
UPDATE range_parted set c = c + 50 WHERE a = 'b' and b > 10 and c >= 96;
79+
80+
UPDATE range_parted set b = 15 WHERE b = 1;
81+
82+
UPDATE range_parted set a = 'b', c = 151 WHERE a = 'a' and c = 200;
83+
84+
UPDATE range_parted set a = 'b', c = 151 WHERE a = 'a' and c = 200;
85+
86+
UPDATE range_parted set a = 'b', c = 150 WHERE a = 'a' and c = 200;
87+
88+
UPDATE range_parted set a = 'b', c = 122 WHERE a = 'a' and c = 200;
89+
90+
UPDATE range_parted set a = 'b', c = 120 WHERE a = 'a' and c = 200;
91+
92+
UPDATE range_parted set a = 'b', c = 112 WHERE a = 'a' and c = 200;
93+
94+
UPDATE range_parted set a = 'b', c = 116 WHERE a = 'a' and c = 200;
95+
96+
UPDATE range_parted set c = c - 50 WHERE c > 97;
97+
98+
update part_def set a = 'd' where a = 'c';
99+
100+
update part_def set a = 'a' where a = 'd';
101+
102+
UPDATE part_a_10_a_20 set a = 'ad' WHERE a = 'a';
103+
104+
UPDATE range_parted set a = 'ad' WHERE a = 'a';
105+
106+
UPDATE range_parted set a = 'bd' WHERE a = 'b';
107+
108+
UPDATE range_parted set a = 'a' WHERE a = 'ad';
109+
110+
UPDATE range_parted set a = 'b' WHERE a = 'bd';
111+
112+
UPDATE list_default set a = 'a' WHERE a = 'd';
113+
114+
UPDATE list_default set a = 'x' WHERE a = 'd';
115+
116+
update utrtest set b = b || b from (values (1), (2)) s(x) where a = s.x
117+
returning *, tableoid::regclass, xmin = pg_current_xact_id()::xid as xmin_ok;
118+
119+
update utrtest set a = 3 - a from (values (1), (2)) s(x) where a = s.x
120+
returning *, tableoid::regclass, xmin = pg_current_xact_id()::xid as xmin_ok;
121+
122+
update utrtest set a = 3 - a from (values (1), (2)) s(x) where a = s.x
123+
returning *, tableoid::regclass;
124+
125+
UPDATE sub_parted set a = 2 WHERE c = 10;
126+
127+
UPDATE list_parted set b = c + a WHERE a = 2;
128+
129+
UPDATE list_parted set c = 70 WHERE b = 1;
130+
131+
UPDATE list_parted set b = 1 WHERE c = 70;
132+
133+
UPDATE list_parted set b = 1 WHERE c = 70;
134+
135+
UPDATE list_parted t1 set a = 2 FROM non_parted t2 WHERE t1.a = t2.id and a = 1;
136+
137+
update hpart1 set a = 3, b=4 where a = 1;
138+
139+
update hash_parted set b = b - 1 where b = 1;
140+
141+
update hash_parted set b = b + 8 where b = 1;

‎crates/pgls_treesitter_grammar/tests/partial_no_errors.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ static PG_REGRESSION_FILES: LazyLock<Vec<&'static str>> = std::sync::LazyLock::n
5252
*/
5353
include_str!("../tests/partial-sqls/insert.sql"),
5454
include_str!("../tests/partial-sqls/copy.sql"),
55+
include_str!("../tests/partial-sqls/update.sql"),
5556
]
5657
});
5758

‎crates/pgls_treesitter_grammar/tests/snapshots/grammar_tests__test_2_sql1.snap‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -8,15 +8,15 @@ program [0..44] 'update auth.users set email = 'my@mail.com';'
88
statement [0..43] 'update auth.users set email = 'my@mail.com''
99
update [0..43] 'update auth.users set email = 'my@mail.com''
1010
keyword_update [0..6] 'update'
11-
relation [7..17] 'auth.users'
11+
update_target [7..17] 'auth.users'
1212
table_reference [7..17] 'auth.users' (@end)
1313
schema_identifier [7..11] 'auth' (@table_reference_1of2)
1414
. [11..12] '.'
1515
table_identifier [12..17] 'users' (@table_reference_2of2)
16-
keyword_set [18..21] 'set'
17-
assignment [22..43] 'email = 'my@mail.com''
18-
column_reference [22..27] 'email' (@left)
19-
any_identifier [22..27] 'email' (@column_reference_1of1)
20-
= [28..29] '='
21-
literal [30..43] ''my@mail.com'' (@right)
16+
update_set_values [18..43] 'set email = 'my@mail.com''
17+
keyword_set [18..21] 'set'
18+
assignment [22..43] 'email = 'my@mail.com'' (@end)
19+
column_identifier [22..27] 'email' (@left)
20+
= [28..29] '='
21+
literal [30..43] ''my@mail.com'' (@end)
2222
; [43..44] ';'

0 commit comments

Comments
 (0)