pytest-querycount 0.3.0__py3-none-any.whl

pytest_querycount/normalize.py

@@ -0,0 +1,97 @@
+"""Reduce a SQL statement to a *fingerprint*.
+
+Two statements share a fingerprint when they differ only in their literal
+values. That is what lets us say "this same query ran 47 times" instead of
+"47 queries ran", which is the difference between a number and a diagnosis.
+
+    >>> fingerprint("SELECT * FROM users WHERE id = 42")
+    'select * from users where id = ?'
+    >>> fingerprint("SELECT * FROM users WHERE id = 43")
+    'select * from users where id = ?'
+"""
+
+from __future__ import annotations
+
+import re
+from functools import lru_cache
+
+_PLACEHOLDER = "?"
+
+# Comments go first: they can contain anything, including quotes that would
+# otherwise unbalance the string matching below.
+_BLOCK_COMMENT = re.compile(r"/\*.*?\*/", re.DOTALL)
+_LINE_COMMENT = re.compile(r"--[^\n]*")
+
+# Postgres dollar quoting ($$body$$ or $tag$body$tag$) must be handled before
+# numeric placeholders, or "$1" inside a body would be mangled.
+_DOLLAR_QUOTED = re.compile(r"\$(\w*)\$.*?\$\1\$", re.DOTALL)
+
+# Single-quoted strings, with '' as the escape for an embedded quote.
+_STRING = re.compile(r"'(?:[^']|'')*'")
+
+# Bind parameters in every style the DBAPIs use. Named styles come before the
+# bare ones so ":name" is not left as a stray colon.
+_PARAM_PYFORMAT = re.compile(r"%\(\w+\)s")
+_PARAM_NAMED = re.compile(r"(?<![:\w]):\w+")  # ":name" but never "::cast"
+_PARAM_NUMERIC = re.compile(r"\$\d+")
+_PARAM_FORMAT = re.compile(r"%s")
+
+# Numbers, once strings are gone so we never touch digits inside a literal.
+# The word boundary keeps identifiers like "table1" intact.
+_NUMBER = re.compile(r"\b\d+\.?\d*(?:[eE][+-]?\d+)?\b")
+
+# Collapse variable-length lists so an IN clause of 3 ids fingerprints the same
+# as one of 300 -- otherwise every batch size looks like a different query.
+_IN_LIST = re.compile(r"\bin\s*\(\s*\?(?:\s*,\s*\?)*\s*\)")
+_VALUES_TUPLE = re.compile(r"\(\s*\?(?:\s*,\s*\?)*\s*\)")
+_VALUES_CLAUSE = re.compile(r"\bvalues\s*((?:\(\s*\?(?:\s*,\s*\?)*\s*\)\s*,?\s*)+)")
+
+_WHITESPACE = re.compile(r"\s+")
+
+
+def _collapse_values(match: re.Match[str]) -> str:
+    """Rewrite a multi-row VALUES clause as a single tuple."""
+    tuples = _VALUES_TUPLE.findall(match.group(1))
+    if len(tuples) <= 1:
+        return match.group(0)
+    first = _VALUES_TUPLE.search(match.group(1))
+    assert first is not None
+    return f"values {first.group(0)}"
+
+
+@lru_cache(maxsize=2048)
+def fingerprint(statement: str) -> str:
+    """Return a stable, literal-free form of ``statement``.
+
+    Cached, because the whole point is that the same statements recur and this
+    runs on every query the suite emits.
+    """
+    sql = _BLOCK_COMMENT.sub(" ", statement)
+    sql = _LINE_COMMENT.sub(" ", sql)
+    sql = _DOLLAR_QUOTED.sub(_PLACEHOLDER, sql)
+    sql = _STRING.sub(_PLACEHOLDER, sql)
+    sql = _PARAM_PYFORMAT.sub(_PLACEHOLDER, sql)
+    sql = _PARAM_NAMED.sub(_PLACEHOLDER, sql)
+    sql = _PARAM_NUMERIC.sub(_PLACEHOLDER, sql)
+    sql = _PARAM_FORMAT.sub(_PLACEHOLDER, sql)
+    sql = _NUMBER.sub(_PLACEHOLDER, sql)
+
+    sql = sql.lower()
+    sql = _IN_LIST.sub(f"in ({_PLACEHOLDER})", sql)
+    sql = _VALUES_CLAUSE.sub(_collapse_values, sql)
+
+    sql = _WHITESPACE.sub(" ", sql).strip()
+    return sql.rstrip(";").strip()
+
+
+def statement_kind(statement: str) -> str:
+    """The leading keyword of ``statement``, lowercased.
+
+    ``"select"``, ``"insert"``, ``"savepoint"``, ... or ``"other"`` when the
+    statement does not start with a word (an empty string, say).
+    """
+    for token in _WHITESPACE.split(_LINE_COMMENT.sub(" ", statement).strip()):
+        cleaned = token.strip("(").lower()
+        if cleaned.isalpha():
+            return cleaned
+    return "other"

