PostgreSQL: support custom enum types

Something I've seen on many Rails + PostgreSQL apps is that the only reason `structure.sql` is used instead of `schema.rb` is to support custom [enum types](https://www.postgresql.org/docs/current/datatype-enum.html). Enum types are great for type integrity and they work well with [`ActiveRecord::Enum`](https://api.rubyonrails.org/classes/ActiveRecord/Enum.html). And schema.rb is much easier to use than structure.sql. It would be great if more PostgreSQL projects could use schema.rb. To enable that, this PR adds native support for PostgreSQL enums in schema.rb.

In migrations, you can now use `create_enum` to add a new enum type, and `t.enum` to add a column:

```ruby
def up
  # note that enums cannot be dropped
  create_enum :mood, ["happy", "sad"]

  change_table :cats do |t|
    t.enum :current_mood, enum_type: "mood", default: "happy", null: false
  end
end
```

The enum definitions and enum columns will be presented in schema.rb, so you can load them into a test database and they'll work correctly.

-------------------------------

It's worth noting again that this is *not* compatible with other database engines. So this will not work with `rails db:system:change` (or the equivalent manual process). My assumption is that this is a fairly unlikely thing to happen midway through a project - as opposed to when the app is first spun up - so it's safe to assume that once you've dug into using enums you probably aren't going to switch databases on a whim.

For what it's worth, the MySQL adapter also supports enums (https://github.com/rails/rails/blob/main/activerecord/test/cases/adapters/mysql2/enum_test.rb) but the syntax doesn't look easy to translate (I also don't have much MySQL experience so don't really want to wade into this).

Author: ghiculescu

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