pytest-querycount 0.3.0__py3-none-any.whl

pytest_querycount/__init__.py

@@ -0,0 +1,31 @@
+"""Fail your tests when they run too many SQL queries."""
+
+from __future__ import annotations
+
+__version__ = "0.3.0"
+
+from pytest_querycount.errors import (
+    DuplicateQueryError,
+    ExplainUnavailableError,
+    NoBackendError,
+    SeqScanError,
+    TooManyQueriesError,
+)
+from pytest_querycount.explain import SeqScan
+from pytest_querycount.normalize import fingerprint
+from pytest_querycount.recorder import Recorder
+from pytest_querycount.records import Duplicate, QueryRecord
+
+__all__ = [
+    "Duplicate",
+    "DuplicateQueryError",
+    "ExplainUnavailableError",
+    "NoBackendError",
+    "QueryRecord",
+    "Recorder",
+    "SeqScan",
+    "SeqScanError",
+    "TooManyQueriesError",
+    "__version__",
+    "fingerprint",
+]

pytest_querycount/backends/__init__.py

@@ -0,0 +1,59 @@
+"""Instrumentation: turning driver activity into :class:`QueryRecord` objects.
+
+A backend is responsible for one database library. It reports queries to every
+recorder currently on the active stack, which is what makes nesting work: the
+per-test recorder and a ``querycount(...)`` block inside that test both see the
+same queries.
+"""
+
+from __future__ import annotations
+
+import contextlib
+from collections.abc import Sequence
+
+from pytest_querycount.recorder import Recorder
+from pytest_querycount.records import QueryRecord
+
+_active: list[Recorder] = []
+_installed: list[str] = []
+
+
+def push(recorder: Recorder) -> None:
+    """Start sending queries to ``recorder``."""
+    _active.append(recorder)
+
+
+def pop(recorder: Recorder) -> None:
+    """Stop sending queries to ``recorder``."""
+    with contextlib.suppress(ValueError):  # a pop without a push is not fatal
+        _active.remove(recorder)
+
+
+def active() -> Sequence[Recorder]:
+    return _active
+
+
+def emit(record: QueryRecord) -> None:
+    """Hand ``record`` to every listening recorder."""
+    for recorder in _active:
+        recorder.add(record)
+
+
+def installed() -> Sequence[str]:
+    """Names of the backends that were successfully instrumented."""
+    return _installed
+
+
+def install_all() -> Sequence[str]:
+    """Instrument every library we support and that is importable.
+
+    Missing libraries are not an error -- most projects use one ORM, not all of
+    them. But *no* backend at all is an error the moment a check runs, since a
+    query budget that cannot see queries would pass silently forever.
+    """
+    from pytest_querycount.backends import sqlalchemy as sqlalchemy_backend
+
+    for backend in (sqlalchemy_backend,):
+        if backend.install() and backend.NAME not in _installed:
+            _installed.append(backend.NAME)
+    return _installed

pytest_querycount/backends/sqlalchemy.py

