在使用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 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)
這個方法提供瞭解決了問題,但稍顯繁瑣。我們可以做一下簡單的封裝,生成一個通用函數:
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 --version psql (PostgreSQL) 9.3.10