FastAPI查询PostgreSQL用户表未找到时返回自定义响应的实现

我在PostgreSQL中搭建了存储用户信息的表，使用以下SQLAlchemy模型和Pydantic模型定义数据结构：

class User(Base):
    __tablename__ = "users"

    id = Column(Integer, primary_key=True, nullable=False, unique=True)
    name = Column(String, nullable=False)
    status = Column(Integer, nullable=False)
    created_at = Column(TIMESTAMP(timezone=True), nullable=False, server_default=text('now()'))

class UserBase(BaseModel):
    name: str
    phone: int
    status: int = 1

参考教程，我已实现通过用户ID查询表，找到用户时返回对应响应模型，未找到则抛出HTTPException：

@router.get("/{id}", response_model=schemas.UserResponse)
def get_user(id: int, db: Session = Depends(get_db)):
    
    user = db.query(models.User).filter(models.User.id == id).first()
    print(user) # checking output
    
    if not user:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, 
            detail=f"User with ID {id} not found.")

    return user

目前该代码可正常返回找到的用户，但我希望在未找到用户时，返回一个status字段为0的对象，而非抛出异常，请问该如何实现？

你可以通过以下步骤实现需求：

1. 调整响应模型的字段可选性
   首先确保schemas.UserResponse允许部分字段为空，比如created_at、name，如果UserResponse是基于UserBase扩展的，可修改为：

   from datetime import datetime
   from typing import Optional
   
   class UserResponse(UserBase):
       id: int
       created_at: Optional[datetime] = None
       class Config:
           orm_mode = True
   

   这样构造返回对象时，可选字段可以设为None或占位值。

2. 修改接口逻辑，替换异常为自定义返回对象
   移除抛出HTTPException的代码，当未查询到用户时，直接构造符合UserResponse结构的对象，将status设为0，同时保持接口返回200状态码：

   @router.get("/{id}", response_model=schemas.UserResponse, status_code=status.HTTP_200_OK)
   def get_user(id: int, db: Session = Depends(get_db)):
       
       user = db.query(models.User).filter(models.User.id == id).first()
       print(user) # checking output
       
       if not user:
           # 构造未找到用户时的返回对象
           return schemas.UserResponse(
               id=id,
               name="",
               phone=0,
               status=0,
               created_at=None
           )
   
       return user
   

3. 关键注意事项

   • 如果UserResponse中有必填字段，必须确保构造的返回对象包含该字段的有效值，比如phone是必填int类型，可设为0作为占位值。
   • 无需再返回404状态码，因为需求是返回一个带status=0的正常响应，保持200状态码即可。

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