pytest-queryguard 0.1.0__py3-none-any.whl

pytest_queryguard/__init__.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+from .api import QueryBudgetExceeded, max_queries
+from .warnings import QueryGuardInternalWarning, QueryGuardWarning
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "QueryBudgetExceeded",
+    "QueryGuardInternalWarning",
+    "QueryGuardWarning",
+    "__version__",
+    "max_queries",
+]

pytest_queryguard/api.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from collections.abc import Iterator
+from contextlib import AbstractContextManager, contextmanager
+
+from .collector import QueryCollector, pop, push
+from .listeners import install
+from .messages import format_budget
+
+
+class QueryBudgetExceeded(AssertionError):
+    """Raised when a block runs more SQL statements than its budget allows."""
+
+
+@contextmanager
+def _measure(limit: int) -> Iterator[QueryCollector]:
+    # The budget failure is raised here; pytest should show the user's with block.
+    __tracebackhide__ = True
+    install()
+    collector = QueryCollector()
+    push(collector)
+    try:
+        yield collector
+    finally:
+        pop(collector)
+    if len(collector.statements) > limit:
+        raise QueryBudgetExceeded(format_budget(limit, collector.statements))
+
+
+def max_queries(limit: int) -> AbstractContextManager[QueryCollector]:
+    """Fail when the block runs more than ``limit`` SQL statements."""
+    if isinstance(limit, bool) or not isinstance(limit, int) or limit < 0:
+        raise ValueError(f"limit must be a non-negative integer, got {limit!r}")
+    return _measure(limit)

pytest_queryguard/callsite.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+import os
+import sys
+import sysconfig
+from collections.abc import Iterator
+from pathlib import Path
+from types import FrameType
+
+from .events import CallSite
+
+
+def _package_dir(name: str) -> str | None:
+    module = sys.modules.get(name)
+    file = getattr(module, "__file__", None)
+    return str(Path(file).resolve().parent) + os.sep if file else None
+
+
+def _compute_prefixes() -> tuple[str, ...]:
+    paths: list[str | None] = [
+        _package_dir("sqlalchemy"),
+        _package_dir("pytest"),
+        _package_dir("_pytest"),
+        _package_dir("pluggy"),
+        str(Path(__file__).resolve().parent) + os.sep,
+    ]
+    for key in ("stdlib", "platstdlib"):
+        path = sysconfig.get_paths().get(key)
+        if path:
+            paths.append(path.rstrip(os.sep) + os.sep)
+            paths.append(str(Path(path).resolve()) + os.sep)
+    return tuple(dict.fromkeys(os.path.normcase(p) for p in paths if p))
+
+
+_prefixes: tuple[str, ...] | None = None
+
+
+def _is_internal(filename: str) -> bool:
+    global _prefixes
+    if filename.startswith("<"):
+        return True
+    if _prefixes is None:
+        _prefixes = _compute_prefixes()
+    if os.path.normcase(filename).startswith(_prefixes):
+        return True
+    parts = Path(filename).parts
+    return "site-packages" in parts or "dist-packages" in parts
+
+
+def _parent_greenlet_frames() -> Iterator[FrameType]:
+    """Frames of the greenlets above the current one, innermost first.
+
+    SQLAlchemy runs the sync code behind an async session in a child greenlet,
+    so the awaiting user code is only reachable through the parent greenlet.
+    """
+    # greenlet comes with SQLAlchemy's async support; only used if already loaded.
+    # Code running inside a SQLAlchemy greenlet has imported it already, while
+    # an import here, after pytester purges sys.modules between runs, could
+    # create a second greenlet module object next to the one SQLAlchemy uses.
+    module = sys.modules.get("greenlet")
+    if module is None:
+        return
+    try:
+        parent = module.getcurrent().parent
+        while parent is not None:
+            frame = parent.gr_frame
+            if frame is not None:
+                yield frame
+            parent = parent.parent
+    except Exception:  # noqa: BLE001  # a plugin bug must never break the user suite
+        return
+
+
+def _first_external(frame: FrameType | None) -> CallSite | None:
+    while frame is not None:
+        filename = frame.f_code.co_filename
+        if not _is_internal(filename):
+            return CallSite(filename, frame.f_lineno, frame.f_code.co_name)
+        frame = frame.f_back
+    return None
+
+
+def find_callsite() -> CallSite | None:
+    """First stack frame outside SQLAlchemy, pytest, stdlib, site-packages and us.
+
+    When the stack ends inside a greenlet (async sessions), the search goes on
+    in the parent greenlets, where the awaiting user code lives.
+    """
+    site = _first_external(sys._getframe(1))
+    if site is not None:
+        return site
+    for frame in _parent_greenlet_frames():
+        site = _first_external(frame)
+        if site is not None:
+            return site
+    return None

pytest_queryguard/collector.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from .events import RelationshipLoadEvent, StatementEvent
+
+
+@dataclass(slots=True, eq=False)
+class QueryCollector:
+    statements: list[StatementEvent] = field(default_factory=list)
+    relationship_loads: list[RelationshipLoadEvent] = field(default_factory=list)
+
+
+# Process-wide on purpose: the code under test may run in another thread.
+_stack: list[QueryCollector] = []
+_generation = 0
+
+
+def generation() -> int:
+    """Changes on every push/pop; lets per-measurement state detect staleness."""
+    return _generation
+
+
+def push(collector: QueryCollector) -> None:
+    global _generation
+    _generation += 1
+    _stack.append(collector)
+
+
+def pop(collector: QueryCollector) -> None:
+    global _generation
+    _generation += 1
+    for i in range(len(_stack) - 1, -1, -1):
+        if _stack[i] is collector:
+            del _stack[i]
+            break
+
+
+def active() -> tuple[QueryCollector, ...]:
+    return tuple(_stack)
+
+
+def _reset_for_tests() -> None:
+    global _generation
+    _generation += 1
+    _stack.clear()

