feat(entity-generator): add enum generation support (#2608)

packages/core/src/naming-strategy/NamingStrategy.ts

@@ -21,7 +21,7 @@ export interface NamingStrategy {
   propertyToColumnName(propertyName: string): string;
 
   /**
-   * Return a column name for a property (used in `EntityGenerator`).
+   * Return a property for a column name (used in `EntityGenerator`).
    */
   columnNameToProperty(columnName: string): string;
 

packages/entity-generator/src/SourceFile.ts

@@ -27,6 +27,7 @@ export class SourceFile {
     });
 
     ret += `export class ${this.meta.className} {\n`;
+    const enumDefinitions: string[] = [];
     Object.values(this.meta.properties).forEach(prop => {
       const decorator = this.getPropertyDecorator(prop, 2);
       const definition = this.getPropertyDefinition(prop, 2);
@@ -38,6 +39,11 @@ export class SourceFile {
       ret += decorator;
       ret += definition;
       ret += '\n';
+
+      if (prop.enum) {
+        const enumClassName = this.namingStrategy.getClassName(this.meta.collection + '_' + prop.fieldNames[0], '_');
+        enumDefinitions.push(this.getEnumClassDefinition(enumClassName, prop.items as string[], 2));
+      }
     });
     ret += '}\n';
 
@@ -47,7 +53,12 @@ export class SourceFile {
       imports.push(`import { ${entity} } from './${entity}';`);
     });
 
-    return `${imports.join('\n')}\n\n${ret}`;
+    ret = `${imports.join('\n')}\n\n${ret}`;
+    if (enumDefinitions.length) {
+      ret += '\n' + enumDefinitions.join('\n');
+    }
+
+    return ret;
   }
 
   getBaseName() {
@@ -75,7 +86,9 @@ export class SourceFile {
 
   private getPropertyDefinition(prop: EntityProperty, padLeft: number): string {
     // string defaults are usually things like SQL functions
-    const useDefault = prop.default != null && typeof prop.default !== 'string';
+    // string defaults can also be enums, for that useDefault should be true.
+    const isEnumOrNonStringDefault = prop.enum || typeof prop.default !== 'string';
+    const useDefault = prop.default != null && isEnumOrNonStringDefault;
     const optional = prop.nullable ? '?' : (useDefault ? '' : '!');
     const ret = `${prop.name}${optional}: ${prop.type}`;
     const padding = ' '.repeat(padLeft);
@@ -84,9 +97,25 @@ export class SourceFile {
       return `${padding + ret};\n`;
     }
 
+    if (prop.enum && typeof prop.default === 'string') {
+      const match = prop.default.match(/^'(.*)'$/);
+      const noQuoteDefault = match?.[1] ?? prop.default;
+      return `${padding}${ret} = ${prop.type}.${noQuoteDefault.toUpperCase()};\n`;
+    }
+
     return `${padding}${ret} = ${prop.default};\n`;
   }
 
+  private getEnumClassDefinition(enumClassName: string, enumValues: string[], padLeft: number): string {
+    const padding = ' '.repeat(padLeft);
+    let ret = `export enum ${enumClassName} {\n`;
+    enumValues.forEach(enumValue => {
+      ret += `${padding}${enumValue.toUpperCase()} = '${enumValue}',\n`;
+    });
+    ret += '}\n';
+    return ret;
+  }
+
   private getPropertyDecorator(prop: EntityProperty, padLeft: number): string {
     const padding = ' '.repeat(padLeft);
     const options = {} as Dictionary;
@@ -99,6 +128,10 @@ export class SourceFile {
       this.getScalarPropertyDecoratorOptions(options, prop);
     }
 
+    if (prop.enum) {
+      options.items = `() => ${prop.type}`;
+    }
+
     this.getCommonDecoratorOptions(options, prop);
     const indexes = this.getPropertyIndexes(prop, options);
     decorator = [...indexes.sort(), decorator].map(d => padding + d).join('\n');
@@ -155,13 +188,17 @@ export class SourceFile {
       options.nullable = true;
     }
 
-    if (prop.default && typeof prop.default === 'string') {
-      if ([`''`, ''].includes(prop.default)) {
-        options.default = `''`;
-      } else if (prop.default.match(/^'.*'$/)) {
-        options.default = prop.default;
+    if (prop.default != null) {
+      if (typeof prop.default === 'string') {
+        if ([`''`, ''].includes(prop.default)) {
+          options.default = `''`;
+        } else if (prop.default.match(/^'.*'$/)) {
+          options.default = prop.default;
+        } else {
+          options.defaultRaw = `\`${prop.default}\``;
+        }
       } else {
-        options.defaultRaw = `\`${prop.default}\``;
+        options.default = prop.default;
       }
     }
   }
@@ -173,6 +210,16 @@ export class SourceFile {
       t = 'datetime';
     }
 
+    if (prop.fieldNames[0] !== this.namingStrategy.propertyToColumnName(prop.name)) {
+      options.fieldName = `'${prop.fieldNames[0]}'`;
+    }
+
+    // for enum properties, we don't need a column type or the property length
+    // in the decorator so return early.
+    if (prop.enum) {
+      return;
+    }
+
     const mappedType1 = this.platform.getMappedType(t);
     const mappedType2 = this.platform.getMappedType(prop.columnTypes[0]);
     const columnType1 = mappedType1.getColumnType({ ...prop, autoincrement: false }, this.platform);
@@ -182,11 +229,7 @@ export class SourceFile {
       options.columnType = this.quote(prop.columnTypes[0]);
     }
 
-    if (prop.fieldNames[0] !== this.namingStrategy.propertyToColumnName(prop.name)) {
-      options.fieldName = `'${prop.fieldNames[0]}'`;
-    }
-
-    if (prop.length && prop.columnTypes[0] !== 'enum') {
+    if (prop.length) {
       options.length = prop.length;
     }
   }
@@ -231,6 +274,10 @@ export class SourceFile {
       return '@PrimaryKey';
     }
 
+    if (prop.enum) {
+      return '@Enum';
+    }
+
     return '@Property';
   }
 

packages/knex/src/schema/DatabaseTable.ts

@@ -228,7 +228,12 @@ export class DatabaseTable {
     return !!this.getPrimaryKey();
   }
 
-  private getPropertyDeclaration(column: Column, namingStrategy: NamingStrategy, schemaHelper: SchemaHelper, compositeFkIndexes: Dictionary<{ keyName: string }>, compositeFkUniques: Dictionary<{ keyName: string }>) {
+  private getPropertyDeclaration(
+    column: Column,
+    namingStrategy: NamingStrategy,
+    schemaHelper: SchemaHelper,
+    compositeFkIndexes: Dictionary<{ keyName: string }>,
+    compositeFkUniques: Dictionary<{ keyName: string }>) {
     const fk = Object.values(this.foreignKeys).find(fk => fk.columnNames.includes(column.name));
     const prop = this.getPropertyName(namingStrategy, column);
     const index = compositeFkIndexes[prop] || this.indexes.find(idx => idx.columnNames[0] === column.name && !idx.composite && !idx.unique && !idx.primary);
@@ -258,6 +263,8 @@ export class DatabaseTable {
       length: column.length,
       index: index ? index.keyName : undefined,
       unique: unique ? unique.keyName : undefined,
+      enum: !!column.enumItems?.length,
+      items: column.enumItems,
       ...fkOptions,
     };
   }
@@ -291,6 +298,12 @@ export class DatabaseTable {
       const parts = fk.referencedTableName.split('.', 2);
       return namingStrategy.getClassName(parts.length > 1 ? parts[1] : parts[0], '_');
     }
+    // If this column is using an enum.
+    if (column.enumItems?.length) {
+      // We will create a new enum name for this type and set it as the property type as well.
+      // The enum name will be a concatenation of the table name and the column name.
+      return namingStrategy.getClassName(this.name + '_' + column.name, '_');
+    }
 
     return column.mappedType?.compareAsType() ?? 'unknown';
   }

packages/knex/src/schema/SchemaHelper.ts

@@ -44,7 +44,7 @@ export abstract class SchemaHelper {
     return this.mapForeignKeys(fks, tableName, schemaName);
   }
 
-  async getEnumDefinitions(connection: AbstractSqlConnection, tableName: string, schemaName?: string): Promise<Dictionary> {
+  async getEnumDefinitions(connection: AbstractSqlConnection, tableName: string, schemaName?: string): Promise<Dictionary<string[]>> {
     return {};
   }
 

packages/mysql-base/src/MySqlSchemaHelper.ts

@@ -1,7 +1,7 @@
 import type { AbstractSqlConnection, Column, Index, Knex, TableDifference } from '@mikro-orm/knex';
 import { SchemaHelper } from '@mikro-orm/knex';
 import type { Dictionary , Type } from '@mikro-orm/core';
-import { StringType, TextType } from '@mikro-orm/core';
+import { EnumType, StringType, TextType } from '@mikro-orm/core';
 
 export class MySqlSchemaHelper extends SchemaHelper {
 
@@ -102,15 +102,15 @@ export class MySqlSchemaHelper extends SchemaHelper {
       + `where k.table_name = '${tableName}' and k.table_schema = database() and c.constraint_schema = database() and k.referenced_column_name is not null`;
   }
 
-  async getEnumDefinitions(connection: AbstractSqlConnection, tableName: string, schemaName?: string): Promise<Dictionary> {
+  async getEnumDefinitions(connection: AbstractSqlConnection, tableName: string, schemaName?: string): Promise<Dictionary<string[]>> {
     const sql =  `select column_name as column_name, column_type as column_type from information_schema.columns
       where data_type = 'enum' and table_name = '${tableName}' and table_schema = database()`;
     const enums = await connection.execute<any[]>(sql);
 
     return enums.reduce((o, item) => {
       o[item.column_name] = item.column_type.match(/enum\((.*)\)/)[1].split(',').map((item: string) => item.match(/'(.*)'/)![1]);
       return o;
-    }, {} as Dictionary<string>);
+    }, {} as Dictionary<string[]>);
   }
 
   async getColumns(connection: AbstractSqlConnection, tableName: string, schemaName?: string): Promise<Column[]> {
@@ -169,7 +169,8 @@ export class MySqlSchemaHelper extends SchemaHelper {
   }
 
   protected wrap(val: string | undefined, type: Type<unknown>): string | undefined {
-    return typeof val === 'string' && val.length > 0 && (type instanceof StringType || type instanceof TextType) ? this.platform.quoteValue(val) : val;
+    return typeof val === 'string' && val.length > 0 && (type instanceof StringType
+      || type instanceof TextType || type instanceof EnumType) ? this.platform.quoteValue(val) : val;
   }
 
 }

packages/postgresql/src/PostgreSqlSchemaHelper.ts

@@ -101,7 +101,7 @@ export class PostgreSqlSchemaHelper extends SchemaHelper {
       order by kcu.table_schema, kcu.table_name, kcu.ordinal_position, kcu.constraint_name`;
   }
 
-  async getEnumDefinitions(connection: AbstractSqlConnection, tableName: string, schemaName: string): Promise<Dictionary> {
+  async getEnumDefinitions(connection: AbstractSqlConnection, tableName: string, schemaName = 'public'): Promise<Dictionary<string[]>> {
     const sql = `select conrelid::regclass as table_from, conname, pg_get_constraintdef(c.oid) as enum_def
       from pg_constraint c join pg_namespace n on n.oid = c.connamespace
       where contype = 'c' and conrelid = '"${schemaName}"."${tableName}"'::regclass order by contype`;
@@ -121,7 +121,7 @@ export class PostgreSqlSchemaHelper extends SchemaHelper {
       }
 
       return o;
-    }, {} as Dictionary<string>);
+    }, {} as Dictionary<string[]>);
   }
 
   createTableColumn(table: Knex.TableBuilder, column: Column, fromTable: DatabaseTable, changedProperties?: Set<string>) {

tests/features/entity-generator/EntityGenerator.test.ts

@@ -89,4 +89,34 @@ describe('EntityGenerator', () => {
     await orm.close(true);
   });
 
+  test('enum with default value [mysql]', async () => {
+    const orm = await initORMMySql('mysql', {}, true);
+    await orm.getSchemaGenerator().dropSchema();
+    await orm.getSchemaGenerator().execute(`
+    create table \`publisher2\` (\`id\` int(10) unsigned not null auto_increment primary key, \`type\` enum('local', 'global') not null default 'local', \`type2\` enum('LOCAL', 'GLOBAL') default 'LOCAL') default character set utf8mb4 engine = InnoDB;
+    `);
+    const generator = orm.getEntityGenerator();
+    const dump = await generator.generate({ save: false, baseDir: './temp/entities' });
+    expect(dump).toMatchSnapshot('mysql-entity-dump-enum-default-value');
+    await orm.getSchemaGenerator().execute(`
+      drop table if exists \`publisher2\`;
+    `);
+    await orm.close(true);
+  });
+
+  test('enum with default value [postgres]', async () => {
+    const orm = await initORMPostgreSql();
+    await orm.getSchemaGenerator().dropSchema();
+    await orm.getSchemaGenerator().execute(`
+    create table "publisher2" ("id" serial primary key, "type" text check ("type" in ('local', 'global')) not null default 'local', "type2" text check ("type2" in ('LOCAL', 'GLOBAL')) default 'LOCAL');
+    `);
+    const generator = orm.getEntityGenerator();
+    const dump = await generator.generate({ save: false, baseDir: './temp/entities' });
+    expect(dump).toMatchSnapshot('postgres-entity-dump-enum-default-value');
+    await orm.getSchemaGenerator().execute(`
+      drop table if exists "publisher2";
+    `);
+    await orm.close(true);
+  });
+
 });