pytest-querycount 0.3.0__tar.gz → 0.4.0__tar.gz

{pytest_querycount-0.3.0 → pytest_querycount-0.4.0}/.github/workflows/ci.yml

@@ -28,15 +28,15 @@ jobs:
           --health-timeout 5s
           --health-retries 10
     steps:
-      - uses: actions/checkout@v5
-      - uses: actions/setup-python@v6
+      - uses: actions/checkout@v7
+      - uses: actions/setup-python@v7
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install
         run: |
           python -m pip install --upgrade pip
           python -m pip install -e ".[dev]"
-          python -m pip install "sqlalchemy==${{ matrix.sqlalchemy }}"
+          python -m pip install "sqlalchemy[asyncio]==${{ matrix.sqlalchemy }}"
       - name: Test
         # Without REQUIRE_PG the plan tests would skip on a broken service and
         # the build would still be green, which is the worst outcome.
@@ -48,8 +48,8 @@ jobs:
   quality:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v5
-      - uses: actions/setup-python@v6
+      - uses: actions/checkout@v7
+      - uses: actions/setup-python@v7
         with:
           python-version: "3.13"
       - run: python -m pip install -e ".[dev]"
@@ -63,8 +63,8 @@ jobs:
   package:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v5
-      - uses: actions/setup-python@v6
+      - uses: actions/checkout@v7
+      - uses: actions/setup-python@v7
         with:
           python-version: "3.13"
       - run: python -m pip install build twine

{pytest_querycount-0.3.0 → pytest_querycount-0.4.0}/.github/workflows/publish.yml

@@ -14,8 +14,8 @@ jobs:
     permissions:
       id-token: write
     steps:
-      - uses: actions/checkout@v5
-      - uses: actions/setup-python@v6
+      - uses: actions/checkout@v7
+      - uses: actions/setup-python@v7
         with:
           python-version: "3.13"
       - run: python -m pip install build

{pytest_querycount-0.3.0 → pytest_querycount-0.4.0}/CHANGELOG.md

@@ -8,9 +8,55 @@ All notable changes to this project are documented here. The format follows
 
 ### Planned
 
-- Django and raw psycopg backends.
+- Raw psycopg backend, for projects using SQL without an ORM.
 - `pytest-xdist` aware reporting.
-- A `--querycount-write-budgets` mode that inserts the markers for you.
+- More plan checks: sorts spilling to disk, nested loops over large sets.
+
+## [0.4.0] - 2026-09-27
+
+### Added
+
+- `--querycount-write-budgets` writes `@pytest.mark.max_queries` markers into
+  your test files using the counts observed in that run, then exits without
+  enforcing them. This is the answer to adopting a budget on a suite that
+  already has five hundred tests: nobody was ever going to read five hundred
+  failures and type five hundred numbers by hand.
+
+  Source editing goes through `ast`, not regular expressions, because "the line
+  above the def" is not something a regular expression finds reliably once
+  decorators, classes, async definitions and multi-line signatures are involved.
+  It matches on qualified names, so two `test_create` methods in different
+  classes are never confused; it leaves a hand-written budget alone; it skips
+  tests that failed, whose query count is whatever they reached before blowing
+  up; and it refuses to touch a file it cannot parse.
+
+- `asyncio` extra, `pytest-querycount[asyncio]`, which brings in
+  `sqlalchemy[asyncio]` and therefore greenlet.
+
+### Fixed
+
+- **Caller attribution was silently lost under asyncio.** The budget still
+  fired, but every failure came without the one piece of information that makes
+  it actionable -- no "from your_file.py:39" line at all.
+
+  SQLAlchemy runs its synchronous internals inside a greenlet spawned per
+  operation, and that greenlet's stack begins at SQLAlchemy's own entry point,
+  so following `f_back` reached only library frames. The awaiting frames live on
+  the *parent* greenlet's stack, reachable through `gr_frame`, and the walk now
+  crosses that boundary. Since the audience for this plugin is FastAPI and
+  Litestar, where async is the default, this was the most important gap it had.
+
+### Verified rather than assumed
+
+The README claimed async engines were instrumented "sync or async" on the
+strength of reasoning alone. They are, and there are now tests for it -- along
+with tests proving that the three async tests fail without the greenlet fix, so
+they guard something real. The plan check works under async PostgreSQL too: the
+EXPLAIN on a raw DBAPI cursor survives being issued from inside the greenlet,
+and the savepoint still restores `enable_seqscan` and keeps the transaction
+usable there.
+
+102 tests, up from 75.
 
 ## [0.3.0] - 2026-09-26
 

