Support for PostgreSQL enum types in Rails 5.1 (including schema dump)

rails_postgres_enum.rb

# Support for PostgreSQL enum types in Rails 5.1 (including schema dump)
module ActiveRecord
  class SchemaDumper
    def dump(stream)
      header(stream)
      extensions(stream)
      enums(stream)
      tables(stream)
      trailer(stream)
      stream
    end

    private

      def enums(stream)
        enum_types = @connection.enum_types
        if enum_types.any?
          stream.puts '  # These are custom enum types that must be created before they can be used in the schema definition'
          enum_types.each do |name, values|
            stream.puts "  create_enum #{name.inspect}, #{values.map(&:inspect)*', '}"
          end
          stream.puts
        end
      end
  end

  module ConnectionAdapters
    class PostgreSQLAdapter
      def enum_types(name = nil)
        where = name.nil? ? '' : " and (t.typname = #{quote(name)})"
        query = <<~SQL
          SELECT t.typname AS enum_name,
                 array_agg(e.enumlabel ORDER BY e.enumsortorder) AS enum_values
          FROM pg_type t
          JOIN pg_enum e ON t.oid = e.enumtypid
          JOIN pg_catalog.pg_namespace n ON n.oid = t.typnamespace
          WHERE (n.nspname = 'public')#{where}
          GROUP BY enum_name
          ORDER BY enum_name
        SQL
        exec_query(query, 'SCHEMA').cast_values.to_h
      end

      def enum_exists?(name)
        enum_names.include?(name.to_s)
      end

      def enum_names
        enum_types.keys
      end

      def enum_values(name)
        enum_types(name).values
      rescue
        []
      end
    end

    module PostgreSQL
      module ColumnDumper
        def schema_type(column)
          column.type.eql?(:enum) ? unquote(column.sql_type).to_sym : super
        end
      end

      module Quoting
        def quote_enum_name(name)
          quote_column_name(name)
        end

        def quote_enum_value(value)
          quote(value)
        end

        def quote_enum_values(values)
          values.map { |value| quote_enum_value(value) }
        end

        def unquote(part)
          if part && part.start_with?('"')
            part[1..-2]
          else
            part
          end
        end
      end

      module SchemaStatements
        def create_enum(name, *values)
          execute "CREATE TYPE #{quote_enum_name(name)} AS ENUM (#{quote_enum_values(values)*', '})"
          load_enum_types(name)
          ColumnMethods.create_enum_column_method(name)
        end

        def drop_enum(name, force: false, if_exists: false, **)
          ColumnMethods.remove_enum_column_method(name)
          sql = %w[DROP TYPE]
          sql += %w[IF EXISTS] if if_exists
          sql << quote_enum_name(name)
          sql << 'CASCADE' if force == :cascade
          execute sql*' '
        end

        def load_enum_types(name)
          oids = ['', '[]'].map { |s| query_value("SELECT #{quote(quote_enum_name(name)+s)}::regtype::oid", 'SCHEMA').to_i }
          load_additional_types(type_map, oids)
        end
      end
    end
  end
end

module PostgreSQL_Enum
  module ColumnMethods
    def create_enum_column_method(enum_type)
      class_eval do
        if !method_defined?(enum_type)
          define_method(enum_type) do |*args, **options|
            args.each { |name| column(name, enum_type, options) }
          end
        end
      end
    end

    def remove_enum_column_method(enum_type)
      class_eval do
        remove_method(enum_type) if method_defined?(enum_type)
      end
    end
  end

  module PostgreSQLAdapter
    def native_database_types
      { enum: { name: 'enum' } }.merge(super)
    end
  end

  module SchemaStatements
    def type_to_sql(type, array: nil, **)
      if enum_exists?(type)
        sql = quote_enum_name(type)
        sql = "#{sql}[]" if array
        sql
      else
        super
      end
    end
  end

  module TypeMapInitializer
    def register_enum_type(row)
      super
      ActiveRecord::ConnectionAdapters::PostgreSQL::ColumnMethods.create_enum_column_method(row['typname'])
    end
  end
end
ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend PostgreSQL_Enum::PostgreSQLAdapter
ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend PostgreSQL_Enum::SchemaStatements
ActiveRecord::ConnectionAdapters::PostgreSQL::ColumnMethods.extend PostgreSQL_Enum::ColumnMethods
ActiveRecord::ConnectionAdapters::PostgreSQL::OID::TypeMapInitializer.prepend PostgreSQL_Enum::TypeMapInitializer