FastAPI WebSocket中如何基于Schema定义响应模型？

假设您有如下Pydantic Schema：

class Item(BaseModel):
    id: int
    name: str

以及一个运行WebSocket的FastAPI服务：

class ConnectionManager:
    def __init__(self, db: Session = Depends(get_db)):
        self.active_connections: list[WebSocket] = []
        self.db = db

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)
    
    async def handler(self, websocket: WebSocket):
        data = await self.receive_json()
        json = self.get_items()
        # 应该发送Item列表
        await websocket.send_json(dict(json))

    def get_items(self) -> List[Item]:
        # 对查询结果应用Schema后返回
        return self.query(Item).all()

WebSocket端点定义如下：

@app.websocket("/ws")
async def ws_endpoint(websocket: WebSocket, manager: ConnectionManager = Depends()):
    await manager.connect(websocket)
    try:
        while True:
            await manager.handle(websocket)
    except WebSocketDisconnect:
        manager.disconnect(websocket)

是否可以为ConnectionManager的get_items等方法定义响应模型，使其无需每次手动转换数据库查询结果即可返回JSON数据？

目前我找到的最接近响应模型的简洁方式是列表推导式：

return [dict(SingleBeacon(**vars(b))) for b in self.db.query(Beacon).all()]

但这种方式扩展性差，且代码不够简洁优雅。

当然可以，有几种更优雅的方式实现自动转换，避免重复的手动转换逻辑：

1. 利用Pydantic的from_orm批量转换 #

Pydantic的BaseModel自带from_orm方法，可直接将ORM模型实例转换为Pydantic模型。我们可以封装通用转换方法提升复用性：

from typing import TypeVar, Generic, List

# 定义泛型类型，适配不同模型
T = TypeVar('T')

class ConnectionManager:
    # ... 保留原有代码 ...

    def get_items(self) -> List[Item]:
        # 注意：区分ORM模型和Pydantic模型，此处假设ORM模型为ItemModel
        db_items = self.db.query(ItemModel).all()
        return [Item.from_orm(item) for item in db_items]
    
    # 通用转换方法，适用于所有模型
    def _convert_to_schema(self, schema_class: Generic[T], db_instances) -> List[T]:
        return [schema_class.from_orm(instance) for instance in db_instances]
    
    # 调用示例：获取其他模型数据
    def get_beacons(self) -> List[SingleBeacon]:
        db_beacons = self.db.query(Beacon).all()
        return self._convert_to_schema(SingleBeacon, db_beacons)

发送JSON时，直接使用Pydantic模型的dict()方法即可：

async def handler(self, websocket: WebSocket):
    data = await websocket.receive_json()
    items = self.get_items()
    await websocket.send_json([item.dict() for item in items])

2. 封装查询+转换的通用方法 #

把查询数据库和转换模型的逻辑合并成一个通用方法，进一步减少重复代码：

class ConnectionManager:
    # ... 保留原有代码 ...

    def query_with_schema(self, orm_model, schema_class: Generic[T]) -> List[T]:
        db_results = self.db.query(orm_model).all()
        return [schema_class.from_orm(result) for result in db_results]
    
    # 使用示例
    def get_items(self) -> List[Item]:
        return self.query_with_schema(ItemModel, Item)
    
    def get_beacons(self) -> List[SingleBeacon]:
        return self.query_with_schema(Beacon, SingleBeacon)

关键注意事项 #

• 严格区分ORM模型和Pydantic模型：比如将数据库ORM模型命名为ItemModel，Pydantic模型命名为Item，避免变量名冲突。
• from_orm比vars()更可靠：它能自动处理ORM的延迟加载属性，无需手动提取实例变量。
• 自定义序列化逻辑：可在Pydantic模型中添加Config类或字段验证器，一次配置即可全局生效，扩展性更强。

通过以上方式，新增查询方法时只需调用通用方法传入对应模型即可，无需重复编写转换逻辑，代码更简洁易维护。

内容的提问来源于stack exchange，提问作者Weed Cookie