pytest_querycount/plugin.py

@@ -0,0 +1,284 @@
+"""pytest integration: options, markers, the fixture, and the summary."""
+
+from __future__ import annotations
+
+from collections.abc import Generator, Sequence
+from typing import Any
+
+import pytest
+
+from pytest_querycount import backends, checks, report
+from pytest_querycount.checks import DEFAULT_DUPLICATE_THRESHOLD, DEFAULT_KINDS, Budget
+from pytest_querycount.recorder import Recorder
+
+RECORDER_KEY = pytest.StashKey[Recorder]()
+_STATS: dict[str, report.TestStats] = {}
+
+
+# -- configuration ---------------------------------------------------------
+
+
+def pytest_addoption(parser: pytest.Parser) -> None:
+    group = parser.getgroup("querycount", "SQL query budgets")
+    group.addoption(
+        "--querycount-report",
+        action="store_true",
+        default=False,
+        help="Print a table of the tests that ran the most SQL queries.",
+    )
+    group.addoption(
+        "--querycount-top",
+        type=int,
+        default=10,
+        metavar="N",
+        help="How many tests the report lists (default: 10).",
+    )
+    group.addoption(
+        "--querycount-max",
+        type=int,
+        default=None,
+        metavar="N",
+        help="Apply a query budget of N to every test that has no explicit one.",
+    )
+    group.addoption(
+        "--querycount-no-seq-scan",
+        action="store_true",
+        default=False,
+        help=(
+            "Apply the missing-index check to every test. Costs one EXPLAIN per "
+            "SELECT, so prefer the marker once you know where to look."
+        ),
+    )
+    parser.addini(
+        "querycount_max",
+        help="Default query budget for tests without an explicit one.",
+        default="",
+    )
+    parser.addini(
+        "querycount_duplicate_threshold",
+        help=(
+            "How many repeats of one query shape count as N+1 "
+            f"(default: {DEFAULT_DUPLICATE_THRESHOLD})."
+        ),
+        default="",
+    )
+
+
+def pytest_configure(config: pytest.Config) -> None:
+    config.addinivalue_line(
+        "markers",
+        "max_queries(n): fail the test if it runs more than n SQL queries.",
+    )
+    config.addinivalue_line(
+        "markers",
+        "no_n_plus_one(threshold=2, kinds=('select',)): fail the test if one "
+        "query shape is executed repeatedly.",
+    )
+    config.addinivalue_line(
+        "markers",
+        "no_seq_scan(ignore=()): fail the test if a query scans a table "
+        "sequentially because no index could serve its filter. PostgreSQL only.",
+    )
+    _STATS.clear()
+    backends.install_all()
+
+
+def pytest_collection_modifyitems(
+    config: pytest.Config,
+    items: list[pytest.Item],
+) -> None:
+    """Reject malformed markers while collecting.
+
+    Validating here rather than mid-test means a typo is a clean usage error
+    instead of an exception raised from inside a hook wrapper.
+    """
+    for item in items:
+        marker = item.get_closest_marker("max_queries")
+        if marker is not None:
+            _max_queries_from_marker(item, marker)
+
+
+# -- per-test recording ----------------------------------------------------
+
+
+@pytest.hookimpl(wrapper=True)
+def pytest_runtest_call(item: pytest.Item) -> Generator[None, object, object]:
+    """Record the call phase, then hold the test to its budget.
+
+    Deliberately wraps only the call phase: fixture setup and teardown run
+    outside it, so schema creation and seeding never consume a budget.
+    """
+    __tracebackhide__ = True
+
+    budget = _budget_for(item)
+    recorder = Recorder(explain=budget.needs_explain)
+    item.stash[RECORDER_KEY] = recorder
+    backends.push(recorder)
+    try:
+        result = yield
+    except BaseException:
+        # The test itself failed. Record what we saw for the report, but let its
+        # exception through untouched: a budget complaint on top of a real
+        # failure buries the thing that actually needs fixing.
+        _record_stats(item, recorder)
+        raise
+    finally:
+        backends.pop(recorder)
+
+    _record_stats(item, recorder)
+    checks.enforce(budget, recorder, label=item.name)
+    return result
+
+
+def _record_stats(item: pytest.Item, recorder: Recorder) -> None:
+    if not recorder.count:
+        return
+    duplicates = recorder.duplicates(threshold=2, kinds=DEFAULT_KINDS)
+    _STATS[item.nodeid] = report.TestStats(
+        nodeid=item.nodeid,
+        count=recorder.count,
+        duration=recorder.total_duration,
+        worst_duplicate=duplicates[0].count if duplicates else 0,
+        worst_sql=duplicates[0].sample.short_sql() if duplicates else None,
+    )
+
+
+def _budget_for(item: pytest.Item) -> Budget:
+    """Resolve the budget: marker first, then --querycount-max, then ini."""
+    budget = Budget(
+        max_queries=_default_max(item.config),
+        duplicate_threshold=_default_threshold(item.config),
+        no_seq_scan=bool(item.config.getoption("querycount_no_seq_scan")),
+    )
+
+    marker = item.get_closest_marker("max_queries")
+    if marker is not None:
+        budget.max_queries = _max_queries_from_marker(item, marker)
+
+    marker = item.get_closest_marker("no_seq_scan")
+    if marker is not None:
+        budget.no_seq_scan = True
+        budget.ignore_tables = tuple(marker.kwargs.get("ignore", ()))
+
+    marker = item.get_closest_marker("no_n_plus_one")
+    if marker is not None:
+        budget.no_n_plus_one = True
+        budget.duplicate_threshold = int(marker.kwargs.get("threshold", budget.duplicate_threshold))
+        if "kinds" in marker.kwargs:
+            kinds = marker.kwargs["kinds"]
+            budget.kinds = None if kinds is None else tuple(kinds)
+
+    return budget
+
+
+def _max_queries_from_marker(item: pytest.Item, marker: pytest.Mark) -> int:
+    if marker.args:
+        value = marker.args[0]
+    elif "n" in marker.kwargs:
+        value = marker.kwargs["n"]
+    else:
+        raise pytest.UsageError(
+            f"{item.nodeid}: @pytest.mark.max_queries needs a number, "
+            "e.g. @pytest.mark.max_queries(3)"
+        )
+    try:
+        return int(value)
+    except (TypeError, ValueError):
+        raise pytest.UsageError(
+            f"{item.nodeid}: @pytest.mark.max_queries({value!r}) is not a number"
+        ) from None
+
+
+def _default_max(config: pytest.Config) -> int | None:
+    from_cli = config.getoption("querycount_max")
+    if from_cli is not None:
+        return int(from_cli)
+    from_ini = str(config.getini("querycount_max") or "").strip()
+    return int(from_ini) if from_ini else None
+
+
+def _default_threshold(config: pytest.Config) -> int:
+    from_ini = str(config.getini("querycount_duplicate_threshold") or "").strip()
+    return int(from_ini) if from_ini else DEFAULT_DUPLICATE_THRESHOLD
+
+
+# -- the fixture -----------------------------------------------------------
+
+
+class _Block:
+    """Context manager returned by ``querycount(...)``.
+
+    ``__enter__`` yields the :class:`Recorder`, so the same object carries the
+    live count during the block and the evidence afterwards.
+    """
+
+    def __init__(self, budget: Budget, label: str) -> None:
+        self._budget = budget
+        self._label = label
+        self.recorder = Recorder(explain=budget.needs_explain)
+
+    def __enter__(self) -> Recorder:
+        backends.push(self.recorder)
+        return self.recorder
+
+    def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
+        __tracebackhide__ = True
+
+        backends.pop(self.recorder)
+        # An exception from inside the block is the more important news.
+        if exc_type is None:
+            checks.enforce(self._budget, self.recorder, label=self._label)
+
+
+@pytest.fixture
+def querycount(request: pytest.FixtureRequest) -> Any:
+    """Record the queries a block of code runs, and optionally budget them.
+
+        def test_detail(client, querycount):
+            with querycount(max_queries=2) as queries:
+                client.get("/users/1")
+            assert queries.duplicates() == []
+
+    Called with no arguments it only observes, which is the way to find out what
+    a budget should be before you commit to one.
+    """
+
+    def factory(
+        max_queries: int | None = None,
+        no_n_plus_one: bool = False,
+        threshold: int = DEFAULT_DUPLICATE_THRESHOLD,
+        kinds: Sequence[str] | None = DEFAULT_KINDS,
+        no_seq_scan: bool = False,
+        ignore: Sequence[str] = (),
+    ) -> _Block:
+        return _Block(
+            Budget(
+                max_queries=max_queries,
+                no_n_plus_one=no_n_plus_one,
+                duplicate_threshold=threshold,
+                kinds=kinds,
+                no_seq_scan=no_seq_scan,
+                ignore_tables=tuple(ignore),
+            ),
+            label="This querycount block",
+        )
+
+    return factory
+
+
+# -- summary ---------------------------------------------------------------
+
+
+def pytest_terminal_summary(
+    terminalreporter: Any,
+    exitstatus: int,
+    config: pytest.Config,
+) -> None:
+    if not config.getoption("querycount_report"):
+        return
+    lines = report.build(_STATS, top=int(config.getoption("querycount_top")))
+    if not lines:
+        return
+    terminalreporter.write_sep("=", "querycount summary")
+    for line in lines:
+        terminalreporter.write_line(line)

