sqlakit 0.13.0__tar.gz → 0.14.0__tar.gz

{sqlakit-0.13.0 → sqlakit-0.14.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlakit
-Version: 0.13.0
+Version: 0.14.0
 Summary: A toolkit for SQLAlchemy applications.
 Keywords: sqlalchemy,database,orm,sql,asyncio
 Author: Anton Ruhlov

{sqlakit-0.13.0 → sqlakit-0.14.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "sqlakit"
-version = "0.13.0"
+version = "0.14.0"
 description = "A toolkit for SQLAlchemy applications."
 readme = "README.md"
 license = "MIT"

{sqlakit-0.13.0 → sqlakit-0.14.0}/pyproject.toml.orig

@@ -1,6 +1,6 @@
 [project]
 name = "sqlakit"
-version = "0.13.0"
+version = "0.14.0"
 description = "A toolkit for SQLAlchemy applications."
 readme = "README.md"
 license = "MIT"

{sqlakit-0.13.0 → sqlakit-0.14.0}/sqlakit/__init__.py

@@ -52,6 +52,7 @@ from .exceptions import (
     UnknownImportPathError,
     UnknownOrderFieldError,
     UnorderedPageError,
+    UnregisteredDatabaseError,
 )
 from .types import (
     DatabaseConfig,
@@ -122,6 +123,7 @@ __all__ = [
     "UnknownImportPathError",
     "UnknownOrderFieldError",
     "UnorderedPageError",
+    "UnregisteredDatabaseError",
     "UrlParts",
     "ValidationArgs",
     "db",

{sqlakit-0.13.0 → sqlakit-0.14.0}/sqlakit/_base.py

@@ -52,6 +52,7 @@ from .exceptions import (
     MissingSessionError,
     RetryNotSupportedError,
     UnknownDatabaseError,
+    UnregisteredDatabaseError,
 )
 
 if TYPE_CHECKING:
@@ -111,6 +112,7 @@ def _elsewhere(frames: tuple[str, ...], skipped: tuple[str, ...]) -> bool:
 
 ConnectionT = TypeVar("ConnectionT")
 SessionT = TypeVar("SessionT")
+ValueT = TypeVar("ValueT")
 
 
 class _Lazy(Generic[ConnectionT]):
@@ -182,18 +184,45 @@ class _Outer(Generic[ConnectionT]):
         connection: What blocks below reuse, with ``join_nested``.
         join_nested: Whether blocks below reuse that connection.
         savepoint: Whether nested blocks run as savepoints.
-        session_savepoint: Whether this block's session needs a savepoint of its
-            own. Two savepoint owners on one connection release each other's out
-            of order, so only one may have it.
-        scope: The block's own scope, whose session the savepoint is for.
+        scope: The block's own scope, whose session the connection is bound to.
+        owner: The scope whose session owns the savepoints of this connection.
+            That session joins with a savepoint of its own, and a block below
+            takes its savepoint through it, so the two are released in the
+            order they were taken. One owner per connection: two of them
+            release each other's savepoints out of order.
 
     """
 
     connection: ConnectionT
     join_nested: bool = True
     savepoint: bool = False
-    session_savepoint: bool = False
     scope: Any = None
+    owner: Any = None
+
+    def savepoint_owner(self) -> Any:  # noqa: ANN401 - a session of either API
+        """Return the session that owns the savepoints here, if one does."""
+        return None if self.owner is None else self.owner.session
+
+
+class _Binding(Generic[ValueT]):
+    """A value bound to a context variable for as long as a block is open.
+
+    The generator `@contextmanager` builds costs more than the two calls it
+    saves, and every block binds two of these.
+    """
+
+    __slots__ = ("_token", "_value", "_var")
+
+    def __init__(self, var: ContextVar[Any], value: ValueT) -> None:
+        self._var = var
+        self._value = value
+
+    def __enter__(self) -> ValueT:
+        self._token = self._var.set(self._value)
+        return self._value
+
+    def __exit__(self, *_: object) -> None:
+        self._var.reset(self._token)
 
 
 class BaseDatabase(Generic[ConnectionT, SessionT]):
@@ -561,12 +590,25 @@ class BaseDatabase(Generic[ConnectionT, SessionT]):
         outer = self._outer.get(None)
         if outer is None or outer.connection is not connection:
             return {}
-        if outer.session_savepoint and outer.scope is self._scope.get(None):
+        if outer.owner is not None and outer.owner is self._scope.get(None):
             return {"join_transaction_mode": "create_savepoint"}
         # Spelled out: SQLAlchemy's default would open a savepoint of its own
         # whenever the connection is already inside one.
         return {"join_transaction_mode": "rollback_only"}
 
+    @staticmethod
+    def _owns_savepoints(outer: _Outer[Any] | None, *, savepoint: bool) -> bool:
+        """Whether this block's session owns the savepoints of its connection.
+
+        A block rolled back at the end gives its session one, so a test or the
+        code under it may commit that session without ending the block. A block
+        that took a savepoint of its own does not: the owner stays the one
+        above, and blocks below take their savepoints through it.
+        """
+        if outer is None:
+            return savepoint
+        return not savepoint and outer.owner is not None
+
     def _plan(self, *, savepoint: bool, rollback: bool) -> tuple[_Outer | None, bool]:
         """Decide what a transaction joins, and whether it is a savepoint.
 
@@ -609,38 +651,34 @@ class BaseDatabase(Generic[ConnectionT, SessionT]):
             return None
         return scope
 
-    @contextmanager
     def _bind(
         self,
         connection: ConnectionT | None,
         checkout: _Lazy[ConnectionT] | None = None,
         *,
         autocommit: bool = False,
-    ) -> Iterator[_Scope[ConnectionT, SessionT]]:
+    ) -> _Binding[_Scope[ConnectionT, SessionT]]:
         """Bind a scope holding ``connection`` to the current context.
 
         Every block gets a scope, and so a session, of its own. Ending that
         session is left to the caller, which knows whether it takes an
         ``await``. A lazy block passes ``checkout`` instead of a connection.
         """
-        scope = _Scope[ConnectionT, SessionT](
+        # Unparameterized: subscripting a generic builds an alias and calls
+        # through it, and every block binds a scope.
+        scope: _Scope[ConnectionT, SessionT] = _Scope(
             connection, checkout=checkout, autocommit=autocommit
         )
-        token = self._scope.set(scope)
-        try:
-            yield scope
-        finally:
-            self._scope.reset(token)
+        return _Binding(self._scope, scope)
 
-    @contextmanager
     def _set_outer(
         self,
         connection: ConnectionT | None,
         *,
         join_nested: bool = True,
         savepoint: bool = False,
-        session_savepoint: bool = False,
-    ) -> Iterator[_Outer[ConnectionT] | None]:
+        owner: Any = None,  # noqa: ANN401 - the scope whose session owns them
+    ) -> _Binding[_Outer[ConnectionT] | None]:
         """Make ``connection`` the outer one for this context. See `_Outer`.
 
         ``None`` leaves this context without an outer transaction at all, as
@@ -652,16 +690,12 @@ class BaseDatabase(Generic[ConnectionT, SessionT]):
                 connection,
                 join_nested=join_nested,
                 savepoint=savepoint,
-                session_savepoint=session_savepoint,
+                owner=owner,
             )
             if connection is not None
             else None
         )
-        token = self._outer.set(outer)
-        try:
-            yield outer
-        finally:
-            self._outer.reset(token)
+        return _Binding(self._outer, outer)
 
 
 DatabaseT = TypeVar("DatabaseT", bound="BaseDatabase[Any, Any]")
@@ -928,8 +962,11 @@ class _DatabaseRegistryMixin(BaseDatabase[Any, Any], Generic[DatabaseT]):
         db._name = alias  # noqa: SLF001
         return db
 
-    def using(self, alias: str) -> _Using:
-        """Return the database under that alias, standing in for the default one.
+    def using(self, target: str | DatabaseT) -> _Using:
+        """Return that database, standing in for the default one.
+
+        Named or handed over, as `recording(using=...)` and a query's `using()`
+        take it.
 
         The block opens on it, and models that live on the default database resolve
         there for as long as it is open:
@@ -945,12 +982,31 @@ class _DatabaseRegistryMixin(BaseDatabase[Any, Any], Generic[DatabaseT]):
 
         Raises:
             UnknownDatabaseError: if nothing is configured under that alias.
+            UnregisteredDatabaseError: if the database is not one this registry
+                holds, since the redirection works by the name it holds it under.
 
         """
+        alias = target if isinstance(target, str) else self._alias_of(target)
         if alias not in self:
             raise UnknownDatabaseError(alias, self.aliases)
         return _Using(self[alias], self._using, alias)
 
+    def _alias_of(self, db: DatabaseT) -> str:
+        """Return the alias this registry holds a database under.
+
+        Raises:
+            UnregisteredDatabaseError: if it holds it under none.
+
+        """
+        # A configured registry is the default database itself, and a
+        # registered one holds it.
+        if db is self or db is self._default:
+            return DEFAULT_ALIAS
+        for alias, held in self._aliased.items():
+            if held is db:
+                return alias
+        raise UnregisteredDatabaseError(self.aliases)
+
     def route(self, *routers: Router | RouterFunction | str) -> None:
         """Say which database a model lives on, for models that do not say it.
 

{sqlakit-0.13.0 → sqlakit-0.14.0}/sqlakit/_db.py

@@ -482,12 +482,23 @@ class Transaction(ContextDecorator, AbstractContextManager["sa.Connection"]):
             owner: Session | None = None
             if outer is not None:
                 connection = outer.connection
+                holder = outer.savepoint_owner()
                 # Without a savepoint the block only takes part in the
                 # transaction around it, which commits it.
-                transaction = connection.begin_nested() if savepoint else None
-                # The block's savepoint isolates it; a session opened inside
-                # must not add a second one on the same connection.
-                session_savepoint = outer.session_savepoint and not savepoint
+                if not savepoint:
+                    transaction = None
+                elif holder is not None:
+                    # The session holding this connection's savepoints takes
+                    # this one too, so a commit of its own releases them in the
+                    # order they were taken.
+                    transaction = holder.begin_nested()
+                    # It opens the savepoint when it next reaches the
+                    # connection, which may be after this block has written
+                    # through a session of its own.
+                    holder.connection()
+                else:
+                    transaction = connection.begin_nested()
+                held_by = outer.owner
             else:
                 borrowed = self.db._scope_to_borrow()  # noqa: SLF001
                 if borrowed is not None:
@@ -496,20 +507,22 @@ class Transaction(ContextDecorator, AbstractContextManager["sa.Connection"]):
                 else:
                     connection = stack.enter_context(self.db.engine.connect())
                     owner, transaction = None, connection.begin()
-                session_savepoint = savepoint
+                held_by = None
             # Unwound in reverse: session, context, transaction, connection.
-            stack.push(self._finish(transaction, owner))
+            stack.push(self._finish(transaction, owner, connection))
             bound = stack.enter_context(
                 self.db._set_outer(  # noqa: SLF001
                     connection,
                     join_nested=self.join_nested,
                     savepoint=savepoint,
-                    session_savepoint=session_savepoint,
+                    owner=held_by,
                 )
             )
             scope = stack.enter_context(self.db._bind(connection))  # noqa: SLF001
             if bound is not None:
                 bound.scope = scope
+                if self.db._owns_savepoints(outer, savepoint=savepoint):  # noqa: SLF001
+                    bound.owner = scope
             stack.push(self._close_session(scope))
         except BaseException:
             stack.close()
@@ -552,6 +565,7 @@ class Transaction(ContextDecorator, AbstractContextManager["sa.Connection"]):
         self,
         transaction: sa.Transaction | None,
         owner: Session | None = None,
+        connection: sa.Connection | None = None,
     ) -> Callable[..., None]:
         """Commit or roll back, unless this block only takes part in another."""
 
@@ -571,6 +585,10 @@ class Transaction(ContextDecorator, AbstractContextManager["sa.Connection"]):
             if transaction is None:
                 return
             if not transaction.is_active:
+                if connection is not None and connection.in_transaction():
+                    # The session of an enclosing block ended this savepoint by
+                    # committing, and the transaction around it holds the work.
+                    return
                 # Rolled back from inside the block. Say so, unless an
                 # exception is already on its way out with the reason.
                 if exc is None:

{sqlakit-0.13.0 → sqlakit-0.14.0}/sqlakit/asyncio/_db.py

@@ -522,12 +522,23 @@ class Transaction(
             owner: AsyncSession | None = None
             if outer is not None:
                 connection = outer.connection
+                holder = outer.savepoint_owner()
                 # Without a savepoint the block only takes part in the
                 # transaction around it, which commits it.
-                transaction = await connection.begin_nested() if savepoint else None
-                # The block's savepoint isolates it; a session opened inside
-                # must not add a second one on the same connection.
-                session_savepoint = outer.session_savepoint and not savepoint
+                if not savepoint:
+                    transaction = None
+                elif holder is not None:
+                    # The session holding this connection's savepoints takes
+                    # this one too, so a commit of its own releases them in the
+                    # order they were taken.
+                    transaction = await holder.begin_nested()
+                    # It opens the savepoint when it next reaches the
+                    # connection, which may be after this block has written
+                    # through a session of its own.
+                    await holder.connection()
+                else:
+                    transaction = await connection.begin_nested()
+                held_by = outer.owner
             else:
                 borrowed = self.db._scope_to_borrow()  # noqa: SLF001
                 if borrowed is not None:
@@ -538,20 +549,22 @@ class Transaction(
                         self.db.engine.connect()
                     )
                     transaction = await connection.begin()
-                session_savepoint = savepoint
+                held_by = None
             # Unwound in reverse: session, context, transaction, connection.
-            stack.push_async_exit(self._finish(transaction, owner))
+            stack.push_async_exit(self._finish(transaction, owner, connection))
             bound = stack.enter_context(
                 self.db._set_outer(  # noqa: SLF001
                     connection,
                     join_nested=self.join_nested,
                     savepoint=savepoint,
-                    session_savepoint=session_savepoint,
+                    owner=held_by,
                 )
             )
             scope = stack.enter_context(self.db._bind(connection))  # noqa: SLF001
             if bound is not None:
                 bound.scope = scope
+                if self.db._owns_savepoints(outer, savepoint=savepoint):  # noqa: SLF001
+                    bound.owner = scope
             stack.push_async_exit(self._close_session(scope))
         except BaseException:
             await stack.aclose()
@@ -594,6 +607,7 @@ class Transaction(
         self,
         transaction: AsyncTransaction | None,
         owner: AsyncSession | None = None,
+        connection: AsyncConnection | None = None,
     ) -> Callable[..., Coroutine[None, None, None]]:
         """Commit or roll back, unless this block only takes part in another."""
 
@@ -613,6 +627,10 @@ class Transaction(
             if transaction is None:
                 return
             if not transaction.is_active:
+                if connection is not None and connection.in_transaction():
+                    # The session of an enclosing block ended this savepoint by
+                    # committing, and the transaction around it holds the work.
+                    return
                 # Rolled back from inside the block. Say so, unless an
                 # exception is already on its way out with the reason.
                 if exc is None:

{sqlakit-0.13.0 → sqlakit-0.14.0}/sqlakit/exceptions.py

@@ -38,6 +38,7 @@ __all__ = [
     "UnknownFieldError",
     "UnknownImportPathError",
     "UnorderedPageError",
+    "UnregisteredDatabaseError",
 ]
 
 
@@ -158,6 +159,17 @@ class UnknownDatabaseError(SQLAKitError, KeyError):
         )
 
 
+class UnregisteredDatabaseError(SQLAKitError, ValueError):
+    """Raised when a registry is handed a database it does not hold."""
+
+    def __init__(self, known: tuple[str, ...] = ()) -> None:
+        super().__init__(
+            "That database is not one this registry holds. Register it with "
+            "`register(alias, db)`, or name one of the databases it has: "
+            f"{', '.join(map(repr, known)) or 'none'}."
+        )
+
+
 class MissingRegistryError(SQLAKitError, ValueError):
     """Raised when a model has nowhere to look a database alias up."""
 

{sqlakit-0.13.0 → sqlakit-0.14.0}/sqlakit/pytest_plugin.py

@@ -42,19 +42,24 @@ import inspect
 import pathlib
 import time
 import warnings
+from collections.abc import Mapping
 from contextlib import AsyncExitStack, ExitStack, contextmanager
 from typing import TYPE_CHECKING, Any
 
 import pytest
 
 from ._debugserver import as_payload, write_report
+from ._recording import Recording, check, require_expectation
 from ._registry import db as importable_db
+from .exceptions import UnknownDatabaseError
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncIterator, Iterator
+    from collections.abc import AsyncIterator, Iterator, Sequence
 
     import sqlalchemy as sa
 
+    from .types import AssertQueries
+
 MARKER = "db"
 SYNC_FIXTURE = "_sqlakit_transaction"
 ASYNC_FIXTURE = "_sqlakit_async_transaction"
@@ -183,7 +188,8 @@ def sqlakit_db(sqlakit_base: Any) -> Any:  # noqa: ANN401
 
     The one the models live on: their registry, which knows every alias, or the
     `Database` they were given in person. Override it for a project with no
-    model layer.
+    model layer, and return a dict or a list for a project on several databases
+    with no registry between them.
     """
     if sqlakit_base is None:
         if not importable_db.is_configured:
@@ -278,6 +284,47 @@ def sqlakit_seed(sqlakit_schema: None) -> None:  # noqa: ARG001 - after the sche
     return
 
 
+@pytest.fixture
+def assert_queries(sqlakit_db: Any) -> AssertQueries:  # noqa: ANN401
+    """Assert what a block asks of the databases under test.
+
+    ```python
+    @pytest.mark.db
+    def test_the_page_costs_two_queries(assert_queries: AssertQueries):
+        with assert_queries(2):
+            User.query.order_by("name").page(limit=10)
+    ```
+
+    The same checks as `sqlakit.testing.assert_queries`: a count, an `at_most`
+    ceiling, and `duplicates=False` to forbid a statement running twice. It
+    watches what `sqlakit_db` returned, so a project whose databases are its
+    own needs no registry for `using` to name one:
+
+    ```python
+    with assert_queries(1, using="replica"):
+        build_report()
+    ```
+    """
+
+    @contextmanager
+    def asserted(
+        count: int | None = None,
+        *,
+        at_most: int | None = None,
+        duplicates: bool = True,
+        using: Any = None,  # noqa: ANN401
+    ) -> Iterator[Recording]:
+        require_expectation(count, at_most, duplicates)
+        recording = Recording()
+        with ExitStack() as stack:
+            for one in _picked(sqlakit_db, _as_asked(using)):
+                stack.enter_context(one.recording(into=recording))
+            yield recording
+        check(recording, count=count, at_most=at_most, duplicates=duplicates)
+
+    return asserted
+
+
 @pytest.fixture
 def _sqlakit_transaction(
     request: pytest.FixtureRequest,
@@ -294,7 +341,8 @@ def _sqlakit_transaction(
                 )
             stack.enter_context(block)
         if _hidden(request):
-            stack.enter_context(sqlakit_db.unbound())
+            for one in _each(sqlakit_db):
+                stack.enter_context(one.unbound())
         with _reported(request, sqlakit_db):
             yield
 
@@ -314,7 +362,8 @@ async def _sqlakit_async_transaction(
             else:
                 stack.enter_context(block)
         if _hidden(request):
-            stack.enter_context(sqlakit_db.unbound())
+            for one in _each(sqlakit_db):
+                stack.enter_context(one.unbound())
         with _reported(request, sqlakit_db):
             yield
 
@@ -331,7 +380,12 @@ def _reported(request: pytest.FixtureRequest, db: Any) -> Iterator[None]:  # noq
         return
     node = request.node
     skip = request.config.getini("sqlakit_skip_queries_from")
-    with db.recording(_named(node), stacks=True, skip_queries_from=skip) as recording:
+    recording = Recording(label=_named(node))
+    with ExitStack() as stack:
+        for one in _each(db):
+            stack.enter_context(
+                one.recording(into=recording, stacks=True, skip_queries_from=skip)
+            )
         yield
     request.config.stash[REPORT].append(
         as_payload(
@@ -384,6 +438,17 @@ def _schema_blocks(
     gets one block per alias. Without it there is one metadata, and one
     database to put it on.
     """
+    databases = _each(db)
+    if len(databases) > 1:
+        if metadata is None:
+            pytest.fail(
+                "`sqlakit_db` returned several databases, and the tables of each "
+                "are the project's to create. Define `sqlakit_schema` with a "
+                "`provisioned_tables()` per database, or `sqlakit_metadata` with "
+                "the tables they share.",
+                pytrace=False,
+            )
+        return [one.provisioned_tables(metadata) for one in databases]
     if base is not None:
         # A database of its own is where the models are pinned, and the alias
         # a registry knows it by is not how the base reaches it.
@@ -403,7 +468,11 @@ def _schema_blocks(
 def _asked_for(request: pytest.FixtureRequest) -> tuple[Any, ...]:
     """Return the databases the marker asks for, by name or in person."""
     marker = request.node.get_closest_marker(MARKER)
-    using = None if marker is None else marker.kwargs.get("using")
+    return _as_asked(None if marker is None else marker.kwargs.get("using"))
+
+
+def _as_asked(using: Any) -> tuple[Any, ...]:  # noqa: ANN401
+    """Return what ``using`` names, as the several it may be."""
     if using is None:
         return ()
     if isinstance(using, str) or not isinstance(using, (list, tuple, set, frozenset)):
@@ -424,6 +493,67 @@ def _hidden(request: pytest.FixtureRequest) -> bool:
     return bool(asked)
 
 
+def _each(db: Any) -> tuple[Any, ...]:  # noqa: ANN401
+    """Return the databases a project handed over, as several or as one.
+
+    `sqlakit_db` returns one database, a registry of them, or several: a list,
+    or a dict naming each, for a project with no registry between them.
+    """
+    if isinstance(db, Mapping):
+        return tuple(db.values())
+    if isinstance(db, (list, tuple)):
+        return tuple(db)
+    return (db,)
+
+
+def _picked(db: Any, using: tuple[Any, ...]) -> list[Any]:  # noqa: ANN401
+    """Return the databases a marker asks for, out of what a project handed over.
+
+    A dict names them itself. A list names them by the alias each database
+    carries, and a registry by the alias it holds one under.
+
+    Raises:
+        UnknownDatabaseError: if an alias belongs to none of them.
+
+    """
+    if not using:
+        return list(_each(db))
+    picked = []
+    for asked in using:
+        if not isinstance(asked, str):
+            picked.append(asked)
+        elif isinstance(db, Mapping):
+            if asked not in db:
+                raise UnknownDatabaseError(asked, tuple(db))
+            picked.append(db[asked])
+        elif isinstance(db, (list, tuple)):
+            picked.append(_of_the_list(db, asked))
+        else:
+            picked.append(db[asked])
+    return picked
+
+
+def _of_the_list(databases: Sequence[Any], asked: str) -> Any:  # noqa: ANN401
+    """Return the database of a list that carries this name.
+
+    Raises:
+        UnknownDatabaseError: if none of them carries it.
+
+    """
+    found = [one for one in databases if asked in one]
+    if not found:
+        aliases = tuple(alias for one in databases for alias in one.aliases)
+        raise UnknownDatabaseError(asked, aliases)
+    if len(found) > 1:
+        pytest.fail(
+            f"`sqlakit_db` returned {len(found)} databases named `{asked}`. "
+            "Name them with `Database(url, alias=...)`, or return a dict "
+            "naming each, so a marker can ask for one.",
+            pytrace=False,
+        )
+    return found[0]
+
+
 def _rolled_back(db: Any, using: tuple[Any, ...]) -> list[Any]:  # noqa: ANN401
     """Return the blocks that undo what a test writes.
 
@@ -431,15 +561,7 @@ def _rolled_back(db: Any, using: tuple[Any, ...]) -> list[Any]:  # noqa: ANN401
     the one they have. Naming one is how a project on several stops paying for
     a connection to each in the tests that read one.
     """
-    if not using:
-        # A registry opens every database it holds, and one database opens
-        # itself. `transactions` is the one a registry has.
-        return [
-            db.transactions(rollback=True)
-            if hasattr(db, "transactions")
-            else db.transaction(rollback=True)
-        ]
-    return [
-        (db[one] if isinstance(one, str) else one).transaction(rollback=True)
-        for one in using
-    ]
+    if not using and len(_each(db)) == 1 and hasattr(db, "transactions"):
+        # A registry opens every database it holds in one block.
+        return [db.transactions(rollback=True)]
+    return [one.transaction(rollback=True) for one in _picked(db, using)]

{sqlakit-0.13.0 → sqlakit-0.14.0}/sqlakit/testing.py

@@ -9,11 +9,12 @@ from . import db as sync_db
 from ._recording import Recording, check, require_expectation
 from .asyncio import db as async_db
 from .exceptions import UnknownDatabaseError
+from .types import AssertQueries
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
 
-__all__ = ["assert_queries"]
+__all__ = ["AssertQueries", "assert_queries"]
 
 
 @contextmanager

{sqlakit-0.13.0 → sqlakit-0.14.0}/sqlakit/types.py

@@ -1,11 +1,12 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Literal, TypeAlias
+from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeAlias
 
 from typing_extensions import TypedDict
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Mapping, Sequence
+    from contextlib import AbstractContextManager
     from pathlib import Path
 
     import sqlalchemy as sa
@@ -13,9 +14,11 @@ if TYPE_CHECKING:
     from sqlalchemy.orm import Query, Session
     from sqlalchemy.pool import Pool
 
+    from ._recording import Recording
     from ._sql import Templates
 
 __all__ = [
+    "AssertQueries",
     "DatabaseConfig",
     "EngineArgs",
     "SessionArgs",
@@ -130,6 +133,32 @@ class QueryStats(TypedDict):
     label: str | None
 
 
+class AssertQueries(Protocol):
+    """The `assert_queries` fixture, for a test that annotates what it asks for.
+
+    Imported from `sqlakit.testing`, where the function of the same name is:
+
+    ```python
+    from sqlakit.testing import AssertQueries
+
+
+    @pytest.mark.db
+    def test_the_page_costs_two_queries(assert_queries: AssertQueries) -> None:
+        with assert_queries(2):
+            User.query.order_by("name").page(limit=10)
+    ```
+    """
+
+    def __call__(
+        self,
+        count: int | None = None,
+        *,
+        at_most: int | None = None,
+        duplicates: bool = True,
+        using: Any = None,  # noqa: ANN401 - an alias, or the database itself
+    ) -> AbstractContextManager[Recording]: ...
+
+
 class DatabaseConfig(UrlParts, total=False):
     """One database in a configuration keyed by alias.