pytest_queryguard/detectors.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import re
+from collections import Counter
+from dataclasses import dataclass
+
+from .collector import QueryCollector
+from .events import CallSite
+
+_STRING = re.compile(r"'(?:[^']|'')*'")
+_NUMBER = re.compile(r"(?<![\w.])\d+(?:\.\d+)?(?![\w.])")
+_PLACEHOLDER = r"""\?|%s|%\(\w+\)s|:\w+|\$\d+|\d+(?:\.\d+)?|'(?:[^']|'')*'"""
+# PostgreSQL drivers may cast each bound value: $1::INTEGER, ?::VARCHAR(20)[].
+_CAST = r"(?:\s*::\s*\w+(?:\s*\(\s*\d+(?:\s*,\s*\d+)*\s*\))?(?:\s*\[\])*)?"
+_ITEM = rf"(?:{_PLACEHOLDER}){_CAST}"
+_IN_LIST = re.compile(
+    rf"\bIN\s*\(\s*{_ITEM}(?:\s*,\s*{_ITEM})*\s*\)",
+    re.IGNORECASE,
+)
+_DOLLAR_PARAM = re.compile(r"\$\d+")
+_WS = re.compile(r"\s+")
+_WRITES = ("INSERT", "UPDATE", "DELETE")
+
+
+def normalize(sql: str) -> str:
+    """Reduce a statement to a shape shared by all its executions."""
+    out = _WS.sub(" ", sql).strip()
+    # Lists first, so every placeholder style collapses to the same text.
+    out = _IN_LIST.sub("IN (...)", out)
+    out = _STRING.sub("?", out)
+    # Numbered parameters ($1, $2) become "?" so they compare equal.
+    out = _DOLLAR_PARAM.sub("?", out)
+    return _NUMBER.sub("?", out)
+
+
+@dataclass(frozen=True, slots=True)
+class RelationshipFinding:
+    relationship: str
+    count: int
+    first: CallSite | None
+    sample_sql: str | None
+
+
+@dataclass(frozen=True, slots=True)
+class RepeatFinding:
+    sql: str
+    count: int
+    first: CallSite | None
+
+
+Finding = RelationshipFinding | RepeatFinding
+
+
+def _relationship_findings(
+    collector: QueryCollector, threshold: int
+) -> list[RelationshipFinding]:
+    parents: dict[str, set[object]] = {}
+    counts: Counter[str] = Counter()
+    first: dict[str, CallSite | None] = {}
+    for load in collector.relationship_loads:
+        if load.parent_key is None:
+            continue
+        name = load.relationship
+        parents.setdefault(name, set()).add(load.parent_key)
+        counts[name] += 1
+        if first.get(name) is None:
+            first[name] = load.callsite
+    samples: dict[str, str] = {}
+    for stmt in collector.statements:
+        if stmt.relationship is not None:
+            samples.setdefault(stmt.relationship, stmt.sql)
+    found = [
+        RelationshipFinding(name, counts[name], first[name], samples.get(name))
+        for name, keys in parents.items()
+        if len(keys) >= threshold
+    ]
+    return sorted(found, key=lambda f: -f.count)
+
+
+def _repeat_findings(collector: QueryCollector, threshold: int) -> list[RepeatFinding]:
+    counts: Counter[str] = Counter()
+    first: dict[str, CallSite | None] = {}
+    for stmt in collector.statements:
+        if stmt.relationship is not None:
+            continue
+        if stmt.sql.lstrip().upper().startswith(_WRITES):
+            continue
+        key = normalize(stmt.sql)
+        counts[key] += 1
+        if first.get(key) is None:
+            first[key] = stmt.callsite
+    found = [
+        RepeatFinding(sql, n, first[sql]) for sql, n in counts.items() if n >= threshold
+    ]
+    return sorted(found, key=lambda f: -f.count)
+
+
+def detect(
+    collector: QueryCollector, *, relationship_threshold: int, repeat_threshold: int
+) -> list[Finding]:
+    findings: list[Finding] = []
+    findings.extend(_relationship_findings(collector, relationship_threshold))
+    findings.extend(_repeat_findings(collector, repeat_threshold))
+    return findings

pytest_queryguard/events.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True, slots=True)
+class CallSite:
+    filename: str
+    lineno: int
+    function: str
+
+
+@dataclass(frozen=True, slots=True)
+class StatementEvent:
+    sql: str
+    callsite: CallSite | None
+    relationship: str | None
+
+
+@dataclass(frozen=True, slots=True)
+class RelationshipLoadEvent:
+    relationship: str
+    parent_key: object
+    callsite: CallSite | None

