Support PREPARE .. AS in PostgreSQL

Author: nene

src/cst/PreparedStatements.ts

@@ -1,6 +1,9 @@
 import { Alias } from "./Alias";
 import { BaseNode, Keyword } from "./Base";
+import { DataType } from "./DataType";
 import { EntityName, Expr, Identifier, ListExpr, ParenExpr } from "./Expr";
+import { AsClause } from "./ProcClause";
+import { Statement } from "./Statement";
 
 export type AllPreparedStatementNodes =
   | AllPreparedStatements
@@ -18,7 +21,8 @@ export interface PrepareStmt extends BaseNode {
   type: "prepare_stmt";
   prepareKw: Keyword<"PREPARE">;
   name: EntityName;
-  source: PrepareFromClause;
+  params?: ParenExpr<ListExpr<DataType>>;
+  source: PrepareFromClause | AsClause<Statement>;
 }
 
 export interface PrepareFromClause extends BaseNode {

src/parser.pegjs

@@ -100,6 +100,7 @@ statement_bigquery
 statement_postgres
   = analyze_stmt
   / explain_stmt
+  / prepare_stmt
   / execute_stmt
   / call_stmt
   / do_stmt
@@ -5526,11 +5527,13 @@ raise_message
  */
 prepare_stmt
   = kw:(PREPARE __) name:(entity_name __)
-    source:prepare_from_clause {
+    params:(paren$list$data_type __)?
+    source:(prepare_from_clause / as_clause$prepareable_statement) {
       return loc({
         type: "prepare_stmt",
         prepareKw: read(kw),
         name: read(name),
+        params: read(params),
         source: source,
       });
     }
@@ -5544,6 +5547,13 @@ prepare_from_clause
     });
   }
 
+prepareable_statement
+  = compound_select_stmt
+  / insert_stmt
+  / update_stmt
+  / delete_stmt
+  / merge_stmt
+
 execute_stmt
   = kw:(EXECUTE __) name:entity_name
     args:(__ (execute_using_clause / paren$list$expr))? {
@@ -7737,6 +7747,7 @@ paren$list$alias$paren$list$column = .
 paren$list$column = .
 paren$list$column_definition = .
 paren$list$create_definition = .
+paren$list$data_type = .
 paren$list$entity_name = .
 paren$list$equals_expr = .
 paren$list$exclusion_param = .
@@ -7804,6 +7815,7 @@ list$column_definition = .
 list$common_table_expr = .
 list$config_parameter_value = .
 list$create_definition = .
+list$data_type = .
 list$delete_clause_table = .
 list$entity_name = .
 list$equals_expr = .
@@ -7902,6 +7914,7 @@ as_clause$__template__
 as_clause$compound_select_stmt = .
 as_clause$func_as_expr_bigquery = .
 as_clause$func_as_expr_postgresql = .
+as_clause$prepareable_statement = .
 as_clause$string_literal = .
 as_clause$expr = .
 /*! as_clause:end */

src/showNode/prepared_statements.ts

@@ -6,7 +6,8 @@ export const preparedStatementsMap: FullTransformMap<
   string,
   AllPreparedStatementNodes
 > = {
-  prepare_stmt: (node) => show([node.prepareKw, node.name, node.source]),
+  prepare_stmt: (node) =>
+    show([node.prepareKw, node.name, node.params, node.source]),
   prepare_from_clause: (node) => show([node.fromKw, node.expr]),
   execute_stmt: (node) => show([node.executeKw, node.name, node.args]),
   execute_immediate_stmt: (node) =>

test/prepared_statements.test.ts

@@ -1,4 +1,4 @@
-import { dialect, parse, testWc, test } from "./test_utils";
+import { dialect, parse, testWc, test, includeAll } from "./test_utils";
 
 describe("prepared statements", () => {
   dialect(["mysql", "mariadb"], () => {
@@ -8,6 +8,33 @@ describe("prepared statements", () => {
     });
   });
 
+  dialect(["postgresql"], () => {
+    it("supports PREPARE .. AS statement", () => {
+      testWc(`PREPARE my_stmt AS SELECT * FROM my_table WHERE id = $1`, {
+        paramTypes: ["$nr"],
+        ...includeAll,
+      });
+      testWc(`PREPARE my_stmt AS UPDATE foo SET bar = $1`, { paramTypes: ["$nr"], ...includeAll });
+      testWc(`PREPARE my_stmt AS DELETE FROM foo WHERE id = $1`, {
+        paramTypes: ["$nr"],
+        ...includeAll,
+      });
+      testWc(`PREPARE my_stmt AS INSERT INTO foo (bar) VALUES ($1)`, {
+        paramTypes: ["$nr"],
+        ...includeAll,
+      });
+      testWc(`PREPARE my_stmt AS MERGE INTO foo USING bar ON x = $1 WHEN MATCHED THEN DELETE`, {
+        paramTypes: ["$nr"],
+        ...includeAll,
+      });
+      testWc(`PREPARE my_stmt AS VALUES ($1), ($2), ($3)`, { paramTypes: ["$nr"], ...includeAll });
+    });
+
+    it("supports PREPARE .. AS with parameters", () => {
+      testWc(`PREPARE my_stmt(INT, TEXT) AS SELECT $1, $2`, { paramTypes: ["$nr"], ...includeAll });
+    });
+  });
+
   dialect(["mysql", "mariadb", "postgresql"], () => {
     it("supports EXECUTE", () => {
       testWc(`EXECUTE my_stmt`);