SQLAlchemy Ltree to Pydantic Model: 两种转换方案

2024-12-28 02:50:03

将 SQLAlchemy ORM 模型（PostgreSQL Ltree 扩展）转换为 Pydantic 模型

在使用 SQLAlchemy 结合 PostgreSQL 的 Ltree 扩展进行数据建模时，可能会遇到如何将 ORM 模型的数据优雅地转换为 Pydantic 模型以便在 API 或其他场景中使用的挑战。本文将深入探讨这个问题，并提供切实可行的解决方案。

问题分析

Pydantic 模型在定义时，类型约束至关重要。它依赖 Python 内置的类型以及 pydantic 自身提供的类型系统。当 SQLAlchemy ORM 模型使用了 PostgreSQL 的特定数据类型，如 LtreeType（通常由 sqlalchemy_utils 库提供），类型不匹配问题就可能出现。具体到本文的问题，Pydantic 模型需要一个 str 类型字段，但 ORM 模型提供了 LtreeType 类型的 path 属性，直接映射时会导致验证错误。

这个问题的根本原因在于： LtreeType 并不直接是 Python 原生或 Pydantic 可以直接理解的数据类型，Pydantic 无法将其作为 str 或者 LtreeType 来验证。当调用 Item.model_validate(db_data, from_attributes=True) 时， Pydantic 会尝试匹配传入的数据结构类型定义。而对于数据库读取出来的 Ltree 实例，它即不是 Python 原生的 str，也不能直接匹配我们 Pydantic 模型中的LtreeType。

解决方案

这里我们将给出两种主要的解决途径。每种方案都有其适用的场景和考量，开发者可以根据具体需求选择合适的方案。

方案一：使用 Pydantic 字段的转换器（Field Serializer）

该方案的主要思路是使用 Pydantic 字段的序列化器（Serializer）特性。在定义 Pydantic 模型时，为特定的字段指定一个转换函数。这个函数负责将 LtreeType 数据转换为 str。

操作步骤：

在 Item 模型中，我们保持 path 字段类型不变（仍然是 LtreeType）。
为 path 字段添加序列化函数，将 Ltree 对象转换为 str 。

代码示例:

from typing import Any
from sqlalchemy.orm import Mapped, mapped_column
from sqlalchemy_utils import LtreeType
from pydantic import BaseModel, ConfigDict, field_serializer
from sqlalchemy import create_engine, text
from sqlalchemy.orm import declarative_base, sessionmaker

# Sqlalchemy Setup
engine = create_engine("postgresql://user:password@localhost:5432/test")
Session = sessionmaker(bind=engine)

Base = declarative_base()
class ItemORM(Base):
    __tablename__ = 'item'

    id: Mapped[int] = mapped_column(primary_key=True)
    label: Mapped[str]
    path: Mapped[str] = mapped_column(LtreeType)


# Pydantic
class Item(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)
    id: int
    label: str
    path: LtreeType

    @field_serializer('path')
    def serialize_path(self, path: LtreeType, _info) -> str:
       return str(path)


# Data
with Session() as session:
    # 示例代码：添加初始化数据和查询
    Base.metadata.create_all(engine)
    session.execute(text("TRUNCATE item"))
    item1 = ItemORM(id=1, label='Test Label', path='A.23')
    session.add(item1)
    session.commit()

    db_data = session.query(ItemORM).first()
    item_dto: Item = Item.model_validate(db_data, from_attributes=True)
    print(item_dto)
    print(type(item_dto.path))

这个方案非常优雅，不需要额外的类型校验和转换步骤，所有的转换都交由Pydantic处理，符合我们Pydantic 的使用原则。输出会显示已经转成字符串后的path值。并且在内部维护的时候，依然是LtreeType的类型。

方案二： Pydantic 模型的路径直接使用字符串

第二种方案直接修改Pydantic 模型字段类型，将LtreeType类型直接修改成字符串str,并在 model_validate 时将 ORM 中的 LtreeType 对象转化为字符串后再传入。

操作步骤:

修改 Pydantic 模型, 将 path 的类型定义修改为 str。
在将数据传递给 model_validate 之前，将ORM对象的 path 值转换为字符串

代码示例:

from typing import Any
from sqlalchemy.orm import Mapped, mapped_column
from sqlalchemy_utils import LtreeType
from pydantic import BaseModel, ConfigDict
from sqlalchemy import create_engine, text
from sqlalchemy.orm import declarative_base, sessionmaker

# Sqlalchemy Setup
engine = create_engine("postgresql://user:password@localhost:5432/test")
Session = sessionmaker(bind=engine)

Base = declarative_base()
class ItemORM(Base):
    __tablename__ = 'item'

    id: Mapped[int] = mapped_column(primary_key=True)
    label: Mapped[str]
    path: Mapped[str] = mapped_column(LtreeType)

# Pydantic
class Item(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)
    id: int
    label: str
    path: str  # path 字段修改为 str 类型

# Data
with Session() as session:
     # 示例代码：添加初始化数据和查询
    Base.metadata.create_all(engine)
    session.execute(text("TRUNCATE item"))
    item1 = ItemORM(id=1, label='Test Label', path='A.23')
    session.add(item1)
    session.commit()
    db_data = session.query(ItemORM).first()

    # 在转换之前把路径转为字符串
    item_data = {**db_data.__dict__, "path": str(db_data.path)}
    item_dto: Item = Item.model_validate(item_data) # 不需要 from_attributes 了，已经转换了字典结构
    print(item_dto)
    print(type(item_dto.path))

在这个方案中，直接在读取数据后转换为目标字符串。这样做逻辑清晰，并且确保在将数据传入到Pydantic模型中验证前是符合规范的数据类型。

选择建议

方案一使用 Pydantic 字段的序列化器（Field Serializer）提供了一种更加符合 Pydantic 原则的方法，允许你在 Pydantic 模型中使用特定的数据类型，而又在数据序列化为字符串时进行处理，对应用程序结构的影响较小，并且更为优雅。如果未来你需要直接使用 Pydantic 进行其他操作（如转换成其他数据格式或者再验证等），它将更为灵活。方案二虽然逻辑简洁明了，但在模型定义层级略显不一致。考虑到在数据交互场景中最终我们需要的 path 大多数字符串格式，这个方案在某些时候也可以满足业务的需求，可根据项目需要灵活选用。

在数据安全方面，需要对外部接收的数据进行验证和清理。

以上两种方案可以有效解决 SQLAlchemy ORM 模型使用 LtreeType 类型与 Pydantic 模型使用 str 类型不匹配的问题，从而确保数据在不同的层之间正确传输和使用。