Altering an Enum field using Alembic

Question

How can I add an element to an Enum field in an alembic migration when using a version of PostgreSQL older than 9.1 (which adds the ALTER TYPE for enums)? This SO question e

Answer 1

Since I got conversion errors and problems with default values, I wrote an even more generalised answer based on the accepted one:

def replace_enum_values(
        name: str,
        old: [str],
        new: [str],
        modify: [(str, str, str)]
):
    """
    Replaces an enum's list of values.

    Args:
        name: Name of the enum
        new: New list of values
        old: Old list of values
        modify: List of tuples of table name
        and column to modify (which actively use the enum).
        Assumes each column has a default val.
    """
    connection = op.get_bind()

    tmp_name = "{}_tmp".format(name)

    # Rename old type
    op.execute(
        "ALTER TYPE {} RENAME TO {};"
        .format(name, tmp_name)
    )

    # Create new type
    lsl = sa.Enum(*new, name=name)
    lsl.create(connection)

    # Replace all usages
    for (table, column) in modify:
        # Get default to re-set later
        default_typed = connection.execute(
            "SELECT column_default "
            "FROM information_schema.columns "
            "WHERE table_name='{table}' "
            "AND column_name='{column}';"
            .format(table=table, column=column)
        ).first()[0]  # type: str

        # Is bracketed already
        default = default_typed[:default_typed.index("::")]

        # Set all now invalid values to default
        connection.execute(
            "UPDATE {table} "
            "SET {column}={default} "
            "WHERE {column} NOT IN {allowed};"
            .format(
                table=table,
                column=column,
                # Invalid: What isn't contained in both new and old
                # Can't just remove what's not in new because we get
                # a type error
                allowed=tuple(set(old).intersection(set(new))),
                default=default
            )
        )

        op.execute(
            "ALTER TABLE {table} "
            # Default needs to be dropped first
            "ALTER COLUMN {column} DROP DEFAULT,"
            # Replace the tpye
            "ALTER COLUMN {column} TYPE {enum_name} USING {column}::text::{enum_name},"
            # Reset default
            "ALTER COLUMN {column} SET DEFAULT {default};"
            .format(
                table=table,
                column=column,
                enum_name=name,
                default=default
            )
        )

    # Remove old type
    op.execute("DROP TYPE {};".format(tmp_name))

This can be called from upgrade / downgrade as such:

replace_enum_values(
    name='enum_name',
    new=["A", "B"],
    old=["A", "C"],
    modify=[('some_table', 'some_column')]
)

All invalidated values will be set to server_default.