pytest_queryguard/listeners.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+import threading
+from typing import Any
+
+from sqlalchemy import event
+from sqlalchemy.engine import Engine
+from sqlalchemy.orm import ORMExecuteState, Session
+
+from . import collector
+from .callsite import find_callsite
+from .events import RelationshipLoadEvent, StatementEvent
+
+_installed = False
+_error: BaseException | None = None
+_local = threading.local()
+
+
+def internal_error() -> BaseException | None:
+    """The first exception swallowed by a listener, if any."""
+    return _error
+
+
+def _set_pending(name: str | None) -> None:
+    _local.pending = None if name is None else (collector.generation(), name)
+
+
+def _take_pending() -> str | None:
+    pending = getattr(_local, "pending", None)
+    _local.pending = None
+    if pending is not None and pending[0] == collector.generation():
+        return str(pending[1])
+    return None
+
+
+def _reset_for_tests() -> None:
+    global _error
+    _error = None
+    _local.pending = None
+
+
+def _on_orm_execute(state: ORMExecuteState) -> None:
+    if _error is not None:
+        return
+    try:
+        targets = collector.active()
+        if not targets:
+            return
+        if not state.is_relationship_load:
+            _set_pending(None)
+            return
+        name = str(state.loader_strategy_path[-1])  # type: ignore[index]
+        parent = state.lazy_loaded_from
+        key: object = None
+        if parent is not None:
+            key = parent.identity_key
+            if key is None:
+                key = id(parent)
+        _set_pending(name)
+        site = find_callsite()
+        load = RelationshipLoadEvent(name, key, site)
+        for target in targets:
+            target.relationship_loads.append(load)
+    except Exception as exc:  # noqa: BLE001  # a plugin bug must never break the user suite
+        _fail(exc)
+
+
+def _on_cursor_execute(
+    conn: Any,
+    cursor: Any,
+    statement: str,
+    parameters: Any,
+    context: Any,
+    executemany: bool,
+) -> None:
+    if _error is not None:
+        return
+    try:
+        targets = collector.active()
+        if not targets:
+            return
+        stmt = StatementEvent(statement, find_callsite(), _take_pending())
+        for target in targets:
+            target.statements.append(stmt)
+    except Exception as exc:  # noqa: BLE001  # a plugin bug must never break the user suite
+        _fail(exc)
+
+
+def _fail(exc: Exception) -> None:
+    global _error
+    if _error is None:
+        _error = exc
+
+
+def record_error(exc: Exception) -> None:
+    """Disable measurement after a bug outside the listeners (first error wins)."""
+    _fail(exc)
+
+
+def swap_error(new: BaseException | None) -> BaseException | None:
+    """Replace the recorded error and return the previous one.
+
+    The plugin clears it for each session and restores the outer session's
+    error afterwards, so in-process nested sessions stay independent.
+    """
+    global _error
+    old, _error = _error, new
+    return old
+
+
+def install() -> None:
+    """Register the class-level listeners once per process."""
+    global _installed
+    if _installed:
+        return
+    event.listen(Session, "do_orm_execute", _on_orm_execute)
+    event.listen(Engine, "before_cursor_execute", _on_cursor_execute)
+    _installed = True

pytest_queryguard/messages.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from .detectors import Finding, RelationshipFinding
+from .events import CallSite, StatementEvent
+
+_SQL_LIMIT = 200
+_REPEAT_FIX = (
+    "fix: fetch all rows in one query (WHERE id IN (...)) instead of in a loop"
+)
+
+
+def _display_path(filename: str, cwd: Path | None) -> str:
+    base = cwd if cwd is not None else Path.cwd()
+    path = Path(filename)
+    try:
+        return path.relative_to(base).as_posix()
+    except ValueError:
+        return filename
+
+
+def _location(site: CallSite, cwd: Path | None) -> str:
+    return f"{_display_path(site.filename, cwd)}:{site.lineno} in {site.function}"
+
+
+def _first_at(site: CallSite | None, cwd: Path | None) -> str:
+    where = "<unknown location>" if site is None else _location(site, cwd)
+    return f"  first at {where}"
+
+
+def _block(finding: Finding, cwd: Path | None) -> list[str]:
+    if isinstance(finding, RelationshipFinding):
+        parent = finding.relationship.split(".", 1)[0]
+        lines = [
+            f"  {finding.relationship} was lazy-loaded {finding.count} times",
+            _first_at(finding.first, cwd),
+        ]
+        if finding.sample_sql is not None:
+            lines.append(f"  sql: {_truncate(finding.sample_sql)}")
+        return [
+            *lines,
+            "  fix: load it with the parent query, e.g.",
+            f"       select({parent}).options(selectinload({finding.relationship}))",
+        ]
+    return [
+        f"  the same statement ran {finding.count} times",
+        _first_at(finding.first, cwd),
+        f"  {_truncate(finding.sql)}",
+        f"  {_REPEAT_FIX}",
+    ]
+
+
+def format_findings(findings: list[Finding], *, cwd: Path | None = None) -> str:
+    lines = ["queryguard: N+1 detected"]
+    for finding in findings:
+        lines.append("")
+        lines.extend(_block(finding, cwd))
+    return "\n".join(lines)
+
+
+def _truncate(sql: str) -> str:
+    sql = " ".join(sql.split())
+    if len(sql) > _SQL_LIMIT:
+        return sql[:_SQL_LIMIT] + "..."
+    return sql
+
+
+def format_budget(
+    limit: int, statements: list[StatementEvent], *, cwd: Path | None = None
+) -> str:
+    lines = [
+        f"queryguard: expected at most {limit} queries, got {len(statements)}",
+        "",
+    ]
+    for number, stmt in enumerate(statements, start=1):
+        line = f"  {number}. {_truncate(stmt.sql)}"
+        if stmt.callsite is not None:
+            line += f"  ({_location(stmt.callsite, cwd)})"
+        lines.append(line)
+    return "\n".join(lines)

pytest_queryguard/plugin.py

