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

Question

我正在使用 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。

Answer 1

您需要手动声明添加的新枚举值，因为很难检测到(尽管并非不可能，但稍后可能会修改此答案)。下面是一个示例迁移，它反射(reflect)了从 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://stackoverflow.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://stackoverflow.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}