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

rails_postgres_enum.rb

# frozen_string_literal: true

# Support for PostgreSQL enum types in Rails 5.2 (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
        return unless 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 |enum_type|
          stream.puts "  create_enum \"#{enum_type.first}\", \"#{enum_type.second.join('", "')}\""
        end
        stream.puts
      end

      def format_colspec(colspec)
        if colspec.is_a? Array
          colspec.reject(&:blank?).map do |value|
            value.is_a?(Hash) ? format_colspec(value) : value
          end.flatten.join(", ")
        else
          colspec.map { |key, value| "#{key}: #{value}" }.join(", ")
        end
      end
  end

  module ConnectionAdapters
    class PostgreSQLAdapter
      def native_database_types
        NATIVE_DATABASE_TYPES.merge(enum: { name: "enum" })
      end

      def enum_types
        query = <<-SQL
          SELECT t.typname AS enum_name,
                 string_agg(e.enumlabel, ' ') AS enum_value
          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'
          GROUP BY enum_name
        SQL
        res =  exec_query(query, 'SCHEMA').cast_values
        res.each do |line|
          line[1] = line[1].split(' ')
        end
        res
      end
    end

    module PostgreSQL
      class SchemaDumper
        def column_spec(column)
          type = schema_type_with_virtual(column)
          spec = prepare_column_options(column)

          if column.type == :enum
            type = 'column'
            enum_type = spec.delete(:enum_type)
            spec = [":#{enum_type}", spec]
          end

          [type, spec]
        end

        def prepare_column_options(column)
          spec = super
          spec[:enum_type] = column.sql_type if column.type == :enum
          spec[:array]     = 'true' if column.respond_to?(:array) && column.array
          spec
        end

        def default_primary_key?(column)
          schema_type(column) == :bigserial || schema_type(column) == :serial
        end
      end
    end

    module PostgreSQL::SchemaStatements
      def create_enum(name, *values)
        execute "CREATE TYPE #{name} AS ENUM (#{ values.map {|v| "'#{v}'" }.join(', ')})"
      end

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