Conditionally set FastAPI response model for route

user_12104222 0
Viewed 2586

I'm trying to return a list of objects of type Company, including only "approved" ones, and with more or less attributes depending on whether the user requesting the list is a superuser or a regular user. This is my code so far:

@router.get("/", response_model=List[schema.CompanyRegularUsers])
def get_companies(db: Session = Depends(get_db), is_superuser: bool = Depends(check_is_superuser)):
    """
    If SU, also include sensitive data.
    """
    if is_superuser:
        return crud.get_companies_admin(db=db)
    return crud.get_companies_user(db=db)
#

The function correctly returns the objects according to request (ie., only is_approved=True companies if a regular request, and both is_approved=True and is_approved=False if requested by a superuser. Problem is, both cases use schema.CompanyRegularUsers, and I'd like to use schema.CompanySuperusers when SU's make the request.

How can I achieve that feature? I.e, is there a way to conditionally set the response_model property of the decorator function?

I've tried using JSONResponse and calling Pydantic's schema.CompanySuperusers.from_orm(), but it won't work with a list of Companies...

pydantic fastapi sqlalchemy orm python
2 Answers

I ended up solving the riddle by returning a custom JSONResponse. It doesn't show up in the automatic documentation, but I think I can tackle that down the road. Code is as follows, in case it helps someone else:

...
from pydantic import parse_obj_as
from fastapi.responses import JSONResponse
from fastapi.encoders import jsonable_encoder
...

@router.get("/", response_model=List[schema.CompanyRegularUsers])
def get_companies(db: Session = Depends(get_db), is_superuser: bool = Depends(check_is_superuser)):
    """
    If SU, also include sensitive data.
    """
    if is_superuser:
        companies = parse_obj_as(List[schema.CompanyAdmin], crud.get_companies_admin(db=db))
        return JSONResponse(jsonable_encoder(companies))
    return crud.get_companies_user(db=db)

So, in the is_admin branch, the path operation calls pydantic's parse_obj_as in order to map the list of objects returned by SQLAlchemy's query as a list of CompanyAdmin objects. Then, it uses jsonable_encoder, the encoder FastAPI uses under the hood for every default response, to serialize the list.

Edit: Typo

user_121042220

You can try to use Union type operator.

Your code would become

from typing import Union

@router.get("/", response_model=List[Union[schema.CompanyRegularUsers, schema.CompanySuperUser]])

this way, you specify as response model a list of either schema.CompanyRegularUsers or schema.CompanySuperUser

Let me know if it works, since I didn't test it

user_121201010
Related