Simplifying PostgreSQL enum migrations with SQLAlchemy and alembic-enums
PostgreSQL enum types migration is not straightforward, and Alembic, the SQLAlchemy migration tool, does not properly support them. I have created a Python package that simplifies the migration of PostgreSQL enums by wrapping a sequence of SQL commands with a Pythonic interface that is easier to understand.
If you are aware of the problem and looking for a solution, you can find the package on PyPI at alembic-enums . The package README provides a step-by-step usage guide and an API overview.
If you want to learn more about the problem, read on.
In PostgreSQL, enums are user-defined types that do not allow the deletion or modification of their fields. To work around this issue, people often create a new enum and reassign fields to a new array.
The StackOverflow question python - How to use Enum with SQLAlchemy and Alembic? - Stack Overflow offers different but similar solutions to this problem.
Additionally, Kamil Kucharski provides a clear and straightforward explanation of the challenge in his blog post Upgrading PostgreSQL’s Enum type with SQLAlchemy using Alembic migration .
The TLDR version of the challenge is as follows (taken from the aforementioned blog post).
Adding an enum element is straightforward:
ALTER TYPE status ADD VALUE 'REJECTED';
Deleting it, however, requires more steps.
ALTER TYPE status RENAME TO status_old;
CREATE TYPE status AS ENUM('STARTED', 'ACCEPTED');
ALTER TABLE transactions
ALTER COLUMN status TYPE status
USING status::text::status;
DROP TYPE status_old;
Things become even more complicated when you want to rename a field. If you want to rename a field from REJECTED to REFUSED, you need to create a temporary ENUM with both fields to update the rows.
ALTER TYPE status RENAME TO status_old
CREATE TYPE status_temp AS
ENUM('STARTED', 'ACCEPTED', 'REJECTED', 'REFUSED');
ALTER TABLE transactions
ALTER COLUMN status TYPE status_temp
USING status::text::status_temp;
UPDATE transactions SET status = 'REFUSED'
WHERE status = 'REJECTED';
CREATE TYPE status AS
ENUM('STARTED', 'ACCEPTED', 'REJECTED');
ALTER TABLE transactions
ALTER COLUMN status TYPE status
USING status::text::status;
DROP TYPE status_old;
DROP TYPE status_temp;
These steps require a lot of mental effort and are too many to be practical. To make it easier and allow myself to iterate faster on my codebase, I wrapped these SQL commands with a Pythonic interface that can be called from an Alembic migration script.
The SQL script above would be transformed into the following Alembic script:
from alembic import op
from alembic_enums import EnumMigration, Column
column = Column("transactions", "status")
enum_migration = EnumMigration(
op=op,
enum_name="status",
old_options=["STARTED", "ACCEPTED", "REJECTED"],
new_options=["STARTED", "ACCEPTED", "REFUSED"],
columns=[column],
)
def upgrade():
with enum_migration.upgrade_ctx():
enum_migration.update_value(column, "REJECTED", "REFUSED")
def downgrade():
with enum_migration.downgrade_ctx():
enum_migration.update_value(column, "REFUSED", "REJECTED")
The package does not automatically detect enum changes or generate migration code, but it provides a helper class to run enum migrations in Alembic migration scripts.
If you find the package helpful, please consider starring the GitHub repository at imankulov/alembic-enums . I would appreciate it if you showed your appreciation for my work.