rbhalla
rbhalla

Reputation: 1189

Excluding pydantic model fields only when returned as part of a FastAPI call

Context

I have a very complex pydantic model with a lot of nested pydantic models. I would like to ensure certain fields are never returned as part of API calls, but I would like those fields present for internal logic.

What I tried

I first tried using pydantic's Field function to specify the exclude flag on the fields I didn't want returned. This worked, however functions in my internal logic had to override this whenever they called .dict() by calling .dict(exclude=None).

Instead, I specified a custom flag return_in_api on the Field, with the goal being to only apply exclusions when FastAPI called .dict(). I tried writing a middleware to call .dict() and pass through my own exclude property based on which nested fields contained return_in_api=False. However FastAPI's middleware was giving me a stream for the response which I didn't want to prematurely resolve.

Instead, I wrote a decorator that called .dict() on the return values of my route handlers with the appropriate exclude value.

Problem

One challenge is that whenever new endpoints get added, the person who added them has to remember to include this decorator, otherwise fields leak.

Ideally I would like to apply this decorator to every route, but doing it through middleware seems to break response streaming.

Upvotes: 5

Views: 11960

Answers (1)

Daniil Fajnberg
Daniil Fajnberg

Reputation: 18663

Excluding fields systematically for all routes

I find it best to work with a concrete albeit super simple example. Let's assume you have the following model:

from pydantic import BaseModel, Field


class Demo(BaseModel):
    foo: str
    bar: str = Field(return_in_api=False)

We want to ensure that bar is never returned in a response, both when the response_model is explicitly provided as an argument to the route decorator and when it is just set as the return annotation for the route handler function. (Assume we do not want to use the built-in exclude parameter for our fields for whatever reason.)

The most reliable way that I found is to subclass fastapi.routing.APIRoute and hook into its __init__ method. By copying a tiny bit of the parent class' code, we can ensure that we will always get the correct response model. Once we have that, it is just a matter of setting up the route's response_model_exclude argument before calling the parent constructor.

Here is what I would suggest:

from collections.abc import Callable
from typing import Any

from fastapi.responses import Response
from fastapi.dependencies.utils import get_typed_return_annotation, lenient_issubclass
from fastapi.routing import APIRoute, Default, DefaultPlaceholder


class CustomAPIRoute(APIRoute):
    def __init__(
        self,
        path: str,
        endpoint: Callable[..., Any],
        *,
        response_model: Any = Default(None),
        **kwargs: Any,
    ) -> None:
        # We need this part to ensure we get the response model,
        # even if it is just set as an annotation on the handler function.
        if isinstance(response_model, DefaultPlaceholder):
            return_annotation = get_typed_return_annotation(endpoint)
            if lenient_issubclass(return_annotation, Response):
                response_model = None
            else:
                response_model = return_annotation
        # Find the fields to exclude:
        if response_model is not None:
            kwargs["response_model_exclude"] = {
                name
                for name, field in response_model.__fields__.items()
                if field.field_info.extra.get("return_in_api") is False
            }
        super().__init__(path, endpoint, response_model=response_model, **kwargs)

We can now set that custom route class on our router (documentation). That way it will be used for all its routes:

from fastapi import FastAPI
# ... import CustomAPIRoute
# ... import Demo

api = FastAPI()
api.router.route_class = CustomAPIRoute


@api.get("/demo1")
async def demo1() -> Demo:
    return Demo(foo="a", bar="b")


@api.get("/demo2", response_model=Demo)
async def demo2() -> dict[str, Any]:
    return {"foo": "x", "bar": "y"}

Trying this simple API example out with uvicorn and GETting the endpoints /demo1 and /demo2 yields the responses {"foo":"a"} and {"foo":"x"} respectively.


Ensuring schema consistency

It is however worth mentioning that (unless we take additional steps) the bar field will still be part of the schema. That means for example the auto-generated OpenAPI documentation for both those endpoints will show bar as a top-level property of the response to expect.

This was not part of your question so I assume you are aware of this and are taking measures to ensure consistency. If not, and for others reading this, you can define a static schema_extra method on the Config of your base model to delete those fields that will never be shown "to the outside" before the schema is returned:

from typing import Any
from pydantic import BaseModel, Field


class CustomBaseModel(BaseModel):
    class Config:
        @staticmethod
        def schema_extra(schema: dict[str, Any]) -> None:
            properties = schema.get("properties", {})
            to_delete = set()
            for name, prop in properties.items():
                if prop.get("return_in_api") is False:
                    to_delete.add(name)
            for name in to_delete:
                del properties[name]


class Demo(CustomBaseModel):
    foo: str
    bar: str = Field(return_in_api=False)

Upvotes: 6

Related Questions