python - 如何在FastAPI中间件中修改响应后的内容长度？

Question

示例代码:

# Here is a minimal reproducible example

import json
from starlette.datastructures import MutableHeaders
from starlette.types import ASGIApp, Receive, Scope, Send, Message
import datetime
import socket
import uvicorn
from fastapi import FastAPI

class MetaDataAdderMiddleware:
    application_generic_urls = ['/openapi.json', '/docs', '/docs/oauth2-redirect', '/redoc']
    
    def __init__(self, app: ASGIApp) -> None:
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:

        start_time = datetime.datetime.utcnow()

        async def send_wrapper(message: Message) -> None:   
         
            if message["type"] == "http.response.body" and len(message["body"]) and not any([scope["path"].startswith(endpoint) for endpoint in MetaDataAdderMiddleware.application_generic_urls]):

                response_body = json.loads(message["body"].decode())

                end_time = datetime.datetime.utcnow()
                response_processing_time_seconds = end_time - start_time
                data = {}
                data["data"] = response_body
                data['metadata'] = {
                    'request_timestamp_utc': start_time,
                    'response_timestamp_utc': end_time,
                    'processing_time_seconds': response_processing_time_seconds,
                    'service_host': socket.gethostname()
                }

                data_to_be_sent_to_user = json.dumps(data, default=str).encode("utf-8")
                message["body"] = data_to_be_sent_to_user

            await send(message)
        await self.app(scope, receive, send_wrapper)


app = FastAPI(
    title="MY DUMMY APP",
)

app.add_middleware(MetaDataAdderMiddleware)

@app.get("/")
async def root():
    return {"message": "Hello World"}

描述: 这是我的用例:FastAPI APP 中的所有端点，无论它们发送什么响应，我都需要使用一些元数据来包装该响应。假设某个端点向我发送以下内容:{"data_key": "data_value"}。但是，用户应该看到，这是最终输出:

{
    "data": {"data_key": "data_value"}
     "metadata": {
          "request_timestamp_utc": "somevalue",
          ...and so on
     }
}

我有一个很大的应用程序和许多路由器。至此我们通过编写中间件就实现了添加Request ID、Authentication、Authorization的功能。

但是，当我点击应用程序的 API 时，在添加上述 MetaDataAdderMiddleware 后，我遇到了以下错误:

ERROR:    Exception in ASGI application
Traceback (most recent call last):
  File "<MY PYTHON PATH>/lib/python3.6/site-packages/uvicorn/protocols/http/httptools_impl.py", line 521, in send
    raise RuntimeError("Response content longer than Content-Length")
RuntimeError: Response content longer than Content-Length

这个错误是足够合乎逻辑的，因为我修改了响应正文，但没有更改内容长度 header 。

这是 send_wrapper 函数级别中的消息和作用域主体的快照，以及我从 httptools_impl.py 级别打印的 header 值: (我编辑了一些字段，以掩盖组织特定的内容)

send_wrapper called
message: {'type': 'http.response.start', 'status': 200, 'headers': [(b'content-length', b'58'), (b'content-type', b'application/json')]}
scope: {'type': 'http', 'asgi': {'version': '3.0', 'spec_version': '2.1'}, 'http_version': '1.1', 'scheme': 'http', 'method': 'POST', 'root_path': '', 'query_string': b'', 'headers': [(b'content-type', b'application/json'), (b'accept', b'*/*'), (b'cache-control', b'no-cache'), (b'accept-encoding', b'gzip, deflate'), (b'content-length', b'238'), (b'connection', b'keep-alive')], 'app': <fastapi.applications.FastAPI object at >, 'fastapi_astack': <contextlib2.AsyncExitStack object at >, 'router': <fastapi.routing.APIRouter object at >, 'endpoint': <function initiate_playbook_execution at >, 'path_params': {}, 'route': <fastapi.routing.APIRoute object at >}
INFO:     - "POST /MYAPI" 200 OK

INSIDE httptools_impl
name: b'content-length' | value: b'58'
self.expected_content_length: 58

