如何在SQLAlchemy和Alembic中使用Enum？

Question

这是我的Post模型：

class Post(Base):
    __tablename__ = 'posts'

    title = db.Column(db.String(120), nullable=False)
    description = db.Column(db.String(2048), nullable=False)

我想在其中添加枚举status。因此，我创建了一个新的Enum：

import enum

class PostStatus(enum.Enum):
    DRAFT='draft'
    APPROVE='approve'
    PUBLISHED='published'

并为模型添加了一个新字段：

class Post(Base):
    ...
    status = db.Column(db.Enum(PostStatus), nullable=False, default=PostStatus.DRAFT.value, server_default=PostStatus.DRAFT.value)

完成后FLASK_APP=server.py flask db migrate，生成了这样的迁移：

def upgrade():
    op.add_column('posts', sa.Column('status', sa.Enum('DRAFT', 'APPROVE', 'PUBLISHED', name='poststatus'), server_default='draft', nullable=False))

尝试升级数据库后，我得到：

sqlalchemy.exc.ProgrammingError: (psycopg2.ProgrammingError) type "poststatus" does not exist
LINE 1: ALTER TABLE posts ADD COLUMN status poststatus DEFAULT 'draf...
                                            ^
 [SQL: "ALTER TABLE posts ADD COLUMN status poststatus DEFAULT 'draft' NOT NULL"]

1. 为什么poststatus没有在数据库级别自动创建类型？在类似的迁移中。
2. 如何server_default正确指定选项？我需要ORM级别的默认值和DB级别的默认值，因为我要更改现有行，因此不应用ORM默认值。
3. 为什么DB中的实际值是'DRAFT'，'APPROVE'，'PUBLISHED'，而不是draft？我认为应该有ENUM值，而不是名称。

先感谢您。

Answer 1

为什么 DB 中的实际值是“DRAFT”、“APPROVE”、“PUBLISHED”，而不是草稿等？我认为应该有 ENUM 值，而不是名称。

正如 Peter Bašista 已经提到的，SQLAlchemy在数据库中使用枚举名称（DRAFT、APPROVE、PUBLISHED）。我认为这是因为枚举值（“draft”、“approve”、...）在 Python 中可以是任意类型，并且不能保证它们是唯一的（除非@unique使用）。

但是，从SQLAlchemy 1.2.3 开始，Enum该类接受一个参数values_callable，该参数可用于在数据库中存储枚举值：

    status = db.Column(
        db.Enum(PostStatus, values_callable=lambda obj: [e.value for e in obj]),
        nullable=False,
        default=PostStatus.DRAFT.value,
        server_default=PostStatus.DRAFT.value
    )

为什么类型 poststatus 没有在 DB 级别自动创建？在类似的迁移中。

我认为基本上您遇到了 alembic 的限制：在某些情况下，它无法正确处理 PostgreSQL 上的枚举。我怀疑您的情况的主要问题是Autogenerate 无法正确处理 postgresql 枚举 #278。

我注意到如果我使用该类型是正确创建的，alembic.op.create_table所以我的解决方法基本上是：

enum_type = SQLEnum(PostStatus, values_callable=lambda enum: [e.value for e in enum])
op.create_table(
    '_dummy',
    sa.Column('id', Integer, primary_key=True),
    sa.Column('status', enum_type)
)
op.drop_table('_dummy')
c_status = Column('status', enum_type, nullable=False)
add_column('posts', c_status)

Answer 2

我只能回答你问题的第三部分。

\n\n

SQLAlchemy 中该类型的文档指出Enum：

\n\n

\n

上面，每个元素的字符串名称，例如 \xe2\x80\x9cone\xe2\x80\x9d、\xe2\x80\x9ctwo\xe2\x80\x9d、\xe2\x80\x9cthird\xe2\x80\x9d 为持久化到数据库；不使用Python 枚举的值（此处表示为整数） ；因此，每个枚举的值可以是任何类型的 Python 对象，无论它是否可持久。

\n

\n\n

因此，通过 SQLAlchemy 的设计，Enum 名称（而不是值）被持久保存到数据库中。

\n

Answer 3

如果您使用的是 PostgreSQL，请使用以下函数示例：

from sqlalchemy.dialects import postgresql
from ... import PostStatus
from alembic import op
import sqlalchemy as sa


def upgrade():
    post_status = postgresql.ENUM(PostStatus, name="status")
    post_status.create(op.get_bind(), checkfirst=True)
    op.add_column('posts', sa.Column('status',  post_status))


def downgrade():
    post_status = postgresql.ENUM(PostStatus, name="status")
    post_status.drop(op.get_bind())

Answer 4

此线程和相关的 StackOverflow 线程诉诸 PostgreSQL 方言特定的类型。然而，在 Alembic 迁移中可以轻松实现通用支持，如下所示。

首先，在要声明自定义 SQLAlchemy Enum 列类型的位置导入 Python 枚举、SQLAlchemy 枚举和SQLAlchemy 声明基。

import enum
from sqlalchemy import Enum
from sqlalchemy.ext.declarative import declarative_base
Base = declarative_base()

我们以OP原来的Python枚举类为例：

class PostStatus(enum.Enum):
    DRAFT='draft'
    APPROVE='approve'
    PUBLISHED='published'

现在我们创建一个SQLAlchemy Enum实例化：

PostStatusType: Enum = Enum(
    PostStatus,
    name="post_status_type",
    create_constraint=True,
    metadata=Base.metadata,
    validate_strings=True,
)

当您运行 Alembicalembic revision --autogenerate -m "Revision Notes"并尝试使用 来应用修订版时alembic upgrade head，您可能会收到有关类型不存在的错误。例如：

...
sqlalchemy.exc.ProgrammingError: (psycopg2.errors.UndefinedObject) type "post_status_type" does not exist
LINE 10:  post_status post_status_type NOT NULL,
...

要解决此问题，请导入 SQLAlchemy Enum 类并将以下内容添加到Alembic 自动生成的修订脚本中的upgrade()和函数中。downgrade()

from myproject.database import PostStatusType
...
def upgrade() -> None:
    PostStatusType.create(op.get_bind(), checkfirst=True)
    ... the remainder of the autogen code...
def downgrade() -> None:
    ...the autogen code...
    PostStatusType.drop(op.get_bind(), checkfirst=True)

最后，请确保sa.Column()使用枚举类型更新表中自动生成的声明，以简单地引用 SQLAlchemy 枚举类型，而不是使用 Alembic 尝试重新声明它。例如在def upgrade() -> None:

op.create_table(
    "my_table",
    sa.Column(
        "post_status",
        PostStatusType,
        nullable=False,
    ),
)