From b1c6db4b8f80b48a45ff9bd3dbba70ae427f8343 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sun, 23 Jun 2019 12:55:16 -0700
Subject: [PATCH] Re-implemented tracing, refs #272

---
 datasette/app.py    |  6 ++--
 datasette/tracer.py | 75 +++++++++++++++++++++++++++++++++++++++++++++
 tests/test_api.py   |  1 -
 3 files changed, 78 insertions(+), 4 deletions(-)

diff --git a/datasette/app.py b/datasette/app.py
index 22584379..4c85a78e 100644
--- a/datasette/app.py
+++ b/datasette/app.py
@@ -37,7 +37,7 @@ from .utils import (
     to_css_class,
 )
 from .utils.asgi import asgi_static, asgi_send_html, asgi_send_json, asgi_send_redirect
-from .tracer import capture_traces, trace
+from .tracer import capture_traces, trace, AsgiTracer
 from .plugins import pm, DEFAULT_PLUGINS
 from .version import __version__
 
@@ -127,7 +127,7 @@ CONFIG_OPTIONS = (
 DEFAULT_CONFIG = {option.name: option.default for option in CONFIG_OPTIONS}
 
 
-async def favicon(scope, recieve, send):
+async def favicon(scope, receive, send):
     await send(
         {
             "type": "http.response.start",
@@ -712,4 +712,4 @@ class Datasette:
         #         if not database.is_mutable:
         #             await database.table_counts(limit=60 * 60 * 1000)
 
-        return app
+        return AsgiTracer(app)
diff --git a/datasette/tracer.py b/datasette/tracer.py
index c6fe0a00..4a46f1e6 100644
--- a/datasette/tracer.py
+++ b/datasette/tracer.py
@@ -1,6 +1,7 @@
 import asyncio
 from contextlib import contextmanager
 import time
+import json
 import traceback
 
 tracers = {}
@@ -53,3 +54,77 @@ def capture_traces(tracer):
     tracers[task_id] = tracer
     yield
     del tracers[task_id]
+
+
+class AsgiTracer:
+    # If the body is larger than this we don't attempt to append the trace
+    max_body_bytes = 1024 * 256  # 256 KB
+
+    def __init__(self, app):
+        self.app = app
+
+    async def __call__(self, scope, receive, send):
+        if b"_trace=1" not in scope.get("query_string", b"").split(b"&"):
+            await self.app(scope, receive, send)
+            return
+        trace_start = time.time()
+        traces = []
+
+        accumulated_body = b""
+        size_limit_exceeded = False
+        response_headers = []
+
+        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)
+                return
+
+            if message["type"] != "http.response.body" or size_limit_exceeded:
+                await send(message)
+                return
+
+            # Accumulate body until the end or until size is exceeded
+            accumulated_body += message["body"]
+            if len(accumulated_body) > self.max_body_bytes:
+                await send(
+                    {
+                        "type": "http.response.body",
+                        "body": accumulated_body,
+                        "more_body": True,
+                    }
+                )
+                size_limit_exceeded = True
+                return
+
+            if not message.get("more_body"):
+                # We have all the body - modify it and send the result
+                # TODO: What to do about Content-Type or other cases?
+                trace_info = {
+                    "request_duration_ms": 1000 * (time.time() - trace_start),
+                    "sum_trace_duration_ms": sum(t["duration_ms"] for t in traces),
+                    "num_traces": len(traces),
+                    "traces": traces,
+                }
+                try:
+                    content_type = [
+                        v.decode("utf8")
+                        for k, v in response_headers
+                        if k.lower() == b"content-type"
+                    ][0]
+                except IndexError:
+                    content_type = ""
+                if "text/html" in content_type and b"</body>" in accumulated_body:
+                    extra = json.dumps(trace_info, indent=2)
+                    extra_html = "<pre>{}</pre></body>".format(extra).encode("utf8")
+                    accumulated_body = accumulated_body.replace(b"</body>", extra_html)
+                elif "json" in content_type and accumulated_body.startswith(b"{"):
+                    data = json.loads(accumulated_body.decode("utf8"))
+                    if "_trace" not in data:
+                        data["_trace"] = trace_info
+                    accumulated_body = json.dumps(data).encode("utf8")
+                await send({"type": "http.response.body", "body": accumulated_body})
+
+        with capture_traces(traces):
+            await self.app(scope, receive, wrapped_send)
diff --git a/tests/test_api.py b/tests/test_api.py
index 96c16175..a32ed5e3 100644
--- a/tests/test_api.py
+++ b/tests/test_api.py
@@ -1613,7 +1613,6 @@ def test_infinity_returned_as_invalid_json_if_requested(app_client):
     ] == response.json
 
 
-@pytest.mark.skip
 def test_trace(app_client):
     response = app_client.get("/fixtures/simple_primary_key.json?_trace=1")
     data = response.json