pythonpostgresqlsqlalchemyalembic

Altering an Enum field using Alembic


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.


Solution

  • 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.