send_wrapper called
message: {'type': 'http.response.body', 'body': b'{"status":true,"stdout":null,"stderr":null,"message":null}'}
scope: {'type': 'http', 'asgi': {'version': '3.0', 'spec_version': '2.1'}, 'http_version': '1.1', 'scheme': 'http', 'method': 'POST', 'root_path': '', 'query_string': b'', 'headers': [(b'content-type', b'application/json'), (b'accept', b'*/*'), (b'cache-control', b'no-cache'), (b'accept-encoding', b'gzip, deflate'), (b'content-length', b'238'), (b'connection', b'keep-alive')], 'app': <fastapi.applications.FastAPI object at >, 'fastapi_astack': <contextlib2.AsyncExitStack object at >, 'router': <fastapi.routing.APIRouter object at >, 'endpoint': <function initiate_playbook_execution at >, 'path_params': {}, 'route': <fastapi.routing.APIRoute object at >}

INSIDE httptools_impl
body: b'{"data": {"status": true, "stdout": null, "stderr": null, "message": null}, "metadata": {"request_timestamp_utc": "BLAH", "response_timestamp_utc": "BLAH", "processing_time_seconds": "0:00:00.469472", "some_field": "some_value"}}'
num_bytes: 286

以下是我更新内容长度的尝试:

• 在更新响应正文后的发送包装函数中，我尝试执行以下操作:

  data_to_be_sent_to_user = json.dumps(data, default=str).encode("utf-8")
  message["body"] = data_to_be_sent_to_user
  headers = MutableHeaders(scope=scope)
  headers["content-length"] = str(len(data_to_be_sent_to_user))
  
  # But this hasn't worked, no change in situation!
  

我该如何继续？

Answer 1

感谢 @MatsLindh 评论，我在这里引用了 Starlette 的 GZipMiddleware 代码库:https://github.com/encode/starlette/blob/fcc4c705ff69182ebd663bc686cb55c242d32683/starlette/middleware/gzip.py#L60

所以想法是，有问题的 content-length 值位于 http.response.start 消息中存在的 header 中。因此，GZipMiddleware 的编写方式是，他们根本没有立即发送第一个 http.response.start 消息。相反，他们还捕获 http.response.body，然后修改响应，找到其长度，然后更新 http.response.start 消息中的长度，然后发送这两条消息的顺序都是正确的。

我能够编写的工作实现，大量借鉴了 GZipMiddleware，如下:

from starlette.types import ASGIApp, Receive, Scope, Send, Message
from starlette.requests import Request
import json
from starlette.datastructures import MutableHeaders


class MetaDataAdderMiddleware:
    application_generic_urls = ['/openapi.json', '/docs', '/docs/oauth2-redirect', '/redoc']
    
    def __init__(
            self, 
            app: ASGIApp
    ) -> None:
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] == "http" and not any([scope["path"].startswith(endpoint) for endpoint in MetaDataAdderMiddleware.application_generic_urls]):
            responder = MetaDataAdderMiddlewareResponder(self.app, self.standard_meta_data, self.additional_custom_information)
            await responder(scope, receive, send)
            return
        await self.app(scope, receive, send)


class MetaDataAdderMiddlewareResponder:

    def __init__(
            self,
            app: ASGIApp,
    ) -> None:
        """
        """
        self.app = app
        self.initial_message: Message = {}

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        self.send = send
        await self.app(scope, receive, self.send_with_meta_response)

    async def send_with_meta_response(self, message: Message):

        message_type = message["type"]
        if message_type == "http.response.start":
            # Don't send the initial message until we've determined how to
            # modify the outgoing headers correctly.
            self.initial_message = message

        elif message_type == "http.response.body":
            response_body = json.loads(message["body"].decode())

            data = {}
            data["data"] = response_body
            data['metadata'] = {
                'field_1': 'value_1',
                'field_2': 'value_2'
            }

            data_to_be_sent_to_user = json.dumps(data, default=str).encode("utf-8")

            headers = MutableHeaders(raw=self.initial_message["headers"])
            headers["Content-Length"] = str(len(data_to_be_sent_to_user))
            message["body"] = data_to_be_sent_to_user

            await self.send(self.initial_message)
            await self.send(message)