{pytest_querycount-0.3.0 → pytest_querycount-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: pytest-querycount
-Version: 0.3.0
+Version: 0.4.0
 Summary: Fail your tests when they run too many SQL queries. Catch N+1 problems in CI, not in production.
 Project-URL: Homepage, https://github.com/ldavidsm/pytest-querycount
 Project-URL: Issues, https://github.com/ldavidsm/pytest-querycount/issues
@@ -44,12 +44,16 @@ Classifier: Topic :: Software Development :: Testing
 Classifier: Typing :: Typed
 Requires-Python: >=3.10
 Requires-Dist: pytest>=8.0
+Provides-Extra: asyncio
+Requires-Dist: sqlalchemy[asyncio]>=2.0; extra == 'asyncio'
 Provides-Extra: dev
+Requires-Dist: aiosqlite>=0.19; extra == 'dev'
 Requires-Dist: mypy>=1.8; extra == 'dev'
 Requires-Dist: psycopg[binary]>=3.1; extra == 'dev'
+Requires-Dist: pytest-asyncio>=0.23; extra == 'dev'
 Requires-Dist: pytest>=8.0; extra == 'dev'
 Requires-Dist: ruff>=0.4; extra == 'dev'
-Requires-Dist: sqlalchemy>=2.0; extra == 'dev'
+Requires-Dist: sqlalchemy[asyncio]>=2.0; extra == 'dev'
 Provides-Extra: postgresql
 Requires-Dist: psycopg[binary]>=3.1; extra == 'postgresql'
 Requires-Dist: sqlalchemy>=2.0; extra == 'postgresql'
@@ -82,7 +86,9 @@ CI run, naming the repeated query and the line of your code that caused it.
 ## Install
 
 ```bash
-pip install "pytest-querycount[sqlalchemy]"
+pip install "pytest-querycount[sqlalchemy]"     # budgets and N+1 detection
+pip install "pytest-querycount[postgresql]"     # and the missing-index check
+pip install "pytest-querycount[asyncio]"        # for async engines
 ```
 
 Requires Python 3.10+, pytest 8+, and SQLAlchemy 2.x. There is nothing to
@@ -206,6 +212,50 @@ Three things this is careful about:
 Needs `pip install "pytest-querycount[postgresql]"`. On SQLite or MySQL the check
 raises rather than passing, because a check that cannot fail is not a check.
 
+### Adopting this on a suite that already exists
+
+The honest problem with query budgets is the first day. Nobody is going to read
+five hundred failures and type five hundred numbers, so in practice a plugin like
+this gets installed, switched on once, and switched off again.
+
+```bash
+pytest --querycount-write-budgets
+```
+
+That runs the suite and writes the markers for you, each holding the count that
+test actually ran:
+
+```diff
++@pytest.mark.max_queries(5)
+ def test_list_shops(db):
+     for shop in db.scalars(select(Shop)):
+         len(shop.items)
+
++@pytest.mark.max_queries(2)
+ def test_list_shops_eagerly(db):
+     ...
+
+ class TestItems:
++    @pytest.mark.max_queries(1)
+     def test_count(self, db):
+         ...
+```
+
+It does not enforce anything on that run -- enforcing a budget while deciding
+what it should be is contradictory. Then you commit the diff and every later run
+holds the suite to it.
+
+The numbers are exact on purpose. A budget with slack in it is not a ratchet: the
+point is that the next query added to that code path turns a test red. Where a
+number looks wrong, that is a finding, not a problem with the tool -- widen it by
+hand and leave a comment saying why.
+
+What it will not do: touch a test that already has a budget, write one for a test
+that failed (its count is whatever it reached before blowing up), or touch a file
+it cannot parse. The editing goes through Python's `ast`, so decorators, classes,
+`async def` and multi-line signatures all land correctly, and two `test_create`
+methods in different classes are never confused for each other.
+
 ### A fixture, for when a marker is too coarse
 
 A marker covers the whole test. When you only care about one block:
@@ -224,6 +274,29 @@ Call it with no arguments to observe without asserting -- that is how you find
 out what the budget should be before committing to one. The object from the
 `with` is the recorder, and it keeps its records after the block ends.
 
+### asyncio
+
+Async engines are instrumented with no extra configuration, and failures still
+name the line of your code that emitted the query.
+
+```python
+@pytest.mark.max_queries(2)
+async def test_authors(session):
+    await list_authors(session)
+```
+
+That second part took work, and is worth knowing about if you are comparing
+tools. SQLAlchemy runs its synchronous internals inside a greenlet spawned per
+operation, so the stack at the moment a query is emitted begins at SQLAlchemy's
+own entry point: your awaiting frames are on a different greenlet's stack, and an
+ordinary stack walk never reaches them. Until 0.4.0 this plugin's async failures
+arrived with no origin at all. It now crosses the greenlet boundary.
+
+The plan check works under async PostgreSQL too, including the savepoint that
+keeps `enable_seqscan` and your transaction intact.
+
+Needs `pip install "pytest-querycount[asyncio]"`, which brings greenlet with it.
+
 ### The report you leave switched on
 
 ```bash
@@ -253,6 +326,7 @@ suite: run it once, look at the rows with a `!`, write budgets for those.
 | `--querycount-top=N` | How many tests the table lists (default 10) |
 | `--querycount-max=N` | Apply a budget of N to every test without an explicit one |
 | `--querycount-no-seq-scan` | Apply the missing-index check to every test |
+| `--querycount-write-budgets` | Write the markers for you, then exit without enforcing. **Modifies your files.** |
 
 `--querycount-max` is how you ratchet: set it just above your current worst
 test, then lower it as you fix things.
@@ -296,11 +370,14 @@ FastAPI, Flask, Litestar -- where nothing is currently maintained.
 
 ## Limitations
 
-- SQLAlchemy 2.x only for now. Django and raw psycopg are on the roadmap.
+- SQLAlchemy 2.x only, sync and async. Raw psycopg is on the roadmap; Django is
+  not, since `assertNumQueries` already covers the budget half there.
 - Under `pytest-xdist` the summary table is per worker, so it will be partial.
   Budgets and N+1 detection are unaffected.
 - `no_seq_scan` is PostgreSQL only, and costs one `EXPLAIN` per SELECT while
   enabled. Budgets and N+1 detection work on any SQLAlchemy backend.
+- `--querycount-write-budgets` edits your files in place. Run it on a clean
+  working tree so the diff is the only thing you have to review.
 - The suggested `CREATE INDEX` is a starting point, not advice. Which columns to
   index, in what order, and whether the index earns its write cost need the whole
   query pattern, not one plan node. A filter over a function call

{pytest_querycount-0.3.0 → pytest_querycount-0.4.0}/README.md

@@ -23,7 +23,9 @@ CI run, naming the repeated query and the line of your code that caused it.
 ## Install
 
 ```bash
-pip install "pytest-querycount[sqlalchemy]"
+pip install "pytest-querycount[sqlalchemy]"     # budgets and N+1 detection
+pip install "pytest-querycount[postgresql]"     # and the missing-index check
+pip install "pytest-querycount[asyncio]"        # for async engines
 ```
 
 Requires Python 3.10+, pytest 8+, and SQLAlchemy 2.x. There is nothing to
@@ -147,6 +149,50 @@ Three things this is careful about:
 Needs `pip install "pytest-querycount[postgresql]"`. On SQLite or MySQL the check
 raises rather than passing, because a check that cannot fail is not a check.
 
+### Adopting this on a suite that already exists
+
+The honest problem with query budgets is the first day. Nobody is going to read
+five hundred failures and type five hundred numbers, so in practice a plugin like
+this gets installed, switched on once, and switched off again.
+
+```bash
+pytest --querycount-write-budgets
+```
+
+That runs the suite and writes the markers for you, each holding the count that
+test actually ran:
+
+```diff
++@pytest.mark.max_queries(5)
+ def test_list_shops(db):
+     for shop in db.scalars(select(Shop)):
+         len(shop.items)
+
++@pytest.mark.max_queries(2)
+ def test_list_shops_eagerly(db):
+     ...
+
+ class TestItems:
++    @pytest.mark.max_queries(1)
+     def test_count(self, db):
+         ...
+```
+
+It does not enforce anything on that run -- enforcing a budget while deciding
+what it should be is contradictory. Then you commit the diff and every later run
+holds the suite to it.
+
+The numbers are exact on purpose. A budget with slack in it is not a ratchet: the
+point is that the next query added to that code path turns a test red. Where a
+number looks wrong, that is a finding, not a problem with the tool -- widen it by
+hand and leave a comment saying why.
+
+What it will not do: touch a test that already has a budget, write one for a test
+that failed (its count is whatever it reached before blowing up), or touch a file
+it cannot parse. The editing goes through Python's `ast`, so decorators, classes,
+`async def` and multi-line signatures all land correctly, and two `test_create`
+methods in different classes are never confused for each other.
+
 ### A fixture, for when a marker is too coarse
 
 A marker covers the whole test. When you only care about one block:
@@ -165,6 +211,29 @@ Call it with no arguments to observe without asserting -- that is how you find
 out what the budget should be before committing to one. The object from the
 `with` is the recorder, and it keeps its records after the block ends.
 
+### asyncio
+
+Async engines are instrumented with no extra configuration, and failures still
+name the line of your code that emitted the query.
+
+```python
+@pytest.mark.max_queries(2)
+async def test_authors(session):
+    await list_authors(session)
+```
+
+That second part took work, and is worth knowing about if you are comparing
+tools. SQLAlchemy runs its synchronous internals inside a greenlet spawned per
+operation, so the stack at the moment a query is emitted begins at SQLAlchemy's
+own entry point: your awaiting frames are on a different greenlet's stack, and an
+ordinary stack walk never reaches them. Until 0.4.0 this plugin's async failures
+arrived with no origin at all. It now crosses the greenlet boundary.
+
+The plan check works under async PostgreSQL too, including the savepoint that
+keeps `enable_seqscan` and your transaction intact.
+
+Needs `pip install "pytest-querycount[asyncio]"`, which brings greenlet with it.
+
 ### The report you leave switched on
 
 ```bash
@@ -194,6 +263,7 @@ suite: run it once, look at the rows with a `!`, write budgets for those.
 | `--querycount-top=N` | How many tests the table lists (default 10) |
 | `--querycount-max=N` | Apply a budget of N to every test without an explicit one |
 | `--querycount-no-seq-scan` | Apply the missing-index check to every test |
+| `--querycount-write-budgets` | Write the markers for you, then exit without enforcing. **Modifies your files.** |
 
 `--querycount-max` is how you ratchet: set it just above your current worst
 test, then lower it as you fix things.
@@ -237,11 +307,14 @@ FastAPI, Flask, Litestar -- where nothing is currently maintained.
 
 ## Limitations
 
-- SQLAlchemy 2.x only for now. Django and raw psycopg are on the roadmap.
+- SQLAlchemy 2.x only, sync and async. Raw psycopg is on the roadmap; Django is
+  not, since `assertNumQueries` already covers the budget half there.
 - Under `pytest-xdist` the summary table is per worker, so it will be partial.
   Budgets and N+1 detection are unaffected.
 - `no_seq_scan` is PostgreSQL only, and costs one `EXPLAIN` per SELECT while
   enabled. Budgets and N+1 detection work on any SQLAlchemy backend.
+- `--querycount-write-budgets` edits your files in place. Run it on a clean
+  working tree so the diff is the only thing you have to review.
 - The suggested `CREATE INDEX` is a starting point, not advice. Which columns to
   index, in what order, and whether the index earns its write cost need the whole
   query pattern, not one plan node. A filter over a function call

{pytest_querycount-0.3.0 → pytest_querycount-0.4.0}/pyproject.toml

@@ -33,11 +33,14 @@ dependencies = ["pytest>=8.0"]
 
 [project.optional-dependencies]
 sqlalchemy = ["sqlalchemy>=2.0"]
+asyncio = ["sqlalchemy[asyncio]>=2.0"]
 postgresql = ["sqlalchemy>=2.0", "psycopg[binary]>=3.1"]
 dev = [
-    "sqlalchemy>=2.0",
+    "sqlalchemy[asyncio]>=2.0",
     "psycopg[binary]>=3.1",
+    "aiosqlite>=0.19",
     "pytest>=8.0",
+    "pytest-asyncio>=0.23",
     "mypy>=1.8",
     "ruff>=0.4",
 ]
@@ -70,3 +73,9 @@ select = ["E", "F", "I", "UP", "B", "SIM", "RUF"]
 [tool.mypy]
 strict = true
 files = ["src/pytest_querycount"]
+
+# greenlet is an optional runtime dependency, arriving only with
+# sqlalchemy[asyncio], and we touch three attributes of it dynamically.
+[[tool.mypy.overrides]]
+module = ["greenlet"]
+ignore_missing_imports = true

{pytest_querycount-0.3.0 → pytest_querycount-0.4.0}/src/pytest_querycount/__init__.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-__version__ = "0.3.0"
+__version__ = "0.4.0"
 
 from pytest_querycount.errors import (
     DuplicateQueryError,

{pytest_querycount-0.3.0 → pytest_querycount-0.4.0}/src/pytest_querycount/plugin.py

@@ -3,17 +3,21 @@
 from __future__ import annotations
 
 from collections.abc import Generator, Sequence
+from pathlib import Path
 from typing import Any
 
 import pytest
 
-from pytest_querycount import backends, checks, report
+from pytest_querycount import backends, checks, report, writer
 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] = {}
 
+# path -> {qualname: observed count}, gathered for --querycount-write-budgets.
+_OBSERVED: dict[Path, dict[str, int]] = {}
+
 
 # -- configuration ---------------------------------------------------------
 
@@ -40,6 +44,16 @@ def pytest_addoption(parser: pytest.Parser) -> None:
         metavar="N",
         help="Apply a query budget of N to every test that has no explicit one.",
     )
