From 7beede0306a5c1b9ff0ab87eba8ba018d9555285 Mon Sep 17 00:00:00 2001
From: Christian De la Hoz <christian.hoz@anixe.pl>
Date: Thu, 2 Sep 2021 20:45:02 +0200
Subject: [PATCH] bug fixes

---
 corpus/create_function/general.txt | 33 ++++++++++++++++++++++++++++++
 corpus/execute_statement.txt       |  2 +-
 corpus/insert_statement/insert.txt | 19 +++++++++++++++++
 grammar.js                         |  9 ++++----
 4 files changed, 58 insertions(+), 5 deletions(-)

diff --git a/corpus/create_function/general.txt b/corpus/create_function/general.txt
index 43f2beb..f96e060 100644
--- a/corpus/create_function/general.txt
+++ b/corpus/create_function/general.txt
@@ -85,12 +85,45 @@ CREATE FUNCTION FOO () RETURNS void AS 'select' LANGUAGE sql VOLATILE;
     (identifier)
     (function_volatility)))
 
+================================================================================
+returns array
+================================================================================
+CREATE FUNCTION FOO () RETURNS text[] AS 'select' LANGUAGE sql;
+--------------------------------------------------------------------------------
+
+(source_file
+  (create_function_statement
+    (function_signature
+      (identifier)
+      (function_parameters))
+    (function_return
+      (identifier))
+    (string)
+    (identifier)))
+
 ================================================================================
 returns setof
 ================================================================================
 CREATE FUNCTION FOO () RETURNS SETOF text AS 'select' LANGUAGE sql;
 --------------------------------------------------------------------------------
 
+(source_file
+  (create_function_statement
+    (function_signature
+      (identifier)
+      (function_parameters))
+    (function_return
+      (return_setof
+        (identifier)))
+    (string)
+    (identifier)))
+
+================================================================================
+returns setof array
+================================================================================
+CREATE FUNCTION FOO () RETURNS SETOF text[] AS 'select' LANGUAGE sql;
+--------------------------------------------------------------------------------
+
 (source_file
   (create_function_statement
     (function_signature
diff --git a/corpus/execute_statement.txt b/corpus/execute_statement.txt
index a4248f9..4b5a149 100644
--- a/corpus/execute_statement.txt
+++ b/corpus/execute_statement.txt
@@ -141,7 +141,7 @@ do $$ begin
 execute format(
     $sql$
         update %1$s
-        set foo = $1
+        set foo = $1, %2$I = $1
     $sql$
     , _tbl_name) using bar;
 end $$;
diff --git a/corpus/insert_statement/insert.txt b/corpus/insert_statement/insert.txt
index 547b90e..45b04ad 100644
--- a/corpus/insert_statement/insert.txt
+++ b/corpus/insert_statement/insert.txt
@@ -78,6 +78,25 @@ insert from select
 insert into foo.my_table select column1 from foo bar;
 --------------------------------------------------------------------------------
 
+(source_file
+  (insert_statement
+    (identifier)
+    (insert_items
+      (select_statement
+        (select_item
+          (identifier))
+        (select_from
+          (from_item
+            (from_table
+              (identifier)
+              (identifier))))))))
+
+================================================================================
+insert from select(1)
+================================================================================
+insert into foo.my_table (select column1 from foo bar);
+--------------------------------------------------------------------------------
+
 (source_file
   (insert_statement
     (identifier)
diff --git a/grammar.js b/grammar.js
index bb1640a..2ea28d9 100644
--- a/grammar.js
+++ b/grammar.js
@@ -114,7 +114,8 @@ module.exports = grammar({
     insert_items: $ => choice(
       seq(kw("default"), kw("values")),
       seq(kw("values"), "(", commaSep($.insert_item), ")"),
-      $.select_statement
+      $.select_statement,
+      seq("(", $.select_statement, ")"),
     ),
     insert_item: $ => choice(
       kw("default"),
@@ -605,12 +606,12 @@ module.exports = grammar({
     function_return: $ => seq(
       kw("returns"),
       choice(
-        $.identifier,
+        $._type,
         $.return_setof,
         $.return_table,
       ),
     ),
-    return_setof: $ => seq(kw("setof"), $.identifier),
+    return_setof: $ => seq(kw("setof"), $._type),
     return_table: $ => seq(kw("table"), "(", commaSep1($.var_declaration), ")"),
 
     function_volatility: $ => choice(
@@ -755,7 +756,7 @@ module.exports = grammar({
     // TODO(chrde): it does not handle nested dollar quotes... perhaps move to an external scanner?
     dollar_quote_string: $ => seq(
       "$", $._identifier, "$",
-        /(([^$]+)|(%\d+\$s)|(\$\d+))+/,
+        /(([^$]+)|(%\d+\$[sI])|(\$\d+))+/,
         //                     ^
         //                     |- matches $1 (execute ... using placeholders)
         //           ^