'FastAPI - How to get the response body in Middleware

Is there any way to get the response content in a middleware? The following code is a copy from 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


Solution 1:[1]

The response body is an iterator, which once it has been iterated through, it cannot be re-iterated again. Thus, you either have to save all the iterated data to a list (or bytes variable) and use that to return a custom Response, or initiate the iterator again. The options below demonstrate both approaches.

Option 1

Save the data to a list and use iterate_in_threadpool to initiate the iterator again, as described here - which is what StreamingResponse uses, as shown 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: If your code uses StreamingResponse, response_body[0] would print out only the first chunk of the response. To get the whole response_body, you should join that list of bytes (chunks), as shown below (.decode() returns a string representation of the bytes object):

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

Note 2: If you have a StreamingResponse streaming a body that wouldn't fit into your server's RAM (for example, a response of 30GB), you may run into memory errors when iterating over the response.body_iterator. This applies to both options listed in this answer.

Option 2

The below demosntrates another approach, where the response body is stored in a bytes object (instead of a list, as shown above), and is used to return a custom Response directly (along with the status_code, headers and media_type of the original response).

@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)

Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source
Solution 1