+    group.addoption(
+        "--querycount-write-budgets",
+        action="store_true",
+        default=False,
+        help=(
+            "Write @pytest.mark.max_queries markers into your test files using "
+            "the counts observed in this run, then exit without enforcing them. "
+            "MODIFIES YOUR FILES -- review the diff afterwards."
+        ),
+    )
     group.addoption(
         "--querycount-no-seq-scan",
         action="store_true",
@@ -80,6 +94,7 @@ def pytest_configure(config: pytest.Config) -> None:
         "sequentially because no index could serve its filter. PostgreSQL only.",
     )
     _STATS.clear()
+    _OBSERVED.clear()
     backends.install_all()
 
 
@@ -126,6 +141,11 @@ def pytest_runtest_call(item: pytest.Item) -> Generator[None, object, object]:
         backends.pop(recorder)
 
     _record_stats(item, recorder)
+
+    if item.config.getoption("querycount_write_budgets"):
+        _observe(item, recorder)
+        return result
+
     checks.enforce(budget, recorder, label=item.name)
     return result
 
@@ -143,6 +163,24 @@ def _record_stats(item: pytest.Item, recorder: Recorder) -> None:
     )
 
 
+def _observe(item: pytest.Item, recorder: Recorder) -> None:
+    """Remember what a passing test ran, so a marker can be written for it.
+
+    Only passing tests: a test that failed part way through ran an arbitrary
+    number of queries, and freezing that number as a budget would be nonsense.
+    """
+    if not recorder.count:
+        return
+    function = getattr(item, "function", None)
+    qualname = getattr(function, "__qualname__", None)
+    if qualname is None:
+        return
+
+    per_file = _OBSERVED.setdefault(Path(str(item.path)), {})
+    # Parametrised cases share one function, so the widest run wins.
+    per_file[qualname] = max(per_file.get(qualname, 0), recorder.count)
+
+
 def _budget_for(item: pytest.Item) -> Budget:
     """Resolve the budget: marker first, then --querycount-max, then ini."""
     budget = Budget(
@@ -274,6 +312,10 @@ def pytest_terminal_summary(
     exitstatus: int,
     config: pytest.Config,
 ) -> None:
+    if config.getoption("querycount_write_budgets"):
+        _write_budgets(terminalreporter)
+        return
+
     if not config.getoption("querycount_report"):
         return
     lines = report.build(_STATS, top=int(config.getoption("querycount_top")))
@@ -282,3 +324,10 @@ def pytest_terminal_summary(
     terminalreporter.write_sep("=", "querycount summary")
     for line in lines:
         terminalreporter.write_line(line)
+
+
+def _write_budgets(terminalreporter: Any) -> None:
+    results = [writer.rewrite_file(path, counts) for path, counts in sorted(_OBSERVED.items())]
+    terminalreporter.write_sep("=", "querycount: budgets written")
+    for line in writer.summarise(results):
+        terminalreporter.write_line(line)

{pytest_querycount-0.3.0 → pytest_querycount-0.4.0}/src/pytest_querycount/recorder.py

@@ -5,8 +5,8 @@ from __future__ import annotations
 import os
 import sys
 from collections import Counter
-from collections.abc import Iterable, Sequence
-from typing import TYPE_CHECKING
+from collections.abc import Iterable, Iterator, Sequence
+from typing import TYPE_CHECKING, Any
 
 from pytest_querycount.records import Duplicate, QueryRecord
 
@@ -26,7 +26,7 @@ _INTERNAL_PATHS = (
     "/threading.py",
 )
 
-_MAX_STACK_DEPTH = 60
+_MAX_STACK_DEPTH = 120
 
 
 def _shorten(path: str) -> str:
@@ -42,19 +42,64 @@ def _shorten(path: str) -> str:
     return path if relative.startswith("..") else relative
 
 
+def _greenlet_module() -> Any:
+    """The greenlet module, or None when it is not installed.
+
+    Imported lazily: greenlet only arrives with ``sqlalchemy[asyncio]``, and a
+    synchronous project should not be made to care that it exists.
+    """
+    try:
+        import greenlet
+    except ImportError:  # pragma: no cover - synchronous install
+        return None
+    return greenlet
+
+
+def _frames(start: FrameType | None) -> Iterator[FrameType]:
+    """Frames from ``start`` outwards, crossing greenlet boundaries.
+
+    Under SQLAlchemy's asyncio support the query is emitted inside a greenlet
+    spawned for it, and that greenlet's stack begins at SQLAlchemy's own entry
+    point -- so following ``f_back`` alone never reaches the caller's code. The
+    awaiting frames live on the *parent* greenlet's stack, which is suspended and
+    reachable through ``gr_frame``.
+
+    Walking that chain is what lets an async test still be told which of its
+    lines emitted the query, rather than being told nothing.
+    """
+    greenlet = _greenlet_module()
+    current = greenlet.getcurrent() if greenlet is not None else None
+
+    frame = start
+    budget = _MAX_STACK_DEPTH
+    while budget > 0:
+        if frame is not None:
+            yield frame
+            frame = frame.f_back
+            budget -= 1
+            continue
+
+        # This greenlet's stack is exhausted. Continue in whoever is waiting on
+        # it. The chain ends at the main greenlet, whose parent is None.
+        if current is None:
+            return
+        current = getattr(current, "parent", None)
+        if current is None:
+            return
+        frame = getattr(current, "gr_frame", None)
+
+
 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)
+        start: 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:
+
+    for frame in _frames(start):
         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