sqlakit 0.1.0__py3-none-any.whl

sqlakit/asyncio/_db.py

@@ -0,0 +1,552 @@
+from __future__ import annotations
+
+import asyncio
+import functools
+import itertools
+import logging
+from contextlib import (
+    AbstractAsyncContextManager,
+    AsyncContextDecorator,
+    AsyncExitStack,
+    asynccontextmanager,
+)
+from functools import cached_property
+from typing import TYPE_CHECKING, Any, Self, TypeVar, cast, overload
+
+import sqlalchemy as sa
+import sqlalchemy.exc
+from sqlalchemy.ext.asyncio import (
+    AsyncConnection,
+    AsyncEngine,
+    AsyncSession,
+    async_sessionmaker,
+    create_async_engine,
+)
+
+from sqlakit._base import (
+    BaseDatabase,
+    BaseRetryingTransaction,
+    default_backoff,
+    fix_sqlite_transactions,
+)
+from sqlakit.exceptions import TransactionRolledBackError
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator, Callable, Coroutine, Sequence
+    from types import TracebackType
+
+    from sqlalchemy.ext.asyncio import AsyncTransaction
+
+    from sqlakit._base import RetryOn, _Scope
+
+    from .orm import Query
+    from .sql import SQL
+
+ModelT = TypeVar("ModelT")
+_FuncT = TypeVar("_FuncT", bound="Callable[..., Any]")
+
+__all__ = ["Database", "RetryingTransaction", "Transaction"]
+
+logger = logging.getLogger("sqlakit")
+
+
+class Database(BaseDatabase[AsyncConnection, AsyncSession]):
+    """The asyncio counterpart of [`sqlakit.Database`][sqlakit.Database].
+
+    Scoping is the same: connections and sessions live in the context, tasks
+    inherit them, transactions nest the same way. Every method that touches the
+    database is awaited.
+    """
+
+    _engine: AsyncEngine | None = None
+    _sessionmaker: async_sessionmaker[AsyncSession] | None = None
+
+    @cached_property
+    def sql(self) -> SQL:
+        """The SQL templates of this database.
+
+        ```python
+        db = Database(DB_URL, templates="app/sql")
+
+        await db.sql("reports/by_team.sql", since=since).typed(TeamReport).all()
+        ```
+
+        Templates need `sqlakit[sql]`, and nothing else here does, so the layer is
+        imported when it is first reached.
+        """
+        # Here rather than at the top, so that `import sqlakit` stays free of
+        # the optional layer and of what it imports.
+        from .sql import SQL  # noqa: PLC0415
+
+        return SQL(self)
+
+    def query(self, model: type[ModelT]) -> Query[ModelT]:
+        """Build a query over a mapped class, on this database.
+
+        ```python
+        await db.query(User).where(User.is_active).order_by("name").page(limit=20)
+        ```
+
+        Any declarative class works, with no model layer under it. A model
+        that has one reaches the same builder as `User.query`, on the database
+        the model belongs to.
+        """
+        # Here rather than at the top: `orm` imports this module.
+        from .orm import Query  # noqa: PLC0415
+
+        return Query(model, self)
+
+    @property
+    def engine(self) -> AsyncEngine:
+        """The underlying engine, created on first access."""
+        if self._engine is None:
+            with self._engine_lock:
+                # Two threads reaching this at once would each build one, and
+                # the pool of the loser would never be disposed.
+                if self._engine is None:
+                    self._engine = self._create_engine()
+        return self._engine
+
+    def _create_engine(self) -> AsyncEngine:
+        engine = create_async_engine(self.url, **self.engine_args)
+        fix_sqlite_transactions(engine.sync_engine)
+        return engine
+
+    def _create_session(self, connection: AsyncConnection) -> AsyncSession:
+        if self._sessionmaker is None:
+            self._sessionmaker = async_sessionmaker(**self.session_args)
+        return self._sessionmaker(
+            bind=connection,
+            **self._session_args_for(connection),
+        )
+
+    @asynccontextmanager
+    async def connect(self) -> AsyncIterator[AsyncConnection]:
+        """Open a connection and bind it, or reuse the one already bound."""
+        outer = self._connection_to_reuse()
+        if outer is not None:
+            async with self._bound(outer) as connection:
+                yield connection
+            return
+        async with self.engine.connect() as opened, self._bound(opened) as connection:
+            yield connection
+
+    @overload
+    def autocommit(self, func: _FuncT) -> _FuncT: ...
+
+    @overload
+    def autocommit(
+        self,
+        func: None = None,
+    ) -> AbstractAsyncContextManager[AsyncConnection]: ...
+
+    def autocommit(
+        self,
+        func: _FuncT | None = None,
+    ) -> _FuncT | AbstractAsyncContextManager[AsyncConnection]:
+        """Run in ``AUTOCOMMIT``, where every statement commits on its own.
+
+        For read-only work, which has nothing to commit, and for statements that
+        cannot run inside a transaction: ``VACUUM``, ``CREATE DATABASE``,
+        ``CREATE INDEX CONCURRENTLY``. Inside a transaction it joins that one, since
+        there is nothing else to commit into.
+
+        A context manager and a decorator, with or without parentheses.
+
+        Args:
+            func: The function to decorate, when used as a bare decorator.
+
+        """
+        autocommit = self._autocommit()
+        if func is not None:
+            return autocommit(func)
+        return autocommit
+
+    @asynccontextmanager
+    async def _autocommit(self) -> AsyncIterator[AsyncConnection]:
+        """Open a connection in ``AUTOCOMMIT`` and bind it, or join the outer one."""
+        outer = self._connection_to_join()
+        if outer is not None:
+            async with self._bound(outer, commit=True) as connection:
+                yield connection
+            return
+        async with self.engine.connect() as opened:
+            await opened.execution_options(isolation_level="AUTOCOMMIT")
+            with self._set_outer(None):
+                async with self._bound(opened, commit=True) as connection:
+                    yield connection
+
+    @overload
+    def transaction(self, func: _FuncT) -> _FuncT: ...
+
+    @overload
+    def transaction(
+        self,
+        func: 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,
+    ) -> Transaction: ...
+
+    @overload
+    def transaction(
+        self,
+        func: 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; this is the main API)
+        self,
+        func: _FuncT | 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 | Transaction | RetryingTransaction:
+        """Run a transaction on a connection bound to the current context.
+
+        Commits when the block exits, rolls back if it raises. Inside another
+        transaction it takes part in that one rather than opening a second
+        connection: the outermost block commits.
+
+        A context manager and a decorator, the latter with or without parentheses:
+
+        ```python
+        async with db.transaction():
+            ...
+
+
+        @db.transaction
+        async def import_users() -> None: ...
+        ```
+
+        Args:
+            func: The function to decorate, when used as a bare decorator.
+            savepoint: Run as a savepoint when nested, so this block can fail and be
+                rolled back on its own, and the blocks below it with it. Off by
+                default: a savepoint costs a round trip.
+            join_nested: Whether blocks below reuse this connection. Turn it off to
+                let them reach the database on their own, seeing nothing of this
+                block and surviving its rollback.
+            rollback: Roll back on the way out rather than commit. Implies
+                ``savepoint``, and is what wraps a test.
+            commit_on_error: Exception types whose escape still commits. The
+                exception propagates; what was written before it stays.
+            retry_on: Exception types, or a predicate over the exception, worth
+                another attempt. Decorator only: retrying re-runs the block, so this
+                returns a [`RetryingTransaction`][sqlakit.asyncio.RetryingTransaction] that
+                type checkers refuse to
+                enter. Only the block that owns the transaction retries.
+            max_retries: How many further attempts ``retry_on`` may buy.
+            backoff: Seconds to wait before attempt ``n``, counted from zero.
+
+        """
+
+        def new_transaction() -> Transaction:
+            return Transaction(
+                self,
+                savepoint=savepoint,
+                join_nested=join_nested,
+                rollback=rollback,
+                commit_on_error=commit_on_error,
+            )
+
+        transaction: Transaction | RetryingTransaction = new_transaction()
+        if retry_on is not None:
+            transaction = RetryingTransaction(
+                new_transaction,
+                retry_on=retry_on,
+                max_retries=max_retries,
+                backoff=backoff,
+            )
+        if func is not None:
+            return transaction(func)
+        return transaction
+
+    @asynccontextmanager
+    async def session_factory(self) -> AsyncIterator[AsyncSession]:
+        """Open a new connection and a session on top of it, and bind both."""
+        async with self.connect():
+            yield self.session
+
+    @asynccontextmanager
+    async def _bound(
+        self,
+        connection: AsyncConnection,
+        *,
+        commit: bool = False,
+    ) -> AsyncIterator[AsyncConnection]:
+        """Bind ``connection`` for the block, ending the session it opened.
+
+        ``commit`` keeps that session's work. Closing a session that isolates
+        itself with a savepoint rolls back to it, discarding what the blocks
+        below committed.
+        """
+        with self._bind(connection) as scope:
+            done = False
+            try:
+                yield connection
+                done = True
+            finally:
+                if scope.session is not None:
+                    if commit and done:
+                        await scope.session.commit()
+                    await scope.session.close()
+
+    @asynccontextmanager
+    async def provisioned_tables(
+        self,
+        metadata: sa.MetaData,
+        *,
+        tables: Sequence[sa.Table] | None = None,
+    ) -> AsyncIterator[None]:
+        """Create these tables here, and drop them when the block ends.
+
+        What a test session opens once, around everything that needs a schema:
+
+        ```python
+        @pytest.fixture(scope="session")
+        async def tables():
+            async with db.provisioned_tables(Model.metadata):
+                yield
+        ```
+
+        Every table of the metadata unless ``tables`` names fewer, which is what a
+        second database wants.
+        """
+        async with self.transaction() as connection:
+            await connection.run_sync(metadata.create_all, tables=tables)
+        try:
+            yield
+        finally:
+            async with self.transaction() as connection:
+                await connection.run_sync(metadata.drop_all, tables=tables)
+
+    async def ping(self) -> bool:
+        """Whether the database answers."""
+        try:
+            async with self.engine.connect() as connection:
+                await connection.execute(sa.text("SELECT 1"))
+        except sa.exc.SQLAlchemyError:
+            return False
+        return True
+
+    async def dispose(self, *, close: bool = True) -> None:
+        """Dispose of the engine and its connection pool."""
+        with self._engine_lock:
+            engine, self._engine = self._engine, None
+        if engine is not None:
+            await engine.dispose(close=close)
+
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(self, *exc_info: object) -> None:
+        await self.dispose()
+
+
+class Transaction(
+    AsyncContextDecorator,
+    AbstractAsyncContextManager["AsyncConnection"],
+):
+    """The transaction an awaited block runs, as an object.
+
+    What [`Database.transaction`][sqlakit.asyncio.Database.transaction] returns.
+
+    Kept as a class rather than a generator so that it also works as a
+    decorator, and so that a single instance can be entered more than once,
+    which is what ``@db.transaction()`` on a coroutine function does.
+    """
+
+    def __init__(
+        self,
+        db: Database,
+        *,
+        savepoint: bool = False,
+        join_nested: bool = True,
+        rollback: bool = False,
+        commit_on_error: type[BaseException]
+        | tuple[type[BaseException], ...]
+        | None = None,
+    ) -> None:
+        self.db = db
+        self.savepoint = savepoint
+        self.join_nested = join_nested
+        self.rollback = rollback
+        self.commit_on_error = commit_on_error
+        self._stacks: list[AsyncExitStack] = []
+
+    def _recreate_cm(self) -> Transaction:
+        """Give every decorated call a transaction of its own.
+
+        `AsyncContextDecorator` reuses one instance for all calls, and two of
+        them running at once would unwind each other's blocks.
+        """
+        return Transaction(
+            self.db,
+            savepoint=self.savepoint,
+            join_nested=self.join_nested,
+            rollback=self.rollback,
+            commit_on_error=self.commit_on_error,
+        )
+
+    async def __aenter__(self) -> AsyncConnection:
+        stack = AsyncExitStack()
+        try:
+            outer, savepoint = self.db._plan(  # noqa: SLF001
+                savepoint=self.savepoint,
+                rollback=self.rollback,
+            )
+            if outer is not None:
+                connection = outer.connection
+                # 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
+            else:
+                connection = await stack.enter_async_context(self.db.engine.connect())
+                transaction = await connection.begin()
+                session_savepoint = savepoint
+            # Unwound in reverse: session, context, transaction, connection.
+            stack.push_async_exit(self._finish(transaction))
+            bound = stack.enter_context(
+                self.db._set_outer(  # noqa: SLF001
+                    connection,
+                    join_nested=self.join_nested,
+                    savepoint=savepoint,
+                    session_savepoint=session_savepoint,
+                )
+            )
+            scope = stack.enter_context(self.db._bind(connection))  # noqa: SLF001
+            if bound is not None:
+                bound.scope = scope
+            stack.push_async_exit(self._close_session(scope))
+        except BaseException:
+            await stack.aclose()
+            raise
+        self._stacks.append(stack)
+        return connection
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        await self._stacks.pop().__aexit__(exc_type, exc, traceback)
+
+    def _close_session(
+        self,
+        scope: _Scope[AsyncConnection, AsyncSession],
+    ) -> Callable[..., Coroutine[None, None, None]]:
+        """Commit the block's session, if it opened one, then close it.
+
+        The block is the unit of work, so what its session holds belongs to the
+        transaction; closing first would roll it back.
+        """
+
+        async def close_session(
+            _exc_type: object,
+            exc: BaseException | None,
+            _traceback: object,
+        ) -> None:
+            if scope.session is None:
+                return
+            if self._keeps(exc):
+                await scope.session.commit()
+            await scope.session.close()
+
+        return close_session
+
+    def _finish(
+        self,
+        transaction: AsyncTransaction | None,
+    ) -> Callable[..., Coroutine[None, None, None]]:
+        """Commit or roll back, unless this block only takes part in another."""
+
+        async def finish(
+            _exc_type: object,
+            exc: BaseException | None,
+            _traceback: object,
+        ) -> None:
+            if transaction is None:
+                return
+            if not transaction.is_active:
+                # Rolled back from inside the block. Say so, unless an
+                # exception is already on its way out with the reason.
+                if exc is None:
+                    raise TransactionRolledBackError
+                return
+            if self._keeps(exc) and not self.rollback:
+                await transaction.commit()
+            else:
+                await transaction.rollback()
+
+        return finish
+
+    def _keeps(self, exc: BaseException | None) -> bool:
+        """Whether the block's work is kept rather than undone."""
+        if exc is None:
+            return True
+        return self.commit_on_error is not None and isinstance(
+            exc, self.commit_on_error
+        )
+
+
+class RetryingTransaction(BaseRetryingTransaction):
+    """A transaction that runs its block again.
+
+    What [`Database.transaction`][sqlakit.asyncio.Database.transaction] returns
+    when it is given ``retry_on``.
+    """
+
+    def __call__(self, func: _FuncT) -> _FuncT:
+        @functools.wraps(func)
+        async def wrapper(*args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
+            transaction = self.transaction()
+            if transaction.db.in_transaction():
+                # Only the block that owns the transaction can restart it: a
+                # retry inside would keep the snapshot that caused the
+                # conflict, and the outer transaction fails anyway.
+                logger.debug(
+                    "%s runs inside another transaction; retrying is up to "
+                    "whoever opened it.",
+                    getattr(func, "__qualname__", func),
+                )
+                async with transaction:
+                    return await func(*args, **kwargs)
+            for attempt in itertools.count():
+                try:
+                    async with self.transaction():
+                        return await func(*args, **kwargs)
+                except Exception as exc:
+                    if not self._retry(exc, attempt=attempt):
+                        raise
+                    await asyncio.sleep(self.backoff(attempt))
+            raise AssertionError  # pragma: no cover - the loop returns or raises
+
+        return cast("_FuncT", wrapper)

sqlakit/asyncio/_registry.py

@@ -0,0 +1,57 @@
+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.
+
+        What a test harness with more than one database needs, in the place 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.is_configured:
+            await super().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."""