From 8a63cdccc7744e6c6969edf47f7c519bf4c25fa6 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Wed, 21 Aug 2024 12:19:18 -0700
Subject: [PATCH] Tracer now catches errors, closes #2405

---
 datasette/database.py |  3 +++
 datasette/tracer.py   | 31 +++++++++++++++++++------------
 tests/test_tracer.py  | 14 ++++++++++++++
 3 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/datasette/database.py b/datasette/database.py
index 8d51befda7..c761dad76b 100644
--- a/datasette/database.py
+++ b/datasette/database.py
@@ -644,6 +644,9 @@ def __init__(self, e, sql, params):
         self.sql = sql
         self.params = params
 
+    def __str__(self):
+        return "QueryInterrupted: {}".format(self.e)
+
 
 class MultipleValues(Exception):
     pass
diff --git a/datasette/tracer.py b/datasette/tracer.py
index 29dd455677..9e66613b99 100644
--- a/datasette/tracer.py
+++ b/datasette/tracer.py
@@ -32,7 +32,7 @@ def trace_child_tasks():
 
 
 @contextmanager
-def trace(type, **kwargs):
+def trace(trace_type, **kwargs):
     assert not TRACE_RESERVED_KEYS.intersection(
         kwargs.keys()
     ), f".trace() keyword parameters cannot include {TRACE_RESERVED_KEYS}"
@@ -45,17 +45,24 @@ def trace(type, **kwargs):
         yield kwargs
         return
     start = time.perf_counter()
-    yield kwargs
-    end = time.perf_counter()
-    trace_info = {
-        "type": type,
-        "start": start,
-        "end": end,
-        "duration_ms": (end - start) * 1000,
-        "traceback": traceback.format_list(traceback.extract_stack(limit=6)[:-3]),
-    }
-    trace_info.update(kwargs)
-    tracer.append(trace_info)
+    captured_error = None
+    try:
+        yield kwargs
+    except Exception as ex:
+        captured_error = ex
+        raise
+    finally:
+        end = time.perf_counter()
+        trace_info = {
+            "type": trace_type,
+            "start": start,
+            "end": end,
+            "duration_ms": (end - start) * 1000,
+            "traceback": traceback.format_list(traceback.extract_stack(limit=6)[:-3]),
+            "error": str(captured_error) if captured_error else None,
+        }
+        trace_info.update(kwargs)
+        tracer.append(trace_info)
 
 
 @contextmanager
diff --git a/tests/test_tracer.py b/tests/test_tracer.py
index 1a4074b004..1e0d700156 100644
--- a/tests/test_tracer.py
+++ b/tests/test_tracer.py
@@ -70,6 +70,20 @@ def test_trace_silently_fails_for_large_page():
         assert "_trace" not in big_response.json
 
 
+def test_trace_query_errors():
+    with make_app_client(settings={"trace_debug": True}) as client:
+        response = client.get(
+            "/fixtures/-/query.json",
+            params={"_trace": 1, "sql": "select * from non_existent_table"},
+        )
+        assert response.status == 400
+
+    data = response.json
+    assert "_trace" in data
+    trace_info = data["_trace"]
+    assert trace_info["traces"][-1]["error"] == "no such table: non_existent_table"
+
+
 def test_trace_parallel_queries():
     with make_app_client(settings={"trace_debug": True}) as client:
         response = client.get("/parallel-queries?_trace=1")