icalialabs / archon-gem Goto Github PK

# frozen_string_literal: true

# Adds migration and schema.rb support to postgres custom enum types, tested on
#
# Usage:

# Require inside `config/application.rb` (inside the Application class):
# ```ruby
# module Artanis
#   class Application < Rails::Application
#     require 'support_for_postgres_enums'
#   end
# end
# ```
#
# Then on migrations:
# ```ruby
# class MyMigration < ActiveRecord::Migration[6.0]
#   def change
#     create_enum :my_custom_type, %w[first_value second_value third_value]
#     add_column :my_table, :my_enum_column, :my_custom_type
#
#     create_table :my_enum_test do |t|
#       t.enum :my_enum_column, type: :my_custom_type
#     end
#   end
# end
# ```
#
## Inspired by:
#  - https://gist.github.com/khamusa/77ada9895db15dbee4e57cb715830a48
#  - https://gist.github.com/clarkdave/5936375

module SupportForPostgresEnums
  def self.adapters
    ActiveRecord::ConnectionAdapters
  end

  def self.postgresql
    adapters::PostgreSQL
  end

  def self.postgresql_adapter
    adapters::PostgreSQLAdapter
  end

  def self.postgresql_oid
    postgresql::OID
  end

  def self.enable
    ActiveRecord::SchemaDumper.send :prepend, OnSchemaDumper

    require 'active_record/connection_adapters/abstract/schema_dumper'
    adapters::SchemaDumper.send :prepend, OnAdapterSchemaDumper

    require 'active_record/migration/command_recorder'
    ActiveRecord::Migration::CommandRecorder.send :prepend, OnCommandRecorder

    require 'active_record/connection_adapters/postgresql/schema_definitions'
    postgresql::TableDefinition.send :prepend, OnTableDefinition

    require 'active_record/connection_adapters/postgresql/oid/enum'

    require 'active_record/connection_adapters/postgresql/oid/type_map_initializer'
    postgresql_oid::TypeMapInitializer.send :prepend, OnTypeMapInitializer

    require 'active_record/connection_adapters/postgresql_adapter'
    postgresql_adapter.send :prepend, self::OnAdapter

    postgresql_adapter::NATIVE_DATABASE_TYPES[:enum] = { name: 'enum' }
  end

  module OnAdapter
    def migration_keys
      super + [:enum_name]
    end

    def prepare_column_options(column, types)
      spec = super(column, types)

      if column.type == :enum
        spec[:enum_name] = column.cast_type.enum_name.inspect
      end

      spec
    end

    def create_enum(name, values)
      return if execute("SELECT 1 FROM pg_type WHERE typname = '#{name}'").any?

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

    def drop_enum(name, _values = nil)
      return if execute("SELECT 1 FROM pg_type WHERE typname = '#{name}'").empty?

      execute "DROP TYPE \"#{name}\""
    end
  end

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

  module OnAdapterSchemaDumper
    private

    def prepare_column_options(column)
      super.tap do |spec|
        if column.type == :enum
          spec[:type] = ":#{column.sql_type_metadata.sql_type}"
        end
      end
    end
  end

  module OnSchemaDumper
    def tables(stream)
      enums(stream)
      super
    end

    private

    def enums(stream)
      statements = []

      defined_enums.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

    def defined_enums
      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 OnTypeMapInitializer
    # We need to know the column name, and the default implementation discards
    # it
    def register_enum_type(row)
      enum_attributes = { type: row['typname'] }
      register row['oid'],
               SupportForPostgresEnums.postgresql_oid::Enum.new(enum_attributes)
    end
  end

  module OnCommandRecorder
    def create_enum(*args, &block)
      record(:create_enum, args, &block)
    end

    def drop_enum(*args, &block)
      record(:drop_enum, args, &block)
    end

    def invert_create_enum(args)
      [:drop_enum, [args.first]]
    end

    def invert_drop_enum(args)
      [:create_enum, args]
    end
  end
end

# Apply the active_record patches:
ActiveSupport.on_load(:active_record) { SupportForPostgresEnums.enable }

# frozen_string_literal: true

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

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

          def type
            :enum
          end
        end
      end
    end
  end
end

require 'archon/support_for_postgres_enums'