Propagate the trace context in database queries if available

This adds the trace context to the SQL statement as a comment, following
the SQLCommenter spec. This allows correlation of application-level
traces with logs and active statements of the database.

synapse/logging/opentracing.py

@@ -850,6 +850,30 @@ def get_active_span_text_map(destination: str | None = None) -> dict[str, str]:
     return carrier
 
 
+@ensure_active_span("get the active span's traceparent", ret=None)
+def get_active_span_traceparent() -> str | None:
+    """
+    Get the W3C Trace Context traceparent string for the currently active span.
+
+    Returns the traceparent in the format ``00-<trace_id>-<span_id>-<flags>``
+    if there is an active Jaeger span, or None otherwise.
+
+    See https://www.w3.org/TR/trace-context/#traceparent-header-field-values
+    """
+    # Jaeger spans expose trace_id, span_id, and flags as int attributes.
+    # Other OpenTracing implementations may not have these.
+    assert opentracing.tracer.active_span is not None
+    ctx = opentracing.tracer.active_span.context
+    trace_id = getattr(ctx, "trace_id", None)
+    span_id = getattr(ctx, "span_id", None)
+    flags = getattr(ctx, "flags", 0)
+
+    if trace_id is None or span_id is None:
+        return None
+
+    return f"00-{trace_id:032x}-{span_id:016x}-{flags:02x}"
+
+
 @ensure_active_span("get the span context as a string.", ret={})
 def active_span_context_as_string() -> str:
     """

synapse/storage/engines/postgres.py

@@ -24,6 +24,7 @@ from typing import TYPE_CHECKING, Any, Mapping, NoReturn, cast
 
 import psycopg2.extensions
 
+from synapse.logging import opentracing
 from synapse.storage.engines._base import (
     AUTO_INCREMENT_PRIMARY_KEYPLACEHOLDER,
     BaseDatabaseEngine,
@@ -40,6 +41,44 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+class _SqlCommenterCursor(psycopg2.extensions.cursor):
+    """A psycopg2 cursor that appends W3C trace context to SQL statements
+    as SQLCommenter comments when OpenTracing is active.
+
+    This propagates the active span's trace context to PostgreSQL, enabling
+    database-side tracing tools to correlate server-side spans with
+    application-level traces.
+
+    See:
+        - https://google.github.io/sqlcommenter/spec/
+        - https://opentelemetry.io/docs/specs/semconv/db/database-spans/#context-propagation
+    """
+
+    def execute(  # type: ignore[override]
+        self,
+        query: str | bytes,
+        vars: Any = None,  # noqa: A002
+    ) -> None:
+        # The traceparent is only added when a trace is actively being
+        # sampled, so untraced queries are not affected.
+        if isinstance(query, str):
+            traceparent = opentracing.get_active_span_traceparent()
+            if traceparent is not None:
+                query = f"{query} /*traceparent='{traceparent}'*/"
+        return super().execute(query, vars)
+
+    def executemany(  # type: ignore[override]
+        self,
+        query: str | bytes,
+        vars_list: Any,  # noqa: A002
+    ) -> None:
+        if isinstance(query, str):
+            traceparent = opentracing.get_active_span_traceparent()
+            if traceparent is not None:
+                query = f"{query} /*traceparent='{traceparent}'*/"
+        return super().executemany(query, vars_list)
+
+
 class PostgresEngine(
     BaseDatabaseEngine[psycopg2.extensions.connection, psycopg2.extensions.cursor]
 ):
@@ -173,6 +212,10 @@ class PostgresEngine(
     def on_new_connection(self, db_conn: "LoggingDatabaseConnection") -> None:
         db_conn.set_isolation_level(self.default_isolation_level)
 
+        # Use a cursor factory that appends W3C trace context to queries
+        # as SQLCommenter comments, propagating spans to the database.
+        db_conn.conn.cursor_factory = _SqlCommenterCursor  # type: ignore[attr-defined]
+
         # Set the bytea output to escape, vs the default of hex
         cursor = db_conn.cursor()
         cursor.execute("SET bytea_output TO escape")