FastAPI - 如何在中间件中获取响应主体

问题描述：

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

@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

解决方案 1：

主体response是一个迭代器，一旦迭代过，就不能再重新迭代。因此，您必须将所有迭代数据保存到list（或bytes变量）并使用它来返回自定义的Response，或者再次启动迭代器。以下选项演示了这两种方法。如果您也想request在里面获取主体middleware，请查看此答案。

选项 1

将数据保存到list并用来iterate_in_threadpool再次启动迭代器，如此处所述—这就是StreamingResponse用途，如此处所示。

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。要获取整个response主体，您应该加入该字节列表（块），如下所示（.decode()返回对象的字符串表示形式bytes）：

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

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

选项 2

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

@app.middleware("http")
async def some_middleware(request: Request, call_next):
    response = await call_next(request)
    chunks = []
    async for chunk in response.body_iterator:
        chunks.append(chunk)
    response_body = b''.join(chunks)
    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)