Support custom postgres enum types on rails migrations and schema.rb

postgres_custom_enums_on_rails_schema_migrations.rb

# frozen_string_literal: true
# Adds migration and schema.rb support to postgres custom enum types, tested on rails 4.2
# The implementation is quite fragile against ActiveRecord version upgrades since it touches many core AR classes.
# The following code should be placed in an initializer
# On migrations:
#
# create_enum(:mood, %w(happy great been_better))
# create_table :person do
#   t.enum :person_mood, enum_name: :mood
# end
#
# Running the above will create a table :person, with a column :person_mood of
# type :mood. This will also be saved on schema.rb so that `rake schema:load`
# works as expected.
#
# To drop an existing enum:
# drop_enum(:mood)
#
# Renaming / updating enums is currently not supported, since it's quite tricky
# (see https://medium.com/pretto/update-postgresql-enums-with-rails-migrations-ceb8a86d8cee)
#
## Inspired by https://gist.github.com/clarkdave/5936375
module PgEnum
  class << self
    def defined_enums(connection)
      query = <<~SQL
        SELECT t.OID, t.typname, t.typtype, array_agg(e.enumlabel) as enumlabels
        FROM pg_type t
        INNER JOIN pg_enum e ON e.enumtypid = t.oid
        WHERE typtype = 'e'
        GROUP BY t.OID, t.typname, t.typtype
      SQL

      connection.select_all(query).each do |enum|
        enum['enumlabels'] = enum['enumlabels'].gsub!(/[{}]/, '').split(',')
      end
    end
  end

  module PostgreSQLAdapter
    extend ActiveSupport::Concern

    included do
      alias_method_chain :migration_keys, :enum_name
      alias_method_chain :prepare_column_options, :enum_name

      ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter::NATIVE_DATABASE_TYPES[:enum] = {
        name: 'enum'
      }
    end

    def create_enum(name, values)
      values = values.map {|v| "'#{v}'"}
      execute "CREATE TYPE #{name} AS ENUM (#{values.join(', ')})"
    end

    def drop_enum(name)
      execute "DROP TYPE #{name}"
    end

    def migration_keys_with_enum_name
      migration_keys_without_enum_name + [:enum_name]
    end

    def prepare_column_options_with_enum_name(column, types)
      spec = prepare_column_options_without_enum_name(column, types)
      spec[:enum_name] = column.cast_type.enum_name.inspect if column.type == :enum
      spec
    end
  end

  # provide support for writing out the 'create_enum' calls in schema.rb
  module SchemaDumper
    extend ActiveSupport::Concern

    included do
      alias_method_chain :tables, :enums
    end

    def tables_with_enums(stream)
      enums(stream)
      tables_without_enums(stream)
    end

    private

    def enums(stream)
      statements = []

      ::PgEnum.defined_enums(@connection).each do |enum|
        enum_name = enum['typname']
        values = enum['enumlabels'].map(&:inspect).join(', ')
        statements << "  create_enum(#{enum_name.inspect}, [#{values}])"
      end

      stream.puts statements.join("\n")
      stream.puts
    end
  end
end

module ActiveRecord
  module ConnectionAdapters
    module PostgreSQL
      module OID # :nodoc:
        class Enum < Type::Value # :nodoc:
          attr_reader :enum_name

          def initialize(options = {})
            @enum_name = options.delete(:enum_name).to_sym
            super
          end
        end
      end

      module ColumnMethods # :nodoc:
        # Enables `t.enum :my_field, enum_name: :my_enum_name` on migrations
        def enum(name, options = {})
          column(name, options.delete(:enum_name), options.except(:enum_name))
        end
      end

      module OID # :nodoc:
        class TypeMapInitializer
          # We need to know the column name, and the default implementation discards it
          def register_enum_type(row)
            register row['oid'], OID::Enum.new(enum_name: row['typname'])
          end
        end
      end
    end

    class PostgreSQLAdapter # :nodoc:
      include PgEnum::PostgreSQLAdapter
    end
  end

  class SchemaDumper
    include PgEnum::SchemaDumper
  end
end