Coverage for app/backend/src/tests/fixtures/query_log.py: 94%
152 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-10 04:35 +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.
87 sql = re.sub(r"'[^']{64,}'", "'...'", sql)
88 return _truncate(re.sub(r"\s+", " ", sql).strip())
91def _shape_id(fingerprint: str) -> str:
92 # Content-addressed, so the three pytest-split nodes agree on ids and merging is a plain dict union. A counter
93 # would be assigned in per-node encounter order and collide across nodes.
94 return hashlib.blake2b(fingerprint.encode(), digest_size=6).hexdigest()
97def _render_example(conn: Any, statement: str, parameters: Any) -> tuple[str, str | None]:
98 """Inline the bound values so the statement can be pasted into psql.
100 psycopg's ClientCursor.mogrify applies the same adaptation and escaping the driver would otherwise do
101 server-side, which is far more faithful than re-implementing literal binding for bytea, arrays and PostGIS
102 geometries. It is purely client-side, so it issues nothing on the connection.
103 """
104 # executemany passes a sequence of parameter sets; one row is enough to have something runnable.
105 if isinstance(parameters, (list, tuple)) and parameters and isinstance(parameters[0], (dict, list, tuple)):
106 parameters = parameters[0]
107 try:
108 params = json.dumps(parameters, default=repr) if parameters else None
109 except TypeError, ValueError:
110 params = None
111 if params is not None:
112 params = _truncate(params)
113 try:
114 cursor = psycopg.ClientCursor(conn.connection.driver_connection)
115 # Truncated past the cap, so the marker is a SQL comment: the result is visibly not runnable rather than
116 # silently invalid.
117 return _truncate(cursor.mogrify(statement, parameters)), params
118 except Exception:
119 # Not worth losing the whole entry over; the fingerprint plus the parameters is still pasteable by hand.
120 return _truncate(statement), params
123# Everything under the backend's src/ is ours; paths are reported relative to it. couchers/ is application code,
124# anything else under src/ is test scaffolding.
125_SRC_ROOT = "/src/"
126_APP_ROOT = "/src/couchers/"
127_SKIP, _APP, _TEST = 0, 1, 2
128# Frames from these are plumbing between our code and the driver, so they never make a useful call site. The
129# recorder's own path is spelled out: a bare "query_log.py" would also swallow test_query_log.py.
130_CALLSITE_SKIP = ("/sqlalchemy/", "/psycopg", "/alembic/", "/fixtures/query_log.py")
133def _frame_kind(code: CodeType) -> int:
134 """_APP, _TEST or _SKIP. Cached by filename, which is all the answer depends on: this runs on every frame of
135 every execution, and the string work is what would otherwise make stack walking too expensive to leave on.
137 Keyed by filename rather than by the code object, because code objects compare equal without regard to
138 co_filename, so two same-bodied functions in different files would share an entry.
139 """
140 filename = code.co_filename
141 known = _frame_cache.get(filename)
142 if known is None:
143 if any(part in filename for part in _CALLSITE_SKIP) or _SRC_ROOT not in filename:
144 known = _SKIP
145 else:
146 known = _APP if _APP_ROOT in filename else _TEST
147 _frame_cache[filename] = known
148 return known
151# How many of our own frames to keep. The innermost is the line that issued the query; the next couple show the
152# chain that got there, which is usually what tells you whether a repeat is a loop.
153_CALLSITE_FRAMES = 3
156def _callsite() -> str:
157 """The innermost few application frames, innermost first, as "path:line in func".
159 Only couchers/ frames: the test and the fixture handler that got here are already implied by the test and span
160 this is recorded under, and including them multiplies the number of distinct call sites for no added meaning.
161 Test frames are used only when a query has no application frame at all, as fixture setup often does not.
162 """
163 frames: list[str] = []
164 fallback = ""
165 frame: FrameType | None = sys._getframe(1)
166 while frame is not None and len(frames) < _CALLSITE_FRAMES:
167 code = frame.f_code
168 kind = _frame_kind(code)
169 if kind != _SKIP:
170 path = code.co_filename.split(_SRC_ROOT, 1)[-1]
171 rendered_frame = f"{path}:{frame.f_lineno} in {code.co_name}"
172 if kind == _APP:
173 frames.append(rendered_frame)
174 elif not fallback:
175 fallback = rendered_frame
176 frame = frame.f_back
177 rendered = " <- ".join(frames) or fallback
178 site_id = _site_ids.get(rendered)
179 if site_id is None:
180 site_id = _shape_id(rendered)
181 _site_ids[rendered] = site_id
182 _sites[site_id] = rendered
183 return site_id
186def _current_span() -> _Span | None:
187 return getattr(_local, "span", None)
190# test_db rebuilds the schema from migrations to diff it against the models. That is schema plumbing rather than an
191# access pattern, and it is already covered by the schema-diff artifact. Note this is not what keeps the real
192# timezone_areas.sql out of the recording: test_migrations is skipped in test:backend anyway, because the backend
193# image has no pg_dump. What bounds the artifact is _MAX_SQL_CHARS.
194_EXCLUDED_MODULES = ("src/tests/test_db.py",)
197def _after_cursor_execute(conn, cursor, statement, parameters, context, executemany):
198 test = _current_test
199 if test is None or test.startswith(_EXCLUDED_MODULES):
200 return
201 fingerprint = _fingerprint(statement)
202 with _lock:
203 shape = _shapes.get(fingerprint)
204 if shape is None or test < shape.first_seen_in:
205 example, params = _render_example(conn, statement, parameters)
206 # The context knows for ORM-issued statements; fall back to reading the statement for the rest.
207 is_write = bool(
208 (context is not None and (context.isinsert or context.isupdate or context.isdelete))
209 or re.match(r"^\s*(INSERT|UPDATE|DELETE)\b", statement, re.IGNORECASE)
210 )
211 _shapes[fingerprint] = _Shape(
212 id=_shape_id(fingerprint),
213 sql=fingerprint,
214 example=example,
215 params=params,
216 write=is_write,
217 first_seen_in=test,
218 )
219 shape = _shapes[fingerprint]
221 span = _current_span()
222 if span is None:
223 # A query outside any RPC or job: fixture setup, or the test body using session_scope() directly.
224 spans = _tests.setdefault(test, [])
225 if spans and spans[-1].kind == "body":
226 span = spans[-1]
227 else:
228 span = _Span(kind="body", name=None)
229 spans.append(span)
230 span.queries.append(shape.id)
231 span.sites.append(_callsite())
234class _SpanScope:
235 """Marks the queries issued inside it as belonging to one RPC call or background job run.
237 The span is registered against the current test on entry, so the recorded order matches call order. It is held
238 in a thread-local because the real-server sessions run handlers on a gRPC executor thread while the test body
239 runs on the main thread; FakeChannel runs them inline, and the same thread-local covers that too.
240 """
242 def __init__(self, kind: str, name: str | None):
243 self._span = _Span(kind=kind, name=name)
244 self._previous: _Span | None = None
246 def __enter__(self) -> _SpanScope:
247 if _current_test is not None:
248 with _lock:
249 _tests.setdefault(_current_test, []).append(self._span)
250 self._previous = _current_span()
251 _local.span = self._span
252 return self
254 def __exit__(self, *exc: object) -> None:
255 _local.span = self._previous
258def span(kind: str, name: str | None) -> Any:
259 """Open a recording span. A no-op unless --query-log is active, so callers need no guard of their own."""
260 if not _enabled:
261 return _NULL_SPAN
262 return _SpanScope(kind, name)
265class _NullSpan:
266 def __enter__(self) -> _NullSpan:
267 return self
269 def __exit__(self, *exc: object) -> None:
270 pass
273_NULL_SPAN = _NullSpan()
276def enable(engine: Engine) -> None:
277 global _enabled
278 _enabled = True
279 event.listen(engine, "after_cursor_execute", _after_cursor_execute)
282def set_current_test(test_id: str | None) -> None:
283 global _current_test
284 _current_test = test_id
285 _local.span = None
288def dump(directory: Path) -> Path:
289 """Write this node's recording. The node suffix keeps the parallel CI jobs from overwriting each other.
291 Gzipped: it is mostly repeated SQL and compresses about fifteen-fold, and this file is carried between CI jobs as
292 an artifact. Read it with `gunzip -c`, or let query_log_report.py merge it.
293 """
294 node = os.environ.get("CI_NODE_INDEX", "local")
295 directory.mkdir(parents=True, exist_ok=True)
296 path = directory / f"data.{node}.json.gz"
297 with _lock:
298 data = {
299 "shapes": {
300 shape.id: {
301 "sql": shape.sql,
302 "example": shape.example,
303 "params": shape.params,
304 "write": shape.write,
305 # The merge step uses this to pick the same `example` the single-node run would have picked.
306 "first_seen_in": shape.first_seen_in,
307 }
308 for shape in _shapes.values()
309 },
310 "sites": dict(sorted(_sites.items())),
311 "tests": {
312 test: [{"kind": s.kind, "name": s.name, "queries": s.queries, "sites": s.sites} for s in spans]
313 for test, spans in sorted(_tests.items())
314 },
315 }
316 # mtime=0 keeps the bytes reproducible, so two identical runs produce byte-identical dumps.
317 path.write_bytes(gzip.compress(json.dumps(data, separators=(",", ":"), sort_keys=True).encode(), mtime=0))
318 return path