remote/datasette
1
0
Fork 0
mirror of https://github.com/simonw/datasette.git synced 2026-09-02 22:54:08 +02:00
Code Issues Projects Releases Packages Wiki Activity

Propagate otel context across the thread boundaries

Browse source 
Spans created on a worker thread resolve their parent from that thread's
ambient context, so without this every span produced below Database came
back as an unparented root, disconnected from the request that caused it.
Carrying the caller's context across each boundary is also what makes the
thread-pool wait visible: db.query covers the full round trip, the new
db.query.execute covers only the work inside the worker, and the gap
between them is the queueing the old tracer folds invisibly into one
number.

- execute_fn()'s executor.submit() and execute_isolated_fn()'s
  run_in_executor() (immutable databases) now run the callable inside a
  contextvars.copy_context(). A *fresh* copy per submit is required:
  concurrently entering one shared Context raises "RuntimeError: cannot
  enter context ... already entered".
- WriteTask carries the otel Context captured on the event loop at enqueue
  time plus an enqueued_at_ns timestamp (both need __slots__ entries, or
  they fail with AttributeError at runtime). _execute_writes attaches that
  context right after the _SHUTDOWN check and detaches it in a finally
  spanning all three execution branches - the write thread is persistent
  and shared, so a leaked token would grow its context stack for every
  write processed afterwards, and a wrong-token detach only logs rather
  than raising.
- New spans: db.query.execute (read worker thread), db.write.queue_wait
  (explicit start/end timestamps, so its duration is the real enqueue ->
  dequeue wait rather than the microseconds spent building the span) and
  db.write.execute (skipped in the conn_exception branch, where fn never
  runs). db.query.execute honours log_sql_errors for the same reason
  db.query does: facet suggestion probes with log_sql_errors=False and
  would otherwise paint two red spans per text column on every table page.
- The write-thread warm-up prepare_connection is left as a documented
  orphan root - no caller context exists that early.

Tests assert actual parent/child span-id relationships in a shared trace,
not just that spans exist, since an unparented root looks identical to a
correct span if you only check presence.

Note that copy_context() copies every ContextVar, not just OTel's, so
Datasette's own context vars (_skip_permission_checks,
_permission_check_cache, _in_datasette_client) now flow into worker
threads where they previously did not.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Alex Garcia 2026-07-30 18:05:34 -07:00
commit 582d79a148
3 changed files with 464 additions and 67 deletions
Download patch file Download diff file Expand all files Collapse all files

228
datasette/database.py
View file