@@ -0,0 +1,290 @@
+from __future__ import annotations
+
+import warnings as std_warnings
+from collections.abc import Generator
+from typing import Any
+
+import pytest
+
+from .collector import QueryCollector, pop, push
+from .detectors import detect
+from .listeners import install, internal_error, record_error, swap_error
+from .messages import format_budget, format_findings
+from .warnings import QueryGuardInternalWarning, QueryGuardWarning
+
+_MODES = ("fail", "warn", "off")
+_internal_reported = pytest.StashKey[bool]()
+_outer_error = pytest.StashKey[BaseException | None]()
+_settings = pytest.StashKey["_Settings"]()
+_report = pytest.StashKey[list[tuple[str, int]]]()
+_REPORT_SIZE = 10
+
+
+class _Settings:
+    def __init__(self, mode: str, relationship: int, repeat: int) -> None:
+        self.mode = mode
+        self.relationship = relationship
+        self.repeat = repeat
+
+
+# pytest 9 reads native [tool.pytest] TOML, where thresholds are real integers
+# and a string-typed option would crash; older versions only know strings.
+_NATIVE_INT = pytest.version_tuple >= (9,)
+_INT_TYPE: dict[str, Any] = {"type": "int"} if _NATIVE_INT else {}
+
+
+def pytest_addoption(parser: pytest.Parser) -> None:
+    group = parser.getgroup("queryguard")
+    group.addoption(
+        "--queryguard-mode",
+        dest="queryguard_mode",
+        default=None,
+        help="N+1 handling: fail, warn or off (overrides queryguard_mode).",
+    )
+    group.addoption(
+        "--queryguard-report",
+        dest="queryguard_report",
+        action="store_true",
+        default=False,
+        help="Print the tests with the most queries at the end of the session.",
+    )
+    parser.addini("queryguard_mode", "fail, warn or off", default="fail")
+    parser.addini(
+        "queryguard_relationship_threshold",
+        "Lazy loads of one relationship treated as N+1 (integer >= 2)",
+        default=2 if _NATIVE_INT else "2",
+        **_INT_TYPE,
+    )
+    parser.addini(
+        "queryguard_repeat_threshold",
+        "Executions of one statement treated as N+1 (integer >= 2)",
+        default=5 if _NATIVE_INT else "5",
+        **_INT_TYPE,
+    )
+
+
+def _threshold(config: pytest.Config, name: str) -> int:
+    detail = ""
+    try:
+        raw: Any = config.getini(name)
+    except (TypeError, ValueError) as exc:
+        raw, detail = None, f": {exc}"
+    try:
+        value = 0 if isinstance(raw, bool) else int(raw)
+    except (TypeError, ValueError):
+        value = 0
+    if value < 2:
+        shown = detail or f", got {raw!r}"
+        raise pytest.UsageError(f"{name} must be an integer >= 2{shown}")
+    return value
+
+
+def _read_settings(config: pytest.Config) -> _Settings:
+    cli_mode = config.getoption("queryguard_mode")
+    if cli_mode is not None:
+        mode, source = cli_mode, "--queryguard-mode"
+    else:
+        source = "queryguard_mode"
+        try:
+            mode = config.getini(source)
+        except (TypeError, ValueError) as exc:
+            # pytest 9 rejects a non-string value from native TOML on its own.
+            raise pytest.UsageError(
+                f"{source} must be one of {', '.join(_MODES)}: {exc}"
+            ) from None
+    if mode not in _MODES:
+        raise pytest.UsageError(
+            f"{source} must be one of {', '.join(_MODES)}, got {mode!r}"
+        )
+    return _Settings(
+        mode,
+        _threshold(config, "queryguard_relationship_threshold"),
+        _threshold(config, "queryguard_repeat_threshold"),
+    )
+
+
+def pytest_configure(config: pytest.Config) -> None:
+    config.addinivalue_line(
+        "markers",
+        "max_queries(n): fail the test if it runs more than n SQL statements.",
+    )
+    config.addinivalue_line(
+        "markers",
+        "allow_nplusone: do not report N+1 query findings for this test.",
+    )
+    config.stash[_outer_error] = swap_error(None)
+    settings = _read_settings(config)
+    config.stash[_settings] = settings
+    if config.getoption("queryguard_report"):
+        config.stash[_report] = []
+    if settings.mode != "off":
+        install()
+
+
+def pytest_unconfigure(config: pytest.Config) -> None:
+    swap_error(config.stash.get(_outer_error, None))
+
+
+def _report_internal_error(config: pytest.Config) -> None:
+    error = internal_error()
+    if error is None or config.stash.get(_internal_reported, False):
+        return
+    config.stash[_internal_reported] = True
+    std_warnings.warn(
+        QueryGuardInternalWarning(
+            f"pytest-queryguard disabled itself after an internal error: {error!r}"
+        ),
+        stacklevel=1,
+    )
+
+
+def _analyse(
+    item: pytest.Item, collector: QueryCollector, settings: _Settings
+) -> str | None:
+    """The report for the finished test, or None; never raises."""
+    if internal_error() is not None or item.get_closest_marker("allow_nplusone"):
+        return None
+    try:
+        findings = detect(
+            collector,
+            relationship_threshold=settings.relationship,
+            repeat_threshold=settings.repeat,
+        )
+        if not findings:
+            return None
+        return format_findings(findings, cwd=item.config.rootpath)
+    except Exception as exc:  # noqa: BLE001  # a plugin bug must never break the user suite
+        record_error(exc)
+        return None
+
+
+def _budget(item: pytest.Item) -> int | None:
+    marker = item.get_closest_marker("max_queries")
+    if marker is None:
+        return None
+    if len(marker.args) > 1 or set(marker.kwargs) - {"n"}:
+        pytest.fail(
+            f"max_queries on {item.nodeid} takes a single argument, the limit",
+            pytrace=False,
+        )
+    if marker.args and "n" in marker.kwargs:
+        pytest.fail(
+            f"max_queries on {item.nodeid} got the limit twice, positionally and as n=",
+            pytrace=False,
+        )
+    limit = marker.args[0] if marker.args else marker.kwargs.get("n")
+    if isinstance(limit, bool) or not isinstance(limit, int) or limit < 0:
+        pytest.fail(
+            f"max_queries on {item.nodeid} needs a non-negative integer "
+            f"as its limit, got {limit!r}",
+            pytrace=False,
+        )
+    return limit
+
+
+def _budget_report(
+    item: pytest.Item, collector: QueryCollector, limit: int
+) -> str | None:
+    """The budget failure text, or None; a bug here never fails the test."""
+    try:
+        if len(collector.statements) <= limit:
+            return None
+        return format_budget(limit, collector.statements, cwd=item.config.rootpath)
+    except Exception as exc:  # noqa: BLE001  # a plugin bug must never break the user suite
+        record_error(exc)
+        return None
+
+
+def _record(item: pytest.Item, collector: QueryCollector, settings: _Settings) -> None:
+    """Remember the test's statement count for the session report."""
+    counts = item.config.stash.get(_report, None)
+    # After an internal error the listeners stop recording, so a count would be wrong.
+    if counts is None or settings.mode == "off" or internal_error() is not None:
+        return
+    counts.append((item.nodeid, len(collector.statements)))
+
+
+def _write_report(
+    terminalreporter: pytest.TerminalReporter, settings: _Settings
+) -> None:
+    terminalreporter.write_sep("=", "queryguard report")
+    # dsession is registered only on the xdist controller when it distributes
+    # tests; the counts live in the workers, so the report would be empty.
+    if terminalreporter.config.pluginmanager.has_plugin("dsession"):
+        terminalreporter.write_line(
+            "queryguard report: not available with pytest-xdist"
+        )
+        return
+    if settings.mode == "off":
+        terminalreporter.write_line(
+            "queryguard report: measurement is off (queryguard_mode = off)"
+        )
+        return
+    if internal_error() is not None:
+        terminalreporter.write_line(
+            "queryguard report: incomplete, measurement stopped after an internal error"
+        )
+    counts = terminalreporter.config.stash.get(_report, [])
+    if not counts:
+        terminalreporter.write_line("no tests were measured")
+        return
+    # sorted() is stable, so tests with equal counts keep the order they ran in.
+    heaviest = sorted(counts, key=lambda entry: entry[1], reverse=True)
+    for nodeid, count in heaviest[:_REPORT_SIZE]:
+        terminalreporter.write_line(f"{count:>5}  {nodeid}")
+
+
+def pytest_terminal_summary(terminalreporter: pytest.TerminalReporter) -> None:
+    config = terminalreporter.config
+    settings = config.stash.get(_settings, None)
+    if settings is None or not config.getoption("queryguard_report"):
+        return
+    try:
+        _write_report(terminalreporter, settings)
+    except Exception as exc:  # noqa: BLE001  # a plugin bug must never break the user suite
+        record_error(exc)
+    _report_internal_error(config)
+
+
+@pytest.hookimpl(wrapper=True)
+def pytest_runtest_call(item: pytest.Item) -> Generator[None, Any, Any]:
+    settings = item.config.stash.get(_settings, None)
+    if settings is None:
+        # pytest_configure did not complete (for example it raised a UsageError
+        # for a bad option), so there is nothing to measure with.
+        return (yield)
+    limit = _budget(item)
+    if settings.mode == "off" and limit is None:
+        return (yield)
+    if limit is not None:
+        install()  # idempotent; budgets also work when the mode is off
+    collector = QueryCollector()
+    push(collector)
+    try:
+        result = yield
+    finally:
+        pop(collector)
+        _report_internal_error(item.config)
+        _record(item, collector, settings)  # failing tests are reported too
+    parts: list[str] = []
+    if limit is not None:
+        budget = _budget_report(item, collector, limit)
+        if budget is not None:
+            parts.append(budget)
+    message = None
+    if settings.mode != "off":
+        message = _analyse(item, collector, settings)
+    _report_internal_error(item.config)
+    if message is not None and (parts or settings.mode == "fail"):
+        parts.append(message)  # one failure: budget first, N+1 after it
+    elif message is not None:
+        path, line, _ = item.reportinfo()
+        std_warnings.warn_explicit(
+            QueryGuardWarning(message),
+            QueryGuardWarning,
+            str(path),
+            (line or 0) + 1,
+        )
+    if parts:
+        pytest.fail("\n\n".join(parts), pytrace=False)
+    return result

