Altering an Enum field using Alembic

In straight SQL, this would work for Postgres, if the order of the things in your enum doesn’t need to be exactly as above:

ALTER TYPE status ADD value 'output_limit_exceeded' after 'timed_out'; 

As of Postgres 9.1 adding a new value to an enum can be done with the ALTER TYPE statement. This is complicated by the fact that it cannot be done in a transaction. However this can be worked around by committing alembic’s transaction see here.

I actually had problems using the older, more verbose, solution because Postgres could not automatically convert the default for the column.

I used a bit simpler approach with less steps than the accepted answer, which I based this on. In this example I will pretend the enum in question is called ‘status_enum’, because in the accepted answer the use of ‘status’ for both the column and enum confused me.

from alembic import op 
import sqlalchemy as sa

name = 'status_enum'
tmp_name = 'tmp_' + name

old_options = ('nonexistent_executable', 'signal', 'success', 'timed_out')
new_options = sorted(old_options + ('output_limit_exceeded',))

new_type = sa.Enum(*new_options, name=name)
old_type = sa.Enum(*old_options, name=name)

tcr = sa.sql.table('testcaseresult',
                   sa.Column('status', new_type, nullable=False))

def upgrade():
    op.execute('ALTER TYPE ' + name + ' RENAME TO ' + tmp_name)

    new_type.create(op.get_bind())
    op.execute('ALTER TABLE testcaseresult ALTER COLUMN status ' +
               'TYPE ' + name + ' USING status::text::' + name)
    op.execute('DROP TYPE ' + tmp_name)


def downgrade():
    # Convert 'output_limit_exceeded' status into 'timed_out'                                                                                                                      
    op.execute(tcr.update().where(tcr.c.status=='output_limit_exceeded')
               .values(status='timed_out'))

    op.execute('ALTER TYPE ' + name + ' RENAME TO ' + tmp_name)

    old_type.create(op.get_bind())
    op.execute('ALTER TABLE testcaseresult ALTER COLUMN status ' +
               'TYPE ' + name + ' USING status::text::' + name)
    op.execute('DROP TYPE ' + tmp_name)

I had the same issue trying to migrate a column type to another. I use the following requirements:

Alembic==0.9.4
SQLAlchemy==1.1.12 

You can provide the argument postgresql_using as a kwarg of alembic.op.alter_column.

from alembic import op
import sqlalchemy as types

op.alter_column(
    table_name='my_table',
    column_name='my_column',
    type_=types.NewType,
    # allows to use postgresql USING
    postgresql_using="my_column::PostgesEquivalentOfNewType",
)

I hope it can help.

This runs without problems:

from alembic import op

def upgrade():
    op.execute("COMMIT")
    op.execute("ALTER TYPE enum_type ADD VALUE 'new_value'")

def downgrade():
    ...

Reference

