Merge pull request rails#41469 from ghiculescu/postgres-enum

PostgreSQL: support custom enum types

Author: rafaelfranca

activerecord/CHANGELOG.md

@@ -1,3 +1,22 @@
+*   PostgreSQL: support custom enum types
+
+    In migrations, use `create_enum` to add a new enum type, and `t.enum` to add a column.
+
+    ```ruby
+    def up
+      create_enum :mood, ["happy", "sad"]
+
+      change_table :cats do |t|
+        t.enum :current_mood, enum_type: "mood", default: "happy", null: false
+      end
+    end
+    ```
+
+    Enums will be presented correctly in `schema.rb`. Note that this is only supported by
+    the PostgreSQL adapter.
+
+    *Alex Ghiculescu*
+
 * Avoid COMMENT statements in PostgreSQL structure dumps
 
     COMMENT statements are now omitted from the output of `db:structure:dump` when using PostgreSQL >= 11.

activerecord/lib/active_record/connection_adapters/abstract_adapter.rb

@@ -454,6 +454,10 @@ def disable_extension(name)
       def enable_extension(name)
       end
 
+      # This is meant to be implemented by the adapters that support custom enum types
+      def create_enum(*) # :nodoc:
+      end
+
       def advisory_locks_enabled? # :nodoc:
         supports_advisory_locks? && @advisory_locks_enabled
       end

activerecord/lib/active_record/connection_adapters/postgresql/column.rb

@@ -32,6 +32,10 @@ def array
         end
         alias :array? :array
 
+        def enum?
+          type == :enum
+        end
+
         def sql_type
           super.delete_suffix("[]")
         end

activerecord/lib/active_record/connection_adapters/postgresql/schema_definitions.rb

@@ -173,11 +173,19 @@ def primary_key(name, type = :primary_key, **options)
         # :method: xml
         # :call-seq: xml(*names, **options)
 
+        ##
+        # :method: timestamptz
+        # :call-seq: timestamptz(*names, **options)
+
+        ##
+        # :method: enum
+        # :call-seq: enum(*names, **options)
+
         included do
           define_column_methods :bigserial, :bit, :bit_varying, :cidr, :citext, :daterange,
             :hstore, :inet, :interval, :int4range, :int8range, :jsonb, :ltree, :macaddr,
             :money, :numrange, :oid, :point, :line, :lseg, :box, :path, :polygon, :circle,
-            :serial, :tsrange, :tstzrange, :tsvector, :uuid, :xml, :timestamptz
+            :serial, :tsrange, :tstzrange, :tsvector, :uuid, :xml, :timestamptz, :enum
         end
       end
 

activerecord/lib/active_record/connection_adapters/postgresql/schema_dumper.rb

@@ -16,6 +16,18 @@ def extensions(stream)
             end
           end
 
+          def types(stream)
+            types = @connection.enum_types
+            if types.any?
+              stream.puts "  # Custom types defined in this database."
+              stream.puts "  # Note that some types may not work with other database engines. Be careful if changing database."
+              types.sort.each do |name, values|
+                stream.puts "  create_enum #{name.inspect}, #{values.split(",").inspect}"
+              end
+              stream.puts
+            end
+          end
+
           def prepare_column_options(column)
             spec = super
             spec[:array] = "true" if column.array?
@@ -26,6 +38,8 @@ def prepare_column_options(column)
               spec = { type: schema_type(column).inspect }.merge!(spec)
             end
 
+            spec[:enum_type] = "\"#{column.sql_type}\"" if column.enum?
+
             spec
           end
 

activerecord/lib/active_record/connection_adapters/postgresql/schema_statements.rb

@@ -542,7 +542,7 @@ def check_constraints(table_name) # :nodoc:
         end
 
         # Maps logical Rails types to PostgreSQL-specific data types.
-        def type_to_sql(type, limit: nil, precision: nil, scale: nil, array: nil, **) # :nodoc:
+        def type_to_sql(type, limit: nil, precision: nil, scale: nil, array: nil, enum_type: nil, **) # :nodoc:
           sql = \
             case type.to_s
             when "binary"
@@ -566,6 +566,10 @@ def type_to_sql(type, limit: nil, precision: nil, scale: nil, array: nil, **) #
               when 5..8; "bigint"
               else raise ArgumentError, "No integer type has byte size #{limit}. Use a numeric with scale 0 instead."
               end
+            when "enum"
+              raise ArgumentError "enum_type is required for enums" if enum_type.nil?
+
+              enum_type
             else
               super
             end

activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb

@@ -164,6 +164,7 @@ def new_client(conn_params)
         money:       { name: "money" },
         interval:    { name: "interval" },
         oid:         { name: "oid" },
+        enum:        {} # special type https://www.postgresql.org/docs/current/datatype-enum.html
       }
 
       OID = PostgreSQL::OID # :nodoc:
@@ -449,6 +450,38 @@ def extensions
         exec_query("SELECT extname FROM pg_extension", "SCHEMA").cast_values
       end
 
