From 34a6b2ac844a0784fae1f36e0243336a48413594 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Wed, 21 Aug 2024 10:58:17 -0700
Subject: [PATCH] Fixed bug with ?_trace=1 and large responses, closes #2404

---
 datasette/tracer.py        |  4 +++-
 datasette/utils/testing.py |  3 +++
 tests/test_tracer.py       | 17 +++++++++++++++++
 3 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/datasette/tracer.py b/datasette/tracer.py
index fc7338b0..29dd4556 100644
--- a/datasette/tracer.py
+++ b/datasette/tracer.py
@@ -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
diff --git a/datasette/utils/testing.py b/datasette/utils/testing.py
index d4990784..1606da05 100644
--- a/datasette/utils/testing.py
+++ b/datasette/utils/testing.py
@@ -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,
diff --git a/tests/test_tracer.py b/tests/test_tracer.py
index ceadee50..1a4074b0 100644
--- a/tests/test_tracer.py
+++ b/tests/test_tracer.py
@@ -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")