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.