@@ -0,0 +1,215 @@
+"""SQLAlchemy 2.x instrumentation.
+
+We listen on the ``Engine`` *class* rather than on individual engines, so every
+engine the suite builds is covered without the project having to register
+anything. Async engines are covered too: ``AsyncEngine`` wraps a sync ``Engine``
+and the cursor events fire on that.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import time
+from typing import Any
+
+from pytest_querycount import backends, explain
+from pytest_querycount.normalize import fingerprint, statement_kind
+from pytest_querycount.recorder import caller_location
+from pytest_querycount.records import QueryRecord
+
+NAME = "sqlalchemy"
+
+# Only PostgreSQL exposes a machine-readable plan we can reason about this way.
+_EXPLAINABLE_DIALECTS = {"postgresql"}
+_SAVEPOINT = "pytest_querycount_explain"
+
+_installed = False
+
+# Fallback timing stack, used when SQLAlchemy hands us no execution context
+# (raw ``exec_driver_sql`` calls, mainly). Durations are informational, so a
+# plain stack is enough and we never let it grow without bound.
+_fallback_starts: list[float] = []
+_CONTEXT_ATTR = "_querycount_start"
+_SCANS_ATTR = "_querycount_scans"
+
+
+def install() -> bool:
+    """Attach the listeners. Returns False when SQLAlchemy is not available."""
+    global _installed
+    if _installed:
+        return True
+    try:
+        from sqlalchemy import event
+        from sqlalchemy.engine import Engine
+    except ImportError:
+        return False
+
+    event.listen(Engine, "before_cursor_execute", _before_cursor_execute)
+    event.listen(Engine, "after_cursor_execute", _after_cursor_execute)
+    _installed = True
+    return True
+
+
+def instrument(engine: Any) -> None:
+    """Attach the listeners to one engine explicitly.
+
+    Only needed for an engine built before the plugin loaded, which in a pytest
+    run should not happen. Kept because "should not happen" is not "cannot".
+    """
+    from sqlalchemy import event
+
+    target = getattr(engine, "sync_engine", engine)
+    event.listen(target, "before_cursor_execute", _before_cursor_execute)
+    event.listen(target, "after_cursor_execute", _after_cursor_execute)
+
+
+def _before_cursor_execute(
+    conn: Any,
+    cursor: Any,
+    statement: str,
+    parameters: Any,
+    context: Any,
+    executemany: bool,
+) -> None:
+    recorders = backends.active()
+    if not recorders:
+        return
+
+    scans: tuple[explain.SeqScan, ...] | None = None
+    if any(recorder.explain for recorder in recorders) and not executemany:
+        scans = _plan_seq_scans(conn, statement, parameters)
+
+    start = time.perf_counter()
+    if context is not None:
+        if scans is not None:
+            # Only reached when EXPLAIN is on, so the context manager's cost here
+            # does not matter the way it does on the per-query timing path.
+            with contextlib.suppress(AttributeError, TypeError):
+                setattr(context, _SCANS_ATTR, scans)
+        try:
+            setattr(context, _CONTEXT_ATTR, start)
+            return
+        except (AttributeError, TypeError):  # pragma: no cover - slotted context
+            pass
+    _fallback_starts.append(start)
+
+
+def _after_cursor_execute(
+    conn: Any,
+    cursor: Any,
+    statement: str,
+    parameters: Any,
+    context: Any,
+    executemany: bool,
+) -> None:
+    if not backends.active():
+        return
+
+    start = _take_start(context)
+    duration = time.perf_counter() - start if start is not None else 0.0
+    scans = _take_scans(context)
+
+    backends.emit(
+        QueryRecord(
+            sql=statement,
+            fingerprint=fingerprint(statement),
+            kind=statement_kind(statement),
+            duration=duration,
+            executemany=bool(executemany),
+            rowcount=_rowcount(cursor),
+            seq_scans=scans or (),
+            explained=scans is not None,
+            location=caller_location(),
+        )
+    )
+
+
+def _take_scans(context: Any) -> tuple[explain.SeqScan, ...] | None:
+    if context is None:
+        return None
+    scans = getattr(context, _SCANS_ATTR, None)
+    if scans is None:
+        return None
+    try:  # noqa: SIM105
+        delattr(context, _SCANS_ATTR)
+    except (AttributeError, TypeError):  # pragma: no cover
+        pass
+    return scans  # type: ignore[no-any-return]
+
+
+def _plan_seq_scans(
+    conn: Any,
+    statement: str,
+    parameters: Any,
+) -> tuple[explain.SeqScan, ...] | None:
+    """Ask PostgreSQL for a plan with sequential scans penalised.
+
+    Returns the filtered sequential scans the planner kept anyway, or None when
+    no plan could be obtained -- a different answer from "found nothing", and
+    the caller must not conflate them.
+
+    Runs on a raw DBAPI cursor rather than through the engine, which keeps the
+    EXPLAIN out of SQLAlchemy's event stream: it is neither counted as one of
+    the test's queries nor able to recurse into this function.
+    """
+    if statement_kind(statement) != "select":
+        # Explaining a write is possible but buys little and risks more.
+        return None
+    if getattr(conn.dialect, "name", None) not in _EXPLAINABLE_DIALECTS:
+        return None
+
+    try:
+        cursor = conn.connection.cursor()
+    except Exception:  # pragma: no cover - pool handed us nothing usable
+        return None
+
+    # The savepoint does double duty: it scopes the enable_seqscan change so we
+    # cannot leak it into the test, and it absorbs a failed EXPLAIN, which would
+    # otherwise abort the transaction the test is running in.
+    guarded = bool(conn.in_transaction())
+    try:
+        if guarded:
+            cursor.execute(f"SAVEPOINT {_SAVEPOINT}")
+        cursor.execute("SET enable_seqscan = off")
+        cursor.execute("EXPLAIN (FORMAT JSON) " + statement, parameters or None)
+        row = cursor.fetchone()
+        return explain.parse(row[0]) if row else ()
+    except Exception:
+        return None
+    finally:
+        try:
+            if guarded:
+                cursor.execute(f"ROLLBACK TO SAVEPOINT {_SAVEPOINT}")
+            else:
+                cursor.execute("RESET enable_seqscan")
+        except Exception:  # pragma: no cover - connection already unusable
+            pass
+        try:  # noqa: SIM105
+            cursor.close()
+        except Exception:  # pragma: no cover
+            pass
+
+
+def _take_start(context: Any) -> float | None:
+    if context is not None:
+        start = getattr(context, _CONTEXT_ATTR, None)
+        if start is not None:
+            # try/except rather than contextlib.suppress: this runs once per
+            # query, and a context manager here is measurable overhead in a
+            # library whose whole job is to sit in that path.
+            try:  # noqa: SIM105
+                delattr(context, _CONTEXT_ATTR)
+            except (AttributeError, TypeError):  # pragma: no cover
+                pass
+            return float(start)
+    if _fallback_starts:
+        return _fallback_starts.pop()
+    return None
+
+
+def _rowcount(cursor: Any) -> int | None:
+    try:
+        count = cursor.rowcount
+    except Exception:  # pragma: no cover - drivers may refuse after DDL
+        return None
+    return count if isinstance(count, int) and count >= 0 else None

pytest_querycount/checks.py

@@ -0,0 +1,166 @@
+"""Turning a recorder plus a budget into a pass or a failure."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+from pytest_querycount import backends
+from pytest_querycount.errors import (
+    DuplicateQueryError,
+    ExplainUnavailableError,
+    NoBackendError,
+    SeqScanError,
+    TooManyQueriesError,
+)
+from pytest_querycount.recorder import Recorder
+from pytest_querycount.records import Duplicate, QueryRecord
+
+DEFAULT_DUPLICATE_THRESHOLD = 2
+DEFAULT_KINDS = ("select",)
+
+
+@dataclass
+class Budget:
+    """What a test is allowed to do to the database."""
+
+    max_queries: int | None = None
+    no_n_plus_one: bool = False
+    duplicate_threshold: int = DEFAULT_DUPLICATE_THRESHOLD
+    kinds: Sequence[str] | None = DEFAULT_KINDS
+    no_seq_scan: bool = False
+    ignore_tables: Sequence[str] = ()
+
+    @property
+    def is_empty(self) -> bool:
+        return self.max_queries is None and not self.no_n_plus_one and not self.no_seq_scan
+
+    @property
+    def needs_explain(self) -> bool:
+        """Whether the backend must pay for a query plan per SELECT."""
+        return self.no_seq_scan
+
+
+def enforce(budget: Budget, recorder: Recorder, label: str = "This test") -> None:
+    """Raise if ``recorder`` violated ``budget``. Cheapest checks first."""
+    # Keep our own frames out of the traceback. The failure message is the
+    # product here; a wall of plugin internals above it is noise.
+    __tracebackhide__ = True
+
+    if budget.is_empty:
+        return
+
+    if not backends.installed():
+        raise NoBackendError(
+            "pytest-querycount has nothing to watch, so this query budget would "
+            "always pass.\n"
+            "Install a supported backend:  pip install 'pytest-querycount[sqlalchemy]'\n"
+            "If your project uses a library we do not support yet, please open an issue."
+        )
+
+    if budget.max_queries is not None and recorder.count > budget.max_queries:
+        raise TooManyQueriesError(_too_many_message(budget, recorder, label))
+
+    if budget.no_n_plus_one:
+        duplicates = recorder.duplicates(
+            threshold=budget.duplicate_threshold,
+            kinds=budget.kinds,
+        )
+        if duplicates:
+            raise DuplicateQueryError(_duplicate_message(duplicates, recorder, label))
+
+    if budget.no_seq_scan:
+        _enforce_no_seq_scan(budget, recorder, label)
+
+
+def _enforce_no_seq_scan(budget: Budget, recorder: Recorder, label: str) -> None:
+    __tracebackhide__ = True
+
+    selects = recorder.of_kind("select")
+    if selects and not recorder.explained_count:
+        # Nothing was explained, so the check saw nothing and would pass no
+        # matter what the queries did. Say so instead.
+        raise ExplainUnavailableError(
+            "no_seq_scan could not obtain a query plan, so it would pass "
+            "regardless of what the queries do.\n"
+            "This check needs PostgreSQL: it reads EXPLAIN (FORMAT JSON), which "
+            "SQLite and MySQL do not provide in a comparable form.\n"
+            "Run these tests against PostgreSQL, or drop the marker."
+        )
+
+    offenders = recorder.seq_scans(ignore=budget.ignore_tables)
+    if offenders:
+        raise SeqScanError(_seq_scan_message(offenders, budget, label))
+
+
+def _seq_scan_message(
+    offenders: Sequence[QueryRecord],
+    budget: Budget,
+    label: str,
+) -> str:
+    skipped = {name.lower() for name in budget.ignore_tables}
+    lines = [
+        f"{label} ran {len(offenders)} quer"
+        f"{'y' if len(offenders) == 1 else 'ies'} that no index could serve.",
+        "",
+        "PostgreSQL still chose a sequential scan with enable_seqscan disabled, "
+        "which means no index covers the filtered columns.",
+        "",
+    ]
+    for record in offenders[:5]:
+        for scan in record.seq_scans:
+            if scan.relation.lower() in skipped:
+                continue
+            lines.append(f"  {scan.describe()}")
+        if record.location:
+            lines.append(f"      from {record.location}")
+        lines.append(f"      {record.short_sql(110)}")
+        lines.append("")
+    if len(offenders) > 5:
+        lines.append(f"  ... and {len(offenders) - 5} more")
+    lines.append(
+        "If a scan is intentional -- a small lookup table read whole -- exclude it "
+        'with @pytest.mark.no_seq_scan(ignore=("table_name",)).'
+    )
+    return "\n".join(lines)
+
+
+def _too_many_message(budget: Budget, recorder: Recorder, label: str) -> str:
+    assert budget.max_queries is not None
+    over = recorder.count - budget.max_queries
+    lines = [
+        f"{label} ran {recorder.count} queries, budget is {budget.max_queries} ({over} over).",
+    ]
+
+    # Lead with the diagnosis. Someone reading a CI log wants the culprit in the
+    # first few lines, not after a list of thirteen SELECTs.
+    duplicates = recorder.duplicates(threshold=2, kinds=budget.kinds)
+    if duplicates:
+        lines += [
+            "",
+            "Repeated query shapes -- most likely where the extra queries come from:",
+            *(f"  {duplicate.describe()}" for duplicate in duplicates[:3]),
+        ]
+
+    lines += ["", recorder.report(limit=12)]
+    return "\n".join(lines)
+
+
+def _duplicate_message(
+    duplicates: Sequence[Duplicate],
+    recorder: Recorder,
+    label: str,
+) -> str:
+    worst = duplicates[0]
+    lines = [
+        f"{label} repeated the same query {worst.count} times ({recorder.count} queries in total).",
+        "",
+        "This is the N+1 pattern. Load the related rows in one go instead -- with "
+        "selectinload() or joinedload() for a relationship, or a single IN query.",
+        "",
+    ]
+    for duplicate in duplicates[:5]:
+        lines.append(f"  {duplicate.describe()}")
+    if len(duplicates) > 5:
+        lines.append(f"  ... and {len(duplicates) - 5} more repeated shapes")
+    return "\n".join(lines)

