How to migrate new Python Enum members in SQLAlchemy for PostgreSQL in semi-automated manner?

Question

I am adding new enum members in an SQLAlchemy migration using Alembic migration tool.

SQLAlchemy uses Python native Enum as:

from enum import Enum

class MyEnum(Enum):
    # Key and value names different to clarify they are separate 
    # concepts in Python
    old = "OLD_VALUE"
    new1 = "NEW1_VALUE"
    new2 = "NEW2_VALUE"

And then:

class MyFantasticModel(Base):

    __tablename__ = "fantasy"

    enum_column = sa.Column(sa.Enum(MyEnum), nullable=False, index=True)

I am aware of the fact that PostgreSQL is difficult what comes to enum migrations, and similar questions have existing low-quality answers. But I still want to do it half automated manner.

Using Python 3.8.

Mikko Ohtamaa · Accepted Answer · 2021-11-28T14:27:55.693

You need to manually declare new enum values that are added, as it would be difficult to detect (though not impossible, will probably retrofit this answer later). Below is an example migration that reflects PSQL keys back from your Enum class.

There is also a downgrade functionality, but it will only work if none of Enum values are being used and is somewhat dangerous.

# revision identifiers, used by Alembic.
from mymodels import MyEnum

revision = ...
down_revision = ....
branch_labels = None
depends_on = None


# By default, SQLAlchemy autogenerates PSQL type
# that is the Enum class name lowercased.
# Because Python enums are created using Python metaclasses,
# you cannot directly read the class name back from the class object (it would return `EnumMeta`).
# You need go through method resolution order table to get the real class instance.
# https://stackoverflow.com/a/54014128/315168
enum_name = MyEnum.mro()[0].__name__.lower()

# PostgreSQL does not have a concept of enum keys and values, only values.
# SQLAlchemy maintains enums in the PostgreSQL by the Python enum key (not value)
# Keys to be added:
enum_keys_to_add = [
    MyEnum.new1.name,
    MyEnum.new2.name,
]

def upgrade():
    for v in enum_keys_to_add:
        # Use ALTER TYPE.
        # See more information https://www.postgresql.org/docs/9.1/sql-altertype.html
        op.execute(f"ALTER TYPE {enum_name} ADD VALUE '{v}'")
    # Note that ALTER TYPE is not effective within the same transaction https://medium.com/makimo-tech-blog/upgrading-postgresqls-enum-type-with-sqlalchemy-using-alembic-migration-881af1e30abe

def downgrade():
    # https://stackoverflow.com/a/39878233/315168
    for v in enum_keys_to_add:
        sql = f"""DELETE FROM pg_enum
            WHERE enumlabel = '{v}'
            AND enumtypid = (
              SELECT oid FROM pg_type WHERE typname = '{enum_name}'
            )"""
        op.execute(sql)

You can also manually inspect the content of your PSQL enum with the following SELECT:

select enum_range(null::myenum);

                                             enum_range                                              
-----------------------------------------------------------------------------------------------------
 {old,new1,new2}

How to migrate new Python Enum members in SQLAlchemy for PostgreSQL in semi-automated manner?

1 Answers1