Better migration of enum types

change_enum_migration.rb

class ChangeEnumMigration < ActiveRecord::Migration
  include MigrationTools

  def change
    migrate_enum :some_enum, [:some_table, :some_column, null: false, default: 'other'],
      keep_values: %w(foo bar),
      new_values: %w(baz qux),
      mappings: {'quux' => 'mux'},
      removed_values: %w(baaaz baaz)
  end
end

migration_tools.rb

# WARNING: don't change methods in this file or old migrations will break
module MigrationTools

  # Migrate enum to new set of values.
  # Params
  #   enum_type: name of the enum type
  #   columns: a list of column definitions that use the enum. Like [table_name, column_name, **column_options]
  # 
  #   keep_values: Values to keep from old enum
  #   new_values: New values of the enum that don't need to be mapped from old values
  #   mappings: mappings of old values to different new values
  #   removed_values: list of removed values (for reverse)
  # NOTE: if removed_values is not complete, cannot be reversed completely
  # NOTE: if column_options is not complete, it will change column, and also inverse won't work.
  def migrate_enum enum_type, *columns, new_values: [], removed_values: [], mappings: {}, **opts
    reversible do |dir|
      dir.up do
        EnumChange.new(columns, enum_type, self, new_values: new_values, mappings: mappings, **opts).perform!
      end
      dir.down do
        EnumChange.new(columns, enum_type, self, new_values: removed_values, mappings: mappings.invert, **opts).perform!
      end
    end
  end

  def create_enum_type enum_type, values
    values = values.map {|v| quote(v)}.join(', ')
    execute "CREATE TYPE #{enum_type} AS ENUM (#{values});"
  end
  
  def drop_enum_type enum_type
    execute("DROP TYPE #{enum_type};")
  end


  class EnumChange < Struct.new(:columns, :enum_type, :keep_values, :new_values, :mappings, :migration_context)
    def initialize columns, enum_type, migration_context, keep_values: [], new_values: [], mappings: {}
      raise ArgumentError, 'no columns' if columns.blank?
      columns = columns.map do |args|
        ChangeColumn.new(migration_context, enum_type, *args)
      end

      super(columns, enum_type, keep_values, new_values, mappings, migration_context)
    end

    def perform!
      # Setup
      rename_type_to_tmp
      create_new_enum_type
      
      # Migrate columns
      columns.each do |column|
        column.rename_column_to_tmp
        column.create_new_column

        if keep_values.any?
          column.migrate_keep_values_to_new_column(quoted_keep_values)
        end

        column.migrate_mappings(mappings)
        column.remove_tmp_column
      end

      # Cleanup
      remove_tmp_type
    end

    # Operations

    def rename_type_to_tmp
      execute("ALTER TYPE #{enum_type} rename to #{tmp_enum_type};")
    end

    def create_new_enum_type
      create_enum_type enum_type, all_target_values
    end

    def remove_tmp_type
      drop_enum_type tmp_enum_type
    end

    # Delegation

    delegate :execute, :create_enum_type, :drop_enum_type, :quote, to: :migration_context

    # Utils

    def tmp_enum_type; "#{enum_type}_old"; end

    def all_target_values
      keep_values + new_values + mappings.values
    end

    def quoted_keep_values
      keep_values.map {|v| quote(v)}.join(', ')
    end

    class ChangeColumn < Struct.new(:migration_context, :enum_type, :table_name, :column_name, :column_options)
      def initialize migration_context, enum_type, table_name, column_name, **column_options
        super(migration_context, enum_type, table_name, column_name, column_options)
      end

      # Operations

      def rename_column_to_tmp
        rename_column table_name, column_name, tmp_column_name
      end

      def create_new_column
        add_column table_name, column_name, enum_type, **column_options
      end

      def migrate_keep_values_to_new_column quoted_keep_values
        execute "UPDATE #{quoted_table_name} SET #{column_name} = #{tmp_column_name}::text::#{enum_type} WHERE #{tmp_column_name} IN (#{quoted_keep_values})"
      end

      def migrate_mappings mappings
        mappings.each do |from, to|
          execute "UPDATE #{quoted_table_name} SET #{column_name} = #{quote(to)} WHERE #{tmp_column_name} = #{quote(from)}"
        end
      end

      def remove_tmp_column
        remove_column table_name, tmp_column_name
      end

      # Utils

      def quoted_table_name
        migration_context.connection.quote_table_name(table_name)
      end

      
      def tmp_column_name; "#{column_name}_old"; end
    
      delegate :execute, :quote, :remove_column, :add_column, :rename_column, to: :migration_context

    end
  end
end