From ef10b8653acd6c7b17effaa1d467031619055805 Mon Sep 17 00:00:00 2001
From: Kononnable <kononnable@gmail.com>
Date: Mon, 23 Sep 2019 20:26:45 +0200
Subject: [PATCH] enum type safety #205

---
 CHANGELOG.md                                  |  1 +
 src/drivers/MysqlDriver.ts                    |  7 ++++-
 src/drivers/PostgresDriver.ts                 | 26 +++++++++----------
 src/models/ColumnInfo.ts                      |  3 ++-
 .../entityTypes/mariadb/entity/Post.ts        |  6 ++---
 .../entityTypes/mysql/entity/Post.ts          |  3 +--
 .../entityTypes/postgres/entity/Post.ts       |  4 +--
 7 files changed, 25 insertions(+), 25 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c73a1c0..f3a3b6d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,7 @@
 ## Unreleased
 
 * change default case conversions for generated files (#196)
+* enum type safety #205
 
 ## 0.3.5
 
diff --git a/src/drivers/MysqlDriver.ts b/src/drivers/MysqlDriver.ts
index 28ee666..9878739 100644
--- a/src/drivers/MysqlDriver.ts
+++ b/src/drivers/MysqlDriver.ts
@@ -165,7 +165,12 @@ export default class MysqlDriver extends AbstractDriver {
                             colInfo.tsType = "string";
                             break;
                         case "enum":
-                            colInfo.tsType = "string";
+                            colInfo.tsType = resp.COLUMN_TYPE.substring(
+                                5,
+                                resp.COLUMN_TYPE.length - 1
+                            )
+                                .replace(/'/gi, '"')
+                                .replace(/","/gi, '" | "');
                             colInfo.options.enum = resp.COLUMN_TYPE.substring(
                                 5,
                                 resp.COLUMN_TYPE.length - 1
diff --git a/src/drivers/PostgresDriver.ts b/src/drivers/PostgresDriver.ts
index 5ed4b6c..9dbb560 100644
--- a/src/drivers/PostgresDriver.ts
+++ b/src/drivers/PostgresDriver.ts
@@ -164,21 +164,16 @@ WHERE "n"."nspname" = table_schema AND "t"."typname"=udt_name
         enumValues: string | null
     ) {
         let ret: {
-            tsType:
-                | "number"
-                | "string"
-                | "boolean"
-                | "Date"
-                | "Buffer"
-                | "object"
-                | "string | object"
-                | "string | string[]"
-                | "any"
-                | null;
+            tsType?: ColumnInfo["tsType"];
             sqlType: string | null;
             isArray: boolean;
             enumValues: string[];
-        } = { tsType: null, sqlType: null, isArray: false, enumValues: [] };
+        } = {
+            tsType: undefined,
+            sqlType: null,
+            isArray: false,
+            enumValues: []
+        };
         ret.sqlType = dataType;
         switch (dataType) {
             case "int2":
@@ -380,19 +375,22 @@ WHERE "n"."nspname" = table_schema AND "t"."typname"=udt_name
                         break;
                     default:
                         if (enumValues) {
+                            ret.tsType = (`"${enumValues
+                                .split(",")
+                                .join('" | "')}"` as never) as string;
                             ret.sqlType = "enum";
                             ret.enumValues = (`"${enumValues
                                 .split(",")
                                 .join('","')}"` as never) as string[];
                         } else {
-                            ret.tsType = null;
+                            ret.tsType = undefined;
                             ret.sqlType = null;
                         }
                         break;
                 }
                 break;
             default:
-                ret.tsType = null;
+                ret.tsType = undefined;
                 ret.sqlType = null;
                 break;
         }
diff --git a/src/models/ColumnInfo.ts b/src/models/ColumnInfo.ts
index 28058cc..b373c99 100644
--- a/src/models/ColumnInfo.ts
+++ b/src/models/ColumnInfo.ts
@@ -15,7 +15,8 @@ export default class ColumnInfo {
         | "object"
         | "string | object"
         | "string | string[]"
-        | "any";
+        | "any"
+        | string;
 
     public relations: RelationInfo[] = [];
 }
diff --git a/test/integration/entityTypes/mariadb/entity/Post.ts b/test/integration/entityTypes/mariadb/entity/Post.ts
index bdd5343..972bd30 100644
--- a/test/integration/entityTypes/mariadb/entity/Post.ts
+++ b/test/integration/entityTypes/mariadb/entity/Post.ts
@@ -2,7 +2,6 @@ import { Entity, PrimaryColumn, Column } from "typeorm";
 
 @Entity("Post")
 export class Post {
-
     @PrimaryColumn()
     id: number;
 
@@ -18,7 +17,7 @@ export class Post {
     @Column("tinyint")
     tinyint: number;
 
-    @Column("tinyint",{width:1})
+    @Column("tinyint", { width: 1 })
     boolean: boolean;
 
     @Column("smallint")
@@ -85,7 +84,7 @@ export class Post {
     longtext: string;
 
     @Column("enum", { enum: ["A", "B", "C"] })
-    enum: string;
+    enum: "A" | "B" | "C";
 
     // In mariaDb Json is recognized as longtext
     // @Column("json")
@@ -117,5 +116,4 @@ export class Post {
 
     @Column("geometrycollection")
     geometrycollection: string;
-
 }
diff --git a/test/integration/entityTypes/mysql/entity/Post.ts b/test/integration/entityTypes/mysql/entity/Post.ts
index 0c49fb5..f76777f 100644
--- a/test/integration/entityTypes/mysql/entity/Post.ts
+++ b/test/integration/entityTypes/mysql/entity/Post.ts
@@ -2,7 +2,6 @@ import { Entity, PrimaryColumn, Column } from "typeorm";
 
 @Entity("Post")
 export class Post {
-
     @PrimaryColumn()
     id: number;
 
@@ -88,7 +87,7 @@ export class Post {
     longtext: string;
 
     @Column("enum", { enum: ["A", "B", "C"] })
-    enum: string;
+    enum: "A" | "B" | "C";
 
     @Column("json")
     json: object;
diff --git a/test/integration/entityTypes/postgres/entity/Post.ts b/test/integration/entityTypes/postgres/entity/Post.ts
index 0dc2de8..1f1cc56 100644
--- a/test/integration/entityTypes/postgres/entity/Post.ts
+++ b/test/integration/entityTypes/postgres/entity/Post.ts
@@ -2,7 +2,6 @@ import { Entity, PrimaryColumn, Column } from "typeorm";
 
 @Entity("Post")
 export class Post {
-
     @PrimaryColumn()
     id: number;
 
@@ -120,7 +119,7 @@ export class Post {
     boolean: boolean;
 
     @Column("enum", { enum: ["A", "B", "C"] })
-    enum: string;
+    enum: "A" | "B" | "C";
 
     @Column("point")
     point: string | object;
@@ -187,5 +186,4 @@ export class Post {
 
     @Column("daterange")
     daterange: string;
-
 }