pytest_querycount/recorder.py

@@ -0,0 +1,214 @@
+"""Collecting queries and answering questions about them."""
+
+from __future__ import annotations
+
+import os
+import sys
+from collections import Counter
+from collections.abc import Iterable, Sequence
+from typing import TYPE_CHECKING
+
+from pytest_querycount.records import Duplicate, QueryRecord
+
+if TYPE_CHECKING:
+    from types import FrameType
+
+# Frames belonging to these files are never the interesting caller: they are the
+# ORM, the driver, the test runner, or us.
+_INTERNAL_PATHS = (
+    "/sqlalchemy/",
+    "/pytest_querycount/",
+    "/_pytest/",
+    "/pluggy/",
+    "/asyncio/",
+    "/greenlet/",
+    "/contextlib.py",
+    "/threading.py",
+)
+
+_MAX_STACK_DEPTH = 60
+
+
+def _shorten(path: str) -> str:
+    """Relative to the working directory when that is shorter, absolute otherwise.
+
+    Failure messages are read in a terminal, and an absolute path to a pytest
+    temp directory pushes the useful part off the edge of the screen.
+    """
+    try:
+        relative = os.path.relpath(path)
+    except ValueError:  # pragma: no cover - different drive on Windows
+        return path
+    return path if relative.startswith("..") else relative
+
+
+def caller_location(skip: int = 1) -> str | None:
+    """``path:lineno`` of the closest frame that is not library machinery."""
+    try:
+        frame: FrameType | None = sys._getframe(skip + 1)
+    except ValueError:  # pragma: no cover - stack shallower than `skip`
+        return None
+    depth = 0
+    while frame is not None and depth < _MAX_STACK_DEPTH:
+        filename = frame.f_code.co_filename.replace(os.sep, "/")
+        if not any(part in filename for part in _INTERNAL_PATHS):
+            return f"{_shorten(frame.f_code.co_filename)}:{frame.f_lineno}"
+        frame = frame.f_back
+        depth += 1
+    return None
+
+
+class Recorder:
+    """Accumulates :class:`QueryRecord` objects and reports on them.
+
+    A recorder only sees queries while it is on the active stack; see
+    :mod:`pytest_querycount.backends`.
+    """
+
+    def __init__(self, explain: bool = False) -> None:
+        self.records: list[QueryRecord] = []
+        self.explain = explain
+        """Whether the backend should obtain a query plan for each SELECT.
+
+        Off unless a check needs it: EXPLAIN means a second round trip per
+        query, which is a real cost to impose on a suite that did not ask."""
+
+    def add(self, record: QueryRecord) -> None:
+        self.records.append(record)
+
+    def clear(self) -> None:
+        self.records.clear()
+
+    # -- basic numbers ----------------------------------------------------
+
+    def __len__(self) -> int:
+        return len(self.records)
+
+    @property
+    def count(self) -> int:
+        """How many statements were executed."""
+        return len(self.records)
+
+    @property
+    def total_duration(self) -> float:
+        """Seconds spent in the driver, summed."""
+        return sum(record.duration for record in self.records)
+
+    def of_kind(self, *kinds: str) -> list[QueryRecord]:
+        wanted = {kind.lower() for kind in kinds}
+        return [record for record in self.records if record.kind in wanted]
+
+    # -- duplicate analysis ----------------------------------------------
+
+    @property
+    def fingerprints(self) -> Counter[str]:
+        return Counter(record.fingerprint for record in self.records)
+
+    def duplicates(
+        self,
+        threshold: int = 2,
+        kinds: Sequence[str] | None = ("select",),
+    ) -> list[Duplicate]:
+        """Query shapes that ran at least ``threshold`` times.
+
+        ``kinds`` restricts the search to certain statement types; ``None``
+        considers all of them. The default looks only at SELECTs, because a
+        repeated SELECT is the N+1 signature, whereas repeated SAVEPOINTs and
+        BEGINs are simply how transactions work.
+        """
+        candidates: Iterable[QueryRecord]
+        candidates = self.records if kinds is None else self.of_kind(*kinds)
+
+        counts: Counter[str] = Counter()
+        samples: dict[str, QueryRecord] = {}
+        locations: dict[str, Counter[str]] = {}
+        for record in candidates:
+            counts[record.fingerprint] += 1
+            samples.setdefault(record.fingerprint, record)
+            if record.location:
+                locations.setdefault(record.fingerprint, Counter())[record.location] += 1
+
+        found = [
+            Duplicate(
+                fingerprint=shape,
+                count=hits,
+                sample=samples[shape],
+                locations=locations.get(shape, Counter()),
+            )
+            for shape, hits in counts.items()
+            if hits >= threshold
+        ]
+        found.sort(key=lambda duplicate: duplicate.count, reverse=True)
+        return found
+
+    # -- query plans ------------------------------------------------------
+
+    @property
+    def explained_count(self) -> int:
+        return sum(1 for record in self.records if record.explained)
+
+    def seq_scans(self, ignore: Sequence[str] = ()) -> list[QueryRecord]:
+        """Records whose plan held a sequential scan no index could serve.
+
+        ``ignore`` names tables to pass over -- a small lookup table is read
+        whole on purpose, and no index would improve it.
+        """
+        skipped = {name.lower() for name in ignore}
+        return [
+            record
+            for record in self.records
+            if any(scan.relation.lower() not in skipped for scan in record.seq_scans)
+        ]
+
+    # -- human output -----------------------------------------------------
+
+    def report(self, width: int = 100, limit: int | None = None, collapse: bool = True) -> str:
+        """The recorded queries, in order, with timings and origins.
+
+        With ``collapse`` a run of identical shapes prints as one line -- twelve
+        copies of the same SELECT say nothing that ``12x`` does not, and they
+        push the useful lines off the screen.
+        """
+        if not self.records:
+            return "no queries recorded"
+
+        groups = self._runs() if collapse else [[record] for record in self.records]
+        shown = groups if limit is None else groups[:limit]
+
+        lines = [
+            f"{self.count} quer{'y' if self.count == 1 else 'ies'} "
+            f"in {self.total_duration * 1000:.1f}ms"
+        ]
+        position = 1
+        for group in shown:
+            first = group[0]
+            elapsed = sum(record.duration for record in group) * 1000
+            if len(group) == 1:
+                label = f"{position}."
+                timing = f"[{elapsed:7.2f}ms]"
+            else:
+                label = f"{position}-{position + len(group) - 1}."
+                timing = f"[{elapsed:7.2f}ms] {len(group)}x"
+            marker = " (executemany)" if first.executemany else ""
+            lines.append(f"  {label:>8} {timing}{marker} {first.short_sql(width)}")
+            if first.location:
+                lines.append(f"  {'':>8} {'':>10}   from {first.location}")
+            position += len(group)
+
+        hidden = self.count - sum(len(group) for group in shown)
+        if hidden > 0:
+            lines.append(f"  ... and {hidden} more")
+        return "\n".join(lines)
+
+    def _runs(self) -> list[list[QueryRecord]]:
+        """Consecutive records sharing a fingerprint, grouped."""
+        runs: list[list[QueryRecord]] = []
+        for record in self.records:
+            if runs and runs[-1][0].fingerprint == record.fingerprint:
+                runs[-1].append(record)
+            else:
+                runs.append([record])
+        return runs
+
+    def __repr__(self) -> str:
+        return f"<Recorder {self.count} queries, {self.total_duration * 1000:.1f}ms>"