@ -1,16 +1,19 @@
import asyncio
import atexit
import contextvars
import inspect
import os
import queue
import sys
import tempfile
import threading
import time
import uuid
from collections import namedtuple
from pathlib import Path
import sqlite_utils
from opentelemetry import context as otel_context_api
from opentelemetry.trace import Status, StatusCode
from .inspect import inspect_hash
@ -351,9 +354,15 @@ class Database:
return _run()
if not write:
# Immutable database - no writes can ever occur, so there is no
# write queue to block; run against a fresh read-only connection
# write queue to block; run against a fresh read-only connection.
# A fresh copy_context() is required per submit (not one shared
# copy reused across calls): concurrent execution of the same
# Context raises "RuntimeError: cannot enter context ... already
# entered". This propagates the caller's otel context (e.g. the
# enclosing db.query span) onto the worker thread.
ctx = contextvars.copy_context()
return await asyncio.get_running_loop().run_in_executor(
self.ds.executor, _run
self.ds.executor, ctx.run, _run
)
# Threaded mode - send to write thread
return await self._send_to_write_thread(fn, isolated_connection=True)
@ -458,8 +467,21 @@ class Database:
task_id = uuid.uuid5(uuid.NAMESPACE_DNS, "datasette.io")
loop = asyncio.get_running_loop()
reply_future = loop.create_future()
# Captured here, on the event loop, at enqueue time: the otel
# Context (carrying the enclosing db.query span, if any) and the
# timestamp used to build the db.write.queue_wait span once this
# task is dequeued on the write thread.
self._write_queue.put(
WriteTask(fn, task_id, loop, reply_future, isolated_connection, transaction)
WriteTask(
fn,
task_id,
loop,
reply_future,
isolated_connection,
transaction,
otel_context_api.get_current(),
time.time_ns(),
)
)
if block:
return await reply_future
@ -473,6 +495,9 @@ class Database:
conn = None
try:
conn = self.connect(write=True)
# This warm-up runs before any write has ever been queued, so
# there is no caller otel context yet to attach - any spans
# created by plugin hooks here are orphans (roots).
self.ds._prepare_connection(conn, self.name)
except Exception as e: # noqa: BLE001
# Stored and re-raised to whoever queues the next write
@ -487,40 +512,79 @@ class Database:
# Best-effort close as the write thread exits
pass
return
exception = None
result = None
if conn_exception is not None:
exception = conn_exception
elif task.isolated_connection:
try:
isolated_connection = self.connect(write=True)
# Restore the caller's otel context (captured on the event loop
# at enqueue time) so spans created while processing this task
# parent correctly to the request that queued it. Must be
# detached below in `finally` - a leaked token silently poisons
# this thread's ambient context for every write processed after
# it, and a *wrong*-token detach only logs a warning rather than
# raising, so this pairing is load-bearing and easy to get wrong
# silently.
token = otel_context_api.attach(task.otel_context)
try:
exception = None
result = None
# Explicit start_time/end_time rather than a `with` block:
# this span's duration is the time the task actually spent
# waiting in the queue (enqueue -> dequeue), not the near-
# zero time spent constructing/ending the span object here.
tracer.start_span(
"db.write.queue_wait", start_time=task.enqueued_at_ns
).end(end_time=time.time_ns())
if conn_exception is not None:
# fn never runs in this branch, so there is nothing to
# wrap in a db.write.execute span.
exception = conn_exception
elif task.isolated_connection:
try:
result = task.fn(isolated_connection)
finally:
isolated_connection.close()
try:
self._all_file_connections.remove(isolated_connection)
except ValueError:
# Was probably a memory connection
pass
except Exception as e: # noqa: BLE001
# Write thread must survive any task failure or the database wedges
sys.stderr.write(f"{e}\n")
sys.stderr.flush()
exception = e
else:
try:
if task.transaction:
with conn:
conn.execute("BEGIN IMMEDIATE")
result = task.fn(conn)
else:
result = task.fn(conn)
except Exception as e: # noqa: BLE001
sys.stderr.write(f"{e}\n")
sys.stderr.flush()
exception = e
_deliver_write_result(task, result, exception)
with tracer.start_as_current_span("db.write.execute") as span:
span.set_attribute(
"datasette.isolated_connection",
task.isolated_connection,
)
span.set_attribute(
"datasette.transaction", task.transaction
)
isolated_connection = self.connect(write=True)
try:
result = task.fn(isolated_connection)
finally:
isolated_connection.close()
try:
self._all_file_connections.remove(
isolated_connection
)
except ValueError:
# Was probably a memory connection
pass
except Exception as e: # noqa: BLE001
# Write thread must survive any task failure or the database wedges
sys.stderr.write(f"{e}\n")
sys.stderr.flush()
exception = e
else:
try:
with tracer.start_as_current_span("db.write.execute") as span:
span.set_attribute(
"datasette.isolated_connection",
task.isolated_connection,
)
span.set_attribute(
"datasette.transaction", task.transaction
)
if task.transaction:
with conn:
conn.execute("BEGIN IMMEDIATE")
result = task.fn(conn)
else:
result = task.fn(conn)
except Exception as e: # noqa: BLE001
sys.stderr.write(f"{e}\n")
sys.stderr.flush()
exception = e
_deliver_write_result(task, result, exception)
finally:
otel_context_api.detach(token)
async def execute_fn(self, fn):
self._check_not_closed()
@ -542,7 +606,13 @@ class Database:
with self._pending_execute_futures_lock:
self._check_not_closed()
future = self.ds.executor.submit(in_thread)
# A fresh copy_context() is required per submit (not one shared
# copy reused across calls): concurrent execution of the same
# Context raises "RuntimeError: cannot enter context ...
# already entered". This propagates the caller's otel context
# (e.g. the enclosing db.query span) onto the worker thread.
ctx = contextvars.copy_context()
future = self.ds.executor.submit(ctx.run, in_thread)
self._pending_execute_futures.add(future)
future.add_done_callback(self._remove_pending_execute_future)
return await asyncio.wrap_future(future)
@ -564,35 +634,51 @@ class Database:
time_limit_ms = custom_time_limit
def sql_operation_in_thread(conn):
with sqlite_timelimit(conn, time_limit_ms):
try:
cursor = conn.cursor()
cursor.execute(sql, params if params is not None else {})
max_returned_rows = self.ds.max_returned_rows
if max_returned_rows == page_size:
max_returned_rows += 1
if max_returned_rows and truncate:
rows = cursor.fetchmany(max_returned_rows + 1)
truncated = len(rows) > max_returned_rows
rows = rows[:max_returned_rows]
else:
rows = cursor.fetchall()
truncated = False
except (sqlite3.OperationalError, sqlite3.DatabaseError) as e:
if e.args == ("interrupted",):
raise QueryInterrupted(e, sql, params)
if log_sql_errors:
sys.stderr.write(
f"ERROR: conn={conn}, sql = {sql!r}, params = {params}: {e}\n"
)
sys.stderr.flush()
raise
# This span is created inside the worker thread. Its parent is
# resolved from the ambient otel context, which was propagated
# onto this thread via copy_context() at the executor.submit()
# boundary in execute_fn() (or run_in_executor() for immutable
# databases) - so it parents correctly to the enclosing
# db.query span despite running on a different thread.
#
# Callers passing log_sql_errors=False are probing and treat a
# failure as an expected answer - see the matching handling on the
# db.query span in execute(). Without this, facet suggestion marks
# two spans per text column as failed on every table page.
with tracer.start_as_current_span(
"db.query.execute",
record_exception=log_sql_errors,
set_status_on_exception=log_sql_errors,
):
with sqlite_timelimit(conn, time_limit_ms):
try:
cursor = conn.cursor()
cursor.execute(sql, params if params is not None else {})
max_returned_rows = self.ds.max_returned_rows
if max_returned_rows == page_size:
max_returned_rows += 1
if max_returned_rows and truncate:
rows = cursor.fetchmany(max_returned_rows + 1)
truncated = len(rows) > max_returned_rows
rows = rows[:max_returned_rows]
else:
rows = cursor.fetchall()
truncated = False
except (sqlite3.OperationalError, sqlite3.DatabaseError) as e:
if e.args == ("interrupted",):
raise QueryInterrupted(e, sql, params)
if log_sql_errors:
sys.stderr.write(
f"ERROR: conn={conn}, sql = {sql!r}, params = {params}: {e}\n"
)
sys.stderr.flush()
raise
if truncate:
return Results(rows, truncated, cursor.description)
if truncate:
return Results(rows, truncated, cursor.description)
else:
return Results(rows, False, cursor.description)
else:
return Results(rows, False, cursor.description)
# SIM117 wants these two context managers merged. They are kept nested
# deliberately: the hand-rolled tracer's wrapper is on its way out, and
@ -924,16 +1010,26 @@ def _apply_write_wrapper(fn, wrapper_factory, track_event):
class WriteTask:
__slots__ = (
"enqueued_at_ns",
"fn",
"isolated_connection",
"loop",
"otel_context",
"reply_future",
"task_id",
"transaction",
)
def __init__(
self, fn, task_id, loop, reply_future, isolated_connection, transaction
self,
fn,
task_id,
loop,
reply_future,
isolated_connection,
transaction,
otel_context,
enqueued_at_ns,
):
self.fn = fn
self.task_id = task_id
@ -941,6 +1037,8 @@ class WriteTask:
self.reply_future = reply_future
self.isolated_connection = isolated_connection
self.transaction = transaction
self.otel_context = otel_context
self.enqueued_at_ns = enqueued_at_ns
def _deliver_write_result(task, result, exception):

