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

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