БЛОГ НА HUSL


Fastapi: Как изменить содержимое ответа, используя middleware

Автор вопроса: Universall

Задача

Хочу написать middleware, который будет добавлять данные к итоговому ответу. Т.е такой ответ:

{
    "user": ...,
    "some_field": ...
}

Должен превратиться в такой:

{
    "some_additional_field": ...,
    "data": {
        "user": ...,
        "some_field": ...
    }
}

Для этого я написал middleware:

class ResponseFormatterMiddleware:
    async def _post_process_response(self, response: Response) -> Response:
        response.body = response.render({
            "some_additional_field": 123,
            "data": response.body
        })

        return response

    async def __call__(self, request: Request, call_next) -> Response:
        response = await call_next(request)
        return await self._post_process_response(response)

И добавил его с помощью:

app.middleware("http")(ResponseFormatterMiddleware())

Но получил ошибку:

AttributeError: '_StreamingResponse' object has no attribute 'body'

После чего я выяснил, что на вход получаю не fastapi.Request, а starlette.middleware.base._CachedRequest, а на после выполнения контроллера я получаю не fastapi.Response, а starlette.middleware.base._StreamingResponse.

Вопрос(ы)

Можно ли сделать сделать +- как в django? Т.е с помощью middleware записывать в request дополнительную информацию, а потом как-либо изменять итоговый ответ.

Можно ли каким-то образом работать с объектами fastapi, а не starlette в middleware?


Источник

Ответы (1 шт):

Автор решения: Universall

На просторах интернета я нашёл пару решений:

Решение 1: В лоб

Достать данные ответа можно с помощью нехитрого выражения:

body: bytes = b"".join([chunk async for chunk in response.body_iterator])

На основе него я смастерил вот такой пример:

import json
from typing import Any

import fastapi
import uvicorn
from starlette.middleware.base import BaseHTTPMiddleware, _StreamingResponse
from starlette.responses import Response

application = fastapi.FastAPI(root_path="/api/")


async def to_fastapi_response(response: _StreamingResponse) -> Response:
    return Response(
        content=await get_streaming_response_body(response),
        status_code=response.status_code,
        headers=response.headers,
        media_type=response.media_type,
        background=response.background
    )


async def get_streaming_response_body(response: _StreamingResponse) -> bytes:
    return b"".join([chunk async for chunk in response.body_iterator])


def update_response_body(response: Response, new_body: Any) -> None:
    response.body = response.render(new_body)  # Изменяем содержимое ответа
    response.headers["content-length"] = str(len(response.body))  # Обновляем длину ответа


class ResponseModifierMiddleware(BaseHTTPMiddleware):
    async def _get_new_body(self, response: Response) -> Any:
        return json.dumps({
            "some_additional_filed": 123,
            "data": json.loads(response.body)
        })

    async def dispatch(self, request, call_next):
        original_response = await call_next(request)
        response = await to_fastapi_response(original_response)  # type: ignore

        if not response.body:
            return response

        update_response_body(
            response=response,  # type: ignore
            new_body=await self._get_new_body(response)  # type: ignore
        )

        return response


@application.get("test/")
async def test_endpoint():
    return {
        "user": {
            "id": 1,
            "username": "test"
        },
        "some_field": 321
    }


application.add_middleware(ResponseModifierMiddleware)
uvicorn.run(application)

Минусы:

  1. Идея пересоздавать объект, как по мне, не очень хорошая
  2. Итоговый (новый) объект всё равно превратится в _StreamingResponse после выполнения middleware

Решение 2: Автоматический декоратор или route_class

https://stackoverflow.com/questions/64115628/get-starlette-request-body-in-the-middleware-context.

Это, конечно, не middleware, но задачу решает. Для fastapi.APIRouter можно указать параметр route_class(Из документации Custom route (path operation) class to be used by this router.).

Он должен наследоваться от fastapi.routing.APIRoute. Метод get_route_handler можно использовать для подмены конечного handler-а, то есть, по сути, для создания декоратора, который будет применён для всех endpoint-ов данного роутера.

Вот немного модифицированный пример из источника:

import json
from typing import Callable, Any

import uvicorn
from fastapi import APIRouter, FastAPI, Request, Response
from fastapi.routing import APIRoute


def update_response_body(response: Response, new_body: Any) -> None:
    response.body = response.render(new_body)
    response.headers["content-length"] = str(len(response.body))


class CustomRoute(APIRoute):
    def get_route_handler(self) -> Callable:
        original_route_handler = super().get_route_handler()

        async def custom_route_handler(request: Request) -> Response:
            response: Response = await original_route_handler(request)

            update_response_body(response, {
                "some_additional_field": 123,
                "data": json.loads(response.body)
            })

            return response

        return custom_route_handler


application = FastAPI(root_path="/api/")
router = APIRouter(route_class=CustomRoute)


@router.get("test/")
async def test_endpoint():
    return {
        "user": {
            "id": 1,
            "username": "test"
        },
        "some_field": 321
    }


application.include_router(router)
uvicorn.run(application)
→ Ссылка

licensed under cc by-sa 3.0 with attribution.