pytest_queryguard/warnings.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+
+class QueryGuardWarning(UserWarning):
+    """An N+1 query pattern found while running in ``warn`` mode."""
+
+
+class QueryGuardInternalWarning(UserWarning):
+    """A bug inside pytest-queryguard; the plugin stops measuring."""

pytest_queryguard-0.1.0.dist-info/METADATA

@@ -0,0 +1,451 @@
+Metadata-Version: 2.5
+Name: pytest-queryguard
+Version: 0.1.0
+Summary: Detect N+1 queries and enforce query budgets for SQLAlchemy 2 in pytest.
+Project-URL: Homepage, https://github.com/MaciejZiel/pytest-queryguard
+Project-URL: Repository, https://github.com/MaciejZiel/pytest-queryguard
+Project-URL: Changelog, https://github.com/MaciejZiel/pytest-queryguard/blob/main/CHANGELOG.md
+Project-URL: Issues, https://github.com/MaciejZiel/pytest-queryguard/issues
+Author: Maciej Zieliński
+License-Expression: MIT
+License-File: LICENSE
+Keywords: fastapi,n+1,pytest,sqlalchemy,testing
+Classifier: Development Status :: 4 - Beta
+Classifier: Framework :: Pytest
+Classifier: Intended Audience :: Developers
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Topic :: Software Development :: Testing
+Classifier: Typing :: Typed
+Requires-Python: >=3.10
+Requires-Dist: pytest>=8
+Requires-Dist: sqlalchemy>=2.0
+Description-Content-Type: text/markdown
+
+# pytest-queryguard
+
+[![CI](https://github.com/MaciejZiel/pytest-queryguard/actions/workflows/ci.yml/badge.svg)](https://github.com/MaciejZiel/pytest-queryguard/actions/workflows/ci.yml)
+
+A pytest plugin that fails tests with N+1 queries and enforces query budgets
+for SQLAlchemy 2, in sync, async and FastAPI code.
+
+## Why
+
+ORM code that looks correct can send one query per row instead of one query
+per page:
+
+```python
+orders = session.scalars(select(Order)).all()  # 1 query
+for order in orders:
+    print(order.customer.name)  # +1 query per order
+```
+
+The tests pass. A test database holds a handful of rows, so 5 extra queries
+cost nothing. In production the same endpoint loads 500 orders, sends 501
+queries and gets slow. Nothing in the test run warned you.
+
+pytest-queryguard watches the SQL each test sends. When it sees the N+1
+pattern, it fails the test and points to the line that caused it, even if
+the test database has only a few rows.
+
+## Before and after
+
+The code below is a small example project. All output in this section is
+copied from a real run.
+
+`app/orders.py`:
+
+```python
+from sqlalchemy import select
+from sqlalchemy.orm import Session
+
+from app.models import Order
+
+
+def list_orders(session: Session) -> list[dict[str, object]]:
+    orders = session.scalars(select(Order)).all()
+    return [{"id": order.id, "customer": order.customer.name} for order in orders]
+```
+
+`app/customers.py`:
+
+```python
+from sqlalchemy.orm import Session
+
+from app.models import Customer
+
+
+def fetch_customers(session: Session, ids: list[int]) -> list[Customer]:
+    customers = []
+    for customer_id in ids:
+        customers.append(session.get(Customer, customer_id))
+    return customers
+
+
+def rename_customers(session: Session, names: dict[int, str]) -> None:
+    for customer_id, name in names.items():
+        session.get(Customer, customer_id).name = name
+        session.flush()
+```
+
+`tests/test_orders.py` (the `session` fixture in `tests/conftest.py`
+creates 50 customers with one order each; `tests/__init__.py` exists, so
+pytest puts the project root on `sys.path` and `app` can be imported):
+
+```python
+import pytest
+
+from app.customers import fetch_customers, rename_customers
+from app.orders import list_orders
+
+
+def test_list_orders(session):
+    assert len(list_orders(session)) == 50
+
+
+def test_fetch_customers(session):
+    assert len(fetch_customers(session, list(range(1, 51)))) == 50
+
+
+@pytest.mark.max_queries(2)
+def test_rename_customers(session):
+    rename_customers(session, {1: "Ada", 2: "Grace"})
+```
+
+All three tests pass without the plugin. With it installed,
+`pytest --queryguard-report` gives the output below. It comes from Python
+3.13 with SQLAlchemy 2.1; Python 3.12 prints the same. Older versions print
+slightly different SQL or function names, for example `in <listcomp>` on
+Python 3.10 and 3.11, or column aliases in the SQL on SQLAlchemy 2.0.
+
+```
+tests/test_orders.py FFF                                                 [100%]
+
+=================================== FAILURES ===================================
+_______________________________ test_list_orders _______________________________
+queryguard: N+1 detected
+
+  Order.customer was lazy-loaded 50 times
+  first at app/orders.py:9 in list_orders
+  sql: SELECT customer.id, customer.name FROM customer WHERE customer.id = ?
+  fix: load it with the parent query, e.g.
+       select(Order).options(selectinload(Order.customer))
+_____________________________ test_fetch_customers _____________________________
+queryguard: N+1 detected
+
+  the same statement ran 50 times
+  first at app/customers.py:9 in fetch_customers
+  SELECT customer.id, customer.name FROM customer WHERE customer.id = ?
+  fix: fetch all rows in one query (WHERE id IN (...)) instead of in a loop
+____________________________ test_rename_customers _____________________________
+queryguard: expected at most 2 queries, got 4
+
+  1. SELECT customer.id, customer.name FROM customer WHERE customer.id = ?  (app/customers.py:15 in rename_customers)
+  2. UPDATE customer SET name=? WHERE customer.id = ?  (app/customers.py:16 in rename_customers)
+  3. SELECT customer.id, customer.name FROM customer WHERE customer.id = ?  (app/customers.py:15 in rename_customers)
+  4. UPDATE customer SET name=? WHERE customer.id = ?  (app/customers.py:16 in rename_customers)
+============================== queryguard report ===============================
+   51  tests/test_orders.py::test_list_orders
+   50  tests/test_orders.py::test_fetch_customers
+    4  tests/test_orders.py::test_rename_customers
+=========================== short test summary info ============================
+FAILED tests/test_orders.py::test_list_orders - Failed: queryguard: N+1 detected
+FAILED tests/test_orders.py::test_fetch_customers - Failed: queryguard: N+1 d...
+FAILED tests/test_orders.py::test_rename_customers - Failed: queryguard: expe...
+============================== 3 failed in 0.07s ===============================
+```
+
+The fixed code loads related rows with the parent query and fetches many
+rows with one `IN` query.
+
+`app/orders.py`, fixed:
+
+```python
+from sqlalchemy import select
+from sqlalchemy.orm import Session, selectinload
+
+from app.models import Order
+
+
+def list_orders(session: Session) -> list[dict[str, object]]:
+    orders = session.scalars(select(Order).options(selectinload(Order.customer))).all()
+    return [{"id": order.id, "customer": order.customer.name} for order in orders]
+```
+
+`app/customers.py`, fixed:
+
+```python
+from collections.abc import Sequence
+
+from sqlalchemy import select
+from sqlalchemy.orm import Session
+
+from app.models import Customer
+
+
+def fetch_customers(session: Session, ids: list[int]) -> Sequence[Customer]:
+    return session.scalars(select(Customer).where(Customer.id.in_(ids))).all()
+
+
+def rename_customers(session: Session, names: dict[int, str]) -> None:
+    customers = session.scalars(select(Customer).where(Customer.id.in_(names)))
+    for customer in customers:
+        customer.name = names[customer.id]
+    session.flush()
+```
+
+The same tests, unchanged:
+
+```
+tests/test_orders.py ...                                                 [100%]
+
+============================== queryguard report ===============================
+    2  tests/test_orders.py::test_list_orders
+    2  tests/test_orders.py::test_rename_customers
+    1  tests/test_orders.py::test_fetch_customers
+============================== 3 passed in 0.04s ===============================
+```
+
+## Install
+
+```bash
+pip install pytest-queryguard
+```
+
+It needs Python 3.10 or newer, pytest 8 or newer and SQLAlchemy 2.0 or
+newer. The plugin registers itself with pytest. There is nothing to import
+or configure.
+
+## Usage
+
+### Automatic N+1 detection
+
+Every test is checked. A test fails when, inside the test function:
+
+- one relationship (for example `Order.customer`) is lazy-loaded for 2 or
+  more different parent objects, or
+- one `SELECT`, with its parameters and literals removed, runs 5 or more
+  times.
+
+Both thresholds can be changed (see [Configuration](#configuration)).
+
+### Query budgets with the `max_queries` marker
+
+```python
+import pytest
+
+
+@pytest.mark.max_queries(3)
+def test_order_page(client):
+    client.get("/orders/1")
+
+
+@pytest.mark.max_queries(n=3)  # the same, with a keyword
+def test_order_page_again(client):
+    client.get("/orders/1")
+```
+
+The test fails when its body runs more SQL statements than the limit. The
+failure lists every statement with the line that issued it. Budgets count
+all statements, writes included, and they always fail, whatever
+`queryguard_mode` is set to. If the test already fails on its own, an
+exceeded budget is not reported, because the original failure wins.
+
+### Query budgets with the `max_queries()` context manager
+
+```python
+from pytest_queryguard import max_queries
+
+
+def test_checkout(service, cart):
+    with max_queries(2):
+        service.checkout(cart)
+```
+
+Only the statements inside the `with` block count. Above the limit it raises
+`QueryBudgetExceeded`, a subclass of `AssertionError`. The context manager
+does not depend on pytest, so it also works in a script or a benchmark. It
+checks the budget only; N+1 detection is done by the plugin around the
+whole test.
+
+### Opting out
+
+```python
+@pytest.mark.allow_nplusone
+def test_legacy_report(session): ...
+```
+
+The test is not checked for N+1. A `max_queries` marker on the same test is
+still enforced.
+
+### Async and FastAPI
+
+Nothing changes for async code. `AsyncSession` and `create_async_engine` run
+on a sync engine underneath, which the plugin already watches. On SQLAlchemy
+2.1 and newer, async support needs `pip install "sqlalchemy[asyncio]"`,
+which provides greenlet. The failure
+points to the `await` in your code, not to SQLAlchemy internals:
+
+```python
+# app/orders.py; the models use SQLAlchemy's AsyncAttrs mixin.
+async def customer_names(session: AsyncSession) -> list[str]:
+    orders = (await session.scalars(select(Order))).all()
+    return [(await order.awaitable_attrs.customer).name for order in orders]
+```
+
+```python
+# tests/test_async.py; main() creates the engine and data, then calls
+# customer_names(session).
+def test_customer_names():
+    asyncio.run(main())
+```
+
+This fails with `Order.customer was lazy-loaded 3 times`, first at the line
+with the `await`. pytest cannot run `async def` tests by itself. With
+pytest-asyncio or anyio you can write the test as `async def` instead of
+calling `asyncio.run`.
+
+FastAPI's `TestClient` runs the application in another thread. Queries from
+any thread are counted, so an N+1 inside an endpoint fails the test that
+called it:
+
+```python
+def test_orders_endpoint(client):
+    assert client.get("/orders").status_code == 200
+```
+
+### Session report
+
+`pytest --queryguard-report` prints, at the end of the run, the 10 tests
+that ran the most queries, failing tests included. See the output in
+[Before and after](#before-and-after).
+
+## Configuration
+
+| Option | Values | Default | Meaning |
+|---|---|---|---|
+| `queryguard_mode` | `fail`, `warn`, `off` | `fail` | `fail` fails a test with N+1. `warn` reports it as a `QueryGuardWarning` and lets the test pass. `off` turns detection off; tests without a budget are not measured. |
+| `queryguard_relationship_threshold` | integer ≥ 2 | `2` | Number of different parent objects for which one relationship must be lazy-loaded to count as N+1. |
+| `queryguard_repeat_threshold` | integer ≥ 2 | `5` | Number of times one normalized `SELECT` must run to count as N+1. |
+
+Any invalid value stops pytest with a usage error.
+
+In `pyproject.toml`:
+
+```toml
+[tool.pytest.ini_options]
+queryguard_mode = "warn"
+queryguard_relationship_threshold = 3
+queryguard_repeat_threshold = 10
+```
+
+With pytest 9 you can use the native TOML table instead:
+
+```toml
+[tool.pytest]
+queryguard_mode = "warn"
+queryguard_relationship_threshold = 3
+queryguard_repeat_threshold = 10
+```
+
+`pytest.ini` and `setup.cfg` work as usual.
+
+The warning classes, `pytest_queryguard.QueryGuardWarning` for N+1 in
+`warn` mode and `pytest_queryguard.QueryGuardInternalWarning` for an
+internal error, work with pytest's usual warning filters. For example, to
+turn N+1 warnings back into failures for one job:
+
+```toml
+[tool.pytest.ini_options]
+queryguard_mode = "warn"
+filterwarnings = ["error::pytest_queryguard.QueryGuardWarning"]
+```
+
+Command-line options:
+
+| Option | Meaning |
+|---|---|
+| `--queryguard-mode=fail\|warn\|off` | Overrides `queryguard_mode` for this run. |
+| `--queryguard-report` | Prints the 10 tests with the most queries at the end of the session. |
+
+## How it works
+
+1. When pytest starts (unless `queryguard_mode = off`; budgets still install
+   the listeners when a test needs them), the plugin adds two class-level
+   SQLAlchemy listeners:
+   `do_orm_execute` on `Session` sees relationship lazy loads, and
+   `before_cursor_execute` on `Engine` sees every SQL statement.
+2. Measurement covers only the test function body, so queries in fixture
+   setup and teardown are not counted.
+3. Each statement is stored with the first stack frame outside SQLAlchemy,
+   pytest, the standard library, `site-packages` and the plugin itself. For
+   async sessions the search continues in the parent greenlets, where the
+   awaiting code runs.
+4. After the test, the recorded events are checked against the budget and
+   the two N+1 rules. A bug inside the plugin disables it with one warning;
+   it never fails your tests.
+
+## Limitations
+
+- SQLAlchemy only. Raw DB-API connections and other ORMs are not seen.
+- Everything in `site-packages` counts as library code. If your application
+  is installed non-editable (for example `pip install .` in a Docker CI
+  image), its frames are skipped too, and call sites point to the test file
+  instead of the application line.
+- One measurement is active per process at a time, so tests in one process
+  must run sequentially. Under pytest-xdist each worker measures its own
+  tests, so detection and budgets work. The report needs a single process,
+  so under pytest-xdist it prints
+  `queryguard report: not available with pytest-xdist`.
+- Tests that use `pytester` to run pytest in-process are also measured by the
+  outer test, which sees the inner run's queries. Mark such tests with
+  `allow_nplusone` or use `pytester.runpytest_subprocess()`.
+- unittest `setUp` and `tearDown` run inside the test call, so their queries
+  are counted, unlike pytest fixtures.
+- A non-strict `xfail` test that has an N+1 is reported as xfailed rather
+  than XPASS, because the N+1 failure looks like the expected failure.
+- `INSERT`, `UPDATE` and `DELETE` are never reported as N+1. At cursor level
+  a unit-of-work flush cannot be told apart from an application write, so
+  all writes are left out. Use a budget to catch write loops.
+
+## Comparison with nplusone
+
+[nplusone](https://github.com/jmcarp/nplusone) solved the same problem
+earlier and also covers Django and Peewee. Its last release, 1.0.0, came
+out in 2018, and it does not support SQLAlchemy 2. pytest-queryguard targets
+SQLAlchemy 2 only, including async sessions, and adds query budgets and a
+fix hint in the failure message. It does not detect unused eager loads,
+which nplusone does.
+
+## Roadmap
+
+- 0.2: missing index detection on PostgreSQL 16 and newer, by running each
+  statement through `EXPLAIN (GENERIC_PLAN)` with `enable_seqscan=off`.
+- Later: Django ORM support.
+
+## Development
+
+The project uses [uv](https://docs.astral.sh/uv/).
+
+```bash
+uv sync
+uv run coverage run -m pytest
+uv run coverage report
+uv run ruff check .
+uv run ruff format --check .
+uv run mypy src
+```
+
+Coverage is measured with `coverage run` rather than `pytest --cov`, because
+`--cov` starts after pytest has imported the plugin and misses its
+import-time lines.
+
+## License
+
+MIT, see the LICENSE file.

pytest_queryguard-0.1.0.dist-info/RECORD

@@ -0,0 +1,16 @@
+pytest_queryguard/__init__.py,sha256=7yzv4shBD-9lRo2AfDrIqEQjI72B7X194M1YlyjYOE4,314
+pytest_queryguard/api.py,sha256=nw4XNTDGGg7IAOEv5WX4C5R0o_AQclD1tDX9QUKcGh8,1143
+pytest_queryguard/callsite.py,sha256=uClQ4Elogn882OwE-SkcZ3dzcrXVWS-ARYoA6wHLynM,3182
+pytest_queryguard/collector.py,sha256=O44eayC8BJzTW8fPbD4nmFNg9KgonM8qyYFv3fgjlMc,1102
+pytest_queryguard/detectors.py,sha256=g0qqb63jgcNTvzZENn7OC4idkHPWA5UFuLDpuuYCNe8,3379
+pytest_queryguard/events.py,sha256=9FJsT2MxD_25DgVO0QRxA3Bakd0VmJiXLUg35h9ehTE,450
+pytest_queryguard/listeners.py,sha256=y6rEFd_ZZPVXf7WxKbnIOkmzzucL0zJlRo0YbDEQfK0,3259
+pytest_queryguard/messages.py,sha256=730nWP83-3ruM7UEjod3qIgeBjGjgSxJPYLxl91Hnow,2514
+pytest_queryguard/plugin.py,sha256=Xq9lUevUAsAmz81CWElWlycCkD58iR8Jr704odNvGxE,10333
+pytest_queryguard/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+pytest_queryguard/warnings.py,sha256=eKFg4bgUygIfNdX4pR4zsWqFFCvGojbb1_1nGbZM95I,262
+pytest_queryguard-0.1.0.dist-info/METADATA,sha256=DA7QHsFdjFHOmQ2P4yv8uIwF9ivk9f0AbLHcFzfDHYg,16074
+pytest_queryguard-0.1.0.dist-info/WHEEL,sha256=W3fkpkm7-wf9vBI5Z-7s0eWkeM-spu78I8Neb98DeEg,87
+pytest_queryguard-0.1.0.dist-info/entry_points.txt,sha256=jjOwpKjkGqp0LohN0BgEaHpnSmZLIZbuG-mN5GHiPFU,49
+pytest_queryguard-0.1.0.dist-info/licenses/LICENSE,sha256=7oE-K45xMCv_h5fb5eRLXTaVfnXwmd8CF03ijwDkGLU,1074
+pytest_queryguard-0.1.0.dist-info/RECORD,,

pytest_queryguard-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.32.4
+Root-Is-Purelib: true
+Tag: py3-none-any

pytest_queryguard-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[pytest11]
+queryguard = pytest_queryguard.plugin

pytest_queryguard-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Maciej Zieliński
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.