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

huangapple go评论89阅读模式
英文:

How should I declare enums in SQLAlchemy using mapped_column (to enable type hinting)?

问题

I am trying to use Enums in SQLAlchemy 2.0 with mapped_column. So far I have the following code (taken from another question):

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)

However, that gives the following error upon the construction of the Campaign class itself.

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

Any hint about how to make this work?

The response from https://stackoverflow.com/questions/28894257/enum-type-in-sqlalchemy-with-postgresql does not apply as I am using version 2 of SQLAlchemy and those answers did not use mapped_column or Mapped types. Also, removing str from CampaignStatus does not help.

英文:

I am trying to use Enums in SQLAlchemy 2.0 with mapped_column. So far I have the following code (taken from another question):

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)

However, that gives the following error upon the construction of the Campaign class itself.

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

Any hint about how to make this work?

The response from https://stackoverflow.com/questions/28894257/enum-type-in-sqlalchemy-with-postgresql does not apply as I am using version 2 of SQLAlchemy and those answers did not use mapped_column or Mapped types. Also, removing str from CampaignStatus does not help.

答案1

得分: 4

以下是您提供的代码部分的中文翻译:

问题的关键在于与__mro__相关的问题导致AttributeError的原因是CampaignStatusType不是一个类,而是一个类型为sqlalchemy.dialects.postgresql.ENUM的实例变量(使用pyright可以验证此点 - 因为它会抱怨Mapped[CampaignStatusType]是一个“非法类型注释:只有类型别名才允许”的错误)。 作为测试,将status的类型注释替换为Mapped[CampaignStatus]确实可以解决问题(并且pyright不报错),但这并不会将列类型连接到所需的postgresql方言的枚举。

因此,在使用方言特定的枚举类型时,唯一的解决方法是使用未注释的构造

status = mapped_column(CampaignStatusType, nullable=False)

然而,如果仍然希望使用类型注释,即任何被Mapped的东西必须是一个类型,并且sqlalchemy.dialects.postgresql.ENUM(导入为pgEnum)是实例CampaignStatusType的基础类型,可能认为以下可能是一个解决方案

# 不要做这个错误的示例,尽管它确实运行
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 typing import get_args
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(
        *get_args(CampaignStatus),
        name="campaignstatus",
        create_constraint=True,
        validate_strings=True,
    ))

请注意,在CampaignStatus上使用typing.get_args并将其展开到此处的Enum,而不是在重复自己的官方示例中所做的。现在包括使用示例:

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'作为输出,显示不经过验证的字符串将不会被插入,而映射到某个枚举的验证字符串将被插入而没有问题。此外,pyright不会产生错误(尽管老实说,这不一定是一个好的度量标准,因为Python中的类型提示仍然相当不成熟,即使在Mapped中放入任何东西,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类型而不需要方言特定的导入,枚举更多或多少会按预期行事。

英文:

The crux of the issue relating to __mro__ causing the AttributeError is that CampaignStatusType is not a class, but rather an instance variable of type sqlalchemy.dialects.postgresql.ENUM (using pyright may verify this - given that it complains about Mapped[CampaignStatusType] being an "Illegal type annotation: variable not allowed unless it is a type alias"). As a test, replacing the type annotation for status with Mapped[CampaignStatus] does resolve the issue (and pyright reports no errors), but that does not hook the column type to the enum with postgresql dialect that is desired.

So the only way around this while using the dialect specific enum type is to use the non-annotated construct:

    status = mapped_column(CampaignStatusType, nullable=False)

However, if type annotation is still desired, i.e. whatever being Mapped must be a type, and that sqlalchemy.dialects.postgresql.ENUM (which was imported as pgEnum) is the underlying type for the instance CampaignStatusType, it may be thought that the following might be a solution

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

While it works, it does NOT actually reflect what will be represented by the data, so DO NOT actually do that. Moreover, it only works because the type annotation is ignored when the specific column type is passed, so putting anything in there will work while having an invalid type.

Now, given that SQLAlchemy is now 2.0 (as the question explicitly want this newer version), perhaps reviewing the documentation and see now native enums should be handled now.

Adapting the examples in the documentation, the following MVCE may now be derived, using all the intended keyword arguments that was passed to the PostgreSQL dialect specific ENUM type passed generic sqlalchemy.Enum instead (aside from metadata=Base.metadata as that's completely superfluous):

from typing import Literal
from typing import get_args
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(
        *get_args(CampaignStatus),
        name="campaignstatus",
        create_constraint=True,
        validate_strings=True,
    ))

Note the use of typing.get_args on CampaignStatus and splat it to the Enum here as opposed to what the official examples have done in repeating themselves. Now to include the usage:

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()

The above will produce failed to insert with 'some_unvalidated_string' as the output, showing that unvalidated strings will not be inserted, while validated strings that are mapped to some enum are inserted without issues. Moreover, pyright will not produce errors (though honestly, this is not necessarily a good metric because type hinting in Python is still fairly half-baked, as pyright did not detect the erroneous example as an error in the very beginning no matter what went inside Mapped, but I digress).

Viewing the newly created entities using 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)

The enum of course cannot be dropped without dropping the campaign table:

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.

So the enum more or less behaves as expected despite only using generic SQLAlchemy types, without needing dialect specific imports.

huangapple
  • 本文由 发表于 2023年5月17日 13:26:42
  • 转载请务必保留本文链接:https://go.coder-hub.com/76268799.html
匿名

发表评论

匿名网友

:?: :razz: :sad: :evil: :!: :smile: :oops: :grin: :eek: :shock: :???: :cool: :lol: :mad: :twisted: :roll: :wink: :idea: :arrow: :neutral: :cry: :mrgreen:

确定