Skip to content

Commit

Permalink
Add enum values option to PostgreSQL enum columns
Browse files Browse the repository at this point in the history
This is done so that you can specify the values of an enum
inline when defining a column. To prevent confusion between
enum type naming, an error would raise if an existing enum type is
being created with different values. If an enum_type is not
specified, the enum type will default to the column name.
Schemas will include the values of the enum type in eachcolumn definition.

    def up
      create_table :cats do |t|
        t.enum :current_mood, enum_type: mood, values: [happy, sad], default: happy, null: false
      end
    end
  • Loading branch information
jenshenny committed Jan 29, 2025
1 parent f132ba6 commit 7a80da9
Show file tree
Hide file tree
Showing 9 changed files with 155 additions and 66 deletions.
17 changes: 17 additions & 0 deletions activerecord/CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -1,3 +1,20 @@
* Add enum values option to PostgreSQL enum columns.

It is now possible to specify the values of an enum type column when
defining a column. If an `enum_type` is not specified, the enum type
name will default to the column name. Schemas will include the values
of the enum type in each column definition.

```ruby
def up
create_table :cats do |t|
t.enum :current_mood, enum_type: "mood", values: ["happy", "sad"], default: "happy", null: false
end
end
```

*Jenny Shen*

* Introduce a before-fork hook in `ActiveSupport::Testing::Parallelization` to clear existing
connections, to avoid fork-safety issues with the mysql2 adapter.

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ def visit_AlterTable(o)
end

def visit_ColumnDefinition(o)
o.sql_type = type_to_sql(o.type, **o.options)
o.sql_type = type_to_sql(o.type, column_name: o.name, **o.options)
column_sql = +"#{quote_column_name(o.name)} #{o.sql_type}"
add_column_options!(column_sql, column_options(o)) unless o.type == :primary_key
column_sql
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,15 +5,30 @@ module ConnectionAdapters
module PostgreSQL
class SchemaCreation < SchemaCreation # :nodoc:
private
delegate :quoted_include_columns_for_index, to: :@conn
delegate :quoted_include_columns_for_index, :create_enum, to: :@conn

def visit_TableDefinition(o)
create_enums(o.columns)
super
end

def visit_AlterTable(o)
create_enums(o.adds.map(&:column))
sql = super
sql << o.constraint_validations.map { |fk| visit_ValidateConstraint fk }.join(" ")
sql << o.exclusion_constraint_adds.map { |con| visit_AddExclusionConstraint con }.join(" ")
sql << o.unique_constraint_adds.map { |con| visit_AddUniqueConstraint con }.join(" ")
end

def create_enums(columns)
enums_to_create = columns.select { |c| c.type == :enum && c.options[:values] }

enums_to_create.each do |c|
enum_type = c.options[:enum_type] || c.name
create_enum(enum_type, c.options[:values])
end
end

def visit_AddForeignKey(o)
super.dup.tap do |sql|
sql << " NOT VALID" unless o.validate?
Expand Down Expand Up @@ -77,7 +92,7 @@ def visit_AddUniqueConstraint(o)

def visit_ChangeColumnDefinition(o)
column = o.column
column.sql_type = type_to_sql(column.type, **column.options)
column.sql_type = type_to_sql(column.type, column_name: column.name, **column.options)
quoted_column_name = quote_column_name(o.name)

change_column_sql = +"ALTER COLUMN #{quoted_column_name} TYPE #{column.sql_type}"
Expand All @@ -91,7 +106,7 @@ def visit_ChangeColumnDefinition(o)
if options[:using]
change_column_sql << " USING #{options[:using]}"
elsif options[:cast_as]
cast_as_type = type_to_sql(options[:cast_as], **options)
cast_as_type = type_to_sql(options[:cast_as], column_name: column.name, **options)
change_column_sql << " USING CAST(#{quoted_column_name} AS #{cast_as_type})"
end

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -282,7 +282,7 @@ def new_column_definition(name, type, **options) # :nodoc:

private
def valid_column_definition_options
super + [:array, :using, :cast_as, :as, :type, :enum_type, :stored]
super + [:array, :using, :cast_as, :as, :type, :enum_type, :stored, :values]
end

def aliased_types(name, fallback)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,11 @@ def prepare_column_options(column)
spec = { type: schema_type(column).inspect }.merge!(spec)
end

spec[:enum_type] = column.sql_type.inspect if column.enum?
if column.enum?
spec[:enum_type] = column.sql_type.inspect unless column.name == column.sql_type
_name, values = @connection.enum_types.find { |name, _values| name == column.sql_type }
spec[:values] = values.inspect
end

spec
end
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -848,7 +848,7 @@ def remove_unique_constraint(table_name, column_name = nil, **options)
end

