如何半自动化完成PostgreSQL下SQLAlchemy的Python Enum新成员迁移？

前置场景说明 #

当前已基于Python原生Enum定义了SQLAlchemy枚举字段，代码结构如下：

枚举定义 #

from enum import Enum

class MyEnum(Enum):
    # Key和value命名做了区分，明确二者在Python中是相互独立的概念
    old = "OLD_VALUE"
    new1 = "NEW1_VALUE"
    new2 = "NEW2_VALUE"

模型定义 #

class MyFantasticModel(Base):
    __tablename__ = "fantasy"
    enum_column = sa.Column(sa.Enum(MyEnum), nullable=False, index=True)

PostgreSQL枚举类型的迁移实现复杂度较高，以下是适配Python 3.8的半自动化实现方案。 #

实现步骤 #

• 生成空白迁移文件
  执行Alembic命令生成自定义迁移文件：
  alembic revision -m "add new values to MyEnum"
• 编写迁移逻辑
  由于PostgreSQL默认不允许在事务中执行ALTER TYPE ... ADD VALUE操作，我们需要单独处理升级、降级逻辑，迁移文件内容参考如下：

  from alembic import op
  from sqlalchemy.dialects import postgresql
  
  # 配置项：根据实际场景修改
  NEW_ENUM_VALUES = ["NEW1_VALUE", "NEW2_VALUE"]
  ENUM_DB_NAME = "myenum"  # 枚举在数据库中的实际名称，默认是Python枚举类名小写
  TARGET_TABLE = "fantasy"
  TARGET_COLUMN = "enum_column"
  
  def upgrade():
      # 开启自动提交，绕过事务限制 *仅PostgreSQL有效*
      op.get_context().connection.autocommit = True
      for val in NEW_ENUM_VALUES:
          op.execute(f"ALTER TYPE {ENUM_DB_NAME} ADD VALUE IF NOT EXISTS '{val}'")
      # 恢复默认事务配置
      op.get_context().connection.autocommit = False
  
  def downgrade():
      # PostgreSQL不支持直接删除枚举值，降级需要替换枚举类型
      tmp_enum_name = f"{ENUM_DB_NAME}_tmp"
      # 降级后保留的旧枚举值列表
      old_enum_values = ["OLD_VALUE"]
      # 1. 创建临时枚举类型
      tmp_enum = postgresql.ENUM(*old_enum_values, name=tmp_enum_name)
      tmp_enum.create(op.get_bind())
      # 2. 转换字段类型，已存在的新值会被转换为对应旧值，无匹配值时会抛出异常
      op.execute(f"""
          ALTER TABLE {TARGET_TABLE}
          ALTER COLUMN {TARGET_COLUMN} TYPE {tmp_enum_name}
          USING {TARGET_COLUMN}::text::{tmp_enum_name}
      """)
      # 3. 删除原枚举类型
      op.execute(f"DROP TYPE {ENUM_DB_NAME}")
      # 4. 临时枚举重命名为原枚举名
      op.execute(f"ALTER TYPE {tmp_enum_name} RENAME TO {ENUM_DB_NAME}")
  

• 执行迁移
  运行迁移命令完成枚举值新增：
  alembic upgrade head

注意事项 #

• PostgreSQL 12及以上版本支持在事务中新增枚举值，可以省略autocommit的调整逻辑
• 降级逻辑执行前请确认所有新增枚举值的记录已被清理或转换，否则会执行失败，生产环境建议提前备份数据
• 如果定义枚举时在sa.Enum中指定了name参数，请将ENUM_DB_NAME替换为你自定义的名称

内容的提问来源于stack exchange，提问作者Mikko Ohtamaa