Coverage for app/backend/src/tests/fixtures/query_log.py: 94%
152 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-24 08:45 +0000
1"""Records every SQL query the suite issues, grouped by test and by the RPC (or background job) that issued it.
3Off unless --query-log is passed, so ordinary local runs pay nothing. CI dumps one JSON file per pytest-split node;
4app/scripts/query_log_report.py merges those, diffs them against develop's and renders the browsable report.
6Each query is stored twice: a fingerprint with the bound parameters replaced by "?" (the stable key used for
7grouping and diffing) and one concrete rendering with the values inlined by psycopg, so it can be copy-pasted
8straight into a psql prompt.
9"""
11import gzip
12import hashlib
13import json
14import os
15import re
16import sys
17import threading
18from dataclasses import dataclass, field
19from pathlib import Path
20from types import CodeType, FrameType
21from typing import Any
23import psycopg
24from sqlalchemy import Engine, event
26# Hard cap on what is stored per statement. Nothing this long is an access pattern worth diffing, and the cap is
27# what bounds the artifact: uncapped, the timezone_areas load alone took a CI node's dump to 495 MB.
28_MAX_SQL_CHARS = 4096
29_TRUNCATION_MARKER = " /* truncated by the query log */"
32@dataclass(slots=True)
33class _Shape:
34 id: str
35 sql: str
36 example: str
37 params: str | None
38 write: bool
39 # The test that first produced this shape. Kept only to make the choice of `example` deterministic.
40 first_seen_in: str
43@dataclass(slots=True)
44class _Span:
45 kind: str
46 name: str | None
47 queries: list[str] = field(default_factory=list)
48 # Parallel to queries: the call site each execution came from. Kept as a separate array so the diff, which reads
49 # only `queries`, cannot be perturbed by line numbers shifting under an unrelated edit.
50 sites: list[str] = field(default_factory=list)
53_lock = threading.Lock()
54_local = threading.local()
56_enabled = False
57_current_test: str | None = None
58_shapes: dict[str, _Shape] = {}
59_tests: dict[str, list[_Span]] = {}
60_sites: dict[str, str] = {}
61_site_ids: dict[str, str] = {}
62_frame_cache: dict[str, int] = {}
65def _truncate(sql: str) -> str:
66 return sql if len(sql) <= _MAX_SQL_CHARS else sql[:_MAX_SQL_CHARS] + _TRUNCATION_MARKER
69def _fingerprint(statement: str) -> str:
70 """The statement with everything that varies with the test data taken out, which is the key we group and diff on.
72 Anything left in here that moves between two identical runs turns the whole report into noise, so each step
73 below is pinned by a test in test_query_log.py.
74 """
75 # sqlcommenter-style comments. Nothing emits these in tests today (tracing is only set up in prod), but a future
76 # change that turns the commenter on would otherwise invalidate every fingerprint at once.
77 sql = re.sub(r"/\*.*?\*/", "", statement, flags=re.DOTALL)
78 # Bind parameters, as psycopg's pyformat paramstyle renders them, plus positional %s for good measure.
79 sql = re.sub(r"%\([^)]*\)s|%s", "?", sql)
80 # An expanded IN (...) list, whose length tracks the test data.
81 sql = re.sub(r"\?(?:\s*,\s*\?)+", "?", sql)
82 # Repeated VALUES tuples from a multi-row insert, whose count tracks the batch size.
83 sql = re.sub(r"\(\?\)(?:\s*,\s*\(\?\))+", "(?)", sql)
84 # Literals inlined into the statement text rather than bound. Bulk resource loads do this: the real
85 # timezone_areas.sql is a few hundred INSERTs each carrying megabytes of WKB hex, so without collapsing them
86 # every row becomes its own multi-megabyte shape. Every literal is matched and the long ones are picked out
87 # afterwards, rather than matching only the long ones: a pattern for "quote, 64 characters, quote" also matches
88 # from one literal's closing quote to the next literal's opening quote, collapsing the statement in between.
89 sql = re.sub(r"'([^']*)'", lambda literal: "'...'" if len(literal.group(1)) >= 64 else literal.group(), sql)
90 return _truncate(re.sub(r"\s+", " ", sql).strip())
93def _shape_id(fingerprint: str) -> str:
94 # Content-addressed, so the three pytest-split nodes agree on ids and merging is a plain dict union. A counter
95 # would be assigned in per-node encounter order and collide across nodes.
96 return hashlib.blake2b(fingerprint.encode(), digest_size=6).hexdigest()
99def _render_example(conn: Any, statement: str, parameters: Any) -> tuple[str, str | None]:
100 """Inline the bound values so the statement can be pasted into psql.
102 psycopg's ClientCursor.mogrify applies the same adaptation and escaping the driver would otherwise do
103 server-side, which is far more faithful than re-implementing literal binding for bytea, arrays and PostGIS
104 geometries. It is purely client-side, so it issues nothing on the connection.
105 """
106 # executemany passes a sequence of parameter sets; one row is enough to have something runnable.
107 if isinstance(parameters, (list, tuple)) and parameters and isinstance(parameters[0], (dict, list, tuple)):
108 parameters = parameters[0]
109 try:
110 params = json.dumps(parameters, default=repr) if parameters else None
111 except TypeError, ValueError:
112 params = None
113 if params is not None:
114 params = _truncate(params)
115 try:
116 cursor = psycopg.ClientCursor(conn.connection.driver_connection)
117 # Truncated past the cap, so the marker is a SQL comment: the result is visibly not runnable rather than
118 # silently invalid.
119 return _truncate(cursor.mogrify(statement, parameters)), params
120 except Exception:
121 # Not worth losing the whole entry over; the fingerprint plus the parameters is still pasteable by hand.
122 return _truncate(statement), params
125# Everything under the backend's src/ is ours; paths are reported relative to it. couchers/ is application code,
126# anything else under src/ is test scaffolding.
127_SRC_ROOT = "/src/"
128_APP_ROOT = "/src/couchers/"
129_SKIP, _APP, _TEST = 0, 1, 2
130# Frames from these are plumbing between our code and the driver, so they never make a useful call site. The
131# recorder's own path is spelled out: a bare "query_log.py" would also swallow test_query_log.py.
132_CALLSITE_SKIP = ("/sqlalchemy/", "/psycopg", "/alembic/", "/fixtures/query_log.py")
135def _frame_kind(code: CodeType) -> int:
136 """_APP, _TEST or _SKIP. Cached by filename, which is all the answer depends on: this runs on every frame of
137 every execution, and the string work is what would otherwise make stack walking too expensive to leave on.
139 Keyed by filename rather than by the code object, because code objects compare equal without regard to
140 co_filename, so two same-bodied functions in different files would share an entry.
141 """
142 filename = code.co_filename
143 known = _frame_cache.get(filename)
144 if known is None:
145 if any(part in filename for part in _CALLSITE_SKIP) or _SRC_ROOT not in filename:
146 known = _SKIP
147 else:
148 known = _APP if _APP_ROOT in filename else _TEST
149 _frame_cache[filename] = known
150 return known
153# How many of our own frames to keep. The innermost is the line that issued the query; the next couple show the
154# chain that got there, which is usually what tells you whether a repeat is a loop.
155_CALLSITE_FRAMES = 3
158def _callsite() -> str:
159 """The innermost few application frames, innermost first, as "path:line in func".
161 Only couchers/ frames: the test and the fixture handler that got here are already implied by the test and span
162 this is recorded under, and including them multiplies the number of distinct call sites for no added meaning.
163 Test frames are used only when a query has no application frame at all, as fixture setup often does not.
164 """
165 frames: list[str] = []
166 fallback = ""
167 frame: FrameType | None = sys._getframe(1)
168 while frame is not None and len(frames) < _CALLSITE_FRAMES:
169 code = frame.f_code
170 kind = _frame_kind(code)
171 if kind != _SKIP:
172 path = code.co_filename.split(_SRC_ROOT, 1)[-1]
173 rendered_frame = f"{path}:{frame.f_lineno} in {code.co_name}"
174 if kind == _APP:
175 frames.append(rendered_frame)
176 elif not fallback:
177 fallback = rendered_frame
178 frame = frame.f_back
179 rendered = " <- ".join(frames) or fallback
180 site_id = _site_ids.get(rendered)
181 if site_id is None:
182 site_id = _shape_id(rendered)
183 _site_ids[rendered] = site_id
184 _sites[site_id] = rendered
185 return site_id
188def _current_span() -> _Span | None:
189 return getattr(_local, "span", None)
192# test_db rebuilds the schema from migrations to diff it against the models. That is schema plumbing rather than an
193# access pattern, and it is already covered by the schema-diff artifact. Note this is not what keeps the real
194# timezone_areas.sql out of the recording: test_migrations is skipped in test:backend anyway, because the backend
195# image has no pg_dump. What bounds the artifact is _MAX_SQL_CHARS.
196_EXCLUDED_MODULES = ("src/tests/test_db.py",)
199def _after_cursor_execute(conn, cursor, statement, parameters, context, executemany):
200 test = _current_test
201 if test is None or test.startswith(_EXCLUDED_MODULES):
202 return
203 fingerprint = _fingerprint(statement)
204 with _lock:
205 shape = _shapes.get(fingerprint)
206 if shape is None or test < shape.first_seen_in:
207 example, params = _render_example(conn, statement, parameters)
208 # The context knows for ORM-issued statements; fall back to reading the statement for the rest.
209 is_write = bool(
210 (context is not None and (context.isinsert or context.isupdate or context.isdelete))
211 or re.match(r"^\s*(INSERT|UPDATE|DELETE)\b", statement, re.IGNORECASE)
212 )
213 _shapes[fingerprint] = _Shape(
214 id=_shape_id(fingerprint),
215 sql=fingerprint,
216 example=example,
217 params=params,
218 write=is_write,
219 first_seen_in=test,
220 )
221 shape = _shapes[fingerprint]
223 span = _current_span()
224 if span is None:
225 # A query outside any RPC or job: fixture setup, or the test body using session_scope() directly.
226 spans = _tests.setdefault(test, [])
227 if spans and spans[-1].kind == "body":
228 span = spans[-1]
229 else:
230 span = _Span(kind="body", name=None)
231 spans.append(span)
232 span.queries.append(shape.id)
233 span.sites.append(_callsite())
236class _SpanScope:
237 """Marks the queries issued inside it as belonging to one RPC call or background job run.
239 The span is registered against the current test on entry, so the recorded order matches call order. It is held
240 in a thread-local because the real-server sessions run handlers on a gRPC executor thread while the test body
241 runs on the main thread; FakeChannel runs them inline, and the same thread-local covers that too.
242 """
244 def __init__(self, kind: str, name: str | None):
245 self._span = _Span(kind=kind, name=name)
246 self._previous: _Span | None = None
248 def __enter__(self) -> _SpanScope:
249 if _current_test is not None:
250 with _lock:
251 _tests.setdefault(_current_test, []).append(self._span)
252 self._previous = _current_span()
253 _local.span = self._span
254 return self
256 def __exit__(self, *exc: object) -> None:
257 _local.span = self._previous
260def span(kind: str, name: str | None) -> Any:
261 """Open a recording span. A no-op unless --query-log is active, so callers need no guard of their own."""
262 if not _enabled:
263 return _NULL_SPAN
264 return _SpanScope(kind, name)
267class _NullSpan:
268 def __enter__(self) -> _NullSpan:
269 return self
271 def __exit__(self, *exc: object) -> None:
272 pass
275_NULL_SPAN = _NullSpan()
278def enable(engine: Engine) -> None:
279 global _enabled
280 _enabled = True
281 event.listen(engine, "after_cursor_execute", _after_cursor_execute)
284def set_current_test(test_id: str | None) -> None:
285 global _current_test
286 _current_test = test_id
287 _local.span = None
290def dump(directory: Path) -> Path:
291 """Write this node's recording. The node suffix keeps the parallel CI jobs from overwriting each other.
293 Gzipped: it is mostly repeated SQL and compresses about fifteen-fold, and this file is carried between CI jobs as
294 an artifact. Read it with `gunzip -c`, or let query_log_report.py merge it.
295 """
296 node = os.environ.get("CI_NODE_INDEX", "local")
297 directory.mkdir(parents=True, exist_ok=True)
298 path = directory / f"data.{node}.json.gz"
299 with _lock:
300 data = {
301 "shapes": {
302 shape.id: {
303 "sql": shape.sql,
304 "example": shape.example,
305 "params": shape.params,
306 "write": shape.write,
307 # The merge step uses this to pick the same `example` the single-node run would have picked.
308 "first_seen_in": shape.first_seen_in,
309 }
310 for shape in _shapes.values()
311 },
312 "sites": dict(sorted(_sites.items())),
313 "tests": {
314 test: [{"kind": s.kind, "name": s.name, "queries": s.queries, "sites": s.sites} for s in spans]
315 for test, spans in sorted(_tests.items())
316 },
317 }
318 # mtime=0 keeps the bytes reproducible, so two identical runs produce byte-identical dumps.
319 path.write_bytes(gzip.compress(json.dumps(data, separators=(",", ":"), sort_keys=True).encode(), mtime=0))
320 return path