109
tests/test_internals_database.py
View file

@ -3,11 +3,13 @@ Tests for the datasette.database.Database class
"""
import asyncio
import threading
import uuid
from types import SimpleNamespace
import pytest
import sqlite_utils
from opentelemetry import context as otel_context_api
from datasette.app import Datasette
from datasette.database import (
@ -1223,3 +1225,110 @@ async def test_database_close_is_idempotent(tmpdir):
# Second call should be a no-op, not raise
db.close()
ds._internal_database.close()
_CONTEXT_LEAK_MARKER_KEY = "otel-context-leak-marker"
@pytest.mark.asyncio
@pytest.mark.parametrize("num_sql_threads", (0, 1))
async def test_write_thread_context_is_detached_between_tasks(
tmp_path, monkeypatch, num_sql_threads
):
"""
The write thread attaches each task's otel Context and must detach it
again before picking up the next task. The thread is persistent and
shared, so a leaked token would grow that thread's context stack for the
rest of the process - and a *wrong*-token detach only logs a warning
rather than raising, so "does it throw" cannot catch either mistake.
Two things are asserted, because neither alone is sufficient:
1. Each task observes the context value that was current on the event
loop when it was queued. This is what fails if the Context is not
carried on WriteTask, or is never attached. It does *not* catch a
missing detach: attach() replaces the current Context wholesale, so a
leftover one from a previous task is simply overwritten.
2. The write thread's attach depth is identical at the same point in
every task. This is what fails if detach is missing - the stack grows
by one per task - and it holds across a task that raises, because the
detach lives in a `finally`.
An otel context value is used rather than a plain contextvars.ContextVar:
a plain var set on the event loop never crosses into the write thread, so
the probe would read None every time and the test could not fail.
"""
name = f"context_leak_test_{num_sql_threads}"
db_path = tmp_path / f"{name}.db"
sqlite3.connect(db_path).close()
ds = Datasette([str(db_path)], settings={"num_sql_threads": num_sql_threads})
db = ds.get_database(name)
await db.execute_write("create table t (id integer primary key)")
write_thread_name = f"_execute_writes for database {name}"
depth = {"value": 0}
real_attach = otel_context_api.attach
real_detach = otel_context_api.detach
def counting_attach(context):
token = real_attach(context)
if threading.current_thread().name == write_thread_name:
depth["value"] += 1
return token
def counting_detach(token):
real_detach(token)
if threading.current_thread().name == write_thread_name:
depth["value"] -= 1
# Patched on the opentelemetry.context module itself, which is what both
# database.py and opentelemetry.trace.use_span() look the functions up on.
monkeypatch.setattr(otel_context_api, "attach", counting_attach)
monkeypatch.setattr(otel_context_api, "detach", counting_detach)
seen_markers = []
seen_depths = []
def probe(conn):
seen_markers.append(otel_context_api.get_value(_CONTEXT_LEAK_MARKER_KEY))
seen_depths.append(depth["value"])
def failing_probe(conn):
probe(conn)
# Exercises the write thread's exception path: the detach still has
# to happen, which is why it lives in a `finally`.
raise ValueError("deliberate failure inside a write task")
try:
for i in range(5):
ctx = otel_context_api.set_value(_CONTEXT_LEAK_MARKER_KEY, f"marker-{i}")
token = real_attach(ctx)
try:
if i == 2:
with pytest.raises(ValueError):
await db.execute_write_fn(failing_probe)
else:
await db.execute_write_fn(probe)
finally:
real_detach(token)
# Sanity check: no marker is active in *this* (event loop) context
# right now, so the final probe is a fair test of the write thread's
# own state rather than something this test forgot to clean up.
assert otel_context_api.get_value(_CONTEXT_LEAK_MARKER_KEY) is None
await db.execute_write_fn(probe)
finally:
db.close()
assert seen_markers == [
"marker-0",
"marker-1",
"marker-2",
"marker-3",
"marker-4",
None,
]
assert len(set(seen_depths)) == 1, (
f"write thread context stack grew across tasks: {seen_depths} - "
"a token was attached without being detached"
)

192
tests/test_telemetry.py
View file

@ -2,12 +2,15 @@ import json
import sqlite3
import subprocess
import sys
import time
import pytest
import sqlite_utils
from opentelemetry.trace import StatusCode
from datasette.app import Datasette
from datasette.telemetry import MAX_SQL_LENGTH, sql_attribute
from datasette.database import Database
from datasette.telemetry import MAX_SQL_LENGTH, sql_attribute, tracer
SECRET_PARAM_VALUE = "SUPER_SECRET_PARAM_VALUE_XYZ_123"
@ -33,6 +36,26 @@ def _spans_for_namespace(otel_spans, namespace):
]
def _children_named(otel_spans, name, parent_span_context):
"""
Finished spans called `name` whose parent really is `parent_span_context`.
Parentage is matched on span id, not on "a span with this name exists" -
a span can exist and still be an unparented root if a thread boundary
dropped the otel context, which is the exact failure these tests exist
to catch.
"""
return [
span
for span in otel_spans.get_finished_spans()
if span.name == name
and span.parent is not None
and span.parent.span_id == parent_span_context.span_id
and span.parent.trace_id == parent_span_context.trace_id
and span.context.trace_id == parent_span_context.trace_id
]
def _all_attribute_values(otel_spans):
"Every attribute value across every finished span, for the 'no leaked param values' test."
values = []
@ -267,3 +290,170 @@ async def test_execute_write_many_records_param_sets_not_rows_returned(otel_span
# calling this a row count would be a lie. Asserted explicitly because the
# attribute really was named datasette.rows_returned at one point.
assert "datasette.rows_returned" not in span.attributes
# --- Context propagation across thread boundaries --------------------------
#
# Every assertion below checks parentage (child.parent.span_id ==
# expected_parent.span_id, in the same trace), not merely that spans exist.
# Spans can exist and still be wrongly parented - or be unparented roots - if
# a thread boundary drops the otel context, which is exactly the failure mode
# these tests exist to prevent.
@pytest.mark.asyncio
async def test_db_query_execute_parents_to_db_query(ds_client, otel_spans):
# execute_fn()'s executor.submit() is thread boundary #1. The
# db.query.execute span is created inside the worker thread; without the
# copy_context() propagation it comes back as an unparented root span
# rather than a child of db.query.
response = await ds_client.get("/fixtures/-/query.json?sql=select+1")
assert response.status_code == 200
query_spans = [
span
for span in _spans_for_namespace(otel_spans, "fixtures")
if span.attributes["db.query.text"] == "select 1"
]
assert query_spans, "expected a db.query span for 'select 1'"
query_span = query_spans[-1]
assert [
span
for span in otel_spans.get_finished_spans()
if span.name == "db.query.execute"
], "expected at least one db.query.execute span"
children = _children_named(otel_spans, "db.query.execute", query_span.context)
assert len(children) == 1, "expected exactly one db.query.execute child of db.query"
# The execute span is strictly contained by the round-trip span, and the
# gap between the two is the thread-pool wait.
assert query_span.start_time <= children[0].start_time
assert children[0].end_time <= query_span.end_time
@pytest.mark.asyncio
async def test_immutable_database_propagates_context(tmp_path, otel_spans):
# Thread boundary #3, the easy one to miss: immutable databases route
# execute_isolated_fn() through loop.run_in_executor() directly rather
# than through the write thread. A span created inside that worker must
# still parent to whatever was current when execute_isolated_fn() was
# awaited, or every immutable-database operation emits orphan roots.
db_path = tmp_path / "t04_immutable.db"
sqlite_utils.Database(str(db_path))["t"].insert({"id": 1}, pk="id")
ds = Datasette()
db = Database(ds, path=str(db_path), is_mutable=False)
ds.add_database(db, name="t04_immutable")
def fn(conn):
with tracer.start_as_current_span("t04-child-in-isolated-worker"):
pass
try:
with tracer.start_as_current_span("t04-parent-on-event-loop") as parent:
parent_context = parent.get_span_context()
await db.execute_isolated_fn(fn)
finally:
ds.remove_database("t04_immutable")
assert [
span
for span in otel_spans.get_finished_spans()
if span.name == "t04-child-in-isolated-worker"
], "expected a span created inside execute_isolated_fn's worker thread"
children = _children_named(
otel_spans, "t04-child-in-isolated-worker", parent_context
)
assert len(children) == 1
@pytest.mark.asyncio
async def test_write_spans_parent_to_db_query(otel_spans):
# Thread boundary #2: WriteTask -> queue.Queue -> the write thread.
# db.write.queue_wait and db.write.execute are both direct children of
# the db.query span that was current on the event loop at enqueue time,
# so they are siblings rather than nested inside one another.
db = Datasette(memory=True).add_memory_database("t04_write_spans")
await db.execute_write("create table docs (id integer primary key)")
query_spans = _spans_for_namespace(otel_spans, "t04_write_spans")
assert query_spans, "expected a db.query span from execute_write()"
query_span = query_spans[-1]
queue_wait_children = _children_named(
otel_spans, "db.write.queue_wait", query_span.context
)
execute_children = _children_named(
otel_spans, "db.write.execute", query_span.context
)
assert len(queue_wait_children) == 1
assert len(execute_children) == 1
execute_span = execute_children[0]
assert execute_span.attributes["datasette.isolated_connection"] is False
assert execute_span.attributes["datasette.transaction"] is True
# Siblings, not parent/child: the queue wait is over by the time the
# write begins.
assert queue_wait_children[0].end_time <= execute_span.start_time
@pytest.mark.asyncio
async def test_write_queue_wait_duration_reflects_real_wait(otel_spans):
# db.write.queue_wait is built from explicit start/end timestamps -
# task.enqueued_at_ns, captured on the event loop, through to the moment
# the write thread dequeued it. If it were a plain `with` block on the
# write thread it would instead measure the microseconds spent building
# the span object, and this assertion would fail.
ds = Datasette(memory=True)
db = ds.add_memory_database("t04_queue_wait")
await db.execute_write("create table docs (id integer primary key)")
def slow_write(conn):
time.sleep(0.1)
# Queue a deliberately slow write without waiting for it, then queue a
# second write immediately behind it: the second task sits in the queue
# for roughly the duration of the first.
_, slow_future = await db._send_to_write_thread(slow_write, block=False)
await db.execute_write("insert into docs (id) values (1)")
await slow_future
query_spans = [
span
for span in _spans_for_namespace(otel_spans, "t04_queue_wait")
if span.attributes["db.query.text"] == "insert into docs (id) values (1)"
]
assert query_spans, "expected a db.query span for the queued-behind insert"
queue_wait_children = _children_named(
otel_spans, "db.write.queue_wait", query_spans[-1].context
)
assert len(queue_wait_children) == 1
duration_ns = queue_wait_children[0].end_time - queue_wait_children[0].start_time
# The slow write sleeps 100ms; anything above 10ms is far beyond the
# microseconds a mis-timestamped span would report.
assert duration_ns > 10_000_000, f"queue wait was only {duration_ns}ns"
@pytest.mark.asyncio
async def test_suppressed_error_does_not_mark_execute_span(ds_client, otel_spans):
"""
The inner db.query.execute span must honour log_sql_errors too.
It is created inside the worker thread, so without record_exception /
set_status_on_exception being passed through it would mark every facet
suggestion probe as failed even though the outer db.query span correctly
reports the failure as suppressed.
"""
db = ds_client.ds.get_database("fixtures")
with pytest.raises(sqlite3.OperationalError):
await db.execute(INVALID_SQL, log_sql_errors=False)
execute_spans = [
span
for span in otel_spans.get_finished_spans()
if span.name == "db.query.execute"
]
assert execute_spans
span = execute_spans[-1]
assert span.status.status_code == StatusCode.UNSET
assert not [event for event in span.events if event.name == "exception"]