postgresql 如何在SQLAlchemy中使用mapped_column声明枚举(以启用类型提示)?

zdwk9cvp  于 2023-05-22  发布在  PostgreSQL
关注(0)|答案(1)|浏览(830)

我尝试在SQLAlchemy 2.0中使用Enumsmapped_column。到目前为止,我有以下代码(取自另一个问题):

from sqlalchemy.dialects.postgresql import ENUM as pgEnum
import enum

class CampaignStatus(str, enum.Enum):
    activated = "activated"
    deactivated = "deactivated"

CampaignStatusType: pgEnum = pgEnum(
    CampaignStatus,
    name="campaignstatus",
    create_constraint=True,
    metadata=Base.metadata,
    validate_strings=True,
)

class Campaign(Base):
    __tablename__ = "campaign"

    id: Mapped[UUID] = mapped_column(primary_key=True, default=uuid4)
    created_at: Mapped[dt.datetime] = mapped_column(default=dt.datetime.now)
    status: Mapped[CampaignStatusType] = mapped_column(nullable=False)

但是,在构造Campaign类本身时会出现以下错误。

Traceback (most recent call last):
  File "<stdin>", line 27, in <module>
    class Campaign(Base):
...
AttributeError: 'ENUM' object has no attribute '__mro__'

有没有什么方法能让它奏效?
来自ENUM type in SQLAlchemy with PostgreSQL的响应不适用,因为我使用的是SQLAlchemy版本2,并且这些答案没有使用mapped_columnMapped类型。此外,从CampaignStatus中删除str也没有帮助。

pengsaosao

pengsaosao1#

导致AttributeError__mro__相关问题的关键在于CampaignStatusType不是一个类,而是一个sqlalchemy.dialects.postgresql.ENUM类型的示例变量(使用pyright可以验证这一点-因为它抱怨Mapped[CampaignStatusType]是一个“非法类型注解:除非是类型别名,否则不允许使用变量“)。作为一个测试,用Mapped[CampaignStatus]替换status的类型注解确实解决了这个问题(pyright没有报告任何错误),但这并没有将列类型与所需的postgresql方言挂钩。
因此,在使用方言特定的枚举类型时,解决这个问题的唯一方法是使用无注解的构造:

status = mapped_column(CampaignStatusType, nullable=False)

然而,如果仍然需要类型注解,即无论Mapped是什么,都必须是一个类型,并且sqlalchemy.dialects.postgresql.ENUM(作为pgEnum导入)是示例CampaignStatusType的基础类型,可以认为以下可能是一个解决方案

# don't do this erroneous example despite it does run
    status: Mapped[sqlalchemy.dialects.postgresql.ENUM] = mapped_column(
        CampaignStatusType,
        nullable=False,
    )

虽然它工作,但它实际上并不反映数据将表示什么,所以不要实际上这样做。此外,它之所以有效,是因为在传递特定列类型时会忽略类型注解,因此在其中放置任何内容都可以在具有无效类型的情况下工作。
现在,考虑到SQLAlchemy现在是2.0(因为问题明确需要这个较新的版本),也许现在应该检查文档并查看本地枚举。
修改文档中的示例,现在可以导出以下MVCE,使用传递给PostgreSQL方言特定ENUM类型的所有预期关键字参数,而不是传递给通用sqlalchemy.Enum(除了metadata=Base.metadata,因为这完全是多余的):

from typing import Literal
from sqlalchemy import Enum
from sqlalchemy.orm import DeclarativeBase
from sqlalchemy.orm import Mapped
from sqlalchemy.orm import mapped_column

CampaignStatus = Literal["activated", "deactivated"]

class Base(DeclarativeBase):
    pass

class Campaign(Base):
    __tablename__ = "campaign"
    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    status: Mapped[CampaignStatus] = mapped_column(Enum(
        "activated", "deactivated",
        name="campaignstatus",
        create_constraint=True,
        validate_strings=True,
    ))

现在加入用法:

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

def main():
    engine = create_engine('postgresql://postgres@localhost/postgres')
    Base.metadata.create_all(engine)
    Session = sessionmaker(bind=engine)
    session = Session()
    session.add(Campaign(status='activated'))
    session.add(Campaign(status='deactivated'))
    session.commit()

    s = 'some_unvalidated_string'
    try:
        session.add(Campaign(status=s))
        session.commit()
    except Exception:
        print("failed to insert with %r" % s)

if __name__ == '__main__':
    main()

上面的代码将生成failed to insert with 'some_unvalidated_string'作为输出,表明未验证的字符串将不会被插入,而Map到某个枚举的已验证字符串将被插入而不会出现问题。此外,pyright不会产生错误(尽管老实说,这不一定是一个好的指标,因为Python中的类型提示仍然相当不成熟,因为pyright在一开始就没有检测到错误的示例,无论Mapped内部发生了什么,但我离题了)。
使用psql查看新创建的实体

postgres=# select * from campaign;
 id |   status    
----+-------------
  1 | activated
  2 | deactivated
(2 rows)
postgres=# \dt campaign;
                                Table "public.campaign"
 Column |      Type      | Collation | Nullable |               Default                
--------+----------------+-----------+----------+--------------------------------------
 id     | integer        |           | not null | nextval('campaign_id_seq'::regclass)
 status | campaignstatus |           | not null | 
Indexes:
    "campaign_pkey" PRIMARY KEY, btree (id)

postgres=# \dT+ campaignstatus;
                                             List of data types
 Schema |      Name      | Internal name  | Size |  Elements   |  Owner   | Access privileges | Description 
--------+----------------+----------------+------+-------------+----------+-------------------+-------------
 public | campaignstatus | campaignstatus | 4    | activated  +| postgres |                   | 
        |                |                |      | deactivated |          |                   | 
(1 row)

当然,在不删除campaign表的情况下,不能删除枚举:

postgres=# drop type campaignstatus;
ERROR:  cannot drop type campaignstatus because other objects depend on it
DETAIL:  column status of table campaign depends on type campaignstatus
HINT:  Use DROP ... CASCADE to drop the dependent objects too.

因此,尽管只使用通用的SQLAlchemy类型,但枚举或多或少地表现为预期的行为,而不需要特定于方言的导入。

相关问题