OpenTelemetry tracing for the database layer (#2862)

PR #2862

tests/conftest.py

@@ -58,6 +58,62 @@ def find_free_port():
         return sock.getsockname()[1]
 
 
+_otel_span_exporter = None
+
+
+@pytest.fixture(scope="session", autouse=True)
+def _otel_provider():
+    """
+    Install a real OTel SDK TracerProvider + InMemorySpanExporter exactly
+    once, before any span is ever created in this process.
+
+    This has to be session-scoped and autouse because
+    `opentelemetry.trace.set_tracer_provider()` is effectively
+    once-per-process: a second call logs a warning and is ignored. So the
+    install must happen exactly once, before anything asserts on spans.
+
+    `datasette.telemetry.tracer` is a module-level `ProxyTracer`. Once a
+    provider exists, the first span it starts resolves a concrete tracer
+    and caches it permanently. It does *not* cache the no-op tracer, so
+    any span started before this fixture runs is merely lost rather than
+    poisoning the tracer for the rest of the process. If the SDK isn't
+    installed, do nothing: core spans stay no-op `NonRecordingSpan`s and
+    the rest of the suite is unaffected.
+    """
+    global _otel_span_exporter
+    try:
+        from opentelemetry import trace as otel_trace
+        from opentelemetry.sdk.trace import TracerProvider
+        from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+        from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
+            InMemorySpanExporter,
+        )
+    except ImportError:
+        return
+    exporter = InMemorySpanExporter()
+    provider = TracerProvider()
+    # SimpleSpanProcessor exports synchronously on span end - no background
+    # batching thread, so assertions immediately after a request never race.
+    provider.add_span_processor(SimpleSpanProcessor(exporter))
+    otel_trace.set_tracer_provider(provider)
+    _otel_span_exporter = exporter
+
+
+@pytest.fixture
+def otel_spans():
+    """
+    Function-scoped access to the finished-spans exporter: clears any spans
+    left over from previous tests, then yields the exporter so a test can
+    call `.get_finished_spans()` after making requests. Skips (rather than
+    fails) if the OTel SDK is not installed.
+    """
+    pytest.importorskip("opentelemetry.sdk")
+    if _otel_span_exporter is None:
+        pytest.skip("OpenTelemetry SDK provider was not installed")
+    _otel_span_exporter.clear()
+    yield _otel_span_exporter
+
+
 @pytest.fixture
 def bare_ds():
     """
@@ -171,6 +227,12 @@ def pytest_collection_modifyitems(config, items):
     move_to_front(items, "test_spatialite_error_if_attempt_to_open_spatialite")
     move_to_front(items, "test_package")
     move_to_front(items, "test_package_with_port")
+    # Same reason: this one shells out to a fresh interpreter. Late in a serial
+    # run the pytest process holds enough threads that the fork half of
+    # subprocess' fork+exec crashes the interpreter on macOS/CPython 3.13
+    # (SIGSEGV/SIGBUS inside _execute_child). Reproduces with any subprocess
+    # call placed there, on an unmodified tree - running it first avoids it.
+    move_to_front(items, "test_datasette_package_never_imports_the_sdk")
 
 
 def move_to_front(items, test_name):

tests/test_internals_database.py

@@ -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 (
@@ -1319,3 +1321,110 @@ async def test_close_releases_memory_connections(num_sql_threads, named):
     for conn in (read_connection, write_connection):
         with pytest.raises(sqlite3.ProgrammingError, match="closed"):
             conn.execute("select 1")
+
+
+_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"
+    )

tests/test_telemetry_registry.py

@@ -0,0 +1,297 @@
+"""
+Two-way conformance between `datasette/telemetry_registry.py` and what
+Datasette actually emits.
+
+This is the test that makes the generated documentation trustworthy. cog
+guarantees the docs match the registry; this guarantees the registry matches
+the code. Without it, both could agree with each other and be wrong.
+
+It checks both directions, and the second one is the one nothing else catches:
+
+- **emitted but not registered** - instrumentation was added without
+  documenting it, so the reference page silently omits it.
+- **registered but never emitted** - the reference page describes a span or
+  attribute that no longer exists, which is worse than omitting it, because a
+  reader will build a dashboard on it.
+
+Both of those directions compare the code against the registry. Neither can
+catch a *rename*, because the call sites now take their names from the
+registry - move `DB_NAMESPACE` to `"db.namespace2"` and code and registry
+still agree with each other, while every existing dashboard breaks. So the
+literal names live here too, spelled out, and are asserted against both the
+registry and the wire. That is the one comparison in this file that is not
+made against a value derived from the registry itself.
+"""
+
+import itertools
+
+import pytest
+import pytest_asyncio
+
+pytest.importorskip("opentelemetry.sdk")
+
+from datasette import telemetry_registry as reg
+from datasette.app import Datasette
+from datasette.database import QueryInterrupted
+from datasette.utils.sqlite import sqlite3
+
+# The names as they appear on the wire, written out rather than read from the
+# registry. If a change to the registry makes one of these fail, that change
+# is renaming something a user's dashboards and saved queries depend on -
+# which is a decision to take deliberately, here, not a line to re-derive.
+EXPECTED_ATTRIBUTES = {
+    "db.query": {
+        "db.system",
+        "db.namespace",
+        "db.query.text",
+        "datasette.callback",
+        "db.operation.name",
+        "datasette.param_count",
+        "datasette.param_sets",
+        "datasette.time_limit_ms",
+        "datasette.rows_returned",
+        "datasette.truncated",
+        "datasette.interrupted",
+        "datasette.sql_error_suppressed",
+        "datasette.executescript",
+        "datasette.executemany",
+    },
+    "db.query.execute": set(),
+    "db.write.queue_wait": set(),
+    "db.write.execute": {
+        "datasette.isolated_connection",
+        "datasette.transaction",
+    },
+    "datasette.startup": set(),
+}
+EXPECTED_SPANS = set(EXPECTED_ATTRIBUTES)
+
+# Named in-memory databases are shared-cache, so two Datasette instances using
+# the same name share one SQLite database - and the second `create table`
+# fails. Every workload below therefore gets its own name.
+_names = itertools.count()
+
+
+def _unique(prefix):
+    return f"{prefix}{next(_names)}"
+
+
+async def exercise():
+    """
+    Drive enough of Datasette to emit every span and attribute the registry
+    claims exists.
+
+    Each call is here because it is the only thing that produces some span or
+    attribute - see the comments. If you add instrumentation on a path this
+    does not reach, add the path rather than loosening the assertions.
+
+    Returns the instance so the caller can close it; startup happens inside
+    so that the `datasette.startup` span lands in the collected set.
+    """
+    name = _unique("registry")
+    ds = Datasette(memory=True)
+    ds.add_memory_database(name)
+    # datasette.startup - and the internal catalog work nested under it
+    await ds.invoke_startup()
+    db = ds.get_database(name)
+
+    # Writes: db.write.queue_wait, db.write.execute, db.query
+    await db.execute_write("create table t (id integer primary key, v text)")
+    # datasette.executemany, datasette.param_sets
+    await db.execute_write_many(
+        "insert into t (id, v) values (?, ?)", [[i, f"v{i}"] for i in range(30)]
+    )
+    # datasette.executescript
+    await db.execute_write_script("create table t2 (id integer); drop table t2;")
+    # datasette.transaction=False - VACUUM cannot run inside a transaction
+    await db.execute_write("vacuum", transaction=False)
+    # datasette.isolated_connection=True
+    await db.execute_isolated_fn(lambda conn: conn.execute("select 1").fetchone())
+
+    # datasette.callback, with named functions so the conformance run sees the
+    # attribute's documented value shape (a qualname, not just "<lambda>")
+    def registry_read_callback(conn):
+        return conn.execute("select count(*) from t").fetchone()
+
+    def registry_write_callback(conn):
+        conn.execute("insert into t (id, v) values (100, 'callback')")
+
+    await db.execute_fn(registry_read_callback)
+    await db.execute_write_fn(registry_write_callback)
+
+    # Reads: db.query.execute, datasette.rows_returned, datasette.truncated,
+    # datasette.param_count, datasette.time_limit_ms
+    await db.execute("select * from t where id > :n", {"n": 5})
+    await db.execute("select * from t", truncate=True)
+
+    # datasette.sql_error_suppressed - the caller is probing and treats
+    # failure as an expected answer
+    with pytest.raises(sqlite3.OperationalError):
+        await db.execute("select nope from t", log_sql_errors=False)
+
+    # datasette.interrupted - only ever set when a query exceeds its time
+    # limit, so the workload has to force one rather than exempt it. An
+    # unbounded recursive CTE cannot finish, so 1ms is always exceeded.
+    with pytest.raises(QueryInterrupted):
+        await db.execute(
+            "with recursive c(x) as (select 0 union all select x+1 from c) "
+            "select * from c",
+            custom_time_limit=1,
+        )
+
+    return ds
+
+
+@pytest_asyncio.fixture
+async def emitted(otel_spans):
+    "Every span name and (span name, attribute key) pair a broad workload emits."
+    # otel_spans has already cleared the exporter, and nothing is cleared
+    # after this point: the workload's own startup emits datasette.startup.
+    ds = await exercise()
+    spans = otel_spans.get_finished_spans()
+    assert spans, "no spans captured - the fixture is not exercising anything"
+    names = set()
+    pairs = set()
+    for span in spans:
+        # str() because span.name is the registry's SpanName instance, and a
+        # set of those would compare equal to literals but read confusingly
+        # in a failure message.
+        names.add(str(span.name))
+        for key in span.attributes or {}:
+            pairs.add((str(span.name), str(key)))
+    ds.close()
+    return {"names": names, "pairs": pairs}
+
+
+def _keys_by_span(pairs):
+    by_span = {}
+    for span_name, key in pairs:
+        by_span.setdefault(span_name, set()).add(key)
+    return by_span
+
+
+@pytest.mark.asyncio
+async def test_workload_emits_exactly_the_expected_names(emitted):
+    """
+    The wire format, pinned to literals.
+
+    Not derived from the registry, so this is what catches a rename that the
+    registry and the call sites make together.
+    """
+    assert emitted["names"] == EXPECTED_SPANS
+    by_span = _keys_by_span(emitted["pairs"])
+    assert {name: by_span.get(name, set()) for name in emitted["names"]} == (
+        EXPECTED_ATTRIBUTES
+    )
+
+
+def test_registry_matches_the_expected_names():
+    "The other half of the rename check: the registry against the same literals."
+    assert {str(span) for span in reg.SPANS} == EXPECTED_SPANS
+    for span in reg.SPANS:
+        assert {str(attribute) for attribute in span.attributes} == EXPECTED_ATTRIBUTES[
+            str(span)
+        ], f"{span} attributes have drifted"
+
+
+@pytest.mark.asyncio
+async def test_every_emitted_span_is_registered(emitted):
+    "A span added without a registry entry would be missing from the docs."
+    unregistered = sorted(
+        name for name in emitted["names"] if reg.span_for(name) is None
+    )
+    assert (
+        not unregistered
+    ), f"these spans are emitted but not in telemetry_registry.SPANS: {unregistered}"
+
+
+@pytest.mark.asyncio
+async def test_every_emitted_attribute_is_registered(emitted):
+    "An attribute added without a registry entry would be missing from the docs."
+    unregistered = sorted(
+        f"{span_name} -> {key}"
+        for span_name, key in emitted["pairs"]
+        if not reg.attribute_allowed(reg.span_for(span_name), key)
+    )
+    assert (
+        not unregistered
+    ), "these span attributes are emitted but not registered: " + ", ".join(
+        unregistered
+    )
+
+
+@pytest.mark.asyncio
+async def test_every_registered_span_is_emitted(emitted):
+    """
+    The direction nothing else catches: the docs must not describe a span that
+    no longer exists.
+    """
+    missing = sorted(
+        str(span)
+        for span in reg.SPANS
+        if not any(reg.span_for(name) is span for name in emitted["names"])
+    )
+    assert not missing, (
+        f"these spans are documented but never emitted by the workload: {missing}. "
+        "Either the instrumentation was removed, or exercise() no longer reaches it."
+    )
+
+
+@pytest.mark.asyncio
+async def test_every_registered_attribute_is_emitted(emitted):
+    """
+    Every registered attribute, optional or not, must actually be set at least
+    once by the workload.
+
+    `optional` describes whether a reader should expect it on every span, not
+    whether the code still sets it - so an attribute deleted from the code but
+    left in the docs has to fail here even when it is marked optional. If a
+    new attribute only appears in some rare case, extend exercise() to reach
+    that case.
+    """
+    by_span = _keys_by_span(emitted["pairs"])
+    missing = []
+    for span in reg.SPANS:
+        emitted_keys = by_span.get(str(span), set())
+        for attribute in span.attributes:
+            if attribute not in emitted_keys:
+                missing.append(f"{span} -> {attribute}")
+    assert not missing, (
+        "these attributes are documented but never emitted by the workload: "
+        + ", ".join(sorted(missing))
+    )
+
+
+def test_registry_has_no_duplicate_names():
+    assert len(set(reg.SPANS)) == len(reg.SPANS)
+    for span in reg.SPANS:
+        assert len(set(span.attributes)) == len(
+            span.attributes
+        ), f"{span} lists an attribute twice"
+
+
+def test_registry_entries_are_documented():
+    "Every entry carries a description - the docs are generated from these."
+    for span in reg.SPANS:
+        assert span.description.strip(), f"{span} has no description"
+        for attribute in span.attributes:
+            assert attribute.description.strip(), f"{span} -> {attribute} has none"
+
+
+def test_registry_entries_are_usable_as_plain_strings():
+    "The str subclassing is the whole reason call sites need no wrapper API."
+    assert isinstance(reg.DB_QUERY, str)
+    assert isinstance(reg.DB_NAMESPACE, str)
+    assert reg.DB_QUERY == "db.query"
+    assert reg.DB_NAMESPACE == "db.namespace"
+    assert f"{reg.DB_QUERY}.execute" == "db.query.execute"
+
+
+def test_span_and_attribute_lookup():
+    assert reg.span_for("db.query") is reg.DB_QUERY
+    assert reg.span_for("datasette.startup") is reg.STARTUP
+    assert reg.span_for("not.a.datasette.span") is None
+    assert reg.attribute_allowed(reg.DB_QUERY, "db.namespace")
+    assert not reg.attribute_allowed(reg.DB_QUERY, "db.namespace.extra")
+    assert not reg.attribute_allowed(reg.DB_QUERY, "datasette.isolated_connection")
+    assert not reg.attribute_allowed(None, "db.namespace")