Fixed bug with ?_trace=1 and large responses, closes #2404

main
Simon Willison 2024-08-21 10:58:17 -07:00
rodzic 9028d7f805
commit 34a6b2ac84
3 zmienionych plików z 23 dodań i 1 usunięć

Wyświetl plik

@ -90,6 +90,7 @@ class AsgiTracer:
async def wrapped_send(message):
nonlocal accumulated_body, size_limit_exceeded, response_headers
if message["type"] == "http.response.start":
response_headers = message["headers"]
await send(message)
@ -102,11 +103,12 @@ class AsgiTracer:
# Accumulate body until the end or until size is exceeded
accumulated_body += message["body"]
if len(accumulated_body) > self.max_body_bytes:
# Send what we have accumulated so far
await send(
{
"type": "http.response.body",
"body": accumulated_body,
"more_body": True,
"more_body": bool(message.get("more_body")),
}
)
size_limit_exceeded = True

Wyświetl plik

@ -62,10 +62,13 @@ class TestClient:
follow_redirects=False,
redirect_count=0,
method="GET",
params=None,
cookies=None,
if_none_match=None,
headers=None,
):
if params:
path += "?" + urlencode(params, doseq=True)
return await self._request(
path=path,
follow_redirects=follow_redirects,

Wyświetl plik

@ -53,6 +53,23 @@ def test_trace(trace_debug):
assert all(isinstance(trace["count"], int) for trace in execute_manys)
def test_trace_silently_fails_for_large_page():
# Max HTML size is 256KB
with make_app_client(settings={"trace_debug": True}) as client:
# Small response should have trace
small_response = client.get("/fixtures/simple_primary_key.json?_trace=1")
assert small_response.status == 200
assert "_trace" in small_response.json
# Big response should not
big_response = client.get(
"/fixtures/-/query.json",
params={"_trace": 1, "sql": "select zeroblob(1024 * 256)"},
)
assert big_response.status == 200
assert "_trace" not in big_response.json
def test_trace_parallel_queries():
with make_app_client(settings={"trace_debug": True}) as client:
response = client.get("/parallel-queries?_trace=1")