sqlakit 0.22.9__tar.gz → 0.22.11__tar.gz

{sqlakit-0.22.9 → sqlakit-0.22.11}/PKG-INFO

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

{sqlakit-0.22.9 → sqlakit-0.22.11}/pyproject.toml

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

{sqlakit-0.22.9 → sqlakit-0.22.11}/pyproject.toml.orig

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

{sqlakit-0.22.9 → sqlakit-0.22.11}/sqlakit/__init__.py

@@ -10,7 +10,7 @@ from ._query import (
     orderable_columns,
 )
 from ._recording import Recording, Statement
-from ._registry import Databases, db
+from ._registry import Databases, autocommit, db, transaction
 from ._routing import Router
 from .exceptions import (
     DEFAULT_ALIAS,
@@ -142,8 +142,10 @@ __all__ = [
     "UnregisteredDatabaseError",
     "UrlParts",
     "ValidationArgs",
+    "autocommit",
     "db",
     "import_models",
     "import_string",
     "orderable_columns",
+    "transaction",
 ]

{sqlakit-0.22.9 → sqlakit-0.22.11}/sqlakit/_base.py

@@ -232,6 +232,9 @@ class BaseDatabase(Generic[ConnectionT, SessionT]):
     """
 
     _engine: Any = None  # narrowed by the subclass
+    # The database to work on instead of this one, or None. Every public member
+    # checks it first, so a database without one pays one attribute read.
+    _override: Any = None
 
     def __init__(
         self,
@@ -299,6 +302,7 @@ class BaseDatabase(Generic[ConnectionT, SessionT]):
         self._listened: Any = None
         self._listening_lock = threading.Lock()
         self._name = alias or DEFAULT_ALIAS
+        self._override = None
 
     def __repr__(self) -> str:
         return f"{type(self).__name__}({self.url.render_as_string()!r})"
@@ -332,6 +336,8 @@ class BaseDatabase(Generic[ConnectionT, SessionT]):
             MissingSessionError: if no connection is bound.
 
         """
+        if self._override is not None:
+            return self._override.session
         try:
             scope = self._scope.get()
         except LookupError:
@@ -384,6 +390,18 @@ class BaseDatabase(Generic[ConnectionT, SessionT]):
         after. `with` is right on either side, awaited or not: it listens, it does
         not run anything.
         """
+        if self._override is not None:
+            with self._override.recording(
+                label,
+                logger=logger,
+                echo=echo,
+                stacks=stacks,
+                skip_queries_from=skip_queries_from,
+                into=into,
+                send_to=send_to,
+            ) as recording:
+                yield recording
+            return
         recording = Recording(label=label) if into is None else into
         self._listen()
         recordings = self._recordings.set((*self._recordings.get(), recording))
@@ -507,6 +525,8 @@ class BaseDatabase(Generic[ConnectionT, SessionT]):
 
         False under ``connect()`` and ``autocommit()``, which open none.
         """
+        if self._override is not None:
+            return self._override.in_transaction()
         return self._outer.get(None) is not None
 
     @property
@@ -542,6 +562,8 @@ class BaseDatabase(Generic[ConnectionT, SessionT]):
         this is False in a block that has only run statements on the
         connection. Reading it opens nothing.
         """
+        if self._override is not None:
+            return self._override.in_session()
         scope = self._scope.get(None)
         return scope is not None and scope.session is not None
 
@@ -563,6 +585,10 @@ class BaseDatabase(Generic[ConnectionT, SessionT]):
         A block opened inside still joins the transaction around it, and rolls
         back with it.
         """
+        if self._override is not None:
+            with self._override.unbound():
+                yield
+            return
         scope = self._scope.get(None)
         if scope is None:
             yield
@@ -617,6 +643,25 @@ class BaseDatabase(Generic[ConnectionT, SessionT]):
         outer = self._outer_to_join()
         return outer, savepoint or rollback or bool(outer and outer.savepoint)
 
+    @contextmanager
+    def override(self, db: Self) -> Iterator[Self]:
+        """Use another database in place of this one until the block ends.
+
+        ```python
+        with Database(TEST_URL) as test_db, db.override(test_db):
+            create_user("ada")  # runs on test_db
+        ```
+
+        Every thread sees it. The database stays open when the block ends, so
+        close it yourself.
+        """
+        before = self._override
+        self._override = None if db is self else db
+        try:
+            yield db
+        finally:
+            self._override = before
+
     def _outer_to_join(self) -> _Outer[ConnectionT] | None:
         """Return the outer transaction new blocks join, if there is one."""
         outer = self._outer.get(None)
@@ -811,15 +856,20 @@ class _DatabaseRegistryMixin(BaseDatabase[Any, Any], Generic[DatabaseT]):
         """Leave everything to [`configure`][sqlakit.Databases.configure]."""
         self._default: DatabaseT | None = None
         self._aliased: dict[str, DatabaseT] = {}
+        self._overridden: dict[str, DatabaseT] = {}
+        self._override = None
         self._routers: tuple[Any, ...] = ()
         self._using: ContextVar[str | None] = ContextVar(
             f"{type(self).__name__}.using", default=None
         )
 
     def __repr__(self) -> str:
-        if not self.is_configured:
+        if self._built_its_own:
+            return super().__repr__()
+        held = self._override
+        if held is None:
             return f"{type(self).__name__}(unconfigured)"
-        return super().__repr__()
+        return f"{type(self).__name__}({held!r})"
 
     def __getitem__(self, alias: str) -> Self | DatabaseT:
         """Return the database configured as ``alias``.
@@ -831,6 +881,8 @@ class _DatabaseRegistryMixin(BaseDatabase[Any, Any], Generic[DatabaseT]):
             UnknownDatabaseError: if nothing is configured under that alias.
 
         """
+        if alias in self._overridden:
+            return self._overridden[alias]
         if alias == DEFAULT_ALIAS:
             return self if self._default is None else self._default
         try:
@@ -839,7 +891,11 @@ class _DatabaseRegistryMixin(BaseDatabase[Any, Any], Generic[DatabaseT]):
             raise UnknownDatabaseError(alias, self.aliases) from None
 
     def __contains__(self, alias: str) -> bool:
-        return alias == DEFAULT_ALIAS or alias in self._aliased
+        return (
+            alias == DEFAULT_ALIAS
+            or alias in self._aliased
+            or alias in self._overridden
+        )
 
     def register(self, alias: str, db: DatabaseT) -> None:
         """Put a database already built under an alias.
@@ -870,6 +926,7 @@ class _DatabaseRegistryMixin(BaseDatabase[Any, Any], Generic[DatabaseT]):
             if self.is_configured:
                 raise DefaultAliasError
             self._default = self._named(alias, db)
+            self._follow_default()
             return
         if alias in self._aliased:
             raise AliasInUseError(alias)
@@ -961,6 +1018,43 @@ class _DatabaseRegistryMixin(BaseDatabase[Any, Any], Generic[DatabaseT]):
         db._name = alias  # noqa: SLF001
         return db
 
+    @contextmanager
+    def override(
+        self,
+        db: DatabaseT,
+        *,
+        alias: str = DEFAULT_ALIAS,
+    ) -> Iterator[DatabaseT]:
+        """Use another database under an alias until the block ends.
+
+        ```python
+        with Database(TEST_URL) as test_db, db.override(test_db):
+            create_user("ada")  # runs on test_db
+        ```
+
+        Every thread sees it. The database stays open when the block ends, so
+        close it yourself.
+        """
+        before = self._overridden.get(alias)
+        named = getattr(db, "_name", None)
+        self._overridden[alias] = self._named(alias, db)
+        self._follow_default()
+        try:
+            yield db
+        finally:
+            if before is None:
+                del self._overridden[alias]
+            else:
+                self._overridden[alias] = before
+            self._follow_default()
+            if named is not None:
+                db._name = named  # noqa: SLF001
+
+    def _follow_default(self) -> None:
+        """Work on the database the default alias leads to, unless it is this one."""
+        held = self._overridden.get(DEFAULT_ALIAS, self._default)
+        self._override = None if held is self else held
+
     def using(self, target: str | DatabaseT) -> _Using:
         """Return that database, standing in for the default one.
 
@@ -997,6 +1091,16 @@ class _DatabaseRegistryMixin(BaseDatabase[Any, Any], Generic[DatabaseT]):
             UnregisteredDatabaseError: if it holds it under none.
 
         """
+        alias = self._held_as(db)
+        if alias is None:
+            raise UnregisteredDatabaseError(self.aliases)
+        return alias
+
+    def _held_as(self, db: object) -> str | None:
+        """Return the alias this registry holds a database under, or None."""
+        for alias, held in self._overridden.items():
+            if held is db:
+                return alias
         # A configured registry is the default database itself, and a
         # registered one holds it.
         if db is self or db is self._default:
@@ -1004,7 +1108,7 @@ class _DatabaseRegistryMixin(BaseDatabase[Any, Any], Generic[DatabaseT]):
         for alias, held in self._aliased.items():
             if held is db:
                 return alias
-        raise UnregisteredDatabaseError(self.aliases)
+        return None
 
     def route(self, *routers: Router | RouterFunction | str) -> None:
         """Say which database a model lives on, for models that do not say it.
@@ -1038,16 +1142,17 @@ class _DatabaseRegistryMixin(BaseDatabase[Any, Any], Generic[DatabaseT]):
         opened with `using()` stands in for the default database.
         """
         placement = self._routed(model) or model.__db__
-        override = self._using.get()
-        if isinstance(placement, str):
-            if override is not None and placement == DEFAULT_ALIAS:
-                placement = override
-            return self[placement]
-        # A model pinned to the database itself follows `using()` as one on the
-        # default alias does, when that database is the one being stood in for.
-        if override is not None and self[DEFAULT_ALIAS] is placement:
-            return self[override]
-        return placement
+        if not isinstance(placement, str):
+            # A model pinned to a database this registry holds follows the
+            # alias it is held under, `using()` and `override()` alike.
+            alias = self._held_as(placement)
+            if alias is None:
+                return placement
+            placement = alias
+        redirected = self._using.get()
+        if redirected is not None and placement == DEFAULT_ALIAS:
+            placement = redirected
+        return self[placement]
 
     def _routed(self, model: type[Any]) -> str | None:
         """Return what the first router says about this model, if anything."""
@@ -1059,13 +1164,18 @@ class _DatabaseRegistryMixin(BaseDatabase[Any, Any], Generic[DatabaseT]):
 
     @property
     def aliases(self) -> tuple[str, ...]:
-        """The aliases configured, the default one first."""
-        return (DEFAULT_ALIAS, *self._aliased)
+        """The aliases configured or overridden, the default one first."""
+        overridden = (
+            alias
+            for alias in self._overridden
+            if alias != DEFAULT_ALIAS and alias not in self._aliased
+        )
+        return (DEFAULT_ALIAS, *self._aliased, *overridden)
 
     @property
     def is_configured(self) -> bool:
         """Whether this registry has a default database to reach."""
-        return "url" in self.__dict__ or self._default is not None
+        return "url" in self.__dict__ or self._override is not None
 
     @property
     def _built_its_own(self) -> bool:
@@ -1073,39 +1183,13 @@ class _DatabaseRegistryMixin(BaseDatabase[Any, Any], Generic[DatabaseT]):
         return "url" in self.__dict__
 
     if not TYPE_CHECKING:
-        # Hidden from type checkers, which keep reading these off `Database`
-        # and its asyncio twin, signatures and all.
-        def _proxy(name: str, *, attribute: bool = False) -> Any:  # noqa: ANN401, N805
-            """Proxy to the database this registry holds, or call its own.
-
-            A registry handed a default proxies to that database. One that
-            `configure` built calls what it inherits, which `super()` reaches.
-            """
-
-            def reach(self: Any, *args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
-                held = self._default
-                found = (
-                    getattr(held, name) if held is not None else getattr(super(), name)
-                )
-                return found if attribute else found(*args, **kwargs)
-
-            return property(reach) if attribute else reach
-
-        connection = _proxy("connection", attribute=True)
-        engine = _proxy("engine", attribute=True)
-        session = _proxy("session", attribute=True)
-        sql = _proxy("sql", attribute=True)
-        assert_queries = _proxy("assert_queries")
-        autocommit = _proxy("autocommit")
-        connect = _proxy("connect")
-        in_session = _proxy("in_session")
-        in_transaction = _proxy("in_transaction")
-        ping = _proxy("ping")
-        provisioned_tables = _proxy("provisioned_tables")
-        query = _proxy("query")
-        session_factory = _proxy("session_factory")
-        transaction = _proxy("transaction")
-        del _proxy
+        # Hidden from type checkers, which read it off `Database`.
+        @property
+        def sql(self) -> Any:  # noqa: ANN401
+            """The templates of the default database, wherever it was built."""
+            if self._built_its_own or self._override is None:
+                return super().sql
+            return self._override.sql
 
     @overload
     def configure(
@@ -1375,6 +1459,17 @@ class BaseRetryingTransaction:
             raise AssertionError  # pragma: no cover
 
 
+def late_bound(source: object, registry: Any) -> Any:  # noqa: ANN401
+    """Return the database ``using=`` points at: the registry, an alias, or a call."""
+    if source is None:
+        return registry
+    if isinstance(source, str):
+        return registry[source]
+    if callable(source):
+        return source()
+    return source
+
+
 def retry_matches(exc: BaseException, retry_on: RetryOn) -> bool:
     """Whether ``retry_on`` claims this exception is worth another attempt."""
     if isinstance(retry_on, type | tuple):

{sqlakit-0.22.9 → sqlakit-0.22.11}/sqlakit/_db.py

@@ -98,6 +98,8 @@ class Database(BaseDatabase[sa.Connection, Session]):
     @property
     def engine(self) -> Engine:
         """The underlying engine, created on first access."""
+        if self._override is not None:
+            return self._override.engine
         if self._engine is None:
             with self._engine_lock:
                 # Two threads reaching this at once would each build one, and
@@ -122,6 +124,8 @@ class Database(BaseDatabase[sa.Connection, Session]):
             MissingConnectionError: if no connection is bound.
 
         """
+        if self._override is not None:
+            return self._override.connection
         return self._reused(self._current_scope())
 
     def _create_session(self, connection: sa.Connection) -> Session:
@@ -149,6 +153,11 @@ class Database(BaseDatabase[sa.Connection, Session]):
     @contextmanager
     def connect(self) -> Iterator[sa.Connection]:
         """Open a connection and bind it, or reuse the one already bound."""
+        elsewhere = self._override
+        if elsewhere is not None:
+            with elsewhere.connect() as connection:
+                yield connection
+            return
         reuse = self._scope_to_reuse()
         if reuse is not None:
             with self._bound(self._reused(reuse)) as connection:
@@ -191,6 +200,11 @@ class Database(BaseDatabase[sa.Connection, Session]):
     @contextmanager
     def _autocommit(self) -> Iterator[sa.Connection]:
         """Open a connection in ``AUTOCOMMIT`` and bind it, or join the outer one."""
+        elsewhere = self._override
+        if elsewhere is not None:
+            with elsewhere.autocommit() as connection:
+                yield connection
+            return
         outer = self._connection_to_join()
         if outer is not None:
             with self._bound(outer, commit=True) as connection:
@@ -321,6 +335,11 @@ class Database(BaseDatabase[sa.Connection, Session]):
         flush, as ``sessionmaker()`` does it. Inside another block it runs on
         the connection already bound.
         """
+        elsewhere = self._override
+        if elsewhere is not None:
+            with elsewhere.session_factory() as session:
+                yield session
+            return
         reuse = self._scope_to_reuse()
         if reuse is not None:
             with self._bound(self._reused(reuse)):
@@ -385,6 +404,10 @@ class Database(BaseDatabase[sa.Connection, Session]):
         Every table of the metadata unless ``tables`` names fewer, as a second
         database wants.
         """
+        if self._override is not None:
+            with self._override.provisioned_tables(metadata, tables=tables):
+                yield
+            return
         with self.transaction() as connection:
             metadata.create_all(connection, tables=tables)
         try:
@@ -395,6 +418,8 @@ class Database(BaseDatabase[sa.Connection, Session]):
 
     def ping(self) -> bool:
         """Whether the database answers."""
+        if self._override is not None:
+            return self._override.ping()
         try:
             with self.engine.connect() as connection:
                 connection.execute(sa.text("SELECT 1"))
@@ -472,6 +497,18 @@ class Transaction(ContextDecorator, AbstractContextManager["sa.Connection"]):
 
     def __enter__(self) -> sa.Connection:
         stack = ExitStack()
+        elsewhere = self.db._override  # noqa: SLF001
+        if elsewhere is not None:
+            connection = stack.enter_context(
+                elsewhere.transaction(
+                    savepoint=self.savepoint,
+                    join_nested=self.join_nested,
+                    rollback=self.rollback,
+                    commit_on_error=self.commit_on_error,
+                )
+            )
+            self._stacks.append(stack)
+            return connection
         try:
             outer, savepoint = self.db._plan(  # noqa: SLF001
                 savepoint=self.savepoint,

sqlakit-0.22.11/sqlakit/_registry.py

@@ -0,0 +1,215 @@
+from __future__ import annotations
+
+import functools
+from contextlib import ExitStack, contextmanager
+from typing import TYPE_CHECKING, Any, TypeVar, cast, overload
+
+from ._base import _DatabaseRegistryMixin, default_backoff, late_bound
+from ._db import Database, RetryingTransaction
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterator
+
+    from ._base import RetryOn
+    from ._db import Transaction
+
+__all__ = ["Databases", "autocommit", "db", "transaction"]
+
+_FuncT = TypeVar("_FuncT", bound="Callable[..., Any]")
+
+
+class Databases(_DatabaseRegistryMixin[Database], Database):
+    """The databases an application talks to, and the default one among them.
+
+    ```python
+    from sqlakit import db
+
+    db.configure(DB_URL)
+
+    db.session  # the default database
+    db["replica"].session  # another, if one was configured
+    ```
+
+    Use it when one database is enough and passing a handle around is not worth
+    it. `Database(url)` is the alternative.
+    """
+
+    _database_class = Database
+
+    @contextmanager
+    def transactions(
+        self,
+        **arguments: Any,  # noqa: ANN401
+    ) -> Iterator[None]:
+        """Open a transaction on every database, not the default one alone.
+
+        A single database needs `transaction(rollback=True)`. A test harness
+        with several needs this:
+
+        ```python
+        with db.transactions(rollback=True):
+            yield
+        ```
+        """
+        with ExitStack() as stack:
+            for alias in self.aliases:
+                stack.enter_context(self[alias].transaction(**arguments))
+            yield
+
+    def dispose(self, *, close: bool = True) -> None:
+        """Dispose of every configured database, not just the default one."""
+        if self._built_its_own:
+            super().dispose(close=close)
+        elif self._default is not None:
+            self._default.dispose(close=close)
+        for db in self._aliased.values():
+            db.dispose(close=close)
+
+
+db = Databases()
+"""The importable registry: one [`Databases`][sqlakit.Databases] for the process.
+
+`db.configure(url)` fills it, and every module then reaches the same connections
+by importing it. Reconfiguring is allowed until something connects, after which
+it raises
+[`DatabaseAlreadyConfiguredError`][sqlakit.DatabaseAlreadyConfiguredError]
+and `dispose()` has to come first.
+"""
+
+
+@overload
+def transaction(func: _FuncT, /) -> _FuncT: ...
+
+
+@overload
+def transaction(
+    *,
+    using: str | Database | Callable[[], Database] | None = None,
+    savepoint: bool = False,
+    join_nested: bool = True,
+    rollback: bool = False,
+    commit_on_error: type[BaseException]
+    | tuple[type[BaseException], ...]
+    | None = None,
+    retry_on: None = None,
+    max_retries: int = 3,
+    backoff: Callable[[int], float] = default_backoff,
+) -> Callable[[_FuncT], _FuncT]: ...
+
+
+@overload
+def transaction(
+    *,
+    using: str | Database | Callable[[], Database] | None = None,
+    savepoint: bool = False,
+    join_nested: bool = True,
+    rollback: bool = False,
+    commit_on_error: type[BaseException]
+    | tuple[type[BaseException], ...]
+    | None = None,
+    retry_on: RetryOn,
+    max_retries: int = 3,
+    backoff: Callable[[int], float] = default_backoff,
+) -> RetryingTransaction: ...
+
+
+def transaction(  # noqa: PLR0913  (all keyword-only, as on `Database`)
+    func: _FuncT | None = None,
+    /,
+    *,
+    using: str | Database | Callable[[], Database] | None = None,
+    savepoint: bool = False,
+    join_nested: bool = True,
+    rollback: bool = False,
+    commit_on_error: type[BaseException]
+    | tuple[type[BaseException], ...]
+    | None = None,
+    retry_on: RetryOn | None = None,
+    max_retries: int = 3,
+    backoff: Callable[[int], float] = default_backoff,
+) -> _FuncT | Callable[[_FuncT], _FuncT] | RetryingTransaction:
+    """Like `Database.transaction`, but find the database on every call.
+
+    ```python
+    @transaction
+    def import_users() -> None: ...
+
+
+    @transaction(using="replica")
+    def rebuild_report() -> None: ...
+
+
+    @transaction(using=lambda: container.resolve(Database))
+    def create_user(name: str) -> None: ...
+    ```
+
+    ``using`` is an alias in [`sqlakit.db`][sqlakit.db], or a function that
+    returns the database. Without it, the default database of `sqlakit.db`.
+    The other arguments are those of `Database.transaction`.
+    """
+
+    def opened() -> Transaction:
+        found: Database = late_bound(using, db)
+        return found.transaction(
+            savepoint=savepoint,
+            join_nested=join_nested,
+            rollback=rollback,
+            commit_on_error=commit_on_error,
+        )
+
+    decorate: Callable[[_FuncT], _FuncT] | RetryingTransaction = (
+        _opening(opened)
+        if retry_on is None
+        else RetryingTransaction(
+            opened, retry_on=retry_on, max_retries=max_retries, backoff=backoff
+        )
+    )
+    if func is not None:
+        return decorate(func)
+    return decorate
+
+
+@overload
+def autocommit(func: _FuncT, /) -> _FuncT: ...
+
+
+@overload
+def autocommit(
+    *,
+    using: str | Database | Callable[[], Database] | None = None,
+) -> Callable[[_FuncT], _FuncT]: ...
+
+
+def autocommit(
+    func: _FuncT | None = None,
+    /,
+    *,
+    using: str | Database | Callable[[], Database] | None = None,
+) -> _FuncT | Callable[[_FuncT], _FuncT]:
+    """Like `Database.autocommit`, but find the database on every call.
+
+    ``using`` works as in [`transaction`][sqlakit.transaction].
+    """
+
+    def opened() -> Any:  # noqa: ANN401
+        found: Database = late_bound(using, db)
+        return found.autocommit()
+
+    decorate: Callable[[_FuncT], _FuncT] = _opening(opened)
+    if func is not None:
+        return decorate(func)
+    return decorate
+
+
+def _opening(opened: Callable[[], Any]) -> Callable[[_FuncT], _FuncT]:
+    """Return a decorator that runs the function inside `opened()`."""
+
+    def decorate(func: _FuncT) -> _FuncT:
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
+            with opened():
+                return func(*args, **kwargs)
+
+        return cast("_FuncT", wrapper)
+
+    return decorate

sqlakit-0.22.11/sqlakit/asyncio/__init__.py

@@ -0,0 +1,12 @@
+from ._db import Database, RetryingTransaction, Transaction
+from ._registry import Databases, autocommit, db, transaction
+
+__all__ = [
+    "Database",
+    "Databases",
+    "RetryingTransaction",
+    "Transaction",
+    "autocommit",
+    "db",
+    "transaction",
+]

{sqlakit-0.22.9 → sqlakit-0.22.11}/sqlakit/asyncio/_db.py

@@ -102,6 +102,8 @@ class Database(BaseDatabase[AsyncConnection, AsyncSession]):
     @property
     def engine(self) -> AsyncEngine:
         """The underlying engine, created on first access."""
+        if self._override is not None:
+            return self._override.engine
         if self._engine is None:
             with self._engine_lock:
                 # Two threads reaching this at once would each build one, and
@@ -127,6 +129,8 @@ class Database(BaseDatabase[AsyncConnection, AsyncSession]):
             MissingConnectionError: if no connection is bound.
 
         """
+        if self._override is not None:
+            return self._override.connection
         scope = self._current_scope()
         if scope.connection is None:
             cell = scope.checkout
@@ -143,6 +147,8 @@ class Database(BaseDatabase[AsyncConnection, AsyncSession]):
 
     async def _aconnection(self) -> AsyncConnection:
         """Return the bound connection, checking it out first if lazy."""
+        if self._override is not None:
+            return await self._override._aconnection()  # noqa: SLF001
         return await self._areused(self._current_scope())
 
     def _create_session(self, connection: AsyncConnection) -> AsyncSession:
@@ -183,6 +189,11 @@ class Database(BaseDatabase[AsyncConnection, AsyncSession]):
     @asynccontextmanager
     async def connect(self) -> AsyncIterator[AsyncConnection]:
         """Open a connection and bind it, or reuse the one already bound."""
+        elsewhere = self._override
+        if elsewhere is not None:
+            async with elsewhere.connect() as connection:
+                yield connection
+            return
         reuse = self._scope_to_reuse()
         if reuse is not None:
             async with self._bound(await self._areused(reuse)) as connection:
@@ -225,6 +236,11 @@ class Database(BaseDatabase[AsyncConnection, AsyncSession]):
     @asynccontextmanager
     async def _autocommit(self) -> AsyncIterator[AsyncConnection]:
         """Open a connection in ``AUTOCOMMIT`` and bind it, or join the outer one."""
+        elsewhere = self._override
+        if elsewhere is not None:
+            async with elsewhere.autocommit() as connection:
+                yield connection
+            return
         outer = self._connection_to_join()
         if outer is not None:
             async with self._bound(outer, commit=True) as connection:
@@ -355,6 +371,11 @@ class Database(BaseDatabase[AsyncConnection, AsyncSession]):
         flush, as ``async_sessionmaker()`` does it. Inside another block it
         runs on the connection already bound.
         """
+        elsewhere = self._override
+        if elsewhere is not None:
+            async with elsewhere.session_factory() as session:
+                yield session
+            return
         reuse = self._scope_to_reuse()
         if reuse is not None:
             async with self._bound(await self._areused(reuse)):
@@ -419,6 +440,10 @@ class Database(BaseDatabase[AsyncConnection, AsyncSession]):
         Every table of the metadata unless ``tables`` names fewer, as a second
         database wants.
         """
+        if self._override is not None:
+            async with self._override.provisioned_tables(metadata, tables=tables):
+                yield
+            return
         async with self.transaction() as connection:
             await connection.run_sync(metadata.create_all, tables=tables)
         try:
@@ -429,6 +454,8 @@ class Database(BaseDatabase[AsyncConnection, AsyncSession]):
 
     async def ping(self) -> bool:
         """Whether the database answers."""
+        if self._override is not None:
+            return await self._override.ping()
         try:
             async with self.engine.connect() as connection:
                 await connection.execute(sa.text("SELECT 1"))
@@ -512,6 +539,18 @@ class Transaction(
 
     async def __aenter__(self) -> AsyncConnection:
         stack = AsyncExitStack()
+        elsewhere = self.db._override  # noqa: SLF001
+        if elsewhere is not None:
+            connection = await stack.enter_async_context(
+                elsewhere.transaction(
+                    savepoint=self.savepoint,
+                    join_nested=self.join_nested,
+                    rollback=self.rollback,
+                    commit_on_error=self.commit_on_error,
+                )
+            )
+            self._stacks.append(stack)
+            return connection
         try:
             outer, savepoint = self.db._plan(  # noqa: SLF001
                 savepoint=self.savepoint,

sqlakit-0.22.11/sqlakit/asyncio/_registry.py

@@ -0,0 +1,189 @@
+from __future__ import annotations
+
+import functools
+from contextlib import AsyncExitStack, asynccontextmanager
+from typing import TYPE_CHECKING, Any, TypeVar, cast, overload
+
+from sqlakit._base import _DatabaseRegistryMixin, default_backoff, late_bound
+
+from ._db import Database, RetryingTransaction
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator, Callable
+
+    from sqlakit._base import RetryOn
+
+    from ._db import Transaction
+
+__all__ = ["Databases", "autocommit", "db", "transaction"]
+
+_FuncT = TypeVar("_FuncT", bound="Callable[..., Any]")
+
+
+class Databases(_DatabaseRegistryMixin[Database], Database):
+    """The databases an application talks to, and the default one among them.
+
+    The asyncio counterpart of [`sqlakit.Databases`][sqlakit.Databases].
+
+    ```python
+    from sqlakit.asyncio import db
+
+    db.configure(DB_URL)
+
+    db.session  # the default database
+    db["replica"].session  # another, if one was configured
+    ```
+    """
+
+    _database_class = Database
+
+    @asynccontextmanager
+    async def transactions(
+        self,
+        **arguments: Any,  # noqa: ANN401
+    ) -> AsyncIterator[None]:
+        """Open a transaction on every database, not the default one alone.
+
+        A test harness with more than one database needs this, where a single
+        database needs `transaction(rollback=True)`.
+        """
+        async with AsyncExitStack() as stack:
+            for alias in self.aliases:
+                await stack.enter_async_context(self[alias].transaction(**arguments))
+            yield
+
+    async def dispose(self, *, close: bool = True) -> None:
+        """Dispose of every configured database, not just the default one."""
+        if self._built_its_own:
+            await super().dispose(close=close)
+        elif self._default is not None:
+            await self._default.dispose(close=close)
+        for db in self._aliased.values():
+            await db.dispose(close=close)
+
+
+db = Databases()
+"""The database an application talks to. Configure it once, use it anywhere."""
+
+
+@overload
+def transaction(func: _FuncT, /) -> _FuncT: ...
+
+
+@overload
+def transaction(
+    *,
+    using: str | Database | Callable[[], Database] | None = None,
+    savepoint: bool = False,
+    join_nested: bool = True,
+    rollback: bool = False,
+    commit_on_error: type[BaseException]
+    | tuple[type[BaseException], ...]
+    | None = None,
+    retry_on: None = None,
+    max_retries: int = 3,
+    backoff: Callable[[int], float] = default_backoff,
+) -> Callable[[_FuncT], _FuncT]: ...
+
+
+@overload
+def transaction(
+    *,
+    using: str | Database | Callable[[], Database] | None = None,
+    savepoint: bool = False,
+    join_nested: bool = True,
+    rollback: bool = False,
+    commit_on_error: type[BaseException]
+    | tuple[type[BaseException], ...]
+    | None = None,
+    retry_on: RetryOn,
+    max_retries: int = 3,
+    backoff: Callable[[int], float] = default_backoff,
+) -> RetryingTransaction: ...
+
+
+def transaction(  # noqa: PLR0913  (all keyword-only, as on `Database`)
+    func: _FuncT | None = None,
+    /,
+    *,
+    using: str | Database | Callable[[], Database] | None = None,
+    savepoint: bool = False,
+    join_nested: bool = True,
+    rollback: bool = False,
+    commit_on_error: type[BaseException]
+    | tuple[type[BaseException], ...]
+    | None = None,
+    retry_on: RetryOn | None = None,
+    max_retries: int = 3,
+    backoff: Callable[[int], float] = default_backoff,
+) -> _FuncT | Callable[[_FuncT], _FuncT] | RetryingTransaction:
+    """Like `Database.transaction`, but find the database on every call.
+
+    The asyncio version of [`sqlakit.transaction`][sqlakit.transaction].
+    """
+
+    def opened() -> Transaction:
+        found: Database = late_bound(using, db)
+        return found.transaction(
+            savepoint=savepoint,
+            join_nested=join_nested,
+            rollback=rollback,
+            commit_on_error=commit_on_error,
+        )
+
+    decorate: Callable[[_FuncT], _FuncT] | RetryingTransaction = (
+        _opening(opened)
+        if retry_on is None
+        else RetryingTransaction(
+            opened, retry_on=retry_on, max_retries=max_retries, backoff=backoff
+        )
+    )
+    if func is not None:
+        return decorate(func)
+    return decorate
+
+
+@overload
+def autocommit(func: _FuncT, /) -> _FuncT: ...
+
+
+@overload
+def autocommit(
+    *,
+    using: str | Database | Callable[[], Database] | None = None,
+) -> Callable[[_FuncT], _FuncT]: ...
+
+
+def autocommit(
+    func: _FuncT | None = None,
+    /,
+    *,
+    using: str | Database | Callable[[], Database] | None = None,
+) -> _FuncT | Callable[[_FuncT], _FuncT]:
+    """Like `Database.autocommit`, but find the database on every call.
+
+    The asyncio version of [`sqlakit.autocommit`][sqlakit.autocommit].
+    """
+
+    def opened() -> Any:  # noqa: ANN401
+        found: Database = late_bound(using, db)
+        return found.autocommit()
+
+    decorate: Callable[[_FuncT], _FuncT] = _opening(opened)
+    if func is not None:
+        return decorate(func)
+    return decorate
+
+
+def _opening(opened: Callable[[], Any]) -> Callable[[_FuncT], _FuncT]:
+    """Return a decorator that runs the function inside `opened()`."""
+
+    def decorate(func: _FuncT) -> _FuncT:
+        @functools.wraps(func)
+        async def wrapper(*args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
+            async with opened():
+                return await func(*args, **kwargs)
+
+        return cast("_FuncT", wrapper)
+
+    return decorate

sqlakit-0.22.9/sqlakit/_registry.py

@@ -1,71 +0,0 @@
-from __future__ import annotations
-
-from contextlib import ExitStack, contextmanager
-from typing import TYPE_CHECKING, Any
-
-from ._base import _DatabaseRegistryMixin
-from ._db import Database
-
-if TYPE_CHECKING:
-    from collections.abc import Iterator
-
-__all__ = ["Databases", "db"]
-
-
-class Databases(_DatabaseRegistryMixin[Database], Database):
-    """The databases an application talks to, and the default one among them.
-
-    ```python
-    from sqlakit import db
-
-    db.configure(DB_URL)
-
-    db.session  # the default database
-    db["replica"].session  # another, if one was configured
-    ```
-
-    Use it when one database is enough and passing a handle around is not worth
-    it. `Database(url)` is the alternative.
-    """
-
-    _database_class = Database
-
-    @contextmanager
-    def transactions(
-        self,
-        **arguments: Any,  # noqa: ANN401
-    ) -> Iterator[None]:
-        """Open a transaction on every database, not the default one alone.
-
-        A single database needs `transaction(rollback=True)`. A test harness
-        with several needs this:
-
-        ```python
-        with db.transactions(rollback=True):
-            yield
-        ```
-        """
-        with ExitStack() as stack:
-            for alias in self.aliases:
-                stack.enter_context(self[alias].transaction(**arguments))
-            yield
-
-    def dispose(self, *, close: bool = True) -> None:
-        """Dispose of every configured database, not just the default one."""
-        if self._built_its_own:
-            super().dispose(close=close)
-        elif self._default is not None:
-            self._default.dispose(close=close)
-        for db in self._aliased.values():
-            db.dispose(close=close)
-
-
-db = Databases()
-"""The importable registry: one [`Databases`][sqlakit.Databases] for the process.
-
-`db.configure(url)` fills it, and every module then reaches the same connections
-by importing it. Reconfiguring is allowed until something connects, after which
-it raises
-[`DatabaseAlreadyConfiguredError`][sqlakit.DatabaseAlreadyConfiguredError]
-and `dispose()` has to come first.
-"""

sqlakit-0.22.9/sqlakit/asyncio/__init__.py

@@ -1,4 +0,0 @@
-from ._db import Database, RetryingTransaction, Transaction
-from ._registry import Databases, db
-
-__all__ = ["Database", "Databases", "RetryingTransaction", "Transaction", "db"]

sqlakit-0.22.9/sqlakit/asyncio/_registry.py

@@ -1,59 +0,0 @@
-from __future__ import annotations
-
-from contextlib import AsyncExitStack, asynccontextmanager
-from typing import TYPE_CHECKING, Any
-
-from sqlakit._base import _DatabaseRegistryMixin
-
-from ._db import Database
-
-if TYPE_CHECKING:
-    from collections.abc import AsyncIterator
-
-__all__ = ["Databases", "db"]
-
-
-class Databases(_DatabaseRegistryMixin[Database], Database):
-    """The databases an application talks to, and the default one among them.
-
-    The asyncio counterpart of [`sqlakit.Databases`][sqlakit.Databases].
-
-    ```python
-    from sqlakit.asyncio import db
-
-    db.configure(DB_URL)
-
-    db.session  # the default database
-    db["replica"].session  # another, if one was configured
-    ```
-    """
-
-    _database_class = Database
-
-    @asynccontextmanager
-    async def transactions(
-        self,
-        **arguments: Any,  # noqa: ANN401
-    ) -> AsyncIterator[None]:
-        """Open a transaction on every database, not the default one alone.
-
-        A test harness with more than one database needs this, where a single
-        database needs `transaction(rollback=True)`.
-        """
-        async with AsyncExitStack() as stack:
-            for alias in self.aliases:
-                await stack.enter_async_context(self[alias].transaction(**arguments))
-            yield
-
-    async def dispose(self, *, close: bool = True) -> None:
-        """Dispose of every configured database, not just the default one."""
-        if self._built_its_own:
-            await super().dispose(close=close)
-        elif self._default is not None:
-            await self._default.dispose(close=close)
-        for db in self._aliased.values():
-            await db.dispose(close=close)
-
-
-db = Databases()
-"""The database an application talks to. Configure it once, use it anywhere."""