pytest_querycount/records.py

@@ -0,0 +1,66 @@
+"""The data the recorder collects."""
+
+from __future__ import annotations
+
+from collections import Counter
+from dataclasses import dataclass, field
+
+from pytest_querycount.explain import SeqScan
+
+
+@dataclass(frozen=True)
+class QueryRecord:
+    """One statement handed to the driver."""
+
+    sql: str
+    """The statement as the driver saw it, literals and placeholders intact."""
+
+    fingerprint: str
+    """``sql`` with its literals replaced -- see :func:`~.normalize.fingerprint`."""
+
+    kind: str
+    """Leading keyword, lowercased: ``"select"``, ``"insert"``, ..."""
+
+    duration: float
+    """Seconds spent inside the driver call."""
+
+    executemany: bool
+    """True when one call carried many parameter sets. Counts as a single query,
+    because it is a single round trip -- which is exactly why it is the fix for
+    a loop of inserts."""
+
+    rowcount: int | None = None
+
+    seq_scans: tuple[SeqScan, ...] = ()
+    """Sequential scans the planner kept even with seq scans penalised, i.e. ones
+    no index could serve. Only populated when a check asked for EXPLAIN."""
+
+    explained: bool = False
+    """Whether a plan was actually obtained. Distinguishes "no problems found"
+    from "we never looked", which must not be reported the same way."""
+
+    location: str | None = None
+    """``path:lineno`` of the nearest frame outside the ORM and the test runner,
+    i.e. the line of *your* code that caused this."""
+
+    def short_sql(self, width: int = 100) -> str:
+        collapsed = " ".join(self.sql.split())
+        if len(collapsed) <= width:
+            return collapsed
+        return collapsed[: width - 1] + "…"
+
+
+@dataclass(frozen=True)
+class Duplicate:
+    """A query shape that ran more than once."""
+
+    fingerprint: str
+    count: int
+    sample: QueryRecord
+    locations: Counter[str] = field(default_factory=Counter)
+
+    def describe(self, width: int = 100) -> str:
+        lines = [f"{self.count}x  {self.sample.short_sql(width)}"]
+        for location, hits in self.locations.most_common(3):
+            lines.append(f"      from {location} ({hits}x)")
+        return "\n".join(lines)