From e2e394b17239236a9f2150b94737c5815323bed6 Mon Sep 17 00:00:00 2001
From: Rene Saarsoo <nene@triin.net>
Date: Sun, 14 Jan 2024 00:05:07 +0200
Subject: [PATCH] Support PostgreSQL CREATE TABLE .. OF type

---
 src/cst/CreateTable.ts                   | 30 +++++++++++++++++-------
 src/parser.pegjs                         | 13 +++++++++-
 src/showNode/create_table.ts             |  2 ++
 test/ddl/create_table_postgresql.test.ts |  5 ++++
 4 files changed, 40 insertions(+), 10 deletions(-)

diff --git a/src/cst/CreateTable.ts b/src/cst/CreateTable.ts
index 4ae4c60f..a47eac3e 100644
--- a/src/cst/CreateTable.ts
+++ b/src/cst/CreateTable.ts
@@ -44,7 +44,8 @@ export type AllCreateTableNodes =
   | CreateTableTablespaceClause
   | CreateTableUsingAccessMethodClause
   | CreateTableWithClause
-  | CreateTableWithoutOidsClause;
+  | CreateTableWithoutOidsClause
+  | CreateTableOfTypeClause;
 
 // CREATE TABLE
 export interface CreateTableStmt extends BaseNode {
@@ -63,6 +64,7 @@ export interface CreateTableStmt extends BaseNode {
   ifNotExistsKw?: [Keyword<"IF">, Keyword<"NOT">, Keyword<"EXISTS">];
   name: EntityName;
   partitionOf?: CreateTablePartitionOfClause;
+  ofType?: CreateTableOfTypeClause;
   columns?: ParenExpr<
     ListExpr<ColumnDefinition | TableConstraint | Constraint<TableConstraint>>
   >;
@@ -218,14 +220,6 @@ export interface CreateTablePartitionByClause extends BaseNode {
   columns: ParenExpr<ListExpr<Expr>>;
 }
 
-// PostgreSQL (we're not including this to PostgresqlCreateTableClause because
-// it comes right after table name, unlike other clauses that come after the column definitions)
-export interface CreateTablePartitionOfClause extends BaseNode {
-  type: "create_table_partition_of_clause";
-  partitionOfKw: [Keyword<"PARTITION">, Keyword<"OF">];
-  table: EntityName;
-}
-
 export interface CreateTablePartitionBoundClause extends BaseNode {
   type: "create_table_partition_bound_clause";
   forValuesKw: [Keyword<"FOR">, Keyword<"VALUES">];
@@ -318,3 +312,21 @@ export interface CreateTableWithoutOidsClause extends BaseNode {
   type: "create_table_without_oids_clause";
   withoutOidsKw: [Keyword<"WITHOUT">, Keyword<"OIDS">];
 }
+
+// PostgreSQL
+//
+// We're not including `PARTITION OF` and `OF type` clauses
+// PostgresqlCreateTableClause because it comes right after table name,
+// unlike other clauses that come after the column definitions.
+
+export interface CreateTablePartitionOfClause extends BaseNode {
+  type: "create_table_partition_of_clause";
+  partitionOfKw: [Keyword<"PARTITION">, Keyword<"OF">];
+  table: EntityName;
+}
+
+export interface CreateTableOfTypeClause extends BaseNode {
+  type: "create_table_of_type_clause";
+  ofKw: Keyword<"OF">;
+  typeName: EntityName;
+}
diff --git a/src/parser.pegjs b/src/parser.pegjs
index fff93205..2d01c93d 100644
--- a/src/parser.pegjs
+++ b/src/parser.pegjs
@@ -1933,6 +1933,7 @@ create_table_stmt
     ifKw:(__ if_not_exists)?
     name:(__ entity_name)
     partitionOf:(__ create_table_partition_of_clause)?
+    ofType:(__ create_table_of_type_clause)?
     columns:(__ paren$list$create_definition)?
     options:(__ table_options)?
     clauses:(__ create_table_clause)*
@@ -1950,6 +1951,7 @@ create_table_stmt
         ifNotExistsKw: read(ifKw),
         name: read(name),
         partitionOf: read(partitionOf),
+        ofType: read(ofType),
         columns: read(columns),
         options: read(options),
         clauses: clauses.map(read),
@@ -1993,7 +1995,7 @@ column_definition
     }
 
 create_table_partition_of_clause
-  = kw:(PARTITION __ OF __) table:entity_name {
+  = kw:(PARTITION __ OF __) table:entity_name &postgres {
     return loc({
       type: "create_table_partition_of_clause",
       partitionOfKw: read(kw),
@@ -2001,6 +2003,15 @@ create_table_partition_of_clause
     });
   }
 
+create_table_of_type_clause
+  = kw:(OF __) typeName:entity_name &postgres {
+    return loc({
+      type: "create_table_of_type_clause",
+      ofKw: read(kw),
+      typeName,
+    });
+  }
+
 create_table_clause
   = as_clause$compound_select_stmt
   / (&bigquery / &mysql) x:create_table_like_clause { return x; }
diff --git a/src/showNode/create_table.ts b/src/showNode/create_table.ts
index 9ecac1ff..b38f8af5 100644
--- a/src/showNode/create_table.ts
+++ b/src/showNode/create_table.ts
@@ -16,6 +16,7 @@ export const createTableMap: FullTransformMap<string, AllCreateTableNodes> = {
       node.ifNotExistsKw,
       node.name,
       node.partitionOf,
+      node.ofType,
       node.columns,
       node.options,
       node.clauses,
@@ -62,4 +63,5 @@ export const createTableMap: FullTransformMap<string, AllCreateTableNodes> = {
     show([node.usingKw, node.method]),
   create_table_with_clause: (node) => show([node.withKw, node.options]),
   create_table_without_oids_clause: (node) => show(node.withoutOidsKw),
+  create_table_of_type_clause: (node) => show([node.ofKw, node.typeName]),
 };
diff --git a/test/ddl/create_table_postgresql.test.ts b/test/ddl/create_table_postgresql.test.ts
index 6db9c77a..48c4384c 100644
--- a/test/ddl/create_table_postgresql.test.ts
+++ b/test/ddl/create_table_postgresql.test.ts
@@ -94,6 +94,11 @@ describe("create table (PostgreSQL)", () => {
       });
     });
 
+    it("supports CREATE TABLE .. OF type", () => {
+      testWc("CREATE TABLE client OF client_type");
+      testWc("CREATE TABLE foo OF schm.my_type");
+    });
+
     it("supports ON COMMIT clause", () => {
       testClauseWc(`ON COMMIT PRESERVE ROWS`);
       testClauseWc(`ON COMMIT DELETE ROWS`);