diff --git a/crates/pgls_completions/src/snapshots/pgls_completions__test_helper__suggests_tables_in_update.snap b/crates/pgls_completions/src/snapshots/pgls_completions__test_helper__suggests_tables_in_update.snap index b64ad3b2d..bde9d523b 100644 --- a/crates/pgls_completions/src/snapshots/pgls_completions__test_helper__suggests_tables_in_update.snap +++ b/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: diff --git a/crates/pgls_tokenizer/src/lib.rs b/crates/pgls_tokenizer/src/lib.rs index 1b316eb01..b791fabe9 100644 --- a/crates/pgls_tokenizer/src/lib.rs +++ b/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] diff --git a/crates/pgls_treesitter_grammar/grammar.js b/crates/pgls_treesitter_grammar/grammar.js index 3a93c597d..e1709fff5 100644 --- a/crates/pgls_treesitter_grammar/grammar.js +++ b/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( diff --git a/crates/pgls_treesitter_grammar/tests/grammar_tests.rs b/crates/pgls_treesitter_grammar/tests/grammar_tests.rs index 98eb1a8c0..3b0d519b2 100644 --- a/crates/pgls_treesitter_grammar/tests/grammar_tests.rs +++ b/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); diff --git a/crates/pgls_treesitter_grammar/tests/partial-sqls/update.sql b/crates/pgls_treesitter_grammar/tests/partial-sqls/update.sql new file mode 100644 index 000000000..037afd495 --- /dev/null +++ b/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; diff --git a/crates/pgls_treesitter_grammar/tests/partial_no_errors.rs b/crates/pgls_treesitter_grammar/tests/partial_no_errors.rs index 33414a131..2b0319d25 100644 --- a/crates/pgls_treesitter_grammar/tests/partial_no_errors.rs +++ b/crates/pgls_treesitter_grammar/tests/partial_no_errors.rs @@ -52,6 +52,7 @@ static PG_REGRESSION_FILES: LazyLock> = 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"), ] }); diff --git a/crates/pgls_treesitter_grammar/tests/snapshots/grammar_tests__test_2_sql1.snap b/crates/pgls_treesitter_grammar/tests/snapshots/grammar_tests__test_2_sql1.snap index a25f91f10..8691db6e4 100644 --- a/crates/pgls_treesitter_grammar/tests/snapshots/grammar_tests__test_2_sql1.snap +++ b/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] ';' diff --git a/crates/pgls_treesitter_grammar/tests/snapshots/grammar_tests__test_2_sql2.snap b/crates/pgls_treesitter_grammar/tests/snapshots/grammar_tests__test_2_sql2.snap index bbe72f15e..e7a1116ab 100644 --- a/crates/pgls_treesitter_grammar/tests/snapshots/grammar_tests__test_2_sql2.snap +++ b/crates/pgls_treesitter_grammar/tests/snapshots/grammar_tests__test_2_sql2.snap @@ -2,23 +2,33 @@ source: crates/pgls_treesitter_grammar/tests/grammar_tests.rs expression: writer --- -update auth.users set users.email = 'my@mail.com'; +update auth.users as u set email = 'my@mail.com' where u.id = 1; ----------------------- -program [0..50] 'update auth.users set users.email = 'my@mail.com';' - statement [0..49] 'update auth.users set users.email = 'my@mail.com'' - update [0..49] 'update auth.users set users.email = 'my@mail.com'' +program [0..64] 'update auth.users as u set email = 'my@mail.com' where u.id = 1;' + statement [0..63] 'update auth.users as u set email = 'my@mail.com' where u.id = 1' + update [0..63] 'update auth.users as u set email = 'my@mail.com' where u.id = 1' keyword_update [0..6] 'update' - relation [7..17] 'auth.users' + update_target [7..22] 'auth.users as u' 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..49] 'users.email = 'my@mail.com'' - column_reference [22..33] 'users.email' (@left) - any_identifier [22..27] 'users' (@column_reference_1of2) - . [27..28] '.' - any_identifier [28..33] 'email' (@column_reference_2of2) - = [34..35] '=' - literal [36..49] ''my@mail.com'' (@right) - ; [49..50] ';' + alias [18..22] 'as u' + keyword_as [18..20] 'as' + any_identifier [21..22] 'u' (@end) + update_set_values [23..48] 'set email = 'my@mail.com'' + keyword_set [23..26] 'set' + assignment [27..48] 'email = 'my@mail.com'' (@end) + column_identifier [27..32] 'email' (@left) + = [33..34] '=' + literal [35..48] ''my@mail.com'' (@end) + where [49..63] 'where u.id = 1' + keyword_where [49..54] 'where' + binary_expression [55..63] 'u.id = 1' (@end) + object_reference [55..59] 'u.id' (@binary_expr_left) + any_identifier [55..56] 'u' (@object_reference_1of2) + . [56..57] '.' + any_identifier [57..59] 'id' (@object_reference_2of2) + = [60..61] '=' (@binary_expr_operator) + literal [62..63] '1' (@end) + ; [63..64] ';' diff --git a/crates/pgls_treesitter_grammar/tests/snapshots/grammar_tests__test_2_sql3.snap b/crates/pgls_treesitter_grammar/tests/snapshots/grammar_tests__test_2_sql3.snap index 51d4ee0ef..4beee2504 100644 --- a/crates/pgls_treesitter_grammar/tests/snapshots/grammar_tests__test_2_sql3.snap +++ b/crates/pgls_treesitter_grammar/tests/snapshots/grammar_tests__test_2_sql3.snap @@ -2,25 +2,31 @@ source: crates/pgls_treesitter_grammar/tests/grammar_tests.rs expression: writer --- -update auth.users set auth.users.email = 'my@mail.com'; +update auth.users set (email, id) = ('my@mail.com', 1); ----------------------- -program [0..55] 'update auth.users set auth.users.email = 'my@mail.com';' - statement [0..54] 'update auth.users set auth.users.email = 'my@mail.com'' - update [0..54] 'update auth.users set auth.users.email = 'my@mail.com'' +program [0..55] 'update auth.users set (email, id) = ('my@mail.com', 1);' + statement [0..54] 'update auth.users set (email, id) = ('my@mail.com', 1)' + update [0..54] 'update auth.users set (email, id) = ('my@mail.com', 1)' 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..54] 'auth.users.email = 'my@mail.com'' - column_reference [22..38] 'auth.users.email' (@left) - schema_identifier [22..26] 'auth' (@column_reference_1of3) - . [26..27] '.' - table_identifier [27..32] 'users' (@column_reference_2of3) - . [32..33] '.' - column_identifier [33..38] 'email' (@column_reference_3of3) - = [39..40] '=' - literal [41..54] ''my@mail.com'' (@right) + update_set_values [18..54] 'set (email, id) = ('my@mail.com', 1)' + keyword_set [18..21] 'set' + assignment [22..54] '(email, id) = ('my@mail.com', 1)' (@end) + lhs_column_list [22..33] '(email, id)' (@left) + ( [22..23] '(' + column_identifier [23..28] 'email' + , [28..29] ',' + column_identifier [30..32] 'id' + ) [32..33] ')' + = [34..35] '=' + rhs_column_list [36..54] '('my@mail.com', 1)' (@end) + ( [36..37] '(' + literal [37..50] ''my@mail.com'' + , [50..51] ',' + literal [52..53] '1' + ) [53..54] ')' (@end) ; [54..55] ';'