如何以半自动方式迁移 SQLAlchemy for PostgreSQL 中的新 Python Enum 成员？

Posted 2023-03-16

【中文标题】如何以半自动方式迁移 SQLAlchemy for PostgreSQL 中的新 Python Enum 成员？

【英文标题】：How to migrate new Python Enum members in SQLAlchemy for PostgreSQL in semi-automated manner?

【发布时间】：2022-01-05 01:18:12

【问题描述】：

我正在使用 Alembic 迁移工具在 SQLAlchemy 迁移中添加新的枚举成员。

SQLAlchemy 使用 Python 原生 Enum 作为：

from enum import Enum

class MyEnum(Enum):
    # Key and value names different to clarify they are separate 
    # concepts in 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 is difficult 涉及到枚举迁移以及类似问题的现有低质量答案这一事实。但我仍然想以半自动化的方式进行。

使用 Python 3.8。

【参考方案1】：

您需要手动声明添加的新枚举值，因为它很难检测到（虽然并非不可能，但稍后可能会修改此答案）。下面是一个示例迁移，它反映了从 Enum 类返回的 PSQL 键。

还有一个降级功能，但只有在没有使用任何Enum 值时才会起作用，而且有点危险。

# revision identifiers, used by Alembic.
from mymodels import MyEnum

revision = ...
down_revision = ....
branch_labels = None
depends_on = None


# By default, SQLAlchemy autogenerates PSQL type
# that is the Enum class name lowercased.
# Because Python enums are created using Python metaclasses,
# you cannot directly read the class name back from the class object (it would return `EnumMeta`).
# You need go through method resolution order table to get the real class instance.
# https://***.com/a/54014128/315168
enum_name = MyEnum.mro()[0].__name__.lower()

# PostgreSQL does not have a concept of enum keys and values, only values.
# SQLAlchemy maintains enums in the PostgreSQL by the Python enum key (not value)
# Keys to be added:
enum_keys_to_add = [
    MyEnum.new1.name,
    MyEnum.new2.name,
]

def upgrade():
    for v in enum_keys_to_add:
        # Use ALTER TYPE.
        # See more information https://www.postgresql.org/docs/9.1/sql-altertype.html
        op.execute(f"ALTER TYPE enum_name ADD VALUE 'v'")
    # Note that ALTER TYPE is not effective within the same transaction https://medium.com/makimo-tech-blog/upgrading-postgresqls-enum-type-with-sqlalchemy-using-alembic-migration-881af1e30abe

def downgrade():
    # https://***.com/a/39878233/315168
    for v in enum_keys_to_add:
        sql = f"""DELETE FROM pg_enum
            WHERE enumlabel = 'v'
            AND enumtypid = (
              SELECT oid FROM pg_type WHERE typname = 'enum_name'
            )"""
        op.execute(sql)

您还可以使用以下SELECT 手动检查 PSQL 枚举的内容：

select enum_range(null::myenum);

                                             enum_range                                              
-----------------------------------------------------------------------------------------------------
 old,new1,new2