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

Question

有没有办法在中间件中获取响应内容？以下代码是 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

Answer 1

response 主体是一个迭代器，一旦被迭代，就不能再被重新迭代。因此，您必须将所有迭代数据保存到 list(或 bytes 变量)并使用它来返回自定义 Response ，或再次启动迭代器。下面的选项演示了这两种方法。如果您还想在 middleware 中获取 request 正文，请查看 this answer .

选项 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

注意 1: 如果您的代码使用 StreamingResponse，response_body[0] 将只打印出第一个 chunk 响应。要获取整个 response_body，您应该加入该字节列表( block )，如下所示(.decode() 返回 bytes< 的字符串表示 对象):

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

注意 2:如果您有一个 StreamingResponse 流式传输不适合您服务器 RAM 的主体(例如，30GB 的响应)，您可以运行迭代 response.body_iterator 时出现内存错误。这适用于此答案中列出的两个选项。

选项 2

下面演示了另一种方法，其中响应正文存储在 bytes 对象中(而不是列表，如上所示)，并用于返回自定义 Response直接(连同原始响应的 status_code、headers 和 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)