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

有没有办法在中间件中获取响应内容？ 以下代码是从here.

@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

The responsebody 是一个迭代器，一旦迭代完成，就无法再次迭代。因此，您要么必须将所有迭代数据保存到list (or bytes变量）并使用它返回自定义Response，或者再次启动迭代器。下面的选项演示了这两种方法。如果您想获得request身体里面的middleware另外，请看一下这个答案.

Option 1

将数据保存到list并使用iterate_in_threadpool再次初始化迭代器，如上所述here- 这是什么StreamingResponse使用，如图here.

from starlette.concurrency import iterate_in_threadpool

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

Note 1:如果您的代码使用StreamingResponse, response_body[0]只会返回第一个chunk of the response。为了得到整个responsebody，您应该加入该字节（块）列表，如下所示（.decode()返回的字符串表示形式bytes目的）：

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

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

Option 2

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

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