I needed to move data while migrating types, including deleting some old types, so I figured I’d write up a more general way of doing this based on the (awesome) accepted answer (https://stackoverflow.com/a/14845740/629272). Hopefully this helps someone else in the same boat!

# This migration will move data from one column to two others based on the type
# for a given row, and modify the type of each row.
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

revision = '000000000001'
down_revision = '000000000000'
branch_labels = None
depends_on = None

# This set of options makes up the old type.
example_types_old = (
    'EXAMPLE_A',
    'EXAMPLE_B',
    'EXAMPLE_C',
)
example_type_enum_old = postgresql.ENUM(*example_types_old, name='exampletype')

# This set of options makes up the new type.
example_types_new = (
    'EXAMPLE_C',
    'EXAMPLE_D',
    'EXAMPLE_E',
)
example_type_enum_new = postgresql.ENUM(*example_types_new, name='exampletype')

# This set of options includes everything from the old and new types.
example_types_tmp = set(example_types_old + example_types_new)
example_type_enum_tmp = postgresql.ENUM(*example_types_tmp, name='_exampletype')

# This is a table view from which we can select and update as necessary. This
# only needs to include the relevant columns which are in either the old or new
# version of the table.
examples_view = sa.Table(
    # Use the name of the actual table so it is modified in the upgrade and
    # downgrade.
    'examples',
    sa.MetaData(),
    sa.Column('id', sa.Integer, primary_key=True),
    # Use the _tmp type so all types are usable.
    sa.Column('example_type', example_type_enum_tmp),
    # This is a column from which the data will be migrated, after which the
    # column will be removed.
    sa.Column('example_old_column', sa.Integer),
    # This is a column to which data from the old column will be added if the
    # type is EXAMPLE_A.
    sa.Column('example_new_column_a', sa.Integer),
    # This is a column to which data from the old column will be added if the
    # type is EXAMPLE_B.
    sa.Column('example_new_column_b', sa.Integer),
)


def upgrade():
    connection = op.get_bind()

    # Add the new column to which data will be migrated.
    example_new_column_a = sa.Column(
        'example_new_column_a',
        sa.Integer,
        nullable=True
    )
    op.add_column('examples', example_new_column_a)

    # Add the new column to which data will be migrated.
    example_new_column_b = sa.Column(
        'example_new_column_b',
        sa.Integer,
        nullable=True
    )
    op.add_column('examples', example_new_column_b)

    # Create the temporary enum and change the example_type column to use the
    # temporary enum.
    # The USING statement automatically maps the old enum to the temporary one.
    example_type_enum_tmp.create(connection, checkfirst=False)
    # Change to the temporary type and map from the old type to the temporary
    # one.
    op.execute('''
        ALTER TABLE examples
            ALTER COLUMN example_type
                TYPE _exampletype
                USING example_type::text::_exampletype
    ''')

    # Move data from example_old_column to example_new_column_a and change its
    # type to EXAMPLE_D if the type is EXAMPLE_A.
    connection.execute(
        examples_view.update().where(
            examples_view.c.example_type == 'EXAMPLE_A'
        ).values(
            example_type='EXAMPLE_D',
            example_new_column_a=examples_view.c.example_old_column,
        )
    )

    # Move data from example_old_column to example_new_column_b and change its
    # type to EXAMPLE_E if the type is EXAMPLE_B.
    connection.execute(
        examples_view.update().where(
            examples_view.c.example_type == 'EXAMPLE_B'
        ).values(
            example_type='EXAMPLE_E',
            example_new_column_b=examples_view.c.example_old_column,
        )
    )

    # Move any remaining data from example_old_column to example_new_column_a
    # and keep its type as EXAMPLE_C.
    connection.execute(
        examples_view.update().where(
            examples_view.c.example_type == 'EXAMPLE_C'
        ).values(
            example_type='EXAMPLE_C',
            example_new_column_a=examples_view.c.example_old_column,
        )
    )

    # Delete the old enum now that the data with the old types have been moved.
    example_type_enum_old.drop(connection, checkfirst=False)

    # Create the new enum and change the example_type column to use the new
    # enum.
    # The USING statement automatically maps the temporary enum to the new one.
    example_type_enum_new.create(connection, checkfirst=False)
    op.execute('''
        ALTER TABLE examples
            ALTER COLUMN example_type
                TYPE exampletype
                USING example_type::text::exampletype
    ''')

    # Delete the temporary enum.
    example_type_enum_tmp.drop(connection, checkfirst=False)

    # Remove the old column.
    op.drop_column('examples', 'example_old_column')


# The downgrade just performs the opposite of all the upgrade operations but in
# reverse.
def downgrade():
    connection = op.get_bind()

    example_old_column = sa.Column(
        'example_old_column',
        sa.Integer,
        nullable=True
    )
    op.add_column('examples', example_old_column)

    example_type_enum_tmp.create(connection, checkfirst=False)
    op.execute('''
        ALTER TABLE examples
            ALTER COLUMN example_type
                TYPE _exampletype
                USING example_type::text::_exampletype
    ''')

    connection.execute(
        examples_view.update().where(
            examples_view.c.example_type == 'EXAMPLE_C'
        ).values(
            example_type='EXAMPLE_C',
            example_old_column=examples_view.c.example_new_column_b,
        )
    )

    connection.execute(
        examples_view.update().where(
            examples_view.c.example_type == 'EXAMPLE_E'
        ).values(
            example_type='EXAMPLE_B',
            example_old_column=examples_view.c.example_new_column_b,
        )
    )

    connection.execute(
        examples_view.update().where(
            examples_view.c.example_type == 'EXAMPLE_D'
        ).values(
            example_type='EXAMPLE_A',
            example_old_column=examples_view.c.example_new_column_a,
        )
    )

    example_type_enum_old.create(connection, checkfirst=False)
    op.execute('''
        ALTER TABLE examples
            ALTER COLUMN example_type
                TYPE exampletype
                USING example_type::text::exampletype
    ''')

    example_type_enum_tmp.drop(connection, checkfirst=False)

    op.drop_column('examples', 'example_new_column_b')
    op.drop_column('examples', 'example_new_column_a')

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.

Found another handy method

op.execute('ALTER TYPE enum_type ADD VALUE new_value')
op.execute('ALTER TYPE enum_type ADD VALUE new_value BEFORE old_value')
op.execute('ALTER TYPE enum_type ADD VALUE new_value AFTER old_value')

This solution is easy to understand and works really well for both upgrade and downgrade. I’ve written this answer in more detailed manner.

Let’s say our enum_type looks like this:

enum_type = ('some_value_1', 'some_value_2')

I want to alter enum_type by adding a new enum, so that it becomes like this:

enum_type = ('some_value_1', 'some_value_2', 'new_value')

This can be done in this way:

from alembic import op


def upgrade():
    op.execute("COMMIT")
    op.execute("ALTER TYPE enum_type ADD VALUE 'new_value'")


def downgrade():
    # Drop 'new_value' from enum_type
    op.execute("ALTER TYPE enum_type RENAME TO enum_type_tmp")

    op.execute("CREATE TYPE enum_type AS ENUM('some_value_1', 'some_value_1')")

    op.execute("DROP TYPE enum_type_tmp")

NOTE: During downgrade, if you’re using enum_type in a table then you can modify the downgrade method as mentioned below:

def downgrade():
    # Drop 'new_value' from enum_type
    op.execute("UPDATE table_name"
               " SET column_name_using_enum_type_value = NULL"
               " WHERE column_name_using_enum_type_value = 'new_value'")    

    op.execute("ALTER TYPE enum_type RENAME TO enum_type_tmp")

    op.execute("CREATE TYPE enum_type AS ENUM('some_value_1', 'some_value_1')")

    op.execute("ALTER TABLE table_name"
               " ALTER COLUMN column_name_using_enum_type_value TYPE enum_type"
               " USING column_name_using_enum_type_value::text::enum_type")

    op.execute("DROP TYPE enum_type_tmp")

This method can be used to update an Enum

def upgrade():
    op.execute("ALTER TYPE categorytype RENAME VALUE 'EXAMPLE_A' TO 'EXAMPLE_B'")


def downgrade():
    op.execute("ALTER TYPE categorytype RENAME VALUE 'EXAMPLE_B' TO 'EXAMPLE_A'")

First alter your column type to VARCHAR().

Then drop your type and create new type with new fields.

and finally alter your column type to your new created type.

def upgrade():
    op.execute(
        '''
        ALTER TABLE your_table ALTER COLUMN your_enum_column TYPE VARCHAR(255);

        DROP TYPE IF EXISTS your_enum_type;

        CREATE TYPE your_enum_type AS ENUM 
            ('value1', 'value2', 'value3', 'value4');

        ALTER TABLE your_table ALTER COLUMN your_enum_column TYPE your_enum_type 
            USING (your_enum_column::your_enum_type);
        '''
    )


def downgrade():
    op.execute(
        '''
        ALTER TABLE your_table ALTER COLUMN your_enum_column TYPE VARCHAR(255);

        DROP TYPE IF EXISTS your_enum_type;

        CREATE TYPE your_enum_type AS ENUM 
            ('value1', 'value2', 'value3');

        ALTER TABLE your_table ALTER COLUMN your_enum_column TYPE your_enum_type 
            USING (your_enum_column::your_enum_type);
        '''
    )

Observation

In order to alleviate the pain when migrating, I always use non-native enums even with PostgreSQL.

Non-native enums are just strings with constraints, if you edit an enum, there are only three scenarios:

1. Rename enum values
2. Remove enum values
3. Add enum values.

For migrations, 2 & 3 are a pair. This is understandable: if you upgrade for adding, then you have to remove when downgrading and vice versa. So let’s categorize them as two types.

Implementation

If it’s renaming, normally I’ll break it into three steps:

1. Drop old constraints
2. Update the rows’ old value to the new value
3. Create new constraints

In alembic, this is done by:

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)
    )

