Automatically migrate enums in alembic

I don't know if people have come up with a good acronym for Python services that compete with MERN or LAMP, but if they have then SQLAlchemy and Alembic are almost certainly included. SQLAlchemy (recently in version 2.0) makes it easy to define ORM schemas for database objects and Alembic keeps everything updated with automatically generated migration files.

If you're using this stack then you probably know the pain that code enums introduce. Declaring an enum requirement in a model is pretty straightforward:

from sqlalchemy import Enum as SqlEnum
enum_field = Column(SqlEnum(MyEnum))

And Alembic will even pick up on the new enum creation:

def upgrade():
    op.add_column('my_table', sa.Column('enum_field', sa.Enum("A", "B", name='myenum'), nullable=True))

So far, so good. Unfortunately when you actually change this enum (as you know does happen) you're out of luck. Alembic ignores this enum value change even when it's outdated from the current database value. So this change:

class MyEnum(Enum):
    A = "A"
    B = "B"

------>

    class MyEnum(Enum):
        A = "A"
        B = "B"
        C = "C"

Creates no diff:

def upgrade():
    pass

And will result in a database error if you actually try to use it.

(builtins.LookupError) C is not among the defined enum values. Enum name: myenum. Possible values: A, Bn[SQL: INSERT INTO invitations...

Spoiler alert: We probably want to use it.

I stumbled upon alembic-autogenerate-enums, which is a neat approach to solve this problem. This lets you make changes to an enum value that's already inserted into the database and have alembic auto-generate the value migration commands:

poetry run alembic revision --autogenerate -m "add new enum value"

This will now result in the following:

def upgrade():
    op.sync_enum_values('public', 'myenum', ['A', 'B'], ['A', 'B', 'C'], [('simple_model', 'enum_field')], False)

def downgrade():
    op.sync_enum_values('public', 'myenum', ['A', 'B', 'C'], ['A', 'B'], [('simple_model', 'enum_field')], True)

Running the upgrade will add value C to the database enum specification without affecting previous values A & B. Downgrades to the previous alembic revision will strip this C value (assuming no existing database objects are using it) and restore state A & B.

The version 0.2.0 that I started using only had support for forward migration of enums (appending new values to the overall definition) but couldn't downgrade() to previous enum revisions. You usually only want to downgrade enum definitions in very limited circumstances, but still, it seemed like a good thing to add for locally testing schema changes. My PR hasn't yet hit pypi but you can grab the latest from master.