create type, for, raise and few fixes

corpus/create_type_statement.txt

@@ -0,0 +1,37 @@
+================================================================================
+empty
+================================================================================
+create type foo;
+--------------------------------------------------------------------------------
+
+(source_file
+  (create_type_statement
+    (identifier)))
+
+================================================================================
+with fields
+================================================================================
+create type foo as ( one text, two bigint[] );
+--------------------------------------------------------------------------------
+
+(source_file
+  (create_type_statement
+    (identifier)
+    (var_declaration
+      (identifier)
+      (identifier))
+    (var_declaration
+      (identifier)
+      (identifier))))
+
+================================================================================
+enum
+================================================================================
+create type foo.bar as enum ('one', 'two');
+--------------------------------------------------------------------------------
+
+(source_file
+  (create_type_statement
+    (identifier)
+    (string)
+    (string)))

corpus/execute_statement.txt

@@ -60,10 +60,11 @@ do $$ begin execute 'command' using 1, foo(bar); end $$;
       (body
         (execute_statement
           (string)
-          (number)
+          (execute_using
-          (function_call
+            (number)
-            (identifier)
+            (function_call
-            (identifier))))
+              (identifier)
+              (identifier)))))
       (dollar_quote))))
 
 ================================================================================
@@ -82,7 +83,8 @@ do $$ begin execute foo() into strict var using 1; end $$;
             (identifier))
           (into
             (identifier))
-          (number)))
+          (execute_using
+            (number))))
       (dollar_quote))))
 
 ================================================================================
@@ -104,3 +106,109 @@ do $$ begin execute 'foo' || 'bar' into _date; end $$;
           (into
             (identifier))))
       (dollar_quote))))
+
+================================================================================
+simple nested dollar quote
+================================================================================
+do $$ begin
+execute format(
+    $sql$
+        select %1$s,
+        select 2,
+    $sql$
+    , _tbl_name) using bar;
+end $$;
+--------------------------------------------------------------------------------
+
+(source_file
+  (do_block
+    (block
+      (dollar_quote)
+      (body
+        (execute_statement
+          (function_call
+            (identifier)
+            (dollar_quote_string)
+            (identifier))
+          (execute_using
+            (identifier))))
+      (dollar_quote))))
+
+================================================================================
+execute format dollar quote
+================================================================================
+do $$ begin
+execute format(
+    $sql$
+        update %1$s
+        set foo = $1
+    $sql$
+    , _tbl_name) using bar;
+end $$;
+--------------------------------------------------------------------------------
+
+(source_file
+  (do_block
+    (block
+      (dollar_quote)
+      (body
+        (execute_statement
+          (function_call
+            (identifier)
+            (dollar_quote_string)
+            (identifier))
+          (execute_using
+            (identifier))))
+      (dollar_quote))))
+
+================================================================================
+complex
+================================================================================
+do $$ begin
+EXECUTE FORMAT(
+    $sql$
+        UPDATE %2$s SET new_val = $1
+        FROM ( SELECT serial AS _serial, created FROM %1$s ) t1
+    $sql$
+    , _tbl_name) USING baz, bar;
+end $$;
+--------------------------------------------------------------------------------
+
+(source_file
+  (do_block
+    (block
+      (dollar_quote)
+      (body
+        (execute_statement
+          (function_call
+            (identifier)
+            (dollar_quote_string)
+            (identifier))
+          (execute_using
+            (identifier)
+            (identifier))))
+      (dollar_quote))))
+
+================================================================================
+complex(1)
+================================================================================
+do $$ begin
+EXECUTE FORMAT($sql$
+foo
+$sql$ , _tbl_name) USING bar.baz;
+end $$;
+--------------------------------------------------------------------------------
+
+(source_file
+  (do_block
+    (block
+      (dollar_quote)
+      (body
+        (execute_statement
+          (function_call
+            (identifier)
+            (dollar_quote_string)
+            (identifier))
+          (execute_using
+            (identifier))))
+      (dollar_quote))))

corpus/plpgsql/block.txt