Let’s forget update_enum_sqls first and just use it as a SQL generator.

If it’s removing, well there are still three steps:

1. Drop old constraints
2. Delete rows with the old value
3. Create new constraints

So basically only the update_enum_sqls might behave differently.

If it’s adding, only two steps:

1. Drop old constraints
2. Create new constraints

Still, we can ignore the update_enum_sqls.

So how to implement it? Not that hard…

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)
    ]

Example

Since we prepared the ingredients, let’s apply:

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", USE WITH CARE!!!
    )

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

The code above can also be found on gist.

This approach is similar to the accepted solution, but with minor differences:

1. it uses op.batch_alter_table instead of op.execute('ALTER TABLE'), so this solution works in both PostgreSQL and SQLite: SQLite doesn’t supports ALTER TABLE but alembic provides a support for it with op.batch_alter_table
2. it doesn’t uses raw SQL

from alembic import op
import sqlalchemy as sa


# Describing of enum
enum_name = "status"
temp_enum_name = f"temp_{enum_name}"
old_values = ("nonexistent_executable", "signal", "success", "timed_out")
new_values = ("output_limit_exceeded", *old_values)
downgrade_to = ("output_limit_exceeded", "timed_out") # on downgrade convert [0] to [1]
old_type = sa.Enum(*old_values, name=enum_name)
new_type = sa.Enum(*new_values, name=enum_name)
temp_type = sa.Enum(*new_values, name=temp_enum_name)

