如何在SQLAlchemy和Alembic中使用枚举？

Question

这是我的Post模型：

class Post(Base):
    __tablename__ = 'posts'

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

我想在此添加Enum 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等？ 我认为应该是枚举值，而不是名称。

谢谢你提前。

- f1nn

Answer 1

为什么数据库中的实际值是'DRAFT'、'APPROVE'、'PUBLISHED'，而不是draft等？我认为应该使用枚举值而不是名称。

正如Peter Bašista所提到的，SQLAlchemy在数据库中使用枚举名称(DRAFT、APPROVE、PUBLISHED)。我认为这样做是因为Python中的枚举值("draft"、"approve"等)可以是任意类型，并且不能保证它们唯一(除非使用@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类型没有在数据库级别自动创建？在类似的迁移中，它是自动创建的。我认为基本上你遇到了alembic的限制：在某些情况下，它无法正确处理PostgreSQL上的枚举类型。我怀疑你的主要问题是Autogenerate doesn't correctly handle postgresql enums #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)

- Felix Schwarz

---

我猜默认值应该传递枚举的名称字段。 - Ahmad

Answer 2

如果你使用的是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())

- Farshid Ashouri

---

谢谢你的回答，Farshid！这正是我需要运行迁移的东西。 - Alessandro M

这种方法容易出错，因为将来你可能会更改“PostStatus”枚举，而此迁移将始终使用最新的枚举更改，而不是你创建此修订时确切的枚举。 - JD Solanki

可以通过不使用 PostStatus 并手动传递枚举值列表来解决，之后您只需从现有类型中添加/删除枚举值即可。 - Krishna

Answer 3

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

SQLAlchemy中Enum类型的文档说明：

以上，每个元素的字符串名称，例如“one”，“two”，“three”，都被持久化到数据库中；Python枚举的值，在此处表示为整数，不会被使用；因此，每个枚举的值可以是任何类型的Python对象，无论它是否可持久化。

因此，Enum的设计是将名称而不是值持久化到数据库中。

- Peter Bašista

Answer 4

这个和相关的 StackOverflow 线程都采用了 PostgreSQL 方言特定的类型。然而，可以通过以下方式轻松地在 Alembic 迁移中实现通用支持。

首先，在您声明自定义的 SQLAlchemy 枚举列类型的位置导入 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枚举实例：

SQLAlchemy枚举

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

当你运行 Alembic 命令 alembic 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枚举类，并将以下内容添加到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,
    ),
)

- Brent

---

我本来想给这个答案点赞，但它的假设是错误的！当你使用PostgreSQL运行SQLAlchemy时，应该使用sqlalchemy.dialects.postgresql.ENUM类型，而不是通常用于SQLAlchemy的相应Enum。 - adir abargil

当尝试再次修改枚举时，这将导致错误。 - adir abargil

Answer 5

如果你想在数据库层面自动创建poststatus类型，你可以使用我的alembic-postgresql-enum库。

安装方法：

pip install alembic-postgresql-enum

在env.py中添加import语句

import alembic_postgresql_enum

下面显示的迁移将自动生成（**）

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    sa.Enum('draft', 'approve', 'published', name='poststatus').create(op.get_bind())
    op.add_column('posts', sa.Column('status', sa.Enum('draft', 'approve', 'published', name='poststatus'), nullable=False))
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column('posts', 'status')
    sa.Enum('draft', 'approve', 'published', name='poststatus').drop(op.get_bind())
    # ### end Alembic commands ###

** 假设使用了 values_callable

- RustyGuard

---

如果您是这个软件包的作者，请遵守自我推广的规定。 - undefined