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 explains the direct process, but I'm not quite sure how best to translate that using alembic.

This is what I have:

new_type = sa.Enum('nonexistent_executable', 'output_limit_exceeded',
                   'signal', 'success', 'timed_out', name='status')
old_type = sa.Enum('nonexistent_executable', 'signal', 'success', 'timed_out',
                   name='status')
tcr = sa.sql.table('testcaseresult',
                   sa.Column('status', new_type, nullable=False))


def upgrade():
    op.alter_column('testcaseresult', u'status', type_=new_type,
                    existing_type=old_type)


def downgrade():
    op.execute(tcr.update().where(tcr.c.status==u'output_limit_exceeded')
               .values(status='timed_out'))
    op.alter_column('testcaseresult', u'status', type_=old_type,
                    existing_type=new_type)

The above unfortunately only produces ALTER TABLE testcaseresult ALTER COLUMN status TYPE status upon upgrade, which essentially does nothing.

Answer 1

I decided to try to follow the postgres approach as directly as possible and came up with the following migration.

from alembic import op
import sqlalchemy as sa

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

old_type = sa.Enum(*old_options, name='status')
new_type = sa.Enum(*new_options, name='status')
tmp_type = sa.Enum(*new_options, name='_status')

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


def upgrade():
    # Create a tempoary "_status" type, convert and drop the "old" type
    tmp_type.create(op.get_bind(), checkfirst=False)
    op.execute('ALTER TABLE testcaseresult ALTER COLUMN status TYPE _status'
               ' USING status::text::_status')
    old_type.drop(op.get_bind(), checkfirst=False)
    # Create and convert to the "new" status type
    new_type.create(op.get_bind(), checkfirst=False)
    op.execute('ALTER TABLE testcaseresult ALTER COLUMN status TYPE status'
               ' USING status::text::status')
    tmp_type.drop(op.get_bind(), checkfirst=False)


def downgrade():
    # Convert 'output_limit_exceeded' status into 'timed_out'
    op.execute(tcr.update().where(tcr.c.status==u'output_limit_exceeded')
               .values(status='timed_out'))
    # Create a tempoary "_status" type, convert and drop the "new" type
    tmp_type.create(op.get_bind(), checkfirst=False)
    op.execute('ALTER TABLE testcaseresult ALTER COLUMN status TYPE _status'
               ' USING status::text::_status')
    new_type.drop(op.get_bind(), checkfirst=False)
    # Create and convert to the "old" status type
    old_type.create(op.get_bind(), checkfirst=False)
    op.execute('ALTER TABLE testcaseresult ALTER COLUMN status TYPE status'
               ' USING status::text::status')
    tmp_type.drop(op.get_bind(), checkfirst=False)

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

Answer 2

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)

Answer 3

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