使SQLModel模型的TSTZRANGE字段兼容FastAPI响应序列化

问题原因 #

写入数据时，Pydantic能自动将Tuple[datetime, datetime]转换为PostgreSQL的TSTZRANGE类型存入数据库，但从数据库刷新对象后，ts_range字段会变成SQLAlchemy的Range对象，而Pydantic默认无法序列化这种对象，导致返回响应时出错。

解决方案 #

方案1：在模型中添加字段验证器（推荐）

通过Pydantic的validator装饰器，在序列化前自动将Range对象转换为Tuple类型，一劳永逸解决所有场景的序列化问题：

from datetime import datetime
from typing import Optional, Tuple

from sqlalchemy.dialects.postgresql import TSTZRANGE, Range
from sqlalchemy import Column
from sqlmodel import Field, SQLModel
from pydantic import validator


class DateRange(SQLModel):
    """Plan rate model"""

    id: Optional[int] = Field(default=None, primary_key=True)

    ts_range: Tuple[datetime, datetime] = Field(nullable=False, sa_column=Column(TSTZRANGE))

    @validator('ts_range', pre=False, always=True)
    def convert_range_to_tuple(cls, value):
        # 如果是SQLAlchemy的Range对象，转换为datetime元组
        if isinstance(value, Range):
            return (value.lower, value.upper)
        return value

方案2：接口内手动转换

在返回对象前，手动将Range对象转为Tuple，适合临时快速解决单个接口的问题：

@router.post(
    "/rates",
    tags=["rates", "create"],
    response_model=DateRange,
)
async def create_rate(
    plan_id: int,
    rate: DateRange,
    db: DB = Depends(get_db),
) -> DateRange:
    """Create new rate"""

    rate = DateRange(**rate.dict())
    db.add(rate)
    try:
        await db.commit()
        await db.refresh(rate)
        # 手动转换Range对象为datetime元组
        rate.ts_range = (rate.ts_range.lower, rate.ts_range.upper)
        return rate
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"Invalid date range [{e}]",
        )

方案3：自定义Pydantic字段类型

如果多个模型都需要使用TSTZRANGE类型，可以自定义一个专属字段类型，统一处理序列化逻辑：

from datetime import datetime
from typing import Tuple, Any
from sqlalchemy.dialects.postgresql import TSTZRANGE, Range
from sqlmodel import Field, SQLModel
from pydantic import GetJsonSchemaHandler
from pydantic.json_schema import JsonSchemaValue
from pydantic_core import core_schema


class TSTZRangeField(Tuple[datetime, datetime]):
    @classmethod
    def __get_pydantic_core_schema__(cls, source_type: Any, handler: GetJsonSchemaHandler) -> core_schema.CoreSchema:
        # 定义序列化逻辑：将Range对象转为datetime元组
        def serialize_range(value: Range) -> Tuple[datetime, datetime]:
            return (value.lower, value.upper)
        
        # 合并基础schema与自定义序列化逻辑
        base_schema = handler(source_type)
        return core_schema.json_serializer_function_override(
            core_schema.no_info_wrap_validator_function(
                lambda v: serialize_range(v) if isinstance(v, Range) else v,
                base_schema
            ),
            serialize_range
        )


# 在模型中使用自定义字段
class DateRange(SQLModel):
    """Plan rate model"""

    id: Optional[int] = Field(default=None, primary_key=True)

    ts_range: TSTZRangeField = Field(nullable=False, sa_column=Column(TSTZRANGE))

内容的提问来源于stack exchange，提问作者Pratik K.