Openstack数据库管理工具alembic更新Enum类型

在使用alembic开发管理数据库时，会遇到一个比较麻烦的问题，就是变更某列的枚举类型，事实上使用sql命令变更相当的简单，一条alter的执行即可：

ALTER TYPE status ADD value 'output_limit_exceeded' after 'timed_out';#删除DROP TYPE status

但这样并不能满足alembic管理的初衷，也无法实现downgrade。

使用google搜索关键字"alembic enum type"，第一个出现的是stackflow的一个帖子Altering an Enum field using Alembic。

from alembic import opimport sqlalchemy as saold_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)

这个方法提供了解决了问题，但稍显繁琐。我们可以做一下简单的封装，生成一个通用函数：

def upgrade_enum(table, column_name, enum_name, old_options, new_options):    old_type = sa.Enum(*old_options, name=enum_name)    new_type = sa.Enum(*new_options, name=enum_name)    tmp_type = sa.Enum(*new_options, name="_" + enum_name)    # Create a temporary type, convert and drop the "old" type    tmp_type.create(op.get_bind(), checkfirst=False)    op.execute(        u'ALTER TABLE {0} ALTER COLUMN {1} TYPE _{2}'        u' USING {1}::text::_{2}'.format(            table,            column_name,            enum_name        )    )    old_type.drop(op.get_bind(), checkfirst=False)    # Create and convert to the "new" type    new_type.create(op.get_bind(), checkfirst=False)    op.execute(        u'ALTER TABLE {0} ALTER COLUMN {1} TYPE {2}'        u' USING {1}::text::{2}'.format(            table,            column_name,            enum_name        )    )    tmp_type.drop(op.get_bind(), checkfirst=False)

这样就可以直接通过传参直接来执行升级或降级操作了。

操作环境：

 alembic-0.8.2-bash-4.2$ psql --versionpsql (PostgreSQL) 9.3.10