pytest_querycount/errors.py

@@ -0,0 +1,39 @@
+"""Failures raised by the plugin.
+
+All of them subclass ``AssertionError`` so pytest renders them as ordinary test
+failures rather than errors.
+"""
+
+from __future__ import annotations
+
+
+class QueryCountError(AssertionError):
+    """Base class for every check failure."""
+
+
+class TooManyQueriesError(QueryCountError):
+    """A test ran more queries than its budget allowed."""
+
+
+class DuplicateQueryError(QueryCountError):
+    """The same query shape was executed repeatedly -- the N+1 signature."""
+
+
+class NoBackendError(QueryCountError):
+    """A check was requested but nothing is instrumented, so it would be a no-op.
+
+    This is deliberately loud: a query budget that silently always passes is
+    worse than no budget at all.
+    """
+
+
+class SeqScanError(QueryCountError):
+    """A query scanned a table sequentially because no index could serve it."""
+
+
+class ExplainUnavailableError(QueryCountError):
+    """A plan check ran where no plan could be obtained.
+
+    Same reasoning as :class:`NoBackendError`: a check that cannot fail is not a
+    check, so we say so rather than pass.
+    """

pytest_querycount/explain.py

@@ -0,0 +1,175 @@
+"""Reading PostgreSQL query plans to find scans no index could serve.
+
+The naive version of this check -- "fail if the plan contains a Seq Scan" --
+does not work, and it is worth saying why, because it shapes everything here.
+
+On a test database of twenty rows PostgreSQL chooses a sequential scan even when
+a perfect index exists, because reading twenty rows is cheaper than descending a
+B-tree. So a Seq Scan tells you nothing about whether an index is missing, and
+thresholding on table size means the check never fires on test data at all.
+
+What we ask instead is whether an index *could* have been used. Run the EXPLAIN
+with ``enable_seqscan`` off, which makes the planner treat sequential scans as
+enormously expensive. If it still picks one, no index can serve that filter --
+and that answer does not depend on how much data the table holds.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from collections.abc import Iterator
+from dataclasses import dataclass
+from typing import Any
+
+_SEQ_SCAN_NODES = {"Seq Scan", "Parallel Seq Scan"}
+
+# Tokens that look like columns in a filter expression but are not.
+_NOT_A_COLUMN = {
+    "and",
+    "or",
+    "not",
+    "any",
+    "all",
+    "true",
+    "false",
+    "null",
+    "is",
+    "like",
+    "ilike",
+    "similar",
+    "between",
+    "in",
+    "case",
+    "when",
+    "then",
+    "else",
+    "end",
+}
+
+_LITERAL = re.compile(r"'(?:[^']|'')*'")
+_CAST = re.compile(r"::\s*[a-z_][a-z0-9_]*(?:\s*\[\s*\])*", re.IGNORECASE)
+_PLACEHOLDER = re.compile(r"\$\d+")
+# PostgreSQL parenthesises a column before casting it, so a varchar filter
+# reads "((city)::text = 'x'::text)". Once the cast is gone the wrapper must go
+# too, or the commonest case of all -- a text column -- yields no column name.
+# The lookbehind is what keeps "lower(email)" from being unwrapped into
+# "lower email", which would suggest an index that does not apply.
+_REDUNDANT_PARENS = re.compile(r"(?<![A-Za-z0-9_\"])\(\s*([a-z_][a-z0-9_]*)\s*\)", re.IGNORECASE)
+
+_COLUMN_BEFORE_OPERATOR = re.compile(
+    r"\b([a-z_][a-z0-9_]*)\s*(?:<=|>=|<>|!=|=|<|>|~~\*?|!~~\*?|@@|\bIS\b|\bLIKE\b)",
+    re.IGNORECASE,
+)
+
+
+@dataclass(frozen=True)
+class SeqScan:
+    """A sequential scan the planner kept even with seq scans penalised."""
+
+    relation: str
+    filter_expression: str
+    columns: tuple[str, ...]
+    """Best-effort column names pulled out of the filter, for the index hint."""
+
+    def suggested_index(self) -> str | None:
+        """A ``CREATE INDEX`` to try, or None when the filter is too complex.
+
+        Deliberately a suggestion and not a promise: which columns to index, in
+        what order, and whether the index earns its write cost are judgement
+        calls that need the whole query pattern, not one plan node.
+        """
+        if not self.columns:
+            return None
+        return f"CREATE INDEX ON {self.relation} ({', '.join(self.columns)});"
+
+    def describe(self) -> str:
+        lines = [f'Seq Scan on "{self.relation}"  Filter: {self.filter_expression}']
+        hint = self.suggested_index()
+        if hint:
+            lines.append(f"      try:  {hint}")
+        return "\n".join(lines)
+
+
+def parse(payload: Any) -> tuple[SeqScan, ...]:
+    """Pull every filtered sequential scan out of an ``EXPLAIN (FORMAT JSON)`` result.
+
+    Accepts the payload however the driver hands it over: already-decoded JSON,
+    or a string still to decode.
+    """
+    document = _decode(payload)
+    if document is None:
+        return ()
+
+    found: list[SeqScan] = []
+    for node in _walk(document):
+        if node.get("Node Type") not in _SEQ_SCAN_NODES:
+            continue
+        expression = node.get("Filter")
+        if not expression:
+            # An unfiltered sequential scan is a deliberate full read of the
+            # table. No index would help, and none is missing.
+            continue
+        relation = node.get("Relation Name") or node.get("Alias") or "?"
+        found.append(
+            SeqScan(
+                relation=str(relation),
+                filter_expression=str(expression),
+                columns=filter_columns(str(expression)),
+            )
+        )
+    return tuple(found)
+
+
+def filter_columns(expression: str) -> tuple[str, ...]:
+    """Column names appearing on the left of a comparison, in order, deduplicated.
+
+    A heuristic on a string PostgreSQL meant for humans, so it gives up rather
+    than guesses: a filter over a function call such as ``lower(email) = $1``
+    yields nothing, and the caller shows the raw expression instead.
+    """
+    cleaned = _LITERAL.sub("''", expression)
+    cleaned = _CAST.sub("", cleaned)
+    cleaned = _PLACEHOLDER.sub("$", cleaned)
+    cleaned = _REDUNDANT_PARENS.sub(r"\1", cleaned)
+
+    seen: dict[str, None] = {}
+    for match in _COLUMN_BEFORE_OPERATOR.finditer(cleaned):
+        name = match.group(1)
+        if name.lower() in _NOT_A_COLUMN:
+            continue
+        # A name immediately followed by "(" is a function, not a column.
+        if cleaned[match.end(1) : match.end(1) + 1] == "(":
+            continue
+        seen.setdefault(name, None)
+    return tuple(seen)
+
+
+def _decode(payload: Any) -> Any:
+    if isinstance(payload, (str, bytes, bytearray)):
+        try:
+            return json.loads(payload)
+        except (ValueError, TypeError):
+            return None
+    return payload
+
+
+def _walk(document: Any) -> Iterator[dict[str, Any]]:
+    """Every plan node in the document, depth first.
+
+    ``EXPLAIN (FORMAT JSON)`` returns a list of statements, each with a "Plan";
+    nodes nest under "Plans". Subplans and CTEs land there too, so one recursive
+    walk reaches all of them.
+    """
+    if isinstance(document, list):
+        for entry in document:
+            yield from _walk(entry)
+        return
+    if not isinstance(document, dict):
+        return
+    if "Plan" in document:
+        yield from _walk(document["Plan"])
+        return
+    yield document
+    for child in document.get("Plans") or ():
+        yield from _walk(child)