# Maps logical Rails types to PostgreSQL-specific data types.
def type_to_sql(type, limit: nil, precision: nil, scale: nil, array: nil, enum_type: nil, **) # :nodoc:
def type_to_sql(type, column_name: nil, limit: nil, precision: nil, scale: nil, array: nil, enum_type: nil, values: nil, **) # :nodoc:
sql = \
case type.to_s
when "binary"
Expand All @@ -873,9 +873,9 @@ def type_to_sql(type, limit: nil, precision: nil, scale: nil, array: nil, enum_t
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?
raise ArgumentError, "enum_type or values is required for enums" if enum_type.nil? && values.nil?

enum_type
enum_type || column_name
else
super
end
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -550,6 +550,11 @@ def create_enum(name, values, **options)
JOIN pg_namespace n ON t.typnamespace = n.oid
WHERE t.typname = #{scope[:name]}
AND n.nspname = #{scope[:schema]}
AND (
SELECT array_agg(e.enumlabel ORDER BY e.enumsortorder)
FROM pg_enum e
WHERE e.enumtypid = t.oid
) = ARRAY[#{sql_values}]::name[]
) THEN
CREATE TYPE #{quote_table_name(name)} AS ENUM (#{sql_values});
END IF;
Expand Down
160 changes: 104 additions & 56 deletions activerecord/test/cases/adapters/postgresql/enum_test.rb
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
require "support/connection_helper"
require "support/schema_dumping_helper"

class PostgresqlEnumTest < ActiveRecord::PostgreSQLTestCase
module PostgresqlEnumSharedTestCases
include ConnectionHelper
include SchemaDumpingHelper

Expand All @@ -19,35 +19,28 @@ class PostgresqlEnum < ActiveRecord::Base
}, prefix: true
end

def setup
@connection = ActiveRecord::Base.lease_connection
@connection.transaction do
@connection.create_enum("mood", ["sad", "ok", "happy"])
@connection.create_table("postgresql_enums") do |t|
t.column :current_mood, :mood
end
end
end

teardown do
def teardown
reset_connection
@connection.drop_table "postgresql_enums", if_exists: true
@connection.drop_enum "mood", if_exists: true
@connection.enum_types.each do |enum_type|
@connection.drop_enum enum_type
end
reset_connection
end

def test_column
PostgresqlEnum.reset_column_information
column = PostgresqlEnum.columns_hash["current_mood"]
assert_equal :enum, column.type
assert_equal "mood", column.sql_type
assert_equal @enum_type, column.sql_type
assert_not_predicate column, :array?

type = PostgresqlEnum.type_for_attribute("current_mood")
assert_not_predicate type, :binary?
end

def test_enum_defaults
@connection.add_column "postgresql_enums", "good_mood", :mood, default: "happy"
@connection.add_column "postgresql_enums", "good_mood", @enum_type, default: "happy"
PostgresqlEnum.reset_column_information

assert_equal "happy", PostgresqlEnum.column_defaults["good_mood"]
Expand Down Expand Up @@ -98,80 +91,55 @@ def test_assigning_enum_to_nil
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_includes output, "# Note that some types may not work with other database engines. Be careful if changing database."

assert_includes output, 'create_enum "mood", ["sad", "ok", "happy"]'

assert_includes output, 't.enum "current_mood", enum_type: "mood"'
assert_includes output, 't.enum "good_mood", default: "happy", null: false, enum_type: "mood"'
end

def test_schema_dump_renamed_enum
@connection.rename_enum :mood, :feeling
@connection.rename_enum @enum_type, :feeling

output = dump_table_schema("postgresql_enums")

assert_includes output, 'create_enum "feeling", ["sad", "ok", "happy"]'

assert_includes output, 't.enum "current_mood", enum_type: "feeling"'
assert_includes output, 't.enum "current_mood", enum_type: "feeling", values: ["sad", "ok", "happy"]'
end

def test_schema_dump_renamed_enum_with_to_option
@connection.rename_enum :mood, to: :feeling
@connection.rename_enum @enum_type, to: :feeling

output = dump_table_schema("postgresql_enums")

assert_includes output, 'create_enum "feeling", ["sad", "ok", "happy"]'

assert_includes output, 't.enum "current_mood", enum_type: "feeling"'
assert_includes output, 't.enum "current_mood", enum_type: "feeling", values: ["sad", "ok", "happy"]'
end

def test_schema_dump_added_enum_value
skip("Adding enum values can not be run in a transaction") if @connection.database_version < 10_00_00

@connection.add_enum_value :mood, :angry, before: :ok
@connection.add_enum_value :mood, :nervous, after: :ok
@connection.add_enum_value :mood, :glad
@connection.add_enum_value @enum_type, :angry, before: :ok
@connection.add_enum_value @enum_type, :nervous, after: :ok
@connection.add_enum_value @enum_type, :glad

assert_nothing_raised do
@connection.add_enum_value :mood, :glad, if_not_exists: true
@connection.add_enum_value :mood, :curious, if_not_exists: true
@connection.add_enum_value @enum_type, :glad, if_not_exists: true
@connection.add_enum_value @enum_type, :curious, if_not_exists: true
end

output = dump_table_schema("postgresql_enums")

assert_includes output, 'create_enum "mood", ["sad", "angry", "ok", "nervous", "happy", "glad", "curious"]'
assert_includes output, "create_enum \"#{@enum_type}\", [\"sad\", \"angry\", \"ok\", \"nervous\", \"happy\", \"glad\", \"curious\"]"
end

