From 669fa21a71c9920923df0f64907c095c0e7d889f Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Wed, 1 May 2019 16:16:15 -0700
Subject: [PATCH] Include request duration in traces

---
 datasette/app.py  | 11 ++++++++---
 tests/test_api.py |  4 ++--
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/datasette/app.py b/datasette/app.py
index d3a8168e..737b5654 100644
--- a/datasette/app.py
+++ b/datasette/app.py
@@ -6,6 +6,7 @@ import json
 import os
 import sys
 import threading
+import time
 import traceback
 import urllib.parse
 from concurrent import futures
@@ -654,6 +655,7 @@ class Datasette:
             async def handle_request(self, request, write_callback, stream_callback):
                 if request.args.get("_trace"):
                     request["traces"] = []
+                    request["trace_start"] = time.time()
                     with capture_traces(request["traces"]):
                         res = await super().handle_request(request, write_callback, stream_callback)
                 else:
@@ -756,7 +758,10 @@ class Datasette:
         async def add_traces_to_response(request, response):
             if request.get("traces") is None:
                 return
-            traces = request["traces"]
+            traces = {
+                "duration": time.time() - request["trace_start"],
+                "queries": 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")
@@ -765,9 +770,9 @@ class Datasette:
                 data = json.loads(response.body.decode("utf8"))
                 if "_traces" not in data:
                     data["_traces"] = {
-                        "num_traces": len(traces),
+                        "num_traces": len(traces["queries"]),
                         "traces": traces,
-                        "duration_sum_ms": sum(t[-1] for t in traces),
+                        "duration_sum_ms": sum(t[-1] for t in traces["queries"]),
                     }
                     response.body = json.dumps(data).encode("utf8")
 
diff --git a/tests/test_api.py b/tests/test_api.py
index 9b75da6e..5043b7c2 100644
--- a/tests/test_api.py
+++ b/tests/test_api.py
@@ -1452,5 +1452,5 @@ def test_trace(app_client):
     traces = data["_traces"]
     assert isinstance(traces["duration_sum_ms"], float)
     assert isinstance(traces["num_traces"], int)
-    assert isinstance(traces["traces"], list)
-    assert len(traces["traces"]) == traces["num_traces"]
+    assert isinstance(traces["traces"], dict)
+    assert len(traces["traces"]["queries"]) == traces["num_traces"]