'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 |