def test_schema_dump_renamed_enum_value
skip("Renaming enum values is only supported in PostgreSQL 10 or later") if @connection.database_version < 10_00_00

@connection.rename_enum_value :mood, from: :ok, to: :okay
@connection.rename_enum_value @enum_type, from: :ok, to: :okay

output = dump_table_schema("postgresql_enums")

assert_includes output, 'create_enum "mood", ["sad", "okay", "happy"]'
assert_includes output, "create_enum \"#{@enum_type}\", [\"sad\", \"okay\", \"happy\"]"
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
def test_create_enum_type_with_different_values
assert_raises ActiveRecord::StatementInvalid do
@connection.create_enum @enum_type, ["mad", "glad"]
end

assert @connection.column_exists?(:postgresql_enums, :best_color, sql_type: "color", default: "blue", null: false)
ensure
$stdout = original
end

def test_drop_enum
Expand Down Expand Up @@ -248,7 +216,7 @@ def test_schema_dump_scoped_to_schemas

output = dump_table_schema("postgresql_enums_in_test_schema")

assert_includes output, 'create_enum "public.mood", ["sad", "ok", "happy"]'
assert_includes output, "create_enum \"public.#{@enum_type}\", [\"sad\", \"ok\", \"happy\"]"
assert_includes output, 'create_enum "mood_in_test_schema", ["sad", "ok", "happy"]'
assert_includes output, 't.enum "current_mood", enum_type: "mood_in_test_schema"'
assert_not_includes output, 'create_enum "other_schema.mood_in_other_schema"'
Expand Down Expand Up @@ -291,3 +259,83 @@ def with_test_schema(name, drop: true)
@connection.schema_cache.clear!
end
end

class PostgresqlEnumTest < ActiveRecord::PostgreSQLTestCase
include PostgresqlEnumSharedTestCases

def setup
@connection = ActiveRecord::Base.lease_connection
@connection.transaction do
@connection.create_enum("mood", ["sad", "ok", "happy"])
@connection.create_table("postgresql_enums") do |t|
t.column :current_mood, :mood
end
end

@enum_type = "mood"
end

def test_schema_dump
@connection.add_column "postgresql_enums", "good_mood", @enum_type, default: "happy", null: false
@connection.add_column "postgresql_enums", "bad_mood", :enum, values: ["angry", "mad", "sad"], default: "sad", null: false
@connection.add_column "postgresql_enums", "party_mood", :enum, values: ["excited", "happy"], default: "happy", enum_type: "fun_mood", null: false

output = dump_table_schema("postgresql_enums")

assert_includes output, "# Note that some types may not work with other database engines. Be careful if changing database."

assert_includes output, "create_enum \"#{@enum_type}\", [\"sad\", \"ok\", \"happy\"]"
assert_includes output, 'create_enum "bad_mood", ["angry", "mad", "sad"]'
assert_includes output, 'create_enum "fun_mood", ["excited", "happy"]'

assert_includes output, "t.enum \"good_mood\", default: \"happy\", null: false, enum_type: \"#{@enum_type}\", values: [\"sad\", \"ok\", \"happy\"]"
assert_includes output, 't.enum "bad_mood", default: "sad", null: false, values: ["angry", "mad", "sad"]'
assert_includes output, 't.enum "party_mood", default: "happy", null: false, enum_type: "fun_mood", values: ["excited", "happy"]'
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
end

class PostgresqlEnumWithValuesAndEnumTypeTest < ActiveRecord::PostgreSQLTestCase
include PostgresqlEnumSharedTestCases

def setup
@connection = ActiveRecord::Base.lease_connection
@connection.transaction do
@connection.create_table("postgresql_enums") do |t|
t.enum :current_mood, values: ["sad", "ok", "happy"], enum_type: "mood"
end
end

@enum_type = "mood"
end
end

class PostgresqlEnumWithValuesTest < ActiveRecord::PostgreSQLTestCase
include PostgresqlEnumSharedTestCases

def setup
@connection = ActiveRecord::Base.lease_connection
@connection.transaction do
@connection.create_table("postgresql_enums") do |t|
t.enum :current_mood, values: ["sad", "ok", "happy"]
end
end

@enum_type = "current_mood"
end
end
2 changes: 1 addition & 1 deletion activerecord/test/cases/migration/invalid_options_test.rb
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ def invalid_add_column_option_exception_message(key)
if current_adapter?(:Mysql2Adapter, :TrilogyAdapter)
default_keys.concat([":auto_increment", ":charset", ":as", ":size", ":unsigned", ":first", ":after", ":type", ":stored"])
elsif current_adapter?(:PostgreSQLAdapter)
default_keys.concat([":array", ":using", ":cast_as", ":as", ":type", ":enum_type", ":stored"])
default_keys.concat([":array", ":using", ":cast_as", ":as", ":type", ":enum_type", ":stored", ":values"])
elsif current_adapter?(:SQLite3Adapter)
default_keys.concat([":as", ":type", ":stored"])
end
Expand Down

0 comments on commit 7a80da9

Please sign in to comment.