src: support enums types in postgres (#10)

* src: support enums in postgres * testdata/ts: sepetate sqlite tests from the rest * npm run build * update js
ariga · Dec 10, 2023 · 35c611b · 35c611b
1 parent bc9df98
commit 35c611b
Show file tree

Hide file tree

Showing 11 changed files with 132 additions and 10 deletions.
diff --git a/build/load.js b/build/load.js
@@ -46,7 +46,8 @@ function loadEntities(dialect,
 // eslint-disable-next-line @typescript-eslint/ban-types
 entities) {
     return __awaiter(this, void 0, void 0, function () {
-        var mockDB, driver, metadataBuilder, entityMetadataValidator, queryRunner, entityMetadatas, _i, entityMetadatas_1, metadata, table, _a, entityMetadatas_2, metadata, table, foreignKeys, _b, entityMetadatas_3, metadata, view, memorySql, queries;
+        var mockDB, driver, metadataBuilder, entityMetadataValidator, queryRunner, entityMetadatas, mockPostgresQueryRunner, _i, entityMetadatas_1, metadata, table, _a, entityMetadatas_2, metadata, table, foreignKeys, _b, entityMetadatas_3, metadata, view, memorySql, queries;
+        var _this = this;
         return __generator(this, function (_c) {
             switch (_c.label) {
                 case 0:
@@ -64,6 +65,27 @@ entities) {
                     entityMetadatas = _c.sent();
                     // validate all created entity metadatas to make sure user created entities are valid and correct
                     entityMetadataValidator.validateMany(entityMetadatas.filter(function (metadata) { return metadata.tableType !== "view"; }), driver);
+                    mockPostgresQueryRunner = function (schema) {
+                        var postgresQueryRunner = queryRunner;
+                        var enumQueries = new Set();
+                        postgresQueryRunner.query = function (query) { return __awaiter(_this, void 0, void 0, function () {
+                            return __generator(this, function (_a) {
+                                if (query.includes("SELECT * FROM current_schema()")) {
+                                    return [2 /*return*/, [{ current_schema: schema }]];
+                                }
+                                if (query.includes('SELECT "n"."nspname", "t"."typname" FROM "pg_type" "t"')) {
+                                    if (enumQueries.has(query)) {
+                                        // mock result for enum that already exists
+                                        return [2 /*return*/, [{}]];
+                                    }
+                                    enumQueries.add(query);
+                                    // mock result for enum that does not exist
+                                    return [2 /*return*/, []];
+                                }
+                                return [2 /*return*/];
+                            });
+                        }); };
+                    };
                     _i = 0, entityMetadatas_1 = entityMetadatas;
                     _c.label = 2;
                 case 2:
@@ -72,6 +94,9 @@ entities) {
                     if (metadata.tableType === "view") {
                         return [3 /*break*/, 4];
                     }
+                    if (dialect === "postgres") {
+                        mockPostgresQueryRunner(metadata.schema);
+                    }
                     table = typeorm_1.Table.create(metadata, driver);
                     return [4 /*yield*/, queryRunner.createTable(table)];
                 case 3:

diff --git a/src/load.ts b/src/load.ts
@@ -2,6 +2,7 @@ import { DataSource, EntitySchema, Table, TableForeignKey } from "typeorm";
 import { ConnectionMetadataBuilder } from "typeorm/connection/ConnectionMetadataBuilder";
 import { EntityMetadataValidator } from "typeorm/metadata-builder/EntityMetadataValidator";
 import { View } from "typeorm/schema-builder/view/View";
+import { PostgresQueryRunner } from "typeorm/driver/postgres/PostgresQueryRunner";
 
 export type Dialect = "mysql" | "postgres" | "mariadb" | "sqlite" | "mssql";
 
@@ -28,11 +29,36 @@ export async function loadEntities(
     driver,
   );
 
+  // mockPostgresQueryRunner used to mock result of queries instead of executing them against the database.
+  const mockPostgresQueryRunner = (schema?: string) => {
+    const postgresQueryRunner = queryRunner as PostgresQueryRunner;
+    const enumQueries = new Set<string>();
+    postgresQueryRunner.query = async (query: string) => {
+      if (query.includes("SELECT * FROM current_schema()")) {
+        return [{ current_schema: schema }];
+      }
+      if (
+        query.includes('SELECT "n"."nspname", "t"."typname" FROM "pg_type" "t"')
+      ) {
+        if (enumQueries.has(query)) {
+          // mock result for enum that already exists
+          return [{}];
+        }
+        enumQueries.add(query);
+        // mock result for enum that does not exist
+        return [];
+      }
+    };
+  };
+
   // create tables statements
   for (const metadata of entityMetadatas) {
     if (metadata.tableType === "view") {
       continue;
     }
+    if (dialect === "postgres") {
+      mockPostgresQueryRunner(metadata.schema);
+    }
     const table = Table.create(metadata, driver);
     await queryRunner.createTable(table);
   }

diff --git a/testdata/ts/atlas-standalone.hcl b/testdata/ts/atlas-standalone.hcl
@@ -15,7 +15,7 @@ data "external_schema" "typeorm" {
     "npx",
     "../..",
     "load",
-    "--path", "./entities",
+    "--path", var.dialect == "sqlite" ? "./entities/sqlite" : "./entities",
     "--dialect", var.dialect,
   ]
 }

diff --git a/testdata/ts/entities/User.ts b/testdata/ts/entities/User.ts
@@ -9,6 +9,12 @@ import {
 } from "typeorm";
 import { Blog } from "./Blog";
 
+export enum UserRole {
+  ADMIN = "admin",
+  EDITOR = "editor",
+  GHOST = "ghost",
+}
+
 @Entity()
 @Unique(["firstName", "lastName"])
 @Check(`"age" > 6`)
@@ -22,6 +28,13 @@ export class User {
   @Column()
   lastName: string;
 
+  @Column({
+    type: "enum",
+    enum: UserRole,
+    default: UserRole.GHOST,
+  })
+  role: UserRole;
+
   @Column()
   @Index("IDX_USER_AGE")
   age: number;

diff --git a/testdata/ts/entities/sqlite/Blog.ts b/testdata/ts/entities/sqlite/Blog.ts
@@ -0,0 +1,14 @@
+import { Column, Entity, ManyToOne, PrimaryGeneratedColumn } from "typeorm";
+import { User } from "./User";
+
+@Entity()
+export class Blog {
+  @PrimaryGeneratedColumn()
+  id: number;
+
+  @Column()
+  title: string;
+
+  @ManyToOne(() => User, (user) => user.blogs)
+  user: User;
+}
diff --git a/testdata/ts/entities/sqlite/User.ts b/testdata/ts/entities/sqlite/User.ts
@@ -0,0 +1,31 @@
+import {
+  Entity,
+  PrimaryGeneratedColumn,
+  Column,
+  OneToMany,
+  Check,
+  Index,
+  Unique,
+} from "typeorm";
+import { Blog } from "./Blog";
+
+@Entity()
+@Unique(["firstName", "lastName"])
+@Check(`"age" > 6`)
+export class User {
+  @PrimaryGeneratedColumn()
+  id: number;
+
+  @Column()
+  firstName: string;
+
+  @Column()
+  lastName: string;
+
+  @Column()
+  @Index("IDX_USER_AGE")
+  age: number;
+
+  @OneToMany(() => Blog, (blog) => blog.user)
+  blogs: Blog[];
+}
diff --git a/testdata/ts/load-entities.ts b/testdata/ts/load-entities.ts
@@ -3,11 +3,20 @@
 import { Dialect, loadEntities } from "../../src/load";
 import { User } from "./entities/User";
 import { Blog } from "./entities/Blog";
+import { User as SqliteUser } from "./entities/sqlite/User";
+import { Blog as SqliteBlog } from "./entities/sqlite/Blog";
 
 // parse the second argument as the dialect
 const dialect = process.argv[2] as Dialect;
 
-// print sql after promise is resolver
-loadEntities(dialect, [User, Blog]).then((sql) => {
-  console.log(sql);
-});
+// sqlite does not support enum, so we need to load different entities
+if (dialect === "sqlite") {
+  loadEntities("sqlite", [SqliteUser, SqliteBlog]).then((sql) => {
+    console.log(sql);
+  });
+} else {
+  // print sql after promise is resolver
+  loadEntities(dialect, [User, Blog]).then((sql) => {
+    console.log(sql);
+  });
+}
diff --git a/...ta/ts/migrations/mysql/20231002071605.sql → ...ta/ts/migrations/mysql/20231210092101.sql b/...ta/ts/migrations/mysql/20231002071605.sql → ...ta/ts/migrations/mysql/20231210092101.sql
@@ -3,6 +3,7 @@ CREATE TABLE `user` (
   `id` int NOT NULL AUTO_INCREMENT,
   `firstName` varchar(255) NOT NULL,
   `lastName` varchar(255) NOT NULL,
+  `role` enum('admin','editor','ghost') NOT NULL DEFAULT "ghost",
   `age` int NOT NULL,
   PRIMARY KEY (`id`),
   UNIQUE INDEX `IDX_c322cd2084cd4b1b2813a90032` (`firstName`, `lastName`),

diff --git a/testdata/ts/migrations/mysql/atlas.sum b/testdata/ts/migrations/mysql/atlas.sum
@@ -1,2 +1,2 @@
-h1:KoVycYqPONHUw7Jz/5GPPdrT8cb9Amgv5PiQtgFpKPQ=
-20231002071605.sql h1:IOSJE0M7CTE5MyKwBf2tD0L/qAZLtfpPb/vXTxa2ThM=
+h1:m9MLvXMU8JlmJzB/PMcRRU4cAYizJCbwLvxTpQ0iyDw=
+20231210092101.sql h1:UBJx3Al66TNzU4eSiuWLaQP+iG+aP81t+kEi4z0m6X8=
diff --git a/...ts/migrations/postgres/20231002071717.sql → ...ts/migrations/postgres/20231210092127.sql b/...ts/migrations/postgres/20231002071717.sql → ...ts/migrations/postgres/20231210092127.sql
@@ -1,8 +1,11 @@
+-- Create enum type "user_role_enum"
+CREATE TYPE "public"."user_role_enum" AS ENUM ('admin', 'editor', 'ghost');
 -- Create "user" table
 CREATE TABLE "public"."user" (
   "id" serial NOT NULL,
   "firstName" character varying NOT NULL,
   "lastName" character varying NOT NULL,
+  "role" "public"."user_role_enum" NOT NULL DEFAULT 'ghost',
   "age" integer NOT NULL,
   PRIMARY KEY ("id"),
   CONSTRAINT "CHK_70c8a9c9c39b98f399c28b8700" CHECK (age > 6)

diff --git a/testdata/ts/migrations/postgres/atlas.sum b/testdata/ts/migrations/postgres/atlas.sum
@@ -1,2 +1,2 @@
-h1:fc6MQIyj84BgIN3Y+ituJkSrSYH0Td2psVyZsaOyVak=
-20231002071717.sql h1:bl0GIRE/KaU87H6fMErK0NK9OU+NSgI+X4YYLhcZ8EQ=
+h1:leQ+fcp10mYR7Y9AI7rO/u/nsjbS+GrYIJkHGyzVmcE=
+20231210092127.sql h1:zKfMfYvki0ZjcyylyA7HfT70n6RPC0X36zHyuCSwN9Q=