from enum import auto, StrEnum
import sqlalchemy as sa
from sqlalchemy_utils import ChoiceType
from sqlmodel import Field, Relationship, SQLModel


class Status(StrEnum):
    STARTED = auto()
    ACCEPTED = auto()
    REJECTED = auto()
    REFUSED = auto()


class Issue(SQLModel, table=True):
    ...
    status: Status = Field(
        sa_column=sa.Column(ChoiceType(Status), nullable=False),
        default=Status.STARTED,
    )

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.