FastAPI - 如何在中间件中获取响应正文

Question

有没有办法在中间件中获取响应内容？以下代码是从此处复制的。

@app.middleware("http")
async def add_process_time_header(request: Request, call_next):
    start_time = time.time()

    response = await call_next(request)

    process_time = time.time() - start_time
    response.headers["X-Process-Time"] = str(process_time)
    return response

Answer 1

主体response是一个迭代器，一旦迭代完成，就无法再次迭代。因此，您要么必须将所有迭代数据保存到 a list（或bytes变量）并使用它返回自定义Response，要么再次启动迭代器。下面的选项演示了这两种方法。如果您也想将尸体request放入其中middleware，请查看此答案。

\n

选项1

\n

将数据保存到 alist并用于iterate_in_threadpool再次启动迭代器，如此处所述-这就是StreamingResponse使用的内容，如此处所示。

\n

from starlette.concurrency import iterate_in_threadpool\n\n@app.middleware("http")\nasync def some_middleware(request: Request, call_next):\n    response = await call_next(request)\n    response_body = [chunk async for chunk in response.body_iterator]\n    response.body_iterator = iterate_in_threadpool(iter(response_body))\n    print(f"response_body={response_body[0].decode()}")\n    return response\n

\n

注 1：如果您的代码使用StreamingResponse,response_body[0]将仅chunk返回response. 要获取整个response主体，您应该加入该字节（块）列表，如下所示（.decode()返回对象的字符串表示形式bytes）：

\n

print(f"response_body={(b\'\'.join(response_body)).decode()}")\n

\n

注意 2：如果您的StreamingResponse流式传输主体不适合服务器的 RAM（例如，30GB 的响应），则在迭代时可能会遇到内存错误response.body_iterator（这适用于列出的两个选项）在此答案中），除非您循环遍历response.body_iterator（如选项 2 所示），但不是将块存储在内存变量中，而是将其存储在磁盘上的某个位置。但是，您随后需要从该磁盘位置检索整个响应数据并将其加载到 RAM 中，以便将其发送回客户端（这可能会进一步延长响应客户端的延迟）\xe2\x80\ x94 在这种情况下，您可以将内容分块加载到 RAM 中并使用，类似于此处、此处以及此处、此处和此处StreamingResponse的演示（在选项 1 中，您可以将迭代器/生成器函数传递给） 。但是，我不建议遵循这种方法，而是让此类端点返回从中间件中排除的大型流响应，如本答案中所述。iterate_in_threadpool

\n

选项2

\n

下面演示了另一种方法，其中响应正文存储在对象中bytes（而不是列表，如上所示），并用于Response直接返回自定义（以及原始响应的status_code,headers和）。media_type

\n

@app.middleware("http")\nasync def some_middleware(request: Request, call_next):\n    response = await call_next(request)\n    response_body = b""\n    async for chunk in response.body_iterator:\n        response_body += chunk\n    print(f"response_body={response_body.decode()}")\n    return Response(content=response_body, status_code=response.status_code, \n        headers=dict(response.headers), media_type=response.media_type)\n

\n