Coverage for src/asgi_profiler/instrument.py: 88%
99 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-22 16:59 +0000
1"""SQL capture.
3Listeners are attached to SQLAlchemy's `Engine` **class**, so every engine in
4the process is covered without the application handing us one. That is what
5makes SQLModel work for free: `sqlmodel.create_engine` returns a SQLAlchemy
6`Engine`, and `sqlmodel.Session` subclasses `sqlalchemy.orm.Session`, so the
7same cursor events fire. Async engines are covered too -- `AsyncEngine` drives
8a sync `Engine` underneath, and that is where the events live.
9"""
11from __future__ import annotations
13import contextlib
14import contextvars
15import threading
16import time
17import traceback
18from collections.abc import Iterator
19from typing import Any
21from sqlalchemy import event
22from sqlalchemy.engine import Engine
24from .models import MAX_SQL_CHARS, Query
26try: # SQLAlchemy pulls greenlet in for its asyncio support
27 import greenlet
28except ImportError: # pragma: no cover - sync-only install
29 greenlet = None
31#: The list being filled for the in-flight request.
32#:
33#: It has to be a *mutable* object rather than a value: SQLAlchemy's sync work
34#: often runs in a worker thread (anyio copies the context into it), so appends
35#: made there must land in an object the request task already holds.
36current_queries: contextvars.ContextVar[list[Query] | None] = contextvars.ContextVar(
37 "asgi_profiler_queries", default=None
38)
40#: `(perf_counter at the start of the request, id of the event loop thread)`.
41#:
42#: The first turns an absolute timestamp into an offset from the start of the
43#: request, which is what makes two statements comparable: whether they
44#: overlapped, how long the gap between them was, whether six of them ran
45#: together or one after another.
46#:
47#: The second is how a blocking call is recognised. See `_blocks_the_loop`.
48current_origin: contextvars.ContextVar[tuple[float, int] | None] = contextvars.ContextVar(
49 "asgi_profiler_origin", default=None
50)
52_IGNORED_FRAME_PARTS = (
53 "/sqlalchemy/",
54 "\\sqlalchemy\\",
55 "/sqlmodel/",
56 "\\sqlmodel\\",
57 "/asgi_profiler/",
58 "\\asgi_profiler\\",
59 "/anyio/",
60 "\\anyio\\",
61 "/asyncio/",
62 "\\asyncio\\",
63 "/greenlet/",
64 "\\greenlet\\",
65 "/gevent/",
66 "\\gevent\\",
67 "/concurrent/futures/",
68 "\\concurrent\\futures\\",
69 "/threading.py",
70 "\\threading.py",
71)
73#: Live capture settings, read by the listeners at event time.
74#:
75#: A dict rather than closure arguments so that a later `install()` can change
76#: them. The listeners are attached to the `Engine` class exactly once, but the
77#: settings they honour must stay reconfigurable -- otherwise the first
78#: `install()` in a process would silently dictate `capture_stacks` for every
79#: application in it, including across tests.
80_settings: dict[str, Any] = {"capture_stacks": True, "stack_depth": 8}
82_installed = False
83_listeners: list[tuple[str, Any]] = []
85#: Cap on how far up the stack to walk before giving up on finding
86#: application frames. Guards against pathological recursion.
87_MAX_STACK_WALK = 250
90def install(*, capture_stacks: bool = True, stack_depth: int = 8) -> None:
91 """Attach the cursor listeners. Safe to call more than once.
93 Calling it again updates `capture_stacks` and `stack_depth` in place; the
94 listeners themselves are only attached the first time.
95 """
96 global _installed
97 _settings["capture_stacks"] = capture_stacks
98 _settings["stack_depth"] = stack_depth
99 if _installed:
100 return
102 # ANN001 missing-type-function-argument, ARG001 unused-function-argument:
103 # SQLAlchemy calls this with a fixed positional signature we do not choose.
104 def _before(conn, cursor, statement, parameters, context, executemany): # noqa: ANN001, ARG001
105 # A stack: a single connection can nest executions.
106 conn.info.setdefault("_profiler_started", []).append(time.perf_counter())
108 # ANN001 missing-type-function-argument, ARG001 unused-function-argument:
109 # same fixed signature as `_before`.
110 def _after(conn, cursor, statement, parameters, context, executemany): # noqa: ANN001, ARG001
111 started = conn.info.get("_profiler_started")
112 if not started: 112 ↛ 113line 112 didn't jump to line 113 because the condition on line 112 was never true
113 return
114 _record(statement, parameters, started.pop(), time.perf_counter())
116 # ANN001 missing-type-function-argument: SQLAlchemy does not export
117 # `ExceptionContext` as a public name to annotate this with.
118 def _error(context) -> None: # noqa: ANN001
119 """A statement that raised never reaches `after_cursor_execute`.
121 Without this the failure is invisible in the UI *and* its start time
122 is stranded on the connection forever, which on a pooled connection
123 that sees repeated errors is an unbounded list.
124 """
125 conn = getattr(context, "connection", None)
126 if conn is None or context.statement is None: 126 ↛ 127line 126 didn't jump to line 127 because the condition on line 126 was never true
127 return # a connect-level failure: no cursor execution to close out
128 started = conn.info.get("_profiler_started")
129 if not started: 129 ↛ 130line 129 didn't jump to line 130 because the condition on line 129 was never true
130 return
131 _record(
132 context.statement,
133 context.parameters,
134 started.pop(),
135 time.perf_counter(),
136 error=_format_error(context.original_exception),
137 )
139 for name, listener in (
140 ("before_cursor_execute", _before),
141 ("after_cursor_execute", _after),
142 ("handle_error", _error),
143 ):
144 event.listen(Engine, name, listener)
145 _listeners.append((name, listener))
147 _installed = True
150def uninstall() -> None:
151 """Detach the listeners. Mainly for test isolation."""
152 global _installed
153 for name, listener in _listeners:
154 with contextlib.suppress(Exception): # already gone
155 event.remove(Engine, name, listener)
156 _listeners.clear()
157 _installed = False
160def _blocks_the_loop(loop_thread: int) -> bool:
161 """Did this statement run *on* the event loop, stopping it?
163 Three situations to tell apart, and only one of them is a bug:
165 * A different thread. Starlette runs `def` endpoints in a worker thread,
166 and `asyncio.to_thread` puts anything else there on request. The loop
167 keeps serving other requests throughout, so this is fine.
168 * The loop thread, inside a greenlet. SQLAlchemy's async bridge runs the
169 driver call in a greenlet it spawned and switches back to the loop while
170 it waits. Also fine, and the common case for `create_async_engine`.
171 * The loop thread, in the main greenlet. Nothing yielded: a synchronous
172 driver call inside an `async def` endpoint, holding the only thread the
173 server has for the entire duration of the query. Every other request
174 in flight is frozen behind it.
176 The greenlet test is on `parent` rather than on any SQLAlchemy internal,
177 because the bridge has been reorganised more than once and an attribute
178 name that moves would turn this into a silent false positive on every
179 async query. A spawned greenlet always has a parent; the main one never
180 does.
182 Under gevent everything runs in a greenlet, so this reports nothing. A
183 detector that misses a real problem is a disappointment; one that invents
184 them is the reason people turn detectors off.
185 """
186 if threading.get_ident() != loop_thread:
187 return False
188 if greenlet is not None and getattr(greenlet.getcurrent(), "parent", None) is not None:
189 return False
190 return True
193def _record(
194 statement: Any,
195 parameters: Any,
196 started: float,
197 ended: float,
198 error: str | None = None,
199) -> None:
200 queries = current_queries.get()
201 if queries is None:
202 return # outside any request: startup, migrations, a shell
204 origin = current_origin.get()
205 queries.append(
206 Query(
207 sql=_normalise(statement),
208 params=_format_params(parameters),
209 duration_ms=(ended - started) * 1000,
210 started_ms=(started - origin[0]) * 1000 if origin else 0.0,
211 blocking=_blocks_the_loop(origin[1]) if origin else False,
212 stack=(_capture_stack(_settings["stack_depth"]) if _settings["capture_stacks"] else []),
213 error=error,
214 )
215 )
218def _normalise(statement: Any) -> str:
219 text = " ".join(str(statement).split())
220 # `max_requests` bounds how many requests are kept, not how big they are.
221 # An ORM bulk insert is a single statement tens of kilobytes long, and one
222 # endpoint doing that can hold hundreds of megabytes in the ring buffer.
223 if len(text) > MAX_SQL_CHARS:
224 return text[:MAX_SQL_CHARS] + f" ... [truncated, {len(text)} chars]"
225 return text
228def _format_params(parameters: Any) -> str:
229 try:
230 text = repr(parameters)
231 except Exception: # a parameter with a hostile __repr__
232 return "<unrepresentable>"
233 return text if len(text) <= 500 else text[:500] + " ..."
236def _format_error(exc: BaseException | None) -> str:
237 if exc is None: 237 ↛ 238line 237 didn't jump to line 238 because the condition on line 237 was never true
238 return "error"
239 text = f"{type(exc).__name__}: {exc}"
240 return text if len(text) <= 500 else text[:500] + " ..."
243def _walk_all_frames() -> Iterator[tuple[str, int, str]]:
244 """Every frame from here outwards, crossing greenlet boundaries.
246 Under an async engine, SQLAlchemy runs the DBAPI call inside a greenlet it
247 spawned. That greenlet's stack is five frames of SQLAlchemy and nothing
248 else -- the application frame that issued the query lives in the *parent*
249 greenlet, which `walk_stack` cannot see. Without this, per-query stacks are
250 silently empty for every async user, which is most of them.
251 """
252 for here, lineno in traceback.walk_stack(None):
253 yield here.f_code.co_filename, lineno, here.f_code.co_name
255 if greenlet is None: 255 ↛ 256line 255 didn't jump to line 256 because the condition on line 255 was never true
256 return
257 current = getattr(greenlet.getcurrent(), "parent", None)
258 while current is not None:
259 frame: Any = getattr(current, "gr_frame", None)
260 while frame is not None: 260 ↛ 263line 260 didn't jump to line 263 because the condition on line 260 was always true
261 yield frame.f_code.co_filename, frame.f_lineno, frame.f_code.co_name
262 frame = frame.f_back
263 current = getattr(current, "parent", None)
266def _capture_stack(depth: int) -> list[str]:
267 """Application frames that led to this query, innermost last.
269 This is the thing that makes an N+1 actionable: it names the line in *your*
270 code that triggered it.
272 Walks outward from the caller and stops as soon as `depth` application
273 frames have been found. `traceback.extract_stack()` would be the obvious
274 call, but it unwinds the *entire* stack and does a `linecache` source
275 lookup per frame, then throws almost all of it away -- which in a real
276 async request stack of 100+ frames costs an order of magnitude more per
277 query than this does.
278 """
279 frames: list[str] = []
280 for seen, (filename, lineno, name) in enumerate(_walk_all_frames()):
281 if seen >= _MAX_STACK_WALK: 281 ↛ 282line 281 didn't jump to line 282 because the condition on line 281 was never true
282 break
283 if any(part in filename for part in _IGNORED_FRAME_PARTS):
284 continue
285 frames.append(f"{filename}:{lineno} in {name}")
286 if len(frames) >= depth:
287 break
288 frames.reverse() # innermost last, matching a traceback
289 return frames