From 58a862cee4a065d463da0a08b87134efb6b6e18d Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sun, 21 Apr 2019 10:41:16 -0700
Subject: [PATCH] ?_trace=1 now adds SQL trace info to JSON/HTML response

Also added documentation for it. Refs #435
---
 datasette/app.py  | 23 ++++++++++++++++++-----
 docs/json_api.rst |  5 +++++
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/datasette/app.py b/datasette/app.py
index f032a36b..a0ed8345 100644
--- a/datasette/app.py
+++ b/datasette/app.py
@@ -624,7 +624,7 @@ class Datasette:
             else:
                 return Results(rows, False, cursor.description)
 
-        with trace("sql", (db_name, sql, params)):
+        with trace("sql", (db_name, sql.strip(), params)):
             results = await self.execute_against_connection_in_thread(
                 db_name, sql_operation_in_thread
             )
@@ -729,10 +729,23 @@ class Datasette:
                 return response.redirect(path)
 
         @app.middleware("response")
-        async def print_traces(request, response):
-            if request.get("traces") is not None:
-                print(json.dumps(request["traces"], indent=2))
-                print("Num traces: {}".format(len(request["traces"])))
+        async def add_traces_to_response(request, response):
+            if request.get("traces") is None:
+                return
+            traces = request["traces"]
+            if "text/html" in response.content_type and b'</body>' in response.body:
+                extra = json.dumps(traces, indent=2)
+                extra_html = "<pre>{}</pre></body>".format(extra).encode("utf8")
+                response.body = response.body.replace(b"</body>", extra_html)
+            elif "json" in response.content_type and response.body.startswith(b"{"):
+                data = json.loads(response.body)
+                if "_traces" not in data:
+                    data["_traces"] = {
+                        "num_traces": len(traces),
+                        "traces": traces,
+                        "duration_sum_ms": sum(t[-1] for t in traces),
+                    }
+                    response.body = json.dumps(data).encode("utf8")
 
         @app.exception(Exception)
         def on_exception(request, exception):
diff --git a/docs/json_api.rst b/docs/json_api.rst
index ef1b4548..af2e53ba 100644
--- a/docs/json_api.rst
+++ b/docs/json_api.rst
@@ -303,6 +303,11 @@ Special table arguments
     Pagination by continuation token - pass the token that was returned in the
     ``"next"`` property by the previous page.
 
+``?_trace=1``
+    Turns on tracing for this page: SQL queries executed during the request will
+    be gathered and included in the response, either in a new ``"_traces"`` key
+    for JSON responses or at the bottom of the page if the response is in HTML.
+
 .. _expand_foreign_keys:
 
 Expanding foreign key references