Add support for native enums on Postgres (#2632)

Reference https://www.postgresql.org/docs/current/static/sql-createtype.html

Closes #394

Signed-off-by: Will Soto <will.soto9@gmail.com>

src/dialects/postgres/schema/columncompiler.js

@@ -30,8 +30,17 @@ assign(ColumnCompiler_PG.prototype, {
 
   // Create the column definition for an enum type.
   // Using method "2" here: http://stackoverflow.com/a/10984951/525714
-  enu(allowed) {
-    return `text check (${this.formatter.wrap(this.args[0])} in ('${allowed.join("', '")}'))`;
+  enu(allowed, options) {
+    options = options || {};
+
+    const values = allowed.join("', '");
+
+    if (options.useNative) {
+      this.tableCompiler.unshiftQuery(`create type "${options.enumName}" as enum ('${values}')`);
+
+      return `"${options.enumName}"`;
+    }
+    return `text check (${this.formatter.wrap(this.args[0])} in ('${values}'))`;
   },
 
   double: 'double precision',

src/schema/columncompiler.js

@@ -25,6 +25,8 @@ ColumnCompiler.prototype.pushQuery = helpers.pushQuery
 
 ColumnCompiler.prototype.pushAdditional = helpers.pushAdditional
 
+ColumnCompiler.prototype.unshiftQuery = helpers.unshiftQuery
+
 ColumnCompiler.prototype._defaultMap = {
   'columnName': function() {
     if (!this.isIncrements) {

src/schema/compiler.js

@@ -1,5 +1,5 @@
 
-import { pushQuery, pushAdditional } from './helpers';
+import { pushQuery, pushAdditional, unshiftQuery } from './helpers';
 
 import { assign, isUndefined } from 'lodash'
 
@@ -20,6 +20,8 @@ assign(SchemaCompiler.prototype, {
 
   pushAdditional: pushAdditional,
 
+  unshiftQuery: unshiftQuery,
+
   createTable: buildTable('create'),
 
   createTableIfNotExists: buildTable('createIfNot'),

src/schema/helpers.js

@@ -34,3 +34,27 @@ export function pushAdditional(fn) {
   fn.call(child, tail(arguments));
   this.sequence.additional = (this.sequence.additional || []).concat(child.sequence);
 }
+
+// Unshift a new query onto the compiled "sequence" stack,
+// creating a new formatter, returning the compiler.
+export function unshiftQuery(query) {
+  if (!query) return;
+  if (isString(query)) {
+    query = {sql: query};
+  }
+  if (!query.bindings) {
+    query.bindings = this.formatter.bindings;
+  }
+  this.sequence.unshift(query);
+
+  let builder;
+  if (this instanceof ColumnCompiler) {
+    builder = this.columnBuilder;
+  } else if (this instanceof TableCompiler) {
+    builder = this.tableBuilder;
+  } else if (this instanceof SchemaCompiler) {
+    builder = this.builder;
+  }
+
+  this.formatter = this.client.formatter(builder);
+}

src/schema/tablecompiler.js

@@ -2,7 +2,7 @@
 
 // Table Compiler
 // -------
-import { pushAdditional, pushQuery } from './helpers';
+import { pushAdditional, pushQuery, unshiftQuery } from './helpers';
 import * as helpers from '../helpers';
 import { groupBy, reduce, map, first, tail, isEmpty, indexOf, isArray, isUndefined } from 'lodash'
 
@@ -23,6 +23,8 @@ TableCompiler.prototype.pushQuery = pushQuery
 
 TableCompiler.prototype.pushAdditional = pushAdditional
 
+TableCompiler.prototype.unshiftQuery = unshiftQuery
+
 // Convert the tableCompiler toSQL
 TableCompiler.prototype.toSQL = function () {
   this[this.method]();

test/integration/schema/index.js

@@ -84,7 +84,7 @@ module.exports = function(knex) {
             const table_name = 'increments_columns_1_test';
             const expected_column = 'id'
             const expected_comment = 'comment_1';
-            
+
             return knex.raw('SELECT c.oid FROM pg_catalog.pg_class c WHERE c.relname = ?',[table_name]).then(function(res) {
               const column_oid = res.rows[0].oid;
 
@@ -106,13 +106,13 @@ module.exports = function(knex) {
             const table_name = 'increments_columns_2_test';
             const expected_column = 'named_2';
             const expected_comment = 'comment_2';
-            
+
             return knex.raw('SELECT c.oid FROM pg_catalog.pg_class c WHERE c.relname = ?',[table_name]).then(function(res) {
               const column_oid = res.rows[0].oid;
-              
+
               return knex.raw('SELECT pg_catalog.col_description(?,?);', [column_oid, '1']).then(function(_res) {
                 const comment = _res.rows[0].col_description;
-                
+
                 return knex.raw('select column_name from INFORMATION_SCHEMA.COLUMNS where table_name = ?;', table_name).then((res) => {
                   const column_name = res.rows[0].column_name;
 
@@ -161,7 +161,7 @@ module.exports = function(knex) {
               TABLE_NAME = ? AND
               COLUMN_NAME = ?
             `
-            
+
             return knex.raw(query,[table_name, expected_column]).then(function(res) {
               const comment = res[0][0].COLUMN_COMMENT
               expect(comment).to.equal(expected_comment)
@@ -183,15 +183,41 @@ module.exports = function(knex) {
               TABLE_NAME = ? AND
               COLUMN_NAME = ?
             `
-            
+
             return knex.raw(query,[table_name, expected_column]).then(function(res) {
               const comment = res[0][0].COLUMN_COMMENT
               expect(comment).to.equal(expected_comment)
             })
-            
+
           });
       });
 
+      describe('enum - postgres', function() {
+        if (!knex || !knex.client || !(/postgres/i.test(knex.client.dialect))) {
+          return Promise.resolve();
+        }
+
+        afterEach(function() {
+          return knex.schema.dropTableIfExists('native_enum_test').raw('DROP TYPE "foo_type"');
+        });
+
+        it('uses native type when useNative is specified', function() {
+          return knex.schema
+            .createTable('native_enum_test', function(table) {
+              table.enum('foo_column', ['a', 'b', 'c'], {
+                useNative: true,
+                enumName: 'foo_type'
+              }).notNull();
+              table.uuid('id').notNull();
+            }).testSql(function(tester) {
+              tester('pg', [
+                'create type "foo_type" as enum (\'a\', \'b\', \'c\')',
+                'create table "native_enum_test" ("foo_column" "foo_type" not null, "id" uuid not null)'
+              ]);
+            });
+        });
+      })
+
       it('Callback function must be supplied', function() {
         expect(function() {
           knex.schema.createTable('callback_must_be_supplied').toString()

test/unit/schema/postgres.js

@@ -541,6 +541,18 @@ describe("PostgreSQL SchemaBuilder", function() {
     expect(tableSql[0].sql).to.equal('alter table "users" add column "foo" text check ("foo" in (\'bar\', \'baz\'))');
   });
 
+  it("adding enum with useNative", function() {
+    tableSql = client.schemaBuilder().table('users', function(table) {
+      table.enu('foo', ['bar', 'baz'], {
+        useNative: true,
+        enumName: 'foo_type'
+      }).notNullable();
+    }).toSQL();
+    equal(2, tableSql.length);
+    expect(tableSql[0].sql).to.equal('create type "foo_type" as enum (\'bar\', \'baz\')')
+    expect(tableSql[1].sql).to.equal('alter table "users" add column "foo" "foo_type" not null');
+  });
+
   it("adding date", function() {
     tableSql = client.schemaBuilder().table('users', function(table) {
       table.date('foo');