Support UPDATE,INSERT,DELETE inside WITH clause (PostgreSQL)

Author: nene

src/cst/Select.ts

@@ -14,7 +14,13 @@ import { FrameClause } from "./WindowFrame";
 import { Literal, StringLiteral } from "./Literal";
 import { MysqlModifier } from "./dialects/Mysql";
 import { ColumnDefinition } from "./CreateTable";
-import { PostgresqlOperator, PostgresqlOperatorExpr } from "./Node";
+import {
+  DeleteStmt,
+  InsertStmt,
+  PostgresqlOperator,
+  PostgresqlOperatorExpr,
+  UpdateStmt,
+} from "./Node";
 
 export type AllSelectNodes =
   | CompoundSelectStmt
@@ -149,7 +155,8 @@ export interface CommonTableExpr extends BaseNode {
   materializedKw?:
     | Keyword<"MATERIALIZED">
     | [Keyword<"NOT">, Keyword<"MATERIALIZED">];
-  expr: ParenExpr<SubSelect>;
+  // PostgreSQL supports UPDATE, DELETE, and INSERT in WITH clause
+  expr: ParenExpr<SubSelect | DeleteStmt | InsertStmt | UpdateStmt>;
   search?: CteSearchClause;
   cycle?: CteCycleClause;
 }

src/parser.pegjs

@@ -191,7 +191,7 @@ common_table_expr
     columns:(__ paren$list$column)?
     asKw:(__ AS)
     materialized:(__ cte_materialized)?
-    select:(__ paren$compound_select_stmt)
+    expr:(__ paren$cte_expr)
     search:(__ cte_search_clause)?
     cycle:(__ cte_cycle_clause)? {
       return loc({
@@ -200,7 +200,7 @@ common_table_expr
         columns: read(columns),
         asKw: read(asKw),
         materializedKw: read(materialized),
-        expr: read(select),
+        expr: read(expr),
         search: read(search),
         cycle: read(cycle),
       });
@@ -209,6 +209,10 @@ common_table_expr
 cte_materialized
   = kws:(NOT __ MATERIALIZED / MATERIALIZED) { return read(kws); }
 
+cte_expr
+  = compound_select_stmt
+  / stmt:(update_stmt / insert_stmt / delete_stmt) &postgres { return stmt; }
+
 cte_search_clause
   = kw:(SEARCH __ (BREADTH / DEPTH) __ FIRST __ BY __) columns:list$ident
     setKw:(__ SET __) result:ident &postgres {
@@ -4944,6 +4948,7 @@ paren$__template__
 
 paren$cast_arg = .
 paren$compound_select_stmt = .
+paren$cte_expr = .
 paren$empty_list = .
 paren$entity_name = .
 paren$expr = .

test/select/with.test.ts

@@ -60,5 +60,17 @@ describe("select WITH", () => {
         SELECT * FROM search_graph
       `);
     });
+
+    it("supports INSERT inside WITH", () => {
+      testWc(`WITH t1 AS (INSERT INTO foo VALUES (1) RETURNING *) SELECT * FROM t1`);
+    });
+
+    it("supports UPDATE inside WITH", () => {
+      testWc(`WITH t1 AS (UPDATE foo SET bar = 1 RETURNING *) SELECT * FROM t1`);
+    });
+
+    it("supports DELETE inside WITH", () => {
+      testWc(`WITH t1 AS (DELETE FROM foo RETURNING *) SELECT * FROM t1`);
+    });
   });
 });