@@ -31,8 +31,8 @@ many declare(s)
 DO $$
 DECLARE one text;
 DECLARE
-    name text;
+    name foo%TYPE;
-    age bigint;
+    age bar%ROWTYPE;
 BEGIN END
 $$;
 --------------------------------------------------------------------------------

corpus/plpgsql/block.txtor_statement.txt

@@ -0,0 +1,123 @@
+================================================================================
+integer
+================================================================================
+do $$ begin
+for foo in 1..10 loop select 1; end loop;
+end $$;
+--------------------------------------------------------------------------------
+
+(source_file
+  (do_block
+    (block
+      (dollar_quote)
+      (body
+        (for_statement
+          (identifier)
+          (number)
+          (number)
+          (select_statement
+            (select_item
+              (number)))))
+      (dollar_quote))))
+
+================================================================================
+integer by step
+================================================================================
+do $$ begin
+for foo in 1..do_something() by 5 loop select 1; end loop;
+end $$;
+--------------------------------------------------------------------------------
+
+(source_file
+  (do_block
+    (block
+      (dollar_quote)
+      (body
+        (for_statement
+          (identifier)
+          (number)
+          (function_call
+            (identifier))
+          (number)
+          (select_statement
+            (select_item
+              (number)))))
+      (dollar_quote))))
+
+================================================================================
+integer reverse
+================================================================================
+do $$ begin
+for foo in reverse 10..1 by 5 loop select 1; end loop;
+end $$;
+--------------------------------------------------------------------------------
+
+(source_file
+  (do_block
+    (block
+      (dollar_quote)
+      (body
+        (for_statement
+          (identifier)
+          (number)
+          (number)
+          (number)
+          (select_statement
+            (select_item
+              (number)))))
+      (dollar_quote))))
+
+================================================================================
+for over query
+================================================================================
+do $$ begin
+for foo, var in select generate_series(1,10) loop select 1; end loop;
+end $$;
+--------------------------------------------------------------------------------
+
+(source_file
+  (do_block
+    (block
+      (dollar_quote)
+      (body
+        (for_statement
+          (identifier)
+          (identifier)
+          (select_statement
+            (select_item
+              (function_call
+                (identifier)
+                (number)
+                (number))))
+          (select_statement
+            (select_item
+              (number)))))
+      (dollar_quote))))
+
+================================================================================
+for over execute
+================================================================================
+do $$ begin
+for foo, var in execute format($sql$ select $1 from %1$s; $sql$, _foo) using bar loop select 1; end loop;
+end $$;
+--------------------------------------------------------------------------------
+
+(source_file
+  (do_block
+    (block
+      (dollar_quote)
+      (body
+        (for_statement
+          (identifier)
+          (identifier)
+          (execute_statement
+            (function_call
+              (identifier)
+              (dollar_quote_string)
+              (identifier))
+            (execute_using
+              (identifier)))
+          (select_statement
+            (select_item
+              (number)))))
+      (dollar_quote))))

corpus/plpgsql/raise_statement.txt

@@ -0,0 +1,70 @@
+================================================================================
+dummy
+================================================================================
+do $$ begin
+raise;
+end $$;
+--------------------------------------------------------------------------------
+
+(source_file
+  (do_block
+    (block
+      (dollar_quote)
+      (body
+        (ERROR))
+      (dollar_quote))))
+
+================================================================================
+basic
+================================================================================
+do $$ begin
+raise 'alarms';
+end $$;
+--------------------------------------------------------------------------------
+
+(source_file
+  (do_block
+    (block
+      (dollar_quote)
+      (body
+        (raise_statement
+          (string)))
+      (dollar_quote))))
+
+================================================================================
+with level
+================================================================================
+do $$ begin
+raise notice 'alarms';
+end $$;
+--------------------------------------------------------------------------------
+
+(source_file
+  (do_block
+    (block
+      (dollar_quote)
+      (body
+        (raise_statement
+          (identifier)
+          (string)))
+      (dollar_quote))))
+
+================================================================================
+with level and args
+================================================================================
+do $$ begin
+raise notice 'alarms %d', _foo, bar;
+end $$;
+--------------------------------------------------------------------------------
+
+(source_file
+  (do_block
+    (block
+      (dollar_quote)
+      (body
+        (raise_statement
+          (identifier)
+          (string)
+          (identifier)
+          (identifier)))
+      (dollar_quote))))

