feat(grammar): partial update statements

crates/pgls_completions/src/snapshots/pgls_completions__test_helper__suggests_tables_in_update.snap

@@ -76,9 +76,9 @@ update public.coos set |
 Results:
 id - public.coos.id (Column)
 name - public.coos.name (Column)
+public - public (Schema)
 information_schema - information_schema (Schema)
 pg_catalog - pg_catalog (Schema)
-pg_toast - pg_toast (Schema)
 
 --------------
 
@@ -107,7 +107,20 @@ information_schema - information_schema (Schema)
 
 update public.coos set name = '|
 update public.coos set name = 'cool' |
+
+Results:
+from - from (Keyword)
+where - where (Keyword)
+
+--------------
+
 update public.coos set name = 'cool' w|
+
+Results:
+where - where (Keyword)
+
+--------------
+
 update public.coos set name = 'cool' where |
 
 Results:

crates/pgls_tokenizer/src/lib.rs

@@ -690,12 +690,12 @@ mod tests {
     #[test]
     fn debug_simple_cast() {
         let result = lex("::test");
-        assert_debug_snapshot!(result, @r###"
+        assert_debug_snapshot!(result, @r#"
         [
             "::" @ DoubleColon,
             "test" @ Ident,
         ]
-        "###);
+        "#);
     }
 
     #[test]

crates/pgls_treesitter_grammar/grammar.js

@@ -2357,7 +2357,7 @@ module.exports = grammar({
 
     _column_indirection: ($) =>
       seq(
-        field("end", $.column_identifier),
+        $.column_identifier,
         repeat(
           choice(
             prec.right($.column_indirection_array_access),
@@ -2389,8 +2389,7 @@ module.exports = grammar({
         ),
       ),
 
-    _set_values: ($) =>
-      partialSeq($.keyword_set, comma_list($.assignment, true)),
+    _set_values: ($) => $.update_set_values,
 
     _column_list: ($) => paren_list(alias($._column, $.column), false),
     _column: ($) =>
@@ -2524,10 +2523,54 @@ module.exports = grammar({
     update: ($) =>
       partialSeq(
         $.keyword_update,
-        optional($.keyword_only),
-        $.relation,
-        $._set_values,
-        optional($.where),
+        $.update_target,
+        $.update_set_values,
+        optional($.update_from),
+        optional(choice($.where, $.where_current_of)),
+      ),
+
+    update_target: ($) =>
+      choice(
+        partialSeq(
+          $.keyword_only,
+          field("end", $.table_reference),
+          optional("*"),
+          optional($.alias),
+        ),
+        seq(field("end", $.table_reference), optional("*"), optional($.alias)),
+      ),
+
+    update_set_values: ($) =>
+      partialSeq($.keyword_set, field("end", comma_list($.assignment, true))),
+
+    update_from: ($) =>
+      partialSeq(
+        $.keyword_from,
+        field(
+          "end",
+          seq(
+            comma_list($.relation, true),
+            repeat(
+              choice(
+                $.join,
+                $.cross_join,
+                $.lateral_join,
+                $.lateral_cross_join,
+              ),
+            ),
+          ),
+        ),
+      ),
+
+    where_current_of: ($) =>
+      prec(
+        1,
+        partialSeq(
+          $.keyword_where,
+          $.keyword_current,
+          $.keyword_of,
+          field("end", $.any_identifier),
+        ),
       ),
 
     table_partition: ($) =>
@@ -2566,10 +2609,40 @@ module.exports = grammar({
       ),
 
     assignment: ($) =>
-      partialSeq(
-        field("left", $.column_reference),
-        "=",
-        field("right", $._expression),
+      choice(
+        partialSeq(
+          field("left", $._column_indirection),
+          "=",
+          field("end", choice($._expression, $.all_fields, $.keyword_default)),
+        ),
+        partialSeq(
+          field("left", $.lhs_column_list),
+          "=",
+          field("end", $.rhs_column_list),
+        ),
+      ),
+
+    lhs_column_list: ($) =>
+      partialSeq("(", comma_list($._column_indirection, false), ")"),
+
+    rhs_column_list: ($) =>
+      choice(
+        wrapped_in_parenthesis(
+          comma_list(
+            choice($._expression, $.all_fields, $.keyword_default),
+            true,
+          ),
+        ),
+        partialSeq(
+          $.keyword_row,
+          wrapped_in_parenthesis(
+            comma_list(
+              choice($._expression, $.all_fields, $.keyword_default),
+              false,
+            ),
+          ),
+        ),
+        $.subquery,
       ),
 
     table_option: ($) =>
@@ -3082,7 +3155,25 @@ module.exports = grammar({
 
     offset: ($) => partialSeq($.keyword_offset, field("end", $.literal)),
 
-    returning: ($) => partialSeq($.keyword_returning, $.select_expression),
+    returning: ($) =>
+      partialSeq(
+        $.keyword_returning,
+        optional($.returning_with),
+        field("end", $.select_expression),
+      ),
+
+    returning_with: ($) =>
+      partialSeq(
+        $.keyword_with,
+        wrapped_in_parenthesis(comma_list($.returning_with_item, true)),
+      ),
+
+    returning_with_item: ($) =>
+      partialSeq(
+        choice($.keyword_old, $.keyword_new),
+        $.keyword_as,
+        field("end", $.any_identifier),
+      ),
 
     grant_statement: ($) =>
       prec.left(

crates/pgls_treesitter_grammar/tests/grammar_tests.rs

@@ -36,8 +36,8 @@ fn test_1() {
 #[test]
 fn test_2() {
     let sql1 = "update auth.users set email = 'my@mail.com';";
-    let sql2 = "update auth.users set users.email = 'my@mail.com';";
-    let sql3 = "update auth.users set auth.users.email = 'my@mail.com';";
+    let sql2 = "update auth.users as u set email = 'my@mail.com' where u.id = 1;";
+    let sql3 = "update auth.users set (email, id) = ('my@mail.com', 1);";
 
     file_snapshot("test_2_sql1", sql1);
     file_snapshot("test_2_sql2", sql2);

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

@@ -0,0 +1,141 @@
+UPDATE update_test SET a = DEFAULT, b = DEFAULT;
+
+UPDATE update_test AS t SET b = 10 WHERE t.a = 10;
+
+UPDATE update_test t SET b = t.b + 10 WHERE t.a = 10;
+
+UPDATE update_test t SET t.b = t.b + 10 WHERE t.a = 10;
+
+UPDATE update_test SET a=v.i FROM (VALUES(100, 20)) AS v(i, j)
+  WHERE update_test.b = v.j;
+
+UPDATE update_test SET a = v.* FROM (VALUES(100, 20)) AS v(i, j)
+  WHERE update_test.b = v.j;
+
+UPDATE update_test SET (c,b,a) = ('bugle', b+11, DEFAULT) WHERE c = 'foo';
+
+UPDATE update_test SET (c,b) = ('car', a+b), a = a + 1 WHERE a = 10;
+
+UPDATE update_test SET (c,b) = ('car', a+b), b = a + 1 WHERE a = 10;
+
+UPDATE update_test
+  SET (b,a) = (select a,b from update_test where b = 41 and c = 'car')
+  WHERE a = 100 AND b = 20;
+
+UPDATE update_test o
+  SET (b,a) = (select a+1,b from update_test i
+               where i.a=o.a and i.b=o.b and i.c is not distinct from o.c);
+
+UPDATE update_test SET (b,a) = (select a+1,b from update_test);
+
+UPDATE update_test SET (b,a) = (select a+1,b from update_test where a = 1000)
+  WHERE a = 11;
+
+UPDATE update_test SET (a,b) = ROW(v.*) FROM (VALUES(21, 100)) AS v(i, j)
+  WHERE update_test.a = v.i;
+
+UPDATE update_test SET (a,b) = (v.*) FROM (VALUES(21, 101)) AS v(i, j)
+  WHERE update_test.a = v.i;
+
+UPDATE update_test AS t SET b = update_test.b + 10 WHERE t.a = 10;
+
+UPDATE update_test SET c = repeat('x', 10000) WHERE c = 'car';
+
+UPDATE update_test t
+  SET (a, b) = (SELECT b, a FROM update_test s WHERE s.a = t.a)
+  WHERE CURRENT_USER = SESSION_USER;
+
+UPDATE part_b_10_b_20 set b = b - 6;
+
+UPDATE part_c_100_200 set c = c - 20, d = c WHERE c = 105;
+
+UPDATE part_b_10_b_20 set a = 'a';
+
+UPDATE range_parted set d = d - 10 WHERE d > 10;
+
+UPDATE range_parted set e = d;
+
+UPDATE part_c_1_100 set c = c + 20 WHERE c = 98;
+
+UPDATE part_b_10_b_20 set c = c + 20 returning c, b, a;
+
+UPDATE part_b_10_b_20 set b = b - 6 WHERE c > 116 returning *;
+
+UPDATE range_parted set b = b - 6 WHERE c > 116 returning a, b + c;
+
+UPDATE upview set c = 199 WHERE b = 4;
+
+UPDATE upview set c = 120 WHERE b = 4;
+
+UPDATE upview set a = 'b', b = 15, c = 120 WHERE b = 4;
+
+UPDATE upview set a = 'b', b = 15 WHERE b = 4;
+
+UPDATE range_parted set c = 95 WHERE a = 'b' and b > 10 and c > 100 returning (range_parted), *;
+
+UPDATE range_parted set c = c + 50 WHERE a = 'b' and b > 10 and c >= 96;
+
+UPDATE range_parted set c = c + 50 WHERE a = 'b' and b > 10 and c >= 96;
+
+UPDATE range_parted set b = 15 WHERE b = 1;
+
+UPDATE range_parted set a = 'b', c = 151 WHERE a = 'a' and c = 200;
+
+UPDATE range_parted set a = 'b', c = 151 WHERE a = 'a' and c = 200;
+
+UPDATE range_parted set a = 'b', c = 150 WHERE a = 'a' and c = 200;
+
+UPDATE range_parted set a = 'b', c = 122 WHERE a = 'a' and c = 200;
+
+UPDATE range_parted set a = 'b', c = 120 WHERE a = 'a' and c = 200;
+
+UPDATE range_parted set a = 'b', c = 112 WHERE a = 'a' and c = 200;
+
+UPDATE range_parted set a = 'b', c = 116 WHERE a = 'a' and c = 200;
+
+UPDATE range_parted set c = c - 50 WHERE c > 97;
+
+update part_def set a = 'd' where a = 'c';
+
+update part_def set a = 'a' where a = 'd';
+
+UPDATE part_a_10_a_20 set a = 'ad' WHERE a = 'a';
+
+UPDATE range_parted set a = 'ad' WHERE a = 'a';
+
+UPDATE range_parted set a = 'bd' WHERE a = 'b';
+
+UPDATE range_parted set a = 'a' WHERE a = 'ad';
+
+UPDATE range_parted set a = 'b' WHERE a = 'bd';
+
+UPDATE list_default set a = 'a' WHERE a = 'd';
+
+UPDATE list_default set a = 'x' WHERE a = 'd';
+
+update utrtest set b = b || b from (values (1), (2)) s(x) where a = s.x
+  returning *, tableoid::regclass, xmin = pg_current_xact_id()::xid as xmin_ok;
+
+update utrtest set a = 3 - a from (values (1), (2)) s(x) where a = s.x
+  returning *, tableoid::regclass, xmin = pg_current_xact_id()::xid as xmin_ok;
+
+update utrtest set a = 3 - a from (values (1), (2)) s(x) where a = s.x
+  returning *, tableoid::regclass;
+
+UPDATE sub_parted set a = 2 WHERE c = 10;
+
+UPDATE list_parted set b = c + a WHERE a = 2;
+
+UPDATE list_parted set c = 70 WHERE b  = 1;
+
+UPDATE list_parted set b = 1 WHERE c = 70;
+
+UPDATE list_parted set b = 1 WHERE c = 70;
+
+UPDATE list_parted t1 set a = 2 FROM non_parted t2 WHERE t1.a = t2.id and a = 1;
+
+update hpart1 set a = 3, b=4 where a = 1;
+
+update hash_parted set b = b - 1 where b = 1;
+
+update hash_parted set b = b + 8 where b = 1;

crates/pgls_treesitter_grammar/tests/partial_no_errors.rs

@@ -52,6 +52,7 @@ static PG_REGRESSION_FILES: LazyLock<Vec<&'static str>> = std::sync::LazyLock::n
          */
         include_str!("../tests/partial-sqls/insert.sql"),
         include_str!("../tests/partial-sqls/copy.sql"),
+        include_str!("../tests/partial-sqls/update.sql"),
     ]
 });
 

crates/pgls_treesitter_grammar/tests/snapshots/grammar_tests__test_2_sql1.snap

@@ -8,15 +8,15 @@ program [0..44] 'update auth.users set email = 'my@mail.com';'
   statement [0..43] 'update auth.users set email = 'my@mail.com''
     update [0..43] 'update auth.users set email = 'my@mail.com''
       keyword_update [0..6] 'update'
-      relation [7..17] 'auth.users'
+      update_target [7..17] 'auth.users'
         table_reference [7..17] 'auth.users' (@end)
           schema_identifier [7..11] 'auth' (@table_reference_1of2)
           . [11..12] '.'
           table_identifier [12..17] 'users' (@table_reference_2of2)
-      keyword_set [18..21] 'set'
-      assignment [22..43] 'email = 'my@mail.com''
-        column_reference [22..27] 'email' (@left)
-          any_identifier [22..27] 'email' (@column_reference_1of1)
-        = [28..29] '='
-        literal [30..43] ''my@mail.com'' (@right)
+      update_set_values [18..43] 'set email = 'my@mail.com''
+        keyword_set [18..21] 'set'
+        assignment [22..43] 'email = 'my@mail.com'' (@end)
+          column_identifier [22..27] 'email' (@left)
+          = [28..29] '='
+          literal [30..43] ''my@mail.com'' (@end)
   ; [43..44] ';'