+      # Returns a list of defined enum types, and their values.
+      def enum_types
+        query = <<~SQL
+          SELECT
+            type.typname AS name,
+            string_agg(enum.enumlabel, ',') AS value
+          FROM pg_enum AS enum
+          JOIN pg_type AS type
+            ON (type.oid = enum.enumtypid)
+          GROUP BY type.typname;
+        SQL
+        exec_query(query, "SCHEMA").cast_values
+      end
+
+      # Given a name and an array of values, creates an enum type.
+      def create_enum(name, values)
+        sql_values = values.map { |s| "'#{s}'" }.join(", ")
+        query = <<~SQL
+          DO $$
+          BEGIN
+              IF NOT EXISTS (
+                SELECT 1 FROM pg_type t
+                WHERE t.typname = '#{name}'
+              ) THEN
+                  CREATE TYPE \"#{name}\" AS ENUM (#{sql_values});
+              END IF;
+          END
+          $$;
+        SQL
+        exec_query(query)
+      end
+
       # Returns the configured supported identifier length supported by PostgreSQL
       def max_identifier_length
         @max_identifier_length ||= query_value("SHOW max_identifier_length", "SCHEMA").to_i

activerecord/lib/active_record/schema_dumper.rb

@@ -47,6 +47,7 @@ def generate_options(config)
     def dump(stream)
       header(stream)
       extensions(stream)
+      types(stream)
       tables(stream)
       trailer(stream)
       stream
@@ -99,6 +100,10 @@ def trailer(stream)
       def extensions(stream)
       end
 
+      # (enum) types are only supported by PostgreSQL
+      def types(stream)
+      end
+
       def tables(stream)
         sorted_tables = @connection.tables.sort
 
@@ -154,6 +159,7 @@ def table(table, stream)
           columns.each do |column|
             raise StandardError, "Unknown type '#{column.sql_type}' for column '#{column.name}'" unless @connection.valid_type?(column.type)
             next if column.name == pk
+
             type, colspec = column_spec(column)
             if type.is_a?(Symbol)
               tbl.print "    t.#{type} #{column.name.inspect}"

activerecord/test/cases/adapters/postgresql/enum_test.rb

@@ -2,20 +2,27 @@
 
 require "cases/helper"
 require "support/connection_helper"
+require "support/schema_dumping_helper"
 
 class PostgresqlEnumTest < ActiveRecord::PostgreSQLTestCase
   include ConnectionHelper
+  include SchemaDumpingHelper
 
   class PostgresqlEnum < ActiveRecord::Base
     self.table_name = "postgresql_enums"
+
+    enum current_mood: {
+      sad: "sad",
+      okay: "ok", # different spelling
+      happy: "happy",
+      aliased_field: "happy"
+    }, _prefix: true
   end
 
   def setup
     @connection = ActiveRecord::Base.connection
     @connection.transaction do
-      @connection.execute <<~SQL
-        CREATE TYPE mood AS ENUM ('sad', 'ok', 'happy');
-      SQL
+      @connection.create_enum("mood", ["sad", "ok", "happy"])
       @connection.create_table("postgresql_enums") do |t|
         t.column :current_mood, :mood
       end
@@ -62,10 +69,9 @@ def test_enum_mapping
   def test_invalid_enum_update
     @connection.execute "INSERT INTO postgresql_enums VALUES (1, 'sad');"
     enum = PostgresqlEnum.first
-    enum.current_mood = "angry"
 
-    assert_raise ActiveRecord::StatementInvalid do
-      enum.save
+    assert_raise ArgumentError do
+      enum.current_mood = "angry"
     end
   end
 
@@ -90,4 +96,47 @@ def test_assigning_enum_to_nil
     assert model.save
     assert_nil model.reload.current_mood
   end
+
+  def test_schema_dump
+    @connection.add_column "postgresql_enums", "good_mood", :mood, default: "happy", null: false
+
+    output = dump_table_schema("postgresql_enums")
+
+    assert output.include?("# Note that some types may not work with other database engines. Be careful if changing database."), output
+
+    assert output.include?('create_enum "mood", ["sad", "ok", "happy"]'), output
+
+    assert output.include?('t.enum "current_mood", enum_type: "mood"'), output
+    assert output.include?('t.enum "good_mood", default: "happy", null: false, enum_type: "mood"'), output
+  end
+
+  def test_schema_load
+    original, $stdout = $stdout, StringIO.new
+
+    ActiveRecord::Schema.define do
+      create_enum :color, ["blue", "green"]
+
+      change_table :postgresql_enums do |t|
+        t.enum :best_color, enum_type: "color", default: "blue", null: false
+      end
+    end
+
+    assert @connection.column_exists?(:postgresql_enums, :best_color, sql_type: "color", default: "blue", null: false)
+  ensure
+    $stdout = original
+  end
+
+  def test_works_with_activerecord_enum
+    model = PostgresqlEnum.create!
+    model.current_mood_okay!
+
+    model = PostgresqlEnum.find(model.id)
+    assert_equal "okay", model.current_mood
+
+    model.current_mood = "happy"
+    model.save!
+
+    model = PostgresqlEnum.find(model.id)
+    assert model.current_mood_happy?
+  end
 end