corpus/plpgsql/return_statement.txt

@@ -37,3 +37,28 @@ $$;
             (select_item
               (number)))))
       (dollar_quote))))
+
+================================================================================
+return query execute
+================================================================================
+DO $$
+BEGIN
+    RETURN QUERY execute format($sql$ select %1$s, $1;$sql$, _foo) using bar;
+END
+$$;
+--------------------------------------------------------------------------------
+
+(source_file
+  (do_block
+    (block
+      (dollar_quote)
+      (body
+        (return_statement
+          (execute_statement
+            (function_call
+              (identifier)
+              (dollar_quote_string)
+              (identifier))
+            (execute_using
+              (identifier)))))
+      (dollar_quote))))

corpus/select_statement/join.txt

@@ -13,8 +13,9 @@ SELECT name FROM products cross join items;
         (from_table
           (identifier))
         (join_item
-          (from_table
+          (from_item
-            (identifier)))))))
+            (from_table
+              (identifier))))))))
 
 ================================================================================
 join on
@@ -32,8 +33,9 @@ SELECT name FROM products join items on products.name = items.name;
           (identifier))
         (join_item
           (join_type)
-          (from_table
+          (from_item
-            (identifier))
+            (from_table
+              (identifier)))
           (join_condition
             (op_expression
               (identifier)
@@ -56,8 +58,9 @@ select name from products join items using(foo);
           (identifier))
         (join_item
           (join_type)
-          (from_table
+          (from_item
-            (identifier))
+            (from_table
+              (identifier)))
           (join_condition
             (identifier)))))))
 
@@ -77,8 +80,9 @@ select name from products natural join items;
           (identifier))
         (join_item
           (join_type)
-          (from_table
+          (from_item
-            (identifier)))))))
+            (from_table
+              (identifier))))))))
 
 ================================================================================
 inner join
@@ -96,8 +100,9 @@ select name from products inner join items on true;
           (identifier))
         (join_item
           (join_type)
-          (from_table
+          (from_item
-            (identifier))
+            (from_table
+              (identifier)))
           (join_condition
             (true)))))))
 
@@ -117,8 +122,9 @@ select name from products left join items on true;
           (identifier))
         (join_item
           (join_type)
-          (from_table
+          (from_item
-            (identifier))
+            (from_table
+              (identifier)))
           (join_condition
             (true)))))))
 
@@ -138,8 +144,9 @@ select name from products right join items on true;
           (identifier))
         (join_item
           (join_type)
-          (from_table
+          (from_item
-            (identifier))
+            (from_table
+              (identifier)))
           (join_condition
             (true)))))))
 
@@ -159,8 +166,9 @@ select name from products full join items on true;
           (identifier))
         (join_item
           (join_type)
-          (from_table
+          (from_item
-            (identifier))
+            (from_table
+              (identifier)))
           (join_condition
             (true)))))))
 
@@ -180,8 +188,9 @@ select name from products left outer join items on true;
           (identifier))
         (join_item
           (join_type)
-          (from_table
+          (from_item
-            (identifier))
+            (from_table
+              (identifier)))
           (join_condition
             (true)))))))
 
@@ -201,8 +210,9 @@ select name from products right outer join items on true;
           (identifier))
         (join_item
           (join_type)
-          (from_table
+          (from_item
-            (identifier))
+            (from_table
+              (identifier)))
           (join_condition
             (true)))))))
 
@@ -222,8 +232,9 @@ select name from products full outer join items on true;
           (identifier))
         (join_item
           (join_type)
-          (from_table
+          (from_item
-            (identifier))
+            (from_table
+              (identifier)))
           (join_condition
             (true)))))))
 
@@ -253,15 +264,18 @@ select name from products
           (identifier))
         (join_item
           (join_type)
-          (from_table
+          (from_item
-            (identifier)))
+            (from_table
+              (identifier))))
         (join_item
-          (from_table
+          (from_item
-            (identifier)))
+            (from_table
+              (identifier))))
         (join_item
           (join_type)
-          (from_table
+          (from_item
-            (identifier))
+            (from_table
+              (identifier)))
           (join_condition
             (op_expression
               (identifier)
@@ -269,43 +283,50 @@ select name from products
               (identifier))))
         (join_item
           (join_type)
-          (from_table
+          (from_item
-            (identifier))
+            (from_table
+              (identifier)))
           (join_condition
             (identifier)))
         (join_item
           (join_type)
-          (from_table
+          (from_item
-            (identifier))
+            (from_table
+              (identifier)))
           (join_condition
             (identifier)))
         (join_item
           (join_type)
-          (from_table
+          (from_item
-            (identifier))
+            (from_table
+              (identifier)))
           (join_condition
             (identifier)))
         (join_item
           (join_type)
-          (from_table
+          (from_item
-            (identifier))
+            (from_table
+              (identifier)))
           (join_condition
             (identifier)))
         (join_item
           (join_type)
-          (from_table
+          (from_item
-            (identifier))
+            (from_table
+              (identifier)))
           (join_condition
             (identifier)))
         (join_item
           (join_type)
-          (from_table
+          (from_item
-            (identifier))
+            (from_table
+              (identifier)))
           (join_condition
             (identifier)))
         (join_item
           (join_type)
-          (from_table
+          (from_item
-            (identifier))
+            (from_table
+              (identifier)))
           (join_condition
             (identifier)))))))

corpus/select_statement/select.txt

@@ -36,6 +36,26 @@ select name into bar from items where true;
       (where_filter
         (true)))))
 
+================================================================================
+into at the end
+================================================================================
+select * from items where true into _bar;
+--------------------------------------------------------------------------------
+
+(source_file
+  (select_statement
+    (select_item
+      (star))
+    (select_from
+      (from_item
+        (from_table
+          (identifier))))
+    (select_where
+      (where_filter
+        (true)))
+    (into
+      (identifier))))
+
 ================================================================================
 into strict
 ================================================================================
@@ -158,9 +178,9 @@ select name limit 1;
 (source_file
   (select_statement
     (select_item
-      (identifier)))
+      (identifier))
-  (ERROR
+    (select_limit
-    (number)))
+      (number))))
 
 ================================================================================
 limit offet
@@ -173,7 +193,8 @@ select name limit 1 offset 5;
     (select_item
       (identifier))
     (select_limit
-      (number)
+      (number))
+    (select_offset
       (number))))
 
 ================================================================================
@@ -186,8 +207,9 @@ select name offset 5 limit 1;
   (select_statement
     (select_item
       (identifier))
+    (select_offset
+      (number))
     (select_limit
-      (number)
       (number))))
 
 ================================================================================
@@ -200,7 +222,8 @@ select name limit all offset 5;
   (select_statement
     (select_item
       (identifier))
-    (select_limit
+    (select_limit)
+    (select_offset
       (number))))
 
 ================================================================================
@@ -213,8 +236,9 @@ select name offset 5 limit all;
   (select_statement
     (select_item
       (identifier))
-    (select_limit
+    (select_offset
-      (number))))
+      (number))
+    (select_limit)))
 
 ================================================================================
 with cte
@@ -241,3 +265,19 @@ select * from foo;
       (from_item
         (from_table
           (identifier))))))
+
+================================================================================
+select alias star
+================================================================================
+SELECT foo.* FROM foo;
+--------------------------------------------------------------------------------
+
+(source_file
+  (select_statement
+    (select_item
+      (identifier)
+      (star))
+    (select_from
+      (from_item
+        (from_table
+          (identifier))))))

grammar.js

@@ -47,6 +47,7 @@ module.exports = grammar({
       $.create_function_statement,
       $.create_table_statement,
       $.create_schema_statement,
+      $.create_type_statement,
       $.select_statement,
       $.insert_statement,
       $.delete_statement,
@@ -58,6 +59,13 @@ module.exports = grammar({
       $.do_block,
     ),
 
+    create_type_statement: $ => seq(
+      kw("create"), kw("type"), $.identifier, optional(choice(
+        seq(kw("as"), kw("enum"), "(", commaSep1($.string), ")"),
+        seq(kw("as"), "(", commaSep1($.var_declaration), ")"),
+      ))
+    ),
+
     // TODO(chrde): update, values
     _with_query_statement: $ => choice(
       $.select_statement,
@@ -384,13 +392,37 @@ module.exports = grammar({
         $._statement,
         $.assign_statement,
         $.return_statement,
+        $.raise_statement,
         $.if_statement,
+        $.for_statement,
         $.execute_statement,
         $.perform_statement,
       ),
       ";",
     ),
 
+    for_statement: $ => seq(
+      kw("for"), commaSep1($.identifier), kw("in"), choice(
+        seq(
+          optional(kw("reverse")),
+          $._value_expression, "..", $._value_expression,
+          optional(seq(kw("by"), $._value_expression))
+        ),
+        $.select_statement,
+        $.execute_statement,
+      ),
+      kw("loop"),
+      repeat($._plpgsql_statement),
+      kw("end"), kw("loop")
+    ),
+
+    // TODO(chrde): https://www.postgresql.org/docs/13/plpgsql-errors-and-messages.html
+    raise_statement: $ => seq(
+      kw("raise"), optional($.identifier),
+      $.string, 
+      optional(seq(",", commaSep($._value_expression)))
+    ),
+
     if_statement: $ => seq(
       kw("if"), $._value_expression, kw("then"), repeat1($._plpgsql_statement),
       repeat(seq(
@@ -405,11 +437,13 @@ module.exports = grammar({
       kw("execute"),
       $._value_expression,
       optional($.into),
-      optional(seq(kw("using"), commaSep1($._value_expression))),
+      optional($.execute_using),
     ),
+    execute_using: $ => seq(kw("using"), commaSep1($._value_expression)),
     assign_statement: $ => seq($.identifier, "=", $._value_expression),
     return_statement: $ => seq(kw("return"), choice(
       seq(kw("query"), $.select_statement),
+      seq(kw("query"), $.execute_statement),
       $._value_expression),
     ),
     perform_statement: $ => seq(kw("perform"), commaSep($.select_item)),
@@ -429,7 +463,8 @@ module.exports = grammar({
       optional($.select_group_by),
       optional($.select_having),
       optional($.select_order_by),
-      optional($.select_limit),
+      optional($._select_limit_offset),
+      optional($.into),
     )),
 
     with_query: $ => seq(kw("with"), commaSep1($.with_query_item)),
@@ -440,14 +475,22 @@ module.exports = grammar({
       optional(choice(kw("materialized"), seq(kw("not"), kw("materialized")))),
       "(", $._with_query_statement, ")"
     ),
-
     into: $ => seq(kw("into"), optional(kw("strict")), commaSep1($.identifier)),
     select_having: $ => seq(kw("having"), $._value_expression),
-    select_limit: $ => choice(
+    _select_limit_offset: $ => choice(
-      seq(kw("limit"), $._value_expression, kw("offset"), $._value_expression),
+      seq($.select_limit, optional($.select_offset)),
-      seq(kw("limit"), kw("all"), kw("offset"), $._value_expression),
+      seq($.select_offset, optional($.select_limit)),
-      seq(kw("offset"), $._value_expression, kw("limit"), kw("all")),
+    ),
-      seq(kw("offset"), $._value_expression, kw("limit"), $._value_expression),
+    select_limit: $ => seq(
+      kw("limit"),
+      choice(
+        kw("all"),
+        $._value_expression
+      )
+    ),
+    select_offset: $ => seq(
+      kw("offset"), $._value_expression,
+      optional(choice(kw("row"), kw("rows")))
     ),
     select_group_by: $ => seq(kw("group"), kw("by"), commaSep1($._value_expression)),
     select_order_by: $ => seq(kw("order"), kw("by"), commaSep1($.order_by_item)),
@@ -460,7 +503,7 @@ module.exports = grammar({
       optional($.identifier)
     ),
     select_from: $ => seq(kw("from"), commaSep1($.from_item)),
-    from_item: $ => seq(
+    from_item: $ => prec.left(seq(
       // TODO(chrde): https://www.postgresql.org/docs/current/sql-select.html
       choice(
         $.from_select,
@@ -468,7 +511,7 @@ module.exports = grammar({
         $.from_function,
       ),
       repeat($.join_item),
-    ),
+    )),
     from_select: $ => seq("(", $.select_statement, ")", optional(kw("as")), $.identifier),
     from_table: $ => seq($.identifier, optional(kw("as")), optional($.identifier)),
     from_function: $ => seq(
@@ -480,11 +523,11 @@ module.exports = grammar({
       )),
     ),
 
-    join_item: $ => choice(
+    join_item: $ => prec.left(choice(
-      seq(kw("natural"), $.join_type, $.from_table),
+      seq(kw("natural"), $.join_type, $.from_item),
-      seq($.join_type, $.from_table, $.join_condition),
+      seq($.join_type, $.from_item, $.join_condition),
-      seq(kw("cross"), kw("join"), $.from_table)
+      seq(kw("cross"), kw("join"), $.from_item)
-    ),
+    )),
     join_condition: $ => choice(
       seq(kw("on"), $._value_expression),
       seq(kw("using"), $._list_of_identifiers)
@@ -579,10 +622,12 @@ module.exports = grammar({
     if_exists: $ => seq(kw("if"), kw("exists")),
     as: $ => seq(kw("as"), $.identifier),
 
-    _type: $ => choice(
+    _type: $ => seq(
-      $.predefined_types,
+      choice($.predefined_types, $.identifier),
-      $.identifier,
+      optional(choice(
-      seq($._type, "[", "]"),
+        repeat1(seq("[", "]")),
+        kw("%rowtype"),
+        kw("%type")))
     ),
 
 // predefined_type:
@@ -639,8 +684,9 @@ module.exports = grammar({
     )),
 
     _value_expression: $ => choice(
-      $.string,
       $.number,
+      $.dollar_quote_string,
+      $.string,
       $.true,
       $.false,
       $.null,
@@ -650,9 +696,24 @@ module.exports = grammar({
       $.function_call,
       $.op_expression,
       $.time_expression,
+      // TODO(chrde): this one feels a bit hacky? perhaps move to identifier regexp
+      seq($.identifier, ".", $.star),
       $.identifier,
     ),
 
+    // TODO(chrde): it does not handle nested dollar quotes... perhaps move to an external scanner?
+    dollar_quote_string: $ => seq(
+      "$", $._identifier, "$",
+        /(([^$]+)|(%\d+\$s)|(\$\d+))+/,
+        //                     ^
+        //                     |- matches $1 (execute ... using placeholders)
+        //           ^
+        //           |- matches %d1s (format placeholders)
+        //  ^
+        //  |- matches anything other than $
+      "$", $._identifier, "$",
+    ),
+
     time_expression: $ => choice(
       seq($.identifier, kw("at"), kw("time"), kw("zone"), $._value_expression),
     ),
@@ -700,6 +761,9 @@ module.exports = grammar({
     star: $ => "*",
     any: $ => /.*/,
     number: $ => /\d+/,
-    identifier: $ => /[a-zA-Z0-9_]+[.a-zA-Z0-9_]*/,
+    identifier: $ => $._identifier,
+    _identifier: $ => /[a-zA-Z0-9_]+(\.?[a-zA-Z0-9_]+)*/,
+    //                                ^
+    //                                |- we dont want to match consecutive dots, e.g: 1..2 consists of 3 tokens
   }
 });