主页 / 问题 / 1576677
Accepted
Universall
Asked:2024-04-18 05:17:05 +0800 CST2024-04-18 05:17:05 +0800 CST

Fastapi：如何使用中间件更改响应内容

  • 772

任务

我想编写将数据添加到最终响应的中间件。也就是说，这个答案：

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

它应该变成这样：

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

为此我编写了中间件：

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, but starlette.middleware.base._CachedRequest，但是执行控制器后我收到的输入不是fastapi.Response, but starlette.middleware.base._StreamingResponse。

问题）

是否可以像 django 中那样做 +- ？也就是说，使用中间件将附加信息写入请求中，然后以某种方式更改最终响应。

是否有可能以某种方式使用对象fastapi而不是starlette中间件？

python

1 个回答

  1. Best Answer

    我在网上找到了几种解决方案：

    解决方案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执行中间件后最终（新）对象仍然会变成

    解决方案 2：自动装饰器或route_class

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

    当然，这不是中间件，但它解决了问题。fastapi.APIRouter您可以为（来自该路由器要使用的route_class文档自定义路由（路径操作）类）指定一个参数。

    它必须继承自fastapi.routing.APIRoute. 该方法get_route_handler可用于替换最终处理程序，即创建一个将应用于给定路由器的所有端点的装饰器。

    这是源代码中稍微修改过的示例：

    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)
    • 0

相关问题