Migrate PostgreSQL non native enum in alembic

migrations.py

import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects import postgresql

# revision identifiers, used by Alembic.
revision = "abcd"
down_revision = "dcba"
branch_labels = None
depends_on = None


def update_enum_sql(table, column, old_value, new_value):
    if new_value is not None:
        return f"UPDATE {table} SET {column} = '{new_value}' where {column} = '{old_value}'"
    else:
        return f"DELETE FROM {table} where {column} = '{old_value}'"


def update_enum_sqls(table, column, olds_to_remove, news_to_add):
    if len(olds_to_remove) != len(news_to_add):
        raise NotImplementedError
    return [
        update_enum_sql(table, column, old, new)
        for old, new in zip(olds_to_remove, news_to_add)
    ]


def update_enum(
    table, column, enum_class_name, target_values, olds_to_remove, news_to_add
):
    op.drop_constraint(f"ck_{table}_{enum_class_name}", table)

    for sql in update_enum_sqls(table, column, olds_to_remove, news_to_add):
        op.execute(sql)

    op.create_check_constraint(
        enum_class_name, table, sa.sql.column(column).in_(target_values)
    )


def upgrade():
    # rename enum
    update_enum(
        "my_table",
        "my_enum",
        "myenumclassname",
        ["NEW", "ENUM", "VALUES"],
        ["OLD"],
        ["NEW"],
    )

    # add enum
    update_enum(
        "my_table",
        "my_enum",
        "myenumclassname",
        ["NEW", "ENUM", "VALUES"],
        [],
        [],
    )


def downgrade():
    # remove enum
    update_enum(
        "my_table",
        "my_enum",
        "myenumclassname",
        ["ENUM", "VALUES"],
        ["NEW"],
        [None],  # this will delete rows with "NEW"
    )

    # edit enum
    update_enum(
        "my_table",
        "my_enum",
        "myenumclassname",
        ["OLD", "ENUM", "VALUES"],
        ["NEW"],
        ["OLD"],
    )