# Describing of table
table_name = "testcaseresult"
column_name = "status"
temp_table = sa.sql.table(
    table_name,
    sa.Column(
        column_name,
        new_type,
        nullable=False
    )
)


def upgrade():
    # temp type to use instead of old one
    temp_type.create(op.get_bind(), checkfirst=False)

    # changing of column type from old enum to new one.
    # SQLite will create temp table for this
    with op.batch_alter_table(table_name) as batch_op:
        batch_op.alter_column(
            column_name,
            existing_type=old_type,
            type_=temp_type,
            existing_nullable=False,
            postgresql_using=f"{column_name}::text::{temp_enum_name}"
        )

    # remove old enum, create new enum
    old_type.drop(op.get_bind(), checkfirst=False)
    new_type.create(op.get_bind(), checkfirst=False)

    # changing of column type from temp enum to new one.
    # SQLite will create temp table for this
    with op.batch_alter_table(table_name) as batch_op:
        batch_op.alter_column(
            column_name,
            existing_type=temp_type,
            type_=new_type,
            existing_nullable=False,
            postgresql_using=f"{column_name}::text::{enum_name}"
        )

    # remove temp enum
    temp_type.drop(op.get_bind(), checkfirst=False)


def downgrade():
    # old enum don't have new value anymore.
    # before downgrading from new enum to old one,
    # we should replace new value from new enum with
    # somewhat of old values from old enum
    op.execute(
        temp_table
        .update()
        .where(
            temp_table.c.status == downgrade_to[0]
        )
        .values(
            status=downgrade_to[1]
        )
    )

    temp_type.create(op.get_bind(), checkfirst=False)

    with op.batch_alter_table(table_name) as batch_op:
        batch_op.alter_column(
            column_name,
            existing_type=new_type,
            type_=temp_type,
            existing_nullable=False,
            postgresql_using=f"{column_name}::text::{temp_enum_name}"
        )

    new_type.drop(op.get_bind(), checkfirst=False)
    old_type.create(op.get_bind(), checkfirst=False)

    with op.batch_alter_table(table_name) as batch_op:
        batch_op.alter_column(
            column_name,
            existing_type=temp_type,
            type_=old_type,
            existing_nullable=False,
            postgresql_using=f"{column_name}::text::{enum_name}"
        )

    temp_type.drop(op.get_bind(), checkfirst=False)

From the accepted solution:

It appears that alembic has no direct support for the USING statement in its alter_table method.

For now alembic supports USING statement in its alter_table method.