diff --git a/synapse/storage/engines/postgres.py b/synapse/storage/engines/postgres.py index d1d880f2a1..34d061e5c6 100644 --- a/synapse/storage/engines/postgres.py +++ b/synapse/storage/engines/postgres.py @@ -21,10 +21,12 @@ import logging from typing import TYPE_CHECKING, Any, Mapping, NoReturn, cast +from urllib.parse import quote import psycopg2.extensions from synapse.logging import opentracing +from synapse.logging.context import current_context from synapse.storage.engines._base import ( AUTO_INCREMENT_PRIMARY_KEYPLACEHOLDER, BaseDatabaseEngine, @@ -40,30 +42,47 @@ 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. +def _build_sqlcommenter_comment() -> str: + """Build a SQLCommenter comment with the logging context and trace context. - This propagates the active span's trace context to PostgreSQL, enabling - database-side tracing tools to correlate server-side spans with - application-level traces. + Returns a ``/*...*/`` comment string to append to a SQL query. The + ``traceparent`` field is only included when a trace is actively being + sampled; the ``log_context`` field is always included. See: - https://google.github.io/sqlcommenter/spec/ - https://opentelemetry.io/docs/specs/semconv/db/database-spans/#context-propagation """ + # Per the SQLCommenter spec, keys are sorted and values are URL-encoded + # then wrapped in single quotes. + pairs: dict[str, str] = { + "log_context": str(current_context()), + } + + traceparent = opentracing.get_active_span_traceparent() + if traceparent is not None: + pairs["traceparent"] = traceparent + + comment = ",".join(f"{k}='{quote(v)}'" for k, v in sorted(pairs.items())) + return f" /*{comment}*/" + + +class _SqlCommenterCursor(psycopg2.extensions.cursor): + """A psycopg2 cursor that appends the logging context and W3C trace context + to SQL statements as SQLCommenter comments. + + This propagates the active span's trace context to PostgreSQL, enabling + database-side tracing tools to correlate server-side spans with + application-level traces. + """ 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}'*/" + query = f"{query}{_build_sqlcommenter_comment()}" return super().execute(query, vars) def executemany( # type: ignore[override] @@ -72,9 +91,7 @@ class _SqlCommenterCursor(psycopg2.extensions.cursor): 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}'*/" + query = f"{query}{_build_sqlcommenter_comment()}" return super().executemany(query, vars_list)