sqlalchemy - 如何在FastAPI中统一返回格式，同时保留Pydantic数据模型？

Question

在FastAPI中，我使用了SQLAlchemy和Pydanic返回数据

@router.get("/1", response_model=User)def read_user(db: Session = Depends(get_db)):    db_user = user_module.get_user(db, user_id="1")    if db_user is None:        raise HTTPException(status_code=404, detail="找不到用户")    return db_user

这种方式可以帮助我规范化返回的模型，但是我想要将所有api统一返回格式为{"code":0,"msg":"success","data":{...}}，使原本返回模型中User的模型置于data中，方便前端管理。

我尝试了使用FastAPI中间件来实现，但是在swagger等文档中无法识别返回模型User；如果重新定义一个通用的Pydantic返回模型嵌套模型，就无法对SQLAlchemy返回的数据模型处理为想要的User模型，请求社区大佬们帮忙解决，万分感谢！

Answer 1

在FastAPI中，你可以通过创建一个自定义的响应类来统一你的API的返回格式。这个自定义的响应类可以继承自FastAPI的Response类，并且可以定义你想要的返回格式。

首先，你需要创建一个新的Pydantic模型来作为你的通用返回格式。这个模型可以包含你想要的任何字段，比如code, msg, 和 data。

from pydantic import BaseModelclass CommonResponse(BaseModel):    code: int    msg: str    data: dict

然后，你可以在每个路由中使用这个新的响应模型。假设你仍然希望User模型是你的返回数据的核心，你可以将它作为CommonResponse模型中的data字段。

from sqlalchemy.orm import Sessionfrom sqlalchemy import create_engine, funcfrom sqlalchemy.ext.declarative import declarative_basefrom sqlalchemy.orm import sessionmakerfrom fastapi import FastAPI, Depends, HTTPExceptionfrom pydantic import BaseModel, ValidationErrorfrom typing import List, Dict, Anyimport datetimeapp = FastAPI()# 定义通用的响应模型class CommonResponse(BaseModel):    code: int    msg: str    data: dict# 创建数据库模型和表Base = declarative_base()class User(Base):    __tablename__ = 'users'    id = Column(Integer, primary_key=True)    name = Column(String)    email = Column(String)    # 其他字段...    # ...# 创建数据库引擎和会话类engine = create_engine('sqlite:///test.db')  # 根据你的数据库配置修改这个URLBase.metadata.create_all(engine)  # 创建表SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)  # 创建会话类@app.get("/1", response_model=CommonResponse)  # 使用通用的响应模型async def read_user(db: Session = Depends(get_db)):  # 使用SQLAlchemy的会话作为依赖项    db_user = User.query.get(1)  # 查询用户ID为1的用户，根据实际情况修改这个查询语句    if db_user is None:  # 如果用户不存在，抛出HTTPException异常        raise HTTPException(status_code=404, detail="找不到用户")    return CommonResponse(code=0, msg="success", data=db_user.dict())  # 将用户对象转换为字典，并作为通用响应模型的`data`字段返回

在这个例子中，我使用了SQLAlchemy来查询数据库，并将查询结果转换为字典。然后，我将这个字典作为通用响应模型的data字段返回。这样，你就可以在你的API中统一返回格式，同时保留Pydantic数据模型。