From 9698a0743bcd65518094f8ba1c6006f2fbacc67e Mon Sep 17 00:00:00 2001 From: d3vyce <44915747+d3vyce@users.noreply.github.com> Date: Thu, 25 Jun 2026 21:11:50 +0200 Subject: [PATCH] chore: rework DB module (#324) --- docs/module/crud.md | 2 +- docs/module/db.md | 151 ++-- docs/module/models.md | 2 +- docs/module/pytest.md | 4 +- docs/reference/db.md | 38 +- docs_src/examples/pagination_search/app.py | 2 + docs_src/examples/pagination_search/db.py | 13 +- src/fastapi_toolsets/__init__.py | 7 +- src/fastapi_toolsets/crud/factory.py | 10 +- src/fastapi_toolsets/db.py | 591 --------------- src/fastapi_toolsets/db/__init__.py | 18 + src/fastapi_toolsets/db/core.py | 315 ++++++++ src/fastapi_toolsets/db/locks.py | 185 +++++ src/fastapi_toolsets/db/m2m.py | 170 +++++ src/fastapi_toolsets/db/testing.py | 69 ++ src/fastapi_toolsets/db/watch.py | 90 +++ src/fastapi_toolsets/fixtures/utils.py | 4 +- src/fastapi_toolsets/pytest/plugin.py | 4 +- src/fastapi_toolsets/pytest/utils.py | 2 +- tests/test_db.py | 814 ++++++++++++++++----- tests/test_example_pagination_search.py | 18 +- tests/test_models.py | 70 +- tests/test_pytest.py | 14 +- 23 files changed, 1662 insertions(+), 931 deletions(-) delete mode 100644 src/fastapi_toolsets/db.py create mode 100644 src/fastapi_toolsets/db/__init__.py create mode 100644 src/fastapi_toolsets/db/core.py create mode 100644 src/fastapi_toolsets/db/locks.py create mode 100644 src/fastapi_toolsets/db/m2m.py create mode 100644 src/fastapi_toolsets/db/testing.py create mode 100644 src/fastapi_toolsets/db/watch.py diff --git a/docs/module/crud.md b/docs/module/crud.md index 588af3f..00da6c6 100644 --- a/docs/module/crud.md +++ b/docs/module/crud.md @@ -167,7 +167,7 @@ user = await UserCrud.update(session, UserUpdate(credits=10), [User.id == user_i ``` !!! warning - `with_for_update` requires an open transaction. Wrap your call in `async with session.begin()` or use the `get_transaction` helper if you are not already inside one. + `with_for_update` requires an open transaction. Wrap your call in `async with session.begin()` or use the `transaction` helper if you are not already inside one. !!! note `NOWAIT` raises `sqlalchemy.exc.OperationalError` immediately if the row is locked rather than waiting. diff --git a/docs/module/db.md b/docs/module/db.md index 230785c..75f757f 100644 --- a/docs/module/db.md +++ b/docs/module/db.md @@ -7,96 +7,137 @@ SQLAlchemy async session management with transactions, table locking, advisory l ## Overview -The `db` module provides helpers to create FastAPI dependencies and context managers for `AsyncSession`, along with utilities for nested transactions, table locks, advisory locks, and polling for row changes. +The `db` module is built around one object, [`Database`](../reference/db.md#fastapi_toolsets.db.Database), which owns the engine and sessionmaker and exposes the FastAPI dependency, a commit-before-response middleware, session/transaction context managers, and table locking. Free helpers cover savepoint-aware transactions, advisory locks, many-to-many association tables, and row-change polling. -## Session dependency +## Setup -Use [`create_db_dependency`](../reference/db.md#fastapi_toolsets.db.create_db_dependency) to create a FastAPI dependency that yields a session and auto-commits on success: +Create one `Database` for your app. Provide a **URL** (the facade builds and disposes the engine) or pass an existing **`engine=`** you own (e.g. for Alembic or `event.listen`). The session factory is built internally with `expire_on_commit=False`. ```python -from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker -from fastapi_toolsets.db import create_db_dependency +from fastapi import Depends, FastAPI +from sqlalchemy.ext.asyncio import AsyncSession -engine = create_async_engine(url="postgresql+asyncpg://...", future=True) -session_maker = async_sessionmaker(bind=engine, expire_on_commit=False) +from fastapi_toolsets.db import Database -get_db = create_db_dependency(session_maker=session_maker) +db = Database("postgresql+asyncpg://postgres:postgres@localhost/app") -@router.get("/users") -async def list_users(session: AsyncSession = Depends(get_db)): +app = FastAPI() +db.install(app) # commit middleware + engine disposal on shutdown + +@app.get("/users") +async def list_users(session: AsyncSession = Depends(db)): ... ``` -## Session context manager +The `Database` instance **is** the dependency: use it directly as `Depends(db)`. The whole request runs as a single transaction (CRUD writes use savepoints under it). -Use [`create_db_context`](../reference/db.md#fastapi_toolsets.db.create_db_context) for sessions outside request handlers (e.g. background tasks, CLI commands): +## Committing before the response + +[`db.install(app)`](../reference/db.md#fastapi_toolsets.db.Database) adds a middleware that commits the request's session when the response starts, after the endpoint returns and before the body is sent. With the middleware installed, the dependency does not commit again. + +The request is committed as a single transaction: + +- **Read-after-write**: a follow-up request sees the write. +- **Atomicity**: multi-write endpoints roll back as a unit on failure. +- **Errors roll back**: on a raised exception the session rolls back and nothing is committed. + +Without `install`, the session commits in the dependency teardown, which runs after the response has been sent. + +!!! warning "Streaming / SSE endpoints" + For a `StreamingResponse` / `EventSourceResponse`, the commit fires at the **start** of the stream. A stream that **writes** must open a short-lived session per write with [`db.session()`](#session-context-manager); the start-time commit will not flush writes made later during the stream. + +## Lifespan + +`db.install(app)` disposes the engine on shutdown, composing around your own lifespan: ```python -from fastapi_toolsets.db import create_db_context +from contextlib import asynccontextmanager -db_context = create_db_context(session_maker=session_maker) +@asynccontextmanager +async def lifespan(app): + await warm_cache() # your startup + yield + await flush_metrics() # your shutdown +app = FastAPI(lifespan=lifespan) +db.install(app) # your shutdown runs first, then the engine is disposed +``` + +If you have no lifespan of your own, [`db.lifespan`](../reference/db.md#fastapi_toolsets.db.Database) works standalone as `FastAPI(lifespan=db.lifespan)`. Engine disposal is idempotent and is a no-op when you passed your own `engine=`. + +## Session context manager + +Use [`db.session()`](../reference/db.md#fastapi_toolsets.db.Database) for sessions outside request handlers (e.g. background tasks, CLI commands). It commits on clean exit and rolls back on exception: + +```python async def seed(): - async with db_context() as session: + async with db.session() as session: ... ``` -## Nested transactions +## Transactions -[`get_transaction`](../reference/db.md#fastapi_toolsets.db.get_transaction) handles savepoints automatically, allowing safe nesting: +[`transaction`](../reference/db.md#fastapi_toolsets.db.transaction) opens a transaction on a session, using a savepoint when one is already open so it nests safely: ```python -from fastapi_toolsets.db import get_transaction +from fastapi_toolsets.db import transaction -async def create_user_with_role(session=session): - async with get_transaction(session=session): +async def create_user_with_role(session): + async with transaction(session): ... - async with get_transaction(session=session): # uses savepoint + async with transaction(session): # uses a savepoint ... ``` +When you have a `Database`, [`db.begin()`](../reference/db.md#fastapi_toolsets.db.Database) opens a session already inside a transaction: + +```python +async with db.begin() as session: + session.add(User(name="ada")) # commits on exit, rolls back on exception +``` + ## Table locking -[`lock_tables`](../reference/db.md#fastapi_toolsets.db.lock_tables) acquires PostgreSQL table-level locks before executing critical sections. It opens a **dedicated session** internally and yields it to the caller, so the lock is guaranteed to be released when the context exits: +[`db.lock_tables`](../reference/db.md#fastapi_toolsets.db.Database) acquires PostgreSQL table-level locks for a critical section. It opens a dedicated session internally and releases the lock when the context exits: ```python -from fastapi_toolsets.db import lock_tables, LockMode +from fastapi_toolsets.db import LockMode -async with lock_tables(session_maker=session_maker, tables=[User], mode=LockMode.EXCLUSIVE) as session: +async with db.lock_tables([User], mode=LockMode.EXCLUSIVE) as session: # No other transaction can modify User until this block exits ... ``` Available lock modes are defined in [`LockMode`](../reference/db.md#fastapi_toolsets.db.LockMode): `ACCESS_SHARE`, `ROW_SHARE`, `ROW_EXCLUSIVE`, `SHARE_UPDATE_EXCLUSIVE`, `SHARE`, `SHARE_ROW_EXCLUSIVE`, `EXCLUSIVE`, `ACCESS_EXCLUSIVE`. -Pass `timeout` to limit how long the lock waits before giving up. On timeout, a [`LockTimeoutError`](../reference/exceptions.md#fastapi_toolsets.exceptions.exceptions.LockTimeoutError) is raised instead of a raw database error: +Pass `timeout` to limit how long the lock waits. On timeout, a [`LockTimeoutError`](../reference/exceptions.md#fastapi_toolsets.exceptions.exceptions.LockTimeoutError) is raised instead of a raw database error: ```python -async with lock_tables(session_maker, [Order], timeout="2s") as session: +async with db.lock_tables([Order], timeout="2s") as session: ... ``` ## Advisory locking -[`advisory_lock`](../reference/db.md#fastapi_toolsets.db.advisory_lock) acquires a PostgreSQL session-level advisory lock. The lock is released explicitly when the context exits, regardless of whether the transaction has committed. +[`advisory_lock`](../reference/db.md#fastapi_toolsets.db.advisory_lock) acquires a PostgreSQL session-level advisory lock on a session you provide. The lock is released when the context exits: ```python from fastapi_toolsets.db import advisory_lock -# Blocking exclusive lock — waits until the lock is free +# Blocking exclusive lock: waits until the lock is free async with advisory_lock(session=session, key=42): ... -# Non-blocking — yields False immediately if already held +# Non-blocking: yields False immediately if already held async with advisory_lock(session=session, key=42, nowait=True) as acquired: if not acquired: raise HTTPException(409, "Resource is locked") -# Blocking with a timeout — raises LockTimeoutError if not acquired in time +# Blocking with a timeout: raises LockTimeoutError if not acquired in time async with advisory_lock(session=session, key=42, timeout="5s"): ... -# Shared — multiple readers allowed simultaneously, blocks exclusive writers +# Shared lock: multiple readers allowed simultaneously, blocks exclusive writers async with advisory_lock(session=session, key=42, shared=True): ... @@ -106,11 +147,11 @@ async with advisory_lock(session=session, key=(1, user_id)): ``` !!! note - Advisory locks use PostgreSQL session-level functions (`pg_advisory_lock` / `pg_advisory_unlock`). The lock is tied to the database connection, not the SQLAlchemy transaction — it is released when the context exits, even if the surrounding transaction is still open. + Advisory locks use PostgreSQL session-level functions (`pg_advisory_lock` / `pg_advisory_unlock`). The lock is tied to the database connection, not the SQLAlchemy transaction, so it is released when the context exits even if the surrounding transaction is still open. ## Row-change polling -[`wait_for_row_change`](../reference/db.md#fastapi_toolsets.db.wait_for_row_change) polls a row until a specific column changes value, useful for waiting on async side effects: +[`wait_for_row_change`](../reference/db.md#fastapi_toolsets.db.wait_for_row_change) polls a row until a specific column changes value: ```python from fastapi_toolsets.db import wait_for_row_change @@ -120,7 +161,7 @@ await wait_for_row_change( session=session, model=Order, pk_value=order_id, - columns=[Order.status], + columns=["status"], interval=1.0, timeout=30.0, ) @@ -128,28 +169,24 @@ await wait_for_row_change( ## Creating a database -!!! info "Added in `v2.1`" - -[`create_database`](../reference/db.md#fastapi_toolsets.db.create_database) creates a database at a given URL. It connects to *server_url* and issues a `CREATE DATABASE` statement: +[`create_database`](../reference/db.md#fastapi_toolsets.db.testing.create_database) (in `fastapi_toolsets.db.testing`) connects to *server_url* and issues a `CREATE DATABASE` statement: ```python -from fastapi_toolsets.db import create_database +from fastapi_toolsets.db.testing import create_database SERVER_URL = "postgresql+asyncpg://postgres:postgres@localhost/postgres" await create_database(db_name="myapp_test", server_url=SERVER_URL) ``` -For test isolation with automatic cleanup, use [`create_worker_database`](../reference/pytest.md#fastapi_toolsets.pytest.utils.create_worker_database) from the `pytest` module instead — it handles drop-before, create, and drop-after automatically. +For test isolation with automatic cleanup, use [`create_worker_database`](../reference/pytest.md#fastapi_toolsets.pytest.utils.create_worker_database) from the `pytest` module, which handles drop-before, create, and drop-after. ## Cleaning up tables -!!! info "Added in `v2.1`" - -[`cleanup_tables`](../reference/db.md#fastapi_toolsets.db.cleanup_tables) truncates all tables: +[`cleanup_tables`](../reference/db.md#fastapi_toolsets.db.testing.cleanup_tables) (in `fastapi_toolsets.db.testing`) truncates all tables: ```python -from fastapi_toolsets.db import cleanup_tables +from fastapi_toolsets.db.testing import cleanup_tables @pytest.fixture(autouse=True) async def clean(db_session): @@ -159,50 +196,50 @@ async def clean(db_session): ## Many-to-Many helpers -SQLAlchemy's ORM collection API triggers lazy-loads when you append to a relationship inside a savepoint (e.g. inside `lock_tables` or a nested `get_transaction`). The three `m2m_*` helpers bypass the ORM collection entirely and issue direct SQL against the association table. +The three `m2m_*` helpers modify a many-to-many association table with direct SQL, without loading the ORM collection. -### `m2m_add` — insert associations +### `m2m_add`: insert associations -[`m2m_add`](../reference/db.md#fastapi_toolsets.db.m2m_add) inserts one or more rows into a secondary table without touching the ORM collection: +[`m2m_add`](../reference/db.md#fastapi_toolsets.db.m2m_add) inserts one or more rows into a secondary table: ```python -from fastapi_toolsets.db import lock_tables, m2m_add +from fastapi_toolsets.db import m2m_add -async with lock_tables(session_maker, [Tag]) as session: +async with db.lock_tables([Tag]) as session: tag = await TagCrud.create(session, TagCreate(name="python")) await m2m_add(session, post, Post.tags, tag) ``` -Pass `ignore_conflicts=True` to silently skip associations that already exist: +Pass `ignore_conflicts=True` to skip associations that already exist: ```python await m2m_add(session, post, Post.tags, tag, ignore_conflicts=True) ``` -### `m2m_remove` — delete associations +### `m2m_remove`: delete associations [`m2m_remove`](../reference/db.md#fastapi_toolsets.db.m2m_remove) deletes specific association rows. Removing a non-existent association is a no-op: ```python -from fastapi_toolsets.db import get_transaction, m2m_remove +from fastapi_toolsets.db import m2m_remove, transaction -async with get_transaction(session): +async with transaction(session): await m2m_remove(session, post, Post.tags, tag1, tag2) ``` -### `m2m_set` — replace the full set +### `m2m_set`: replace the full set -[`m2m_set`](../reference/db.md#fastapi_toolsets.db.m2m_set) atomically replaces all associations: it deletes every existing row for the owner instance then inserts the new set. Passing no related instances clears the association entirely: +[`m2m_set`](../reference/db.md#fastapi_toolsets.db.m2m_set) replaces all associations: it deletes every existing row for the owner instance then inserts the new set. Passing no related instances clears the association: ```python -from fastapi_toolsets.db import get_transaction, m2m_set +from fastapi_toolsets.db import m2m_set, transaction # Replace all tags -async with get_transaction(session): +async with transaction(session): await m2m_set(session, post, Post.tags, tag_a, tag_b) # Clear all tags -async with get_transaction(session): +async with transaction(session): await m2m_set(session, post, Post.tags) ``` diff --git a/docs/module/models.md b/docs/module/models.md index eae194c..a5c86cb 100644 --- a/docs/module/models.md +++ b/docs/module/models.md @@ -134,7 +134,7 @@ SessionLocal = async_sessionmaker(engine, expire_on_commit=False, class_=EventSe ``` !!! info "Callbacks fire on `session.commit()` only — not on savepoints." - Savepoints created by [`get_transaction`](db.md) or `begin_nested()` do **not** + Savepoints created by [`transaction`](db.md) or `begin_nested()` do **not** trigger callbacks. All events accumulated across flushes are dispatched once when the outermost `commit()` is called. diff --git a/docs/module/pytest.md b/docs/module/pytest.md index cf6e448..b0c22af 100644 --- a/docs/module/pytest.md +++ b/docs/module/pytest.md @@ -107,10 +107,10 @@ url = worker_database_url("postgresql+asyncpg://user:pass@localhost/myapp", defa ## Manual table cleanup -[`cleanup_tables`](../reference/db.md#fastapi_toolsets.db.cleanup_tables) truncates all tables in a single statement and can be called directly when you need more control: +[`cleanup_tables`](../reference/db.md#fastapi_toolsets.db.testing.cleanup_tables) truncates all tables in a single statement and can be called directly when you need more control: ```python -from fastapi_toolsets.db import cleanup_tables +from fastapi_toolsets.pytest import cleanup_tables @pytest.fixture(autouse=True) async def clean(db_session): diff --git a/docs/reference/db.md b/docs/reference/db.md index 66b3204..144fd97 100644 --- a/docs/reference/db.md +++ b/docs/reference/db.md @@ -1,46 +1,48 @@ # `db` -Here's the reference for all database session utilities, transaction helpers, and locking functions. +Here's the reference for the `Database` facade, the transaction helper, locking +functions, many-to-many helpers, and row-watching utilities. You can import them directly from `fastapi_toolsets.db`: ```python from fastapi_toolsets.db import ( + Database, LockMode, advisory_lock, - cleanup_tables, - create_database, - create_db_dependency, - create_db_context, - get_transaction, lock_tables, m2m_add, m2m_remove, m2m_set, + transaction, wait_for_row_change, ) ``` +## ::: fastapi_toolsets.db.Database + +## ::: fastapi_toolsets.db.transaction + ## ::: fastapi_toolsets.db.LockMode -## ::: fastapi_toolsets.db.create_db_dependency - -## ::: fastapi_toolsets.db.create_db_context - -## ::: fastapi_toolsets.db.get_transaction - ## ::: fastapi_toolsets.db.lock_tables ## ::: fastapi_toolsets.db.advisory_lock -## ::: fastapi_toolsets.db.wait_for_row_change - -## ::: fastapi_toolsets.db.create_database - -## ::: fastapi_toolsets.db.cleanup_tables - ## ::: fastapi_toolsets.db.m2m_add ## ::: fastapi_toolsets.db.m2m_remove ## ::: fastapi_toolsets.db.m2m_set + +## ::: fastapi_toolsets.db.wait_for_row_change + +Admin and test helpers live in `fastapi_toolsets.db.testing`: + +```python +from fastapi_toolsets.db.testing import cleanup_tables, create_database +``` + +## ::: fastapi_toolsets.db.testing.create_database + +## ::: fastapi_toolsets.db.testing.cleanup_tables diff --git a/docs_src/examples/pagination_search/app.py b/docs_src/examples/pagination_search/app.py index 8a6348f..56c096f 100644 --- a/docs_src/examples/pagination_search/app.py +++ b/docs_src/examples/pagination_search/app.py @@ -2,8 +2,10 @@ from fastapi import FastAPI from fastapi_toolsets.exceptions import init_exceptions_handlers +from .db import db from .routes import router app = FastAPI() +db.install(app=app) init_exceptions_handlers(app=app) app.include_router(router=router) diff --git a/docs_src/examples/pagination_search/db.py b/docs_src/examples/pagination_search/db.py index 8826de6..6c78899 100644 --- a/docs_src/examples/pagination_search/db.py +++ b/docs_src/examples/pagination_search/db.py @@ -1,17 +1,14 @@ from typing import Annotated from fastapi import Depends -from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine +from sqlalchemy.ext.asyncio import AsyncSession -from fastapi_toolsets.db import create_db_context, create_db_dependency +from fastapi_toolsets.db import Database DATABASE_URL = "postgresql+asyncpg://postgres:postgres@localhost:5432/postgres" -engine = create_async_engine(url=DATABASE_URL, future=True) -async_session_maker = async_sessionmaker(bind=engine, expire_on_commit=False) +db = Database(url=DATABASE_URL) -get_db = create_db_dependency(session_maker=async_session_maker) -get_db_context = create_db_context(session_maker=async_session_maker) +get_db = db - -SessionDep = Annotated[AsyncSession, Depends(get_db)] +SessionDep = Annotated[AsyncSession, Depends(db)] diff --git a/src/fastapi_toolsets/__init__.py b/src/fastapi_toolsets/__init__.py index 3170c73..466c8b1 100644 --- a/src/fastapi_toolsets/__init__.py +++ b/src/fastapi_toolsets/__init__.py @@ -7,16 +7,19 @@ Example usage: from fastapi import FastAPI, Depends from fastapi_toolsets.exceptions import init_exceptions_handlers from fastapi_toolsets.crud import CrudFactory - from fastapi_toolsets.db import create_db_dependency + from fastapi_toolsets.db import Database from fastapi_toolsets.schemas import Response + db = Database("postgresql+asyncpg://postgres:postgres@localhost/app") + app = FastAPI() + db.install(app) init_exceptions_handlers(app) UserCrud = CrudFactory(User) @app.get("/users/{user_id}", response_model=Response[dict]) - async def get_user(user_id: int, session = Depends(get_db)): + async def get_user(user_id: int, session = Depends(db)): user = await UserCrud.get(session, [User.id == user_id]) return Response(data={"user": user.username}, message="Success") """ diff --git a/src/fastapi_toolsets/crud/factory.py b/src/fastapi_toolsets/crud/factory.py index 38bf53f..08f18ca 100644 --- a/src/fastapi_toolsets/crud/factory.py +++ b/src/fastapi_toolsets/crud/factory.py @@ -22,7 +22,7 @@ from sqlalchemy.orm import DeclarativeBase, QueryableAttribute, selectinload from sqlalchemy.sql.base import ExecutableOption from sqlalchemy.sql.roles import WhereHavingRole -from ..db import get_transaction +from ..db import transaction from ..exceptions import InvalidOrderFieldError, NotFoundError from ..schemas import ( CursorPaginatedResponse, @@ -716,7 +716,7 @@ class AsyncCrud(Generic[ModelType]): Returns: Created model instance, or ``Response[schema]`` when ``schema`` is given. """ - async with get_transaction(session): + async with transaction(session): m2m_exclude = cls._m2m_schema_fields() data = ( obj.model_dump(exclude=m2m_exclude) if m2m_exclude else obj.model_dump() @@ -1067,7 +1067,7 @@ class AsyncCrud(Generic[ModelType]): Raises: NotFoundError: If no record found """ - async with get_transaction(session): + async with transaction(session): m2m_exclude = cls._m2m_schema_fields() # Eagerly load M2M relationships that will be updated so that @@ -1127,7 +1127,7 @@ class AsyncCrud(Generic[ModelType]): Returns: Model instance """ - async with get_transaction(session): + async with transaction(session): values = obj.model_dump(exclude_unset=True) q = insert(cls.model).values(**values) if set_: @@ -1189,7 +1189,7 @@ class AsyncCrud(Generic[ModelType]): Returns: ``None``, or ``Response[None]`` when ``return_response=True``. """ - async with get_transaction(session): + async with transaction(session): result = await session.execute(select(cls.model).where(and_(*filters))) objects = result.scalars().all() for obj in objects: diff --git a/src/fastapi_toolsets/db.py b/src/fastapi_toolsets/db.py deleted file mode 100644 index 9ec30e1..0000000 --- a/src/fastapi_toolsets/db.py +++ /dev/null @@ -1,591 +0,0 @@ -"""Database utilities: sessions, transactions, and locks.""" - -import asyncio -from collections.abc import AsyncGenerator, Callable -from contextlib import AbstractAsyncContextManager, asynccontextmanager -from enum import Enum -from typing import Any, TypeVar, cast - -import asyncpg -from sqlalchemy import Table, delete, text, tuple_ -from sqlalchemy import exc as sa_exc -from sqlalchemy.dialects.postgresql import insert as pg_insert -from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine -from sqlalchemy.orm import DeclarativeBase, QueryableAttribute -from sqlalchemy.orm.relationships import RelationshipProperty - -from .exceptions import LockTimeoutError, NotFoundError, PoolExhaustedError - - -def _is_lock_not_available(e: sa_exc.DBAPIError) -> bool: - return e.orig is not None and isinstance( - e.orig.__cause__, asyncpg.exceptions.LockNotAvailableError - ) - - -__all__ = [ - "LockMode", - "advisory_lock", - "cleanup_tables", - "create_database", - "create_db_context", - "create_db_dependency", - "get_transaction", - "lock_tables", - "m2m_add", - "m2m_remove", - "m2m_set", - "wait_for_row_change", -] - - -_SessionT = TypeVar("_SessionT", bound=AsyncSession) - - -def create_db_dependency( - session_maker: async_sessionmaker[_SessionT], -) -> Callable[[], AsyncGenerator[_SessionT, None]]: - """Create a FastAPI dependency for database sessions. - - Creates a dependency function that yields a session and auto-commits - if a transaction is active when the request completes. - - Args: - session_maker: Async session factory from create_session_factory() - - Returns: - An async generator function usable with FastAPI's Depends() - - Example: - ```python - from fastapi import Depends - from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker - from fastapi_toolsets.db import create_db_dependency - - engine = create_async_engine("postgresql+asyncpg://...") - SessionLocal = async_sessionmaker(engine, expire_on_commit=False) - get_db = create_db_dependency(SessionLocal) - - @app.get("/users") - async def list_users(session: AsyncSession = Depends(get_db)): - ... - ``` - """ - - async def get_db() -> AsyncGenerator[_SessionT, None]: - async with session_maker() as session: - try: - await session.connection() - except sa_exc.TimeoutError as e: - raise PoolExhaustedError() from e - yield session - if session.in_transaction(): - await session.commit() - - return get_db - - -def create_db_context( - session_maker: async_sessionmaker[_SessionT], -) -> Callable[[], AbstractAsyncContextManager[_SessionT]]: - """Create a context manager for database sessions. - - Creates a context manager for use outside of FastAPI request handlers, - such as in background tasks, CLI commands, or tests. - - Args: - session_maker: Async session factory from create_session_factory() - - Returns: - An async context manager function - - Example: - ```python - from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker - from fastapi_toolsets.db import create_db_context - - engine = create_async_engine("postgresql+asyncpg://...") - SessionLocal = async_sessionmaker(engine, expire_on_commit=False) - get_db_context = create_db_context(SessionLocal) - - async def background_task(): - async with get_db_context() as session: - user = await UserCrud.get(session, [User.id == 1]) - ... - ``` - """ - get_db = create_db_dependency(session_maker) - return asynccontextmanager(get_db) - - -@asynccontextmanager -async def get_transaction( - session: AsyncSession, -) -> AsyncGenerator[AsyncSession, None]: - """Get a transaction context, handling nested transactions. - - If already in a transaction, creates a savepoint (nested transaction). - Otherwise, starts a new transaction. - - Args: - session: AsyncSession instance - - Yields: - The session within the transaction context - - Example: - ```python - async with get_transaction(session): - session.add(model) - # Auto-commits on exit, rolls back on exception - ``` - """ - if session.in_transaction(): - async with session.begin_nested(): - yield session - else: - async with session.begin(): - yield session - - -class LockMode(str, Enum): - """PostgreSQL table lock modes. - - See: https://www.postgresql.org/docs/current/explicit-locking.html - """ - - ACCESS_SHARE = "ACCESS SHARE" - ROW_SHARE = "ROW SHARE" - ROW_EXCLUSIVE = "ROW EXCLUSIVE" - SHARE_UPDATE_EXCLUSIVE = "SHARE UPDATE EXCLUSIVE" - SHARE = "SHARE" - SHARE_ROW_EXCLUSIVE = "SHARE ROW EXCLUSIVE" - EXCLUSIVE = "EXCLUSIVE" - ACCESS_EXCLUSIVE = "ACCESS EXCLUSIVE" - - -def lock_tables( - session_maker: async_sessionmaker[_SessionT], - tables: list[type[DeclarativeBase]], - *, - mode: LockMode = LockMode.SHARE_UPDATE_EXCLUSIVE, - timeout: str = "5s", -) -> AbstractAsyncContextManager[_SessionT]: - """Lock PostgreSQL tables for the duration of a transaction. - - Args: - session_maker: Async session factory used to create the dedicated - session. - tables: List of SQLAlchemy model classes to lock. - mode: Lock mode (default: SHARE UPDATE EXCLUSIVE). - timeout: Lock timeout (default: "5s"). - - Yields: - The dedicated session, open within the locked transaction. - - Raises: - SQLAlchemyError: If the lock cannot be acquired within *timeout*. - - Example: - ```python - from fastapi_toolsets.db import lock_tables, LockMode - - async with lock_tables(session_maker, [User, Account]) as session: - # Tables are locked; changes are committed when the context exits. - user = await UserCrud.get(session, [User.id == 1]) - user.balance += 100 - - # With custom lock mode - async with lock_tables(session_maker, [Order], mode=LockMode.EXCLUSIVE) as session: - await process_order(session, order_id) - ``` - """ - table_names = ",".join(table.__tablename__ for table in tables) - - @asynccontextmanager - async def _lock() -> AsyncGenerator[_SessionT, None]: - async with session_maker() as session: - try: - await session.execute(text(f"SET LOCAL lock_timeout='{timeout}'")) - await session.execute(text(f"LOCK {table_names} IN {mode.value} MODE")) - yield session - await session.commit() - except sa_exc.TimeoutError as e: - await session.rollback() - raise PoolExhaustedError( - f"Connection pool exhausted while locking '{table_names}'. " - ) from e - except sa_exc.DBAPIError as e: - await session.rollback() - if _is_lock_not_available(e): - raise LockTimeoutError( - f"Lock on '{table_names}' could not be acquired within {timeout}." - ) from e - raise # pragma: no cover - except BaseException: - await session.rollback() - raise - - return _lock() - - -@asynccontextmanager -async def advisory_lock( - session: AsyncSession, - key: int | tuple[int, int], - *, - shared: bool = False, - nowait: bool = False, - timeout: str | None = None, -) -> AsyncGenerator[bool, None]: - """Acquire a PostgreSQL session-level advisory lock. - - Args: - session: AsyncSession instance. - key: Lock key — a single ``int`` (bigint) or a ``(int, int)`` pair for namespacing. - shared: Acquire a shared lock (multiple holders allowed). Default is exclusive. - nowait: Return ``False`` immediately if the lock is unavailable instead of waiting. - timeout: Maximum wait time (e.g. ``"5s"``, ``"500ms"``). Raises ``DBAPIError`` - if exceeded. Ignored when *nowait* is ``True``. - - Yields: - ``True`` if the lock was acquired, ``False`` if *nowait* is ``True`` and the lock - is already held. - - Raises: - LockTimeoutError: If *timeout* is set and the lock cannot be acquired in time. - - Example: - ```python - from fastapi_toolsets.db import advisory_lock - - async with advisory_lock(session, 42): - ... - - async with advisory_lock(session, 42, nowait=True) as acquired: - if not acquired: - raise HTTPException(409, "Resource is locked") - - async with advisory_lock(session, 42, timeout="5s"): - ... - - async with advisory_lock(session, (1, user_id), shared=True): - ... - ``` - """ - suffix = "_shared" if shared else "" - acquire_fn = f"{'pg_try_advisory_lock' if nowait else 'pg_advisory_lock'}{suffix}" - release_fn = f"pg_advisory_unlock{suffix}" - - if isinstance(key, tuple): - k1, k2 = key - args = "CAST(:k1 AS integer), CAST(:k2 AS integer)" - params: dict[str, int] = {"k1": k1, "k2": k2} - else: - args = ":k" - params = {"k": key} - - acquire_sql = text(f"SELECT {acquire_fn}({args})") - release_sql = text(f"SELECT {release_fn}({args})") - - if timeout is not None and not nowait: - await session.execute(text(f"SET LOCAL lock_timeout='{timeout}'")) - - try: - result = await session.execute(acquire_sql, params) - except sa_exc.DBAPIError as e: - if _is_lock_not_available(e): - raise LockTimeoutError( - f"Advisory lock {key!r} could not be acquired within {timeout}." - ) from e - raise # pragma: no cover - acquired = result.scalar() if nowait else True - try: - yield acquired - finally: - if acquired: - await session.execute(release_sql, params) - - -async def create_database( - db_name: str, - *, - server_url: str, -) -> None: - """Create a database. - - Connects to *server_url* using ``AUTOCOMMIT`` isolation and issues a - ``CREATE DATABASE`` statement for *db_name*. - - Args: - db_name: Name of the database to create. - server_url: URL used for server-level DDL (must point to an existing - database on the same server). - - Example: - ```python - from fastapi_toolsets.db import create_database - - SERVER_URL = "postgresql+asyncpg://postgres:postgres@localhost/postgres" - await create_database("myapp_test", server_url=SERVER_URL) - ``` - """ - engine = create_async_engine(server_url, isolation_level="AUTOCOMMIT") - try: - async with engine.connect() as conn: - await conn.execute(text(f"CREATE DATABASE {db_name}")) - finally: - await engine.dispose() - - -async def cleanup_tables( - session: AsyncSession, - base: type[DeclarativeBase], -) -> None: - """Truncate all tables for fast between-test cleanup. - - Executes a single ``TRUNCATE … RESTART IDENTITY CASCADE`` statement - across every table in *base*'s metadata, which is significantly faster - than dropping and re-creating tables between tests. - - This is a no-op when the metadata contains no tables. - - Args: - session: An active async database session. - base: SQLAlchemy DeclarativeBase class containing model metadata. - - Example: - ```python - @pytest.fixture - async def db_session(worker_db_url): - async with create_db_session(worker_db_url, Base) as session: - yield session - await cleanup_tables(session, Base) - ``` - """ - tables = base.metadata.sorted_tables - if not tables: - return - - table_names = ", ".join(f'"{t.name}"' for t in tables) - await session.execute(text(f"TRUNCATE {table_names} RESTART IDENTITY CASCADE")) - await session.commit() - - -_M = TypeVar("_M", bound=DeclarativeBase) - - -async def wait_for_row_change( - session: AsyncSession, - model: type[_M], - pk_value: Any, - *, - columns: list[str] | None = None, - interval: float = 0.5, - timeout: float | None = None, -) -> _M: - """Poll a database row until a change is detected. - - Queries the row every ``interval`` seconds and returns the model instance - once a change is detected in any column (or only the specified ``columns``). - - Args: - session: AsyncSession instance - model: SQLAlchemy model class - pk_value: Primary key value of the row to watch - columns: Optional list of column names to watch. If None, all columns - are watched. - interval: Polling interval in seconds (default: 0.5) - timeout: Maximum time to wait in seconds. None means wait forever. - - Returns: - The refreshed model instance with updated values - - Raises: - NotFoundError: If the row does not exist or is deleted during polling - TimeoutError: If timeout expires before a change is detected - - Example: - ```python - from fastapi_toolsets.db import wait_for_row_change - - # Wait for any column to change - updated = await wait_for_row_change(session, User, user_id) - - # Watch specific columns with a timeout - updated = await wait_for_row_change( - session, User, user_id, - columns=["status", "email"], - interval=1.0, - timeout=30.0, - ) - ``` - """ - instance = await session.get(model, pk_value) - if instance is None: - raise NotFoundError(f"{model.__name__} with pk={pk_value!r} not found") - - if columns is not None: - watch_cols = columns - else: - watch_cols = [attr.key for attr in model.__mapper__.column_attrs] - - initial = {col: getattr(instance, col) for col in watch_cols} - - elapsed = 0.0 - while True: - await asyncio.sleep(interval) - elapsed += interval - - if timeout is not None and elapsed >= timeout: - raise TimeoutError( - f"No change detected on {model.__name__} " - f"with pk={pk_value!r} within {timeout}s" - ) - - session.expunge(instance) - instance = await session.get(model, pk_value) - - if instance is None: - raise NotFoundError(f"{model.__name__} with pk={pk_value!r} was deleted") - - current = {col: getattr(instance, col) for col in watch_cols} - if current != initial: - return instance - - -def _m2m_prop(rel_attr: QueryableAttribute) -> RelationshipProperty: # type: ignore[type-arg] - """Return the validated M2M RelationshipProperty for *rel_attr*. - - Raises TypeError if *rel_attr* is not a Many-to-Many relationship. - """ - prop = rel_attr.property - if not isinstance(prop, RelationshipProperty) or prop.secondary is None: - raise TypeError( - f"m2m helpers require a Many-to-Many relationship attribute, " - f"got {rel_attr!r}. Use a relationship with a secondary table." - ) - return prop - - -async def m2m_add( - session: AsyncSession, - instance: DeclarativeBase, - rel_attr: QueryableAttribute, - *related: DeclarativeBase, - ignore_conflicts: bool = False, -) -> None: - """Insert rows into a Many-to-Many association table without loading the ORM collection. - - Args: - session: DB async session. - instance: The "owner" side model instance (e.g. the ``A`` in ``A.b_list``). - rel_attr: The M2M relationship attribute on the model class (e.g. ``A.b_list``). - *related: One or more related instances to associate with ``instance``. - ignore_conflicts: When ``True``, silently skip rows that already exist - in the association table (``ON CONFLICT DO NOTHING``). - - Raises: - TypeError: If ``rel_attr`` is not a Many-to-Many relationship. - """ - prop = _m2m_prop(rel_attr) - if not related: - return - - secondary = cast(Table, prop.secondary) - assert secondary is not None # guaranteed by _m2m_prop - sync_pairs = prop.secondary_synchronize_pairs - assert sync_pairs is not None # set whenever secondary is set - - # synchronize_pairs: [(parent_col, assoc_col), ...] - # secondary_synchronize_pairs: [(related_col, assoc_col), ...] - rows: list[dict[str, Any]] = [] - for rel_instance in related: - row: dict[str, Any] = {} - for parent_col, assoc_col in prop.synchronize_pairs: - row[assoc_col.name] = getattr(instance, cast(str, parent_col.key)) - for related_col, assoc_col in sync_pairs: - row[assoc_col.name] = getattr(rel_instance, cast(str, related_col.key)) - rows.append(row) - - stmt = pg_insert(secondary).values(rows) - if ignore_conflicts: - stmt = stmt.on_conflict_do_nothing() - await session.execute(stmt) - - -async def m2m_remove( - session: AsyncSession, - instance: DeclarativeBase, - rel_attr: QueryableAttribute, - *related: DeclarativeBase, -) -> None: - """Remove rows from a Many-to-Many association table without loading the ORM collection. - - Args: - session: DB async session. - instance: The "owner" side model instance (e.g. the ``A`` in ``A.b_list``). - rel_attr: The M2M relationship attribute on the model class (e.g. ``A.b_list``). - *related: One or more related instances to disassociate from ``instance``. - - Raises: - TypeError: If ``rel_attr`` is not a Many-to-Many relationship. - """ - prop = _m2m_prop(rel_attr) - if not related: - return - - secondary = cast(Table, prop.secondary) - assert secondary is not None # guaranteed by _m2m_prop - related_pairs = prop.secondary_synchronize_pairs - assert related_pairs is not None # set whenever secondary is set - - parent_where = [ - assoc_col == getattr(instance, cast(str, parent_col.key)) - for parent_col, assoc_col in prop.synchronize_pairs - ] - - if len(related_pairs) == 1: - related_col, assoc_col = related_pairs[0] - related_values = [getattr(r, cast(str, related_col.key)) for r in related] - related_where = assoc_col.in_(related_values) - else: - assoc_cols = [ac for _, ac in related_pairs] - rel_cols = [rc for rc, _ in related_pairs] - related_values_t = [ - tuple(getattr(r, cast(str, rc.key)) for rc in rel_cols) for r in related - ] - related_where = tuple_(*assoc_cols).in_(related_values_t) - - await session.execute(delete(secondary).where(*parent_where, related_where)) - - -async def m2m_set( - session: AsyncSession, - instance: DeclarativeBase, - rel_attr: QueryableAttribute, - *related: DeclarativeBase, -) -> None: - """Replace the entire Many-to-Many association set atomically. - - Args: - session: DB async session. - instance: The "owner" side model instance (e.g. the ``A`` in ``A.b_list``). - rel_attr: The M2M relationship attribute on the model class (e.g. ``A.b_list``). - *related: The new complete set of related instances. - - Raises: - TypeError: If ``rel_attr`` is not a Many-to-Many relationship. - """ - prop = _m2m_prop(rel_attr) - secondary = cast(Table, prop.secondary) - assert secondary is not None # guaranteed by _m2m_prop - - parent_where = [ - assoc_col == getattr(instance, cast(str, parent_col.key)) - for parent_col, assoc_col in prop.synchronize_pairs - ] - await session.execute(delete(secondary).where(*parent_where)) - - if related: - await m2m_add(session, instance, rel_attr, *related) diff --git a/src/fastapi_toolsets/db/__init__.py b/src/fastapi_toolsets/db/__init__.py new file mode 100644 index 0000000..e41f90f --- /dev/null +++ b/src/fastapi_toolsets/db/__init__.py @@ -0,0 +1,18 @@ +"""Database package: the ``Database`` facade plus PostgreSQL power-tools.""" + +from .core import Database, transaction +from .locks import LockMode, advisory_lock, lock_tables +from .m2m import m2m_add, m2m_remove, m2m_set +from .watch import wait_for_row_change + +__all__ = [ + "Database", + "LockMode", + "advisory_lock", + "lock_tables", + "m2m_add", + "m2m_remove", + "m2m_set", + "transaction", + "wait_for_row_change", +] diff --git a/src/fastapi_toolsets/db/core.py b/src/fastapi_toolsets/db/core.py new file mode 100644 index 0000000..6feb413 --- /dev/null +++ b/src/fastapi_toolsets/db/core.py @@ -0,0 +1,315 @@ +"""The ``Database`` facade: session lifecycle, dependency, middleware, transactions.""" + +from collections.abc import AsyncGenerator +from contextlib import AbstractAsyncContextManager, asynccontextmanager +from typing import Any + +from sqlalchemy import exc as sa_exc +from sqlalchemy.ext.asyncio import ( + AsyncEngine, + AsyncSession, + async_sessionmaker, + create_async_engine, +) +from sqlalchemy.orm import DeclarativeBase +from starlette.requests import Request +from starlette.types import ASGIApp, Message, Receive, Scope, Send + +from ..exceptions import PoolExhaustedError +from .locks import LockMode, lock_tables + + +@asynccontextmanager +async def transaction( + session: AsyncSession, +) -> AsyncGenerator[AsyncSession, None]: + """Run a block inside a savepoint-aware transaction. + + If *session* is already in a transaction, a nested transaction (savepoint) + is opened so the block can roll back independently. Otherwise a top-level + transaction is started. Commits on clean exit, rolls back on exception. + + Args: + session: AsyncSession instance. + + Yields: + The session within the transaction context. + + Example: + ```python + from fastapi_toolsets.db import transaction + + async with transaction(session): + session.add(model) + ``` + """ + if session.in_transaction(): + async with session.begin_nested(): + yield session + else: + async with session.begin(): + yield session + + +class _CommitOnResponseMiddleware: + """Commit the request's DB session before the response is sent.""" + + def __init__(self, app: ASGIApp, *, state_attr: str) -> None: + self.app = app + self.state_attr = state_attr + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] != "http": + await self.app(scope, receive, send) + return + + async def send_wrapper(message: Message) -> None: + if message["type"] == "http.response.start": + # ``scope["state"]`` is the same dict ``request.state`` writes + # to, so this is the session stashed by the dependency. + state = scope.get("state") + session = state.get(self.state_attr) if state else None + if session is not None and session.in_transaction(): + await session.commit() + await send(message) + + await self.app(scope, receive, send_wrapper) + + +class Database: + """One object that owns the engine, sessions, dependency, and middleware. + + Provide exactly one of *url* (the facade builds and disposes the engine) or + *engine* (an engine you own, e.g. for Alembic or ``event.listen``, left + untouched). + + Args: + url: Database connection URL (e.g. ``"postgresql+asyncpg://..."``). + engine: An existing :class:`AsyncEngine` to reuse instead of *url*. + session_class: Session class for the sessionmaker (e.g. ``EventSession``). + expire_on_commit: Expire attributes after commit. Defaults to ``False``. + autoflush: Autoflush the session before queries. Defaults to ``True``. + **engine_options: Extra keyword arguments forwarded to + :func:`create_async_engine` (URL mode only, e.g. ``pool_size``, + ``echo``, ``connect_args``). + + Raises: + TypeError: If neither or both of *url* and *engine* are given, or if + *engine_options* are passed together with *engine*. + + Example: + ```python + from fastapi import Depends, FastAPI + from fastapi_toolsets.db import Database + + db = Database("postgresql+asyncpg://postgres:postgres@localhost/app") + + app = FastAPI() + db.install(app) + + @app.get("/users/{user_id}") + async def get_user(user_id: int, session=Depends(db)): + return await UserCrud.get(session, [User.id == user_id]) + ``` + """ + + def __init__( + self, + url: str | None = None, + *, + engine: AsyncEngine | None = None, + session_class: type[AsyncSession] = AsyncSession, + expire_on_commit: bool = False, + autoflush: bool = True, + **engine_options: Any, + ) -> None: + if (url is None) == (engine is None): + raise TypeError( + "Database requires exactly one of 'url' or 'engine' " + "(got both or neither)." + ) + if engine is not None and engine_options: + raise TypeError( + "engine_options are only valid in URL mode; configure the " + "engine you pass via 'engine=' yourself." + ) + + if engine is not None: + self._owns_engine = False + self.engine: AsyncEngine = engine + else: + assert url is not None # guaranteed by the XOR check above + self._owns_engine = True + self.engine = create_async_engine(url, **engine_options) + self._sessionmaker: async_sessionmaker[AsyncSession] = async_sessionmaker( + self.engine, + class_=session_class, + expire_on_commit=expire_on_commit, + autoflush=autoflush, + ) + # Private, per-instance state attribute; cannot collide with another + # Database or be mismatched against the middleware. + self._state_attr = f"_ft_db_session_{id(self):x}" + self._middleware_installed = False + self._disposed = False + + async def _dispose(self) -> None: + """Dispose the engine once, only if we own it (idempotent).""" + if self._owns_engine and not self._disposed: + self._disposed = True + await self.engine.dispose() + + @asynccontextmanager + async def lifespan(self, app: Any) -> AsyncGenerator[None, None]: + """Dispose the engine on shutdown; use as ``FastAPI(lifespan=db.lifespan)``. + + Args: + app: The ASGI application (unused; required by the lifespan protocol). + + Yields: + Control to the application for its lifetime. + + Example: + ```python + app = FastAPI(lifespan=db.lifespan) + ``` + """ + try: + yield + finally: + await self._dispose() + + def install(self, app: Any) -> None: + """Wire the commit middleware and engine disposal onto *app*. + + Args: + app: The FastAPI/Starlette application to wire. + + Example: + ```python + @asynccontextmanager + async def lifespan(app): + ... # your startup + yield + ... # your shutdown + + app = FastAPI(lifespan=lifespan) + db.install(app) + ``` + """ + app.add_middleware(_CommitOnResponseMiddleware, state_attr=self._state_attr) + self._middleware_installed = True + + inner_lifespan = app.router.lifespan_context + + @asynccontextmanager + async def _composed(app_: Any) -> AsyncGenerator[None, None]: + async with self.lifespan(app_): + async with inner_lifespan(app_): + yield + + app.router.lifespan_context = _composed + + @asynccontextmanager + async def _open(self) -> AsyncGenerator[AsyncSession, None]: + """Open a session and eagerly acquire a connection (fail-fast on pool).""" + async with self._sessionmaker() as session: + try: + await session.connection() + except sa_exc.TimeoutError as e: + raise PoolExhaustedError() from e + yield session + + async def __call__(self, request: Request) -> AsyncGenerator[AsyncSession, None]: + """FastAPI dependency: yield a session and commit once at the right time. + + Args: + request: The incoming request (injected by FastAPI). + + Yields: + An AsyncSession for the duration of the request. + + Example: + ```python + @app.get("/users/{user_id}") + async def get_user(user_id: int, session=Depends(db)): + return await UserCrud.get(session, [User.id == user_id]) + ``` + """ + async with self._open() as session: + setattr(request.state, self._state_attr, session) + yield session + if not self._middleware_installed and session.in_transaction(): + await session.commit() + + @asynccontextmanager + async def session(self) -> AsyncGenerator[AsyncSession, None]: + """Open a session outside request handlers (background tasks, CLI, tests). + + Commits on clean exit, rolls back on exception. + + Yields: + An AsyncSession ready for database operations. + + Example: + ```python + async with db.session() as session: + user = await UserCrud.get(session, [User.id == 1]) + ``` + """ + async with self._open() as session: + yield session + if session.in_transaction(): + await session.commit() + + @asynccontextmanager + async def begin(self) -> AsyncGenerator[AsyncSession, None]: + """Open a session already inside a transaction (sugar for the common case). + + Equivalent to ``session()`` + :func:`transaction`. Commits on clean exit, + rolls back on exception. + + Yields: + An AsyncSession open within a transaction. + + Example: + ```python + async with db.begin() as session: + session.add(User(name="ada")) + ``` + """ + async with self.session() as session, transaction(session): + yield session + + def lock_tables( + self, + tables: list[type[DeclarativeBase]], + *, + mode: LockMode = LockMode.SHARE_UPDATE_EXCLUSIVE, + timeout: str = "5s", + ) -> AbstractAsyncContextManager[AsyncSession]: + """Lock PostgreSQL tables for the duration of a dedicated transaction. + + Opens its own session from the facade's sessionmaker, changes are + committed when the context exits. + + Args: + tables: List of SQLAlchemy model classes to lock. + mode: Lock mode (default: ``SHARE UPDATE EXCLUSIVE``). + timeout: Lock timeout (default: ``"5s"``). + + Yields: + The dedicated session, open within the locked transaction. + + Raises: + LockTimeoutError: If the lock cannot be acquired within *timeout*. + PoolExhaustedError: If the connection pool is exhausted. + + Example: + ```python + async with db.lock_tables([User, Account]) as session: + user = await UserCrud.get(session, [User.id == 1]) + user.balance += 100 + ``` + """ + return lock_tables(self._sessionmaker, tables, mode=mode, timeout=timeout) diff --git a/src/fastapi_toolsets/db/locks.py b/src/fastapi_toolsets/db/locks.py new file mode 100644 index 0000000..56efa0b --- /dev/null +++ b/src/fastapi_toolsets/db/locks.py @@ -0,0 +1,185 @@ +"""PostgreSQL locking helpers: table locks and advisory locks.""" + +from collections.abc import AsyncGenerator +from contextlib import AbstractAsyncContextManager, asynccontextmanager +from enum import Enum +from typing import TypeVar + +import asyncpg +from sqlalchemy import exc as sa_exc +from sqlalchemy import text +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker +from sqlalchemy.orm import DeclarativeBase + +from ..exceptions import LockTimeoutError, PoolExhaustedError + +_SessionT = TypeVar("_SessionT", bound=AsyncSession) + + +def _is_lock_not_available(e: sa_exc.DBAPIError) -> bool: + return e.orig is not None and isinstance( + e.orig.__cause__, asyncpg.exceptions.LockNotAvailableError + ) + + +class LockMode(str, Enum): + """PostgreSQL table lock modes. + + See: https://www.postgresql.org/docs/current/explicit-locking.html + """ + + ACCESS_SHARE = "ACCESS SHARE" + ROW_SHARE = "ROW SHARE" + ROW_EXCLUSIVE = "ROW EXCLUSIVE" + SHARE_UPDATE_EXCLUSIVE = "SHARE UPDATE EXCLUSIVE" + SHARE = "SHARE" + SHARE_ROW_EXCLUSIVE = "SHARE ROW EXCLUSIVE" + EXCLUSIVE = "EXCLUSIVE" + ACCESS_EXCLUSIVE = "ACCESS EXCLUSIVE" + + +def lock_tables( + session_maker: async_sessionmaker[_SessionT], + tables: list[type[DeclarativeBase]], + *, + mode: LockMode = LockMode.SHARE_UPDATE_EXCLUSIVE, + timeout: str = "5s", +) -> AbstractAsyncContextManager[_SessionT]: + """Lock PostgreSQL tables for the duration of a transaction. + + Prefer the method on a :class:`Database` instance; use this + directly only when you manage your own session factory. + + Args: + session_maker: Async session factory used to create the dedicated + session. + tables: List of SQLAlchemy model classes to lock. + mode: Lock mode (default: SHARE UPDATE EXCLUSIVE). + timeout: Lock timeout (default: "5s"). + + Yields: + The dedicated session, open within the locked transaction. + + Raises: + LockTimeoutError: If the lock cannot be acquired within *timeout*. + PoolExhaustedError: If the connection pool is exhausted. + + Example: + ```python + from fastapi_toolsets.db import lock_tables + + async with lock_tables(session_maker, [User, Account]) as session: + user = await UserCrud.get(session, [User.id == 1]) + user.balance += 100 + ``` + """ + table_names = ",".join(table.__tablename__ for table in tables) + + @asynccontextmanager + async def _lock() -> AsyncGenerator[_SessionT, None]: + async with session_maker() as session: + try: + await session.execute(text(f"SET LOCAL lock_timeout='{timeout}'")) + await session.execute(text(f"LOCK {table_names} IN {mode.value} MODE")) + yield session + await session.commit() + except sa_exc.TimeoutError as e: + await session.rollback() + raise PoolExhaustedError( + f"Connection pool exhausted while locking '{table_names}'. " + ) from e + except sa_exc.DBAPIError as e: + await session.rollback() + if _is_lock_not_available(e): + raise LockTimeoutError( + f"Lock on '{table_names}' could not be acquired within {timeout}." + ) from e + raise # pragma: no cover + except BaseException: + await session.rollback() + raise + + return _lock() + + +@asynccontextmanager +async def advisory_lock( + session: AsyncSession, + key: int | tuple[int, int], + *, + shared: bool = False, + nowait: bool = False, + timeout: str | None = None, +) -> AsyncGenerator[bool, None]: + """Acquire a PostgreSQL session-level advisory lock. + + Args: + session: AsyncSession instance. + key: Lock key, either a single ``int`` (bigint) or a ``(int, int)`` pair for namespacing. + shared: Acquire a shared lock (multiple holders allowed). Default is exclusive. + nowait: Return ``False`` immediately if the lock is unavailable instead of waiting. + timeout: Maximum wait time (e.g. ``"5s"``, ``"500ms"``). Raises ``DBAPIError`` + if exceeded. Ignored when *nowait* is ``True``. + + Yields: + ``True`` if the lock was acquired, ``False`` if *nowait* is ``True`` and the lock + is already held. + + Raises: + LockTimeoutError: If *timeout* is set and the lock cannot be acquired in time. + + Example: + ```python + from fastapi_toolsets.db import advisory_lock + + async with advisory_lock(session, 42): + ... + + async with advisory_lock(session, 42, nowait=True) as acquired: + if not acquired: + raise HTTPException(409, "Resource is locked") + + async with advisory_lock(session, 42, timeout="5s"): + ... + + async with advisory_lock(session, (1, user_id), shared=True): + ... + ``` + """ + suffix = "_shared" if shared else "" + acquire_fn = f"{'pg_try_advisory_lock' if nowait else 'pg_advisory_lock'}{suffix}" + release_fn = f"pg_advisory_unlock{suffix}" + + if isinstance(key, tuple): + k1, k2 = key + args = "CAST(:k1 AS integer), CAST(:k2 AS integer)" + params: dict[str, int] = {"k1": k1, "k2": k2} + else: + args = ":k" + params = {"k": key} + + acquire_sql = text(f"SELECT {acquire_fn}({args})") + release_sql = text(f"SELECT {release_fn}({args})") + + # Lock management runs raw SQL on the caller's session. Guard it with + # ``no_autoflush`` so acquiring or releasing the lock never flushes the + # caller's pending ORM changes; SQLAlchemy 2.1 autoflushes on raw + # ``text()`` too, where 2.0 did not. + try: + with session.no_autoflush: + if timeout is not None and not nowait: + await session.execute(text(f"SET LOCAL lock_timeout='{timeout}'")) + result = await session.execute(acquire_sql, params) + except sa_exc.DBAPIError as e: + if _is_lock_not_available(e): + raise LockTimeoutError( + f"Advisory lock {key!r} could not be acquired within {timeout}." + ) from e + raise # pragma: no cover + acquired = result.scalar() if nowait else True + try: + yield acquired + finally: + if acquired: + with session.no_autoflush: + await session.execute(release_sql, params) diff --git a/src/fastapi_toolsets/db/m2m.py b/src/fastapi_toolsets/db/m2m.py new file mode 100644 index 0000000..b4bc979 --- /dev/null +++ b/src/fastapi_toolsets/db/m2m.py @@ -0,0 +1,170 @@ +"""Many-to-Many association-table helpers (direct, without loading collections).""" + +from typing import Any, TypeVar, cast + +from sqlalchemy import ColumnElement, Table, delete, tuple_ +from sqlalchemy.dialects.postgresql import insert as pg_insert +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.orm import DeclarativeBase, QueryableAttribute +from sqlalchemy.orm.relationships import RelationshipProperty + +_M = TypeVar("_M", bound=DeclarativeBase) + + +def _m2m_prop(rel_attr: QueryableAttribute) -> tuple[RelationshipProperty, Table]: # type: ignore[type-arg] + """Return the validated M2M RelationshipProperty and its secondary table. + + Raises TypeError if *rel_attr* is not a Many-to-Many relationship. + """ + prop = rel_attr.property + if not isinstance(prop, RelationshipProperty) or prop.secondary is None: + raise TypeError( + f"m2m helpers require a Many-to-Many relationship attribute, " + f"got {rel_attr!r}. Use a relationship with a secondary table." + ) + return prop, cast(Table, prop.secondary) + + +def _parent_where( + prop: RelationshipProperty, # type: ignore[type-arg] + instance: DeclarativeBase, +) -> list[ColumnElement[bool]]: + """Build the WHERE clauses matching the owner side of *instance*.""" + return [ + assoc_col == getattr(instance, cast(str, parent_col.key)) + for parent_col, assoc_col in prop.synchronize_pairs + ] + + +async def m2m_add( + session: AsyncSession, + instance: DeclarativeBase, + rel_attr: QueryableAttribute, + *related: DeclarativeBase, + ignore_conflicts: bool = False, +) -> None: + """Insert rows into a Many-to-Many association table without loading the ORM collection. + + Args: + session: DB async session. + instance: The "owner" side model instance (e.g. the ``A`` in ``A.b_list``). + rel_attr: The M2M relationship attribute on the model class (e.g. ``A.b_list``). + *related: One or more related instances to associate with ``instance``. + ignore_conflicts: When ``True``, silently skip rows that already exist + in the association table (``ON CONFLICT DO NOTHING``). + + Raises: + TypeError: If ``rel_attr`` is not a Many-to-Many relationship. + + Example: + ```python + from fastapi_toolsets.db import m2m_add, transaction + + async with transaction(session): + await m2m_add(session, post, Post.tags, tag1, tag2) + ``` + """ + prop, secondary = _m2m_prop(rel_attr) + if not related: + return + + sync_pairs = prop.secondary_synchronize_pairs + assert sync_pairs is not None # set whenever secondary is set + + # synchronize_pairs: [(parent_col, assoc_col), ...] + # secondary_synchronize_pairs: [(related_col, assoc_col), ...] + rows: list[dict[str, Any]] = [] + for rel_instance in related: + row: dict[str, Any] = {} + for parent_col, assoc_col in prop.synchronize_pairs: + row[assoc_col.name] = getattr(instance, cast(str, parent_col.key)) + for related_col, assoc_col in sync_pairs: + row[assoc_col.name] = getattr(rel_instance, cast(str, related_col.key)) + rows.append(row) + + stmt = pg_insert(secondary).values(rows) + if ignore_conflicts: + stmt = stmt.on_conflict_do_nothing() + await session.execute(stmt) + + +async def m2m_remove( + session: AsyncSession, + instance: DeclarativeBase, + rel_attr: QueryableAttribute, + *related: DeclarativeBase, +) -> None: + """Remove rows from a Many-to-Many association table without loading the ORM collection. + + Args: + session: DB async session. + instance: The "owner" side model instance (e.g. the ``A`` in ``A.b_list``). + rel_attr: The M2M relationship attribute on the model class (e.g. ``A.b_list``). + *related: One or more related instances to disassociate from ``instance``. + + Raises: + TypeError: If ``rel_attr`` is not a Many-to-Many relationship. + + Example: + ```python + from fastapi_toolsets.db import m2m_remove, transaction + + async with transaction(session): + await m2m_remove(session, post, Post.tags, tag1) + ``` + """ + prop, secondary = _m2m_prop(rel_attr) + if not related: + return + + related_pairs = prop.secondary_synchronize_pairs + assert related_pairs is not None # set whenever secondary is set + + parent_where = _parent_where(prop, instance) + + if len(related_pairs) == 1: + related_col, assoc_col = related_pairs[0] + related_values = [getattr(r, cast(str, related_col.key)) for r in related] + related_where = assoc_col.in_(related_values) + else: + assoc_cols = [ac for _, ac in related_pairs] + rel_cols = [rc for rc, _ in related_pairs] + related_values_t = [ + tuple(getattr(r, cast(str, rc.key)) for rc in rel_cols) for r in related + ] + related_where = tuple_(*assoc_cols).in_(related_values_t) + + await session.execute(delete(secondary).where(*parent_where, related_where)) + + +async def m2m_set( + session: AsyncSession, + instance: DeclarativeBase, + rel_attr: QueryableAttribute, + *related: DeclarativeBase, +) -> None: + """Replace the entire Many-to-Many association set atomically. + + Args: + session: DB async session. + instance: The "owner" side model instance (e.g. the ``A`` in ``A.b_list``). + rel_attr: The M2M relationship attribute on the model class (e.g. ``A.b_list``). + *related: The new complete set of related instances. + + Raises: + TypeError: If ``rel_attr`` is not a Many-to-Many relationship. + + Example: + ```python + from fastapi_toolsets.db import m2m_set, transaction + + async with transaction(session): + await m2m_set(session, post, Post.tags, tag1, tag2) # replaces all + ``` + """ + prop, secondary = _m2m_prop(rel_attr) + + await session.execute(delete(secondary).where(*_parent_where(prop, instance))) + + if related: + await m2m_add(session, instance, rel_attr, *related) diff --git a/src/fastapi_toolsets/db/testing.py b/src/fastapi_toolsets/db/testing.py new file mode 100644 index 0000000..64c0d84 --- /dev/null +++ b/src/fastapi_toolsets/db/testing.py @@ -0,0 +1,69 @@ +"""Database admin and test helpers: DDL and truncation.""" + +from sqlalchemy import text +from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine +from sqlalchemy.orm import DeclarativeBase + + +async def create_database( + db_name: str, + *, + server_url: str, +) -> None: + """Create a database. + + Connects to *server_url* using ``AUTOCOMMIT`` isolation and issues a + ``CREATE DATABASE`` statement for *db_name*. + + Args: + db_name: Name of the database to create. + server_url: URL used for server-level DDL (must point to an existing + database on the same server). + + Example: + ```python + from fastapi_toolsets.db.testing import create_database + + SERVER_URL = "postgresql+asyncpg://postgres:postgres@localhost/postgres" + await create_database("myapp_test", server_url=SERVER_URL) + ``` + """ + engine = create_async_engine(server_url, isolation_level="AUTOCOMMIT") + try: + async with engine.connect() as conn: + await conn.execute(text(f"CREATE DATABASE {db_name}")) + finally: + await engine.dispose() + + +async def cleanup_tables( + session: AsyncSession, + base: type[DeclarativeBase], +) -> None: + """Truncate all tables for fast between-test cleanup. + + Executes a single ``TRUNCATE … RESTART IDENTITY CASCADE`` statement + across every table in *base*'s metadata. + + This is a no-op when the metadata contains no tables. + + Args: + session: An active async database session. + base: SQLAlchemy DeclarativeBase class containing model metadata. + + Example: + ```python + @pytest.fixture + async def db_session(worker_db_url): + async with create_db_session(worker_db_url, Base) as session: + yield session + await cleanup_tables(session, Base) + ``` + """ + tables = base.metadata.sorted_tables + if not tables: + return + + table_names = ", ".join(f'"{t.name}"' for t in tables) + await session.execute(text(f"TRUNCATE {table_names} RESTART IDENTITY CASCADE")) + await session.commit() diff --git a/src/fastapi_toolsets/db/watch.py b/src/fastapi_toolsets/db/watch.py new file mode 100644 index 0000000..514da01 --- /dev/null +++ b/src/fastapi_toolsets/db/watch.py @@ -0,0 +1,90 @@ +"""Row-watching helpers: poll a database row until it changes.""" + +import asyncio +from typing import Any, TypeVar + +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.orm import DeclarativeBase + +from ..exceptions import NotFoundError + +_M = TypeVar("_M", bound=DeclarativeBase) + + +async def wait_for_row_change( + session: AsyncSession, + model: type[_M], + pk_value: Any, + *, + columns: list[str] | None = None, + interval: float = 0.5, + timeout: float | None = None, +) -> _M: + """Poll a database row until a change is detected. + + Queries the row every ``interval`` seconds and returns the model instance + once a change is detected in any column (or only the specified ``columns``). + + Args: + session: AsyncSession instance. + model: SQLAlchemy model class. + pk_value: Primary key value of the row to watch. + columns: Optional list of column names to watch. If None, all columns + are watched. + interval: Polling interval in seconds (default: 0.5). + timeout: Maximum time to wait in seconds. None means wait forever. + + Returns: + The refreshed model instance with updated values. + + Raises: + NotFoundError: If the row does not exist or is deleted during polling. + TimeoutError: If timeout expires before a change is detected. + + Example: + ```python + from fastapi_toolsets.db import wait_for_row_change + + # Wait for any column to change + updated = await wait_for_row_change(session, User, user_id) + + # Watch specific columns with a timeout + updated = await wait_for_row_change( + session, User, user_id, + columns=["status", "email"], + interval=1.0, + timeout=30.0, + ) + ``` + """ + instance = await session.get(model, pk_value) + if instance is None: + raise NotFoundError(f"{model.__name__} with pk={pk_value!r} not found") + + if columns is not None: + watch_cols = columns + else: + watch_cols = [attr.key for attr in model.__mapper__.column_attrs] + + initial = {col: getattr(instance, col) for col in watch_cols} + + elapsed = 0.0 + while True: + await asyncio.sleep(interval) + elapsed += interval + + if timeout is not None and elapsed >= timeout: + raise TimeoutError( + f"No change detected on {model.__name__} " + f"with pk={pk_value!r} within {timeout}s" + ) + + session.expunge(instance) + instance = await session.get(model, pk_value) + + if instance is None: + raise NotFoundError(f"{model.__name__} with pk={pk_value!r} was deleted") + + current = {col: getattr(instance, col) for col in watch_cols} + if current != initial: + return instance diff --git a/src/fastapi_toolsets/fixtures/utils.py b/src/fastapi_toolsets/fixtures/utils.py index 48dadfb..d25a978 100644 --- a/src/fastapi_toolsets/fixtures/utils.py +++ b/src/fastapi_toolsets/fixtures/utils.py @@ -9,7 +9,7 @@ from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import DeclarativeBase -from ..db import get_transaction +from ..db import transaction from ..logger import get_logger from ..types import ModelType from .enum import LoadStrategy @@ -229,7 +229,7 @@ async def _load_ordered( model_name = type(instances[0]).__name__ loaded: list[DeclarativeBase] = [] - async with get_transaction(session): + async with transaction(session): for model_cls, group in _group_by_type(instances): match strategy: case LoadStrategy.INSERT: diff --git a/src/fastapi_toolsets/pytest/plugin.py b/src/fastapi_toolsets/pytest/plugin.py index 2aba2f9..1b2b5eb 100644 --- a/src/fastapi_toolsets/pytest/plugin.py +++ b/src/fastapi_toolsets/pytest/plugin.py @@ -9,7 +9,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import DeclarativeBase, selectinload from sqlalchemy.orm.interfaces import ExecutableOption, ORMOption -from ..db import get_transaction +from ..db import transaction from ..fixtures import FixtureRegistry, LoadStrategy @@ -106,7 +106,7 @@ def _create_fixture_function( loaded: list[DeclarativeBase] = [] - async with get_transaction(session): + async with transaction(session): for instance in instances: if strategy == LoadStrategy.INSERT: session.add(instance) diff --git a/src/fastapi_toolsets/pytest/utils.py b/src/fastapi_toolsets/pytest/utils.py index 94f61bd..11ed032 100644 --- a/src/fastapi_toolsets/pytest/utils.py +++ b/src/fastapi_toolsets/pytest/utils.py @@ -15,7 +15,7 @@ from sqlalchemy.ext.asyncio import ( ) from sqlalchemy.orm import DeclarativeBase -from ..db import cleanup_tables, create_database +from ..db.testing import cleanup_tables, create_database from ..models.watched import EventSession diff --git a/tests/test_db.py b/tests/test_db.py index b4bec40..4852a19 100644 --- a/tests/test_db.py +++ b/tests/test_db.py @@ -1,9 +1,14 @@ -"""Tests for fastapi_toolsets.db module.""" +"""Tests for fastapi_toolsets.db module (v5 ``Database`` facade).""" import asyncio import uuid +from contextlib import asynccontextmanager +from unittest.mock import AsyncMock, patch import pytest +from fastapi import Depends, FastAPI +from fastapi.responses import StreamingResponse +from httpx import ASGITransport, AsyncClient from sqlalchemy import ( Column, ForeignKey, @@ -16,7 +21,12 @@ from sqlalchemy import ( ) from sqlalchemy.engine import make_url from sqlalchemy.exc import IntegrityError -from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine +from sqlalchemy.ext.asyncio import ( + AsyncEngine, + AsyncSession, + async_sessionmaker, + create_async_engine, +) from sqlalchemy.orm import ( DeclarativeBase, Mapped, @@ -24,21 +34,21 @@ from sqlalchemy.orm import ( relationship, selectinload, ) +from starlette.requests import Request from fastapi_toolsets.db import ( + Database, LockMode, advisory_lock, - cleanup_tables, - create_database, - create_db_context, - create_db_dependency, - get_transaction, lock_tables, m2m_add, m2m_remove, m2m_set, + transaction, wait_for_row_change, ) +from fastapi_toolsets.db.core import _CommitOnResponseMiddleware +from fastapi_toolsets.db.testing import cleanup_tables, create_database from fastapi_toolsets.exceptions import ( LockTimeoutError, NotFoundError, @@ -46,153 +56,309 @@ from fastapi_toolsets.exceptions import ( ) from fastapi_toolsets.pytest import create_db_session -from .conftest import DATABASE_URL, Base, Post, Role, RoleCrud, Tag, User, UserCrud +from .conftest import ( + DATABASE_URL, + Base, + Post, + Role, + RoleCreate, + RoleCrud, + Tag, + User, + UserCrud, +) -class TestCreateDbDependency: - """Tests for create_db_dependency.""" +def _make_request() -> Request: + """Minimal ASGI HTTP request for exercising the Database dependency directly.""" + return Request({"type": "http", "headers": []}) + + +class TestDatabaseConstruction: + """Construction contract: provide exactly one of url / engine.""" + + def test_requires_url_or_engine(self): + """Neither url nor engine raises TypeError.""" + with pytest.raises(TypeError): + Database() @pytest.mark.anyio - async def test_yields_session(self): - """Dependency yields a valid session.""" - engine = create_async_engine(DATABASE_URL, echo=False) - session_factory = async_sessionmaker(engine, expire_on_commit=False) - get_db = create_db_dependency(session_factory) + async def test_both_url_and_engine_raises(self, engine): + """Both url and engine raises TypeError.""" + with pytest.raises(TypeError): + Database(DATABASE_URL, engine=engine) - async for session in get_db(): + @pytest.mark.anyio + async def test_engine_options_with_engine_raises(self, engine): + """engine_options are rejected in engine= mode.""" + with pytest.raises(TypeError): + Database(engine=engine, pool_size=5) + + @pytest.mark.anyio + async def test_url_mode_owns_engine(self): + """URL mode builds and owns the engine.""" + db = Database(DATABASE_URL) + try: + assert db._owns_engine is True + assert db.engine is not None + finally: + await db.engine.dispose() + + @pytest.mark.anyio + async def test_engine_mode_borrows_engine(self, engine): + """engine= mode reuses the given engine and does not own it.""" + db = Database(engine=engine) + assert db._owns_engine is False + assert db.engine is engine + + @pytest.mark.anyio + async def test_distinct_instances_use_distinct_state_attrs(self, engine): + """Two Database instances never share a request-state attribute.""" + a = Database(engine=engine) + b = Database(engine=engine) + assert a._state_attr != b._state_attr + + @pytest.mark.anyio + async def test_lifespan_disposes_owned_engine(self): + """The lifespan disposes the engine it built (URL mode).""" + db = Database(DATABASE_URL) + # ``AsyncEngine.dispose`` is read-only on the instance, so patch the class. + with patch.object(AsyncEngine, "dispose", new=AsyncMock()) as disposed: + async with db.lifespan(None): + pass + disposed.assert_awaited_once() + await db.engine.dispose() + + @pytest.mark.anyio + async def test_lifespan_skips_borrowed_engine(self): + """The lifespan leaves a borrowed engine untouched (engine= mode).""" + eng = create_async_engine(DATABASE_URL, echo=False) + db = Database(engine=eng) + with patch.object(AsyncEngine, "dispose", new=AsyncMock()) as disposed: + async with db.lifespan(None): + pass + disposed.assert_not_awaited() + await eng.dispose() + + +class TestLifespanComposition: + """``install`` composes engine disposal around the app's own lifespan.""" + + @pytest.mark.anyio + async def test_install_composes_user_lifespan(self): + """A user-defined lifespan runs, and the engine is disposed after it.""" + events: list[str] = [] + + @asynccontextmanager + async def user_lifespan(app): + events.append("startup") + yield + events.append("shutdown") + + db = Database(DATABASE_URL) + app = FastAPI(lifespan=user_lifespan) + db.install(app) + + with patch.object(AsyncEngine, "dispose", new=AsyncMock()) as disposed: + async with app.router.lifespan_context(app): + assert events == ["startup"] + disposed.assert_not_awaited() + # User shutdown runs, then the engine is disposed. + assert events == ["startup", "shutdown"] + disposed.assert_awaited_once() + await db.engine.dispose() + + @pytest.mark.anyio + async def test_install_disposes_without_user_lifespan(self): + """``install`` disposes the engine even when the app has no custom lifespan.""" + db = Database(DATABASE_URL) + app = FastAPI() + db.install(app) + + with patch.object(AsyncEngine, "dispose", new=AsyncMock()) as disposed: + async with app.router.lifespan_context(app): + disposed.assert_not_awaited() + disposed.assert_awaited_once() + await db.engine.dispose() + + @pytest.mark.anyio + async def test_disposal_is_idempotent(self): + """Combining ``lifespan=db.lifespan`` with ``install`` disposes only once.""" + db = Database(DATABASE_URL) + app = FastAPI(lifespan=db.lifespan) + db.install(app) + + with patch.object(AsyncEngine, "dispose", new=AsyncMock()) as disposed: + async with app.router.lifespan_context(app): + pass + disposed.assert_awaited_once() + await db.engine.dispose() + + @pytest.mark.anyio + async def test_install_skips_disposal_for_borrowed_engine(self, engine): + """``install`` never disposes an engine it does not own.""" + db = Database(engine=engine) + app = FastAPI() + db.install(app) + + with patch.object(AsyncEngine, "dispose", new=AsyncMock()) as disposed: + async with app.router.lifespan_context(app): + pass + disposed.assert_not_awaited() + + +class TestDatabaseDependency: + """Tests for the FastAPI dependency (``Depends(db)`` / ``db.__call__``).""" + + @pytest.mark.anyio + async def test_yields_session(self, engine): + """Dependency yields a valid session.""" + db = Database(engine=engine) + async for session in db(_make_request()): assert isinstance(session, AsyncSession) break - await engine.dispose() + @pytest.mark.anyio + async def test_auto_commits_transaction(self, engine, session_maker): + """Without middleware, the dependency commits an open transaction on exit.""" + db = Database(engine=engine) + + async for session in db(_make_request()): + role = Role(name="test_role_dep") + session.add(role) + await session.flush() + + async with session_maker() as verify: + result = await RoleCrud.first(verify, [Role.name == "test_role_dep"]) + assert result is not None @pytest.mark.anyio - async def test_auto_commits_transaction(self): - """Dependency auto-commits if transaction is active.""" - engine = create_async_engine(DATABASE_URL, echo=False) - session_factory = async_sessionmaker(engine, expire_on_commit=False) - - async with engine.begin() as conn: - await conn.run_sync(Base.metadata.create_all) - - try: - get_db = create_db_dependency(session_factory) - - async for session in get_db(): - role = Role(name="test_role_dep") - session.add(role) - await session.flush() - - async with session_factory() as verify_session: - result = await RoleCrud.first( - verify_session, [Role.name == "test_role_dep"] - ) - assert result is not None - finally: - async with engine.begin() as conn: - await conn.run_sync(Base.metadata.drop_all) - await engine.dispose() - - @pytest.mark.anyio - async def test_in_transaction_on_yield(self): + async def test_in_transaction_on_yield(self, engine): """Session is already in a transaction when the endpoint body starts.""" - engine = create_async_engine(DATABASE_URL, echo=False) - session_factory = async_sessionmaker(engine, expire_on_commit=False) - get_db = create_db_dependency(session_factory) - - async for session in get_db(): + db = Database(engine=engine) + async for session in db(_make_request()): assert session.in_transaction() break - await engine.dispose() - @pytest.mark.anyio - async def test_no_commit_when_not_in_transaction(self): - """Dependency skips commit if the session is no longer in a transaction on exit.""" - engine = create_async_engine(DATABASE_URL, echo=False) - session_factory = async_sessionmaker(engine, expire_on_commit=False) - get_db = create_db_dependency(session_factory) - - async for session in get_db(): - # Manually commit — session exits the transaction + async def test_no_commit_when_not_in_transaction(self, engine): + """Dependency skips commit if the session left its transaction on exit.""" + db = Database(engine=engine) + async for session in db(_make_request()): await session.commit() assert not session.in_transaction() - # The dependency's post-yield path must not call commit again (no error) - - await engine.dispose() + # The post-yield path must not call commit again (no error). @pytest.mark.anyio - async def test_data_inside_lock_is_committed(self): - """Changes made inside lock_tables are committed when the context exits.""" - engine = create_async_engine(DATABASE_URL, echo=False) - session_factory = async_sessionmaker(engine, expire_on_commit=False) - - async with engine.begin() as conn: - await conn.run_sync(Base.metadata.create_all) - - try: - async with lock_tables(session_factory, [Role]) as session: - role = Role(name="lock_committed") - session.add(role) - - async with session_factory() as verify: - result = await RoleCrud.first(verify, [Role.name == "lock_committed"]) - assert result is not None - finally: - async with engine.begin() as conn: - await conn.run_sync(Base.metadata.drop_all) - await engine.dispose() - - -class TestCreateDbContext: - """Tests for create_db_context.""" + async def test_stashes_session_on_request_state(self, engine): + """Dependency exposes the session on request.state for the commit middleware.""" + db = Database(engine=engine) + request = _make_request() + async for session in db(request): + assert getattr(request.state, db._state_attr) is session + break @pytest.mark.anyio - async def test_context_manager_yields_session(self): + async def test_skips_commit_when_middleware_installed(self, engine, session_maker): + """With ``install()``, the dependency must NOT commit — the middleware owns it. + + Here no middleware actually runs (we call the dependency directly), so the + open transaction is rolled back on session close and nothing persists. + """ + db = Database(engine=engine) + db.install(FastAPI()) + + async for session in db(_make_request()): + role = Role(name="mw_owns_commit") + session.add(role) + await session.flush() + + async with session_maker() as verify: + result = await RoleCrud.first(verify, [Role.name == "mw_owns_commit"]) + assert result is None + + +class TestDatabaseSession: + """Tests for ``db.session()`` (sessions outside request handlers).""" + + @pytest.mark.anyio + async def test_context_manager_yields_session(self, engine): """Context manager yields a valid session.""" - engine = create_async_engine(DATABASE_URL, echo=False) - session_factory = async_sessionmaker(engine, expire_on_commit=False) - get_db_context = create_db_context(session_factory) - - async with get_db_context() as session: + db = Database(engine=engine) + async with db.session() as session: assert isinstance(session, AsyncSession) - await engine.dispose() + @pytest.mark.anyio + async def test_context_manager_commits(self, engine, session_maker): + """Context manager commits on exit.""" + db = Database(engine=engine) + + async with db.session() as session: + role = Role(name="context_role") + session.add(role) + await session.flush() + + async with session_maker() as verify: + result = await RoleCrud.first(verify, [Role.name == "context_role"]) + assert result is not None @pytest.mark.anyio - async def test_context_manager_commits(self): - """Context manager commits on exit.""" - engine = create_async_engine(DATABASE_URL, echo=False) - session_factory = async_sessionmaker(engine, expire_on_commit=False) - - async with engine.begin() as conn: - await conn.run_sync(Base.metadata.create_all) + async def test_no_commit_when_not_in_transaction(self, engine): + """Context skips commit if the session left its transaction on exit.""" + db = Database(engine=engine) + async with db.session() as session: + await session.commit() + assert not session.in_transaction() + @pytest.mark.anyio + async def test_pool_exhausted_raises_pool_exhausted_error(self): + """PoolExhaustedError is raised when the pool is exhausted on session entry.""" + db = Database(DATABASE_URL, pool_size=1, max_overflow=0, pool_timeout=0.1) try: - get_db_context = create_db_context(session_factory) - - async with get_db_context() as session: - role = Role(name="context_role") - session.add(role) - await session.flush() - - async with session_factory() as verify_session: - result = await RoleCrud.first( - verify_session, [Role.name == "context_role"] - ) - assert result is not None + async with db.session(): # checks out the single available connection + with pytest.raises(PoolExhaustedError): + async with db.session(): + pass finally: - async with engine.begin() as conn: - await conn.run_sync(Base.metadata.drop_all) - await engine.dispose() + await db.engine.dispose() -class TestGetTransaction: - """Tests for get_transaction context manager.""" +class TestDatabaseBegin: + """Tests for ``db.begin()`` (open a session already in a transaction).""" + + @pytest.mark.anyio + async def test_commits_on_success(self, engine, session_maker): + """The block commits when it exits cleanly.""" + db = Database(engine=engine) + async with db.begin() as session: + session.add(Role(name="begin_role")) + + async with session_maker() as verify: + result = await RoleCrud.first(verify, [Role.name == "begin_role"]) + assert result is not None + + @pytest.mark.anyio + async def test_rolls_back_on_exception(self, engine, session_maker): + """The block rolls back on exception.""" + db = Database(engine=engine) + with pytest.raises(ValueError): + async with db.begin() as session: + session.add(Role(name="begin_rollback_role")) + await session.flush() + raise ValueError("Simulated error") + + async with session_maker() as verify: + result = await RoleCrud.first(verify, [Role.name == "begin_rollback_role"]) + assert result is None + + +class TestTransaction: + """Tests for the ``transaction`` context manager (savepoint-aware primitive).""" @pytest.mark.anyio async def test_starts_transaction(self, db_session: AsyncSession): - """get_transaction starts a new transaction.""" - async with get_transaction(db_session): + """transaction starts a new transaction.""" + async with transaction(db_session): role = Role(name="tx_role") db_session.add(role) @@ -202,12 +368,12 @@ class TestGetTransaction: @pytest.mark.anyio async def test_nested_transaction_uses_savepoint(self, db_session: AsyncSession): """Nested transactions use savepoints.""" - async with get_transaction(db_session): + async with transaction(db_session): role1 = Role(name="outer_role") db_session.add(role1) await db_session.flush() - async with get_transaction(db_session): + async with transaction(db_session): role2 = Role(name="inner_role") db_session.add(role2) @@ -220,7 +386,7 @@ class TestGetTransaction: async def test_rollback_on_exception(self, db_session: AsyncSession): """Transaction rolls back on exception.""" try: - async with get_transaction(db_session): + async with transaction(db_session): role = Role(name="rollback_role") db_session.add(role) await db_session.flush() @@ -234,13 +400,13 @@ class TestGetTransaction: @pytest.mark.anyio async def test_nested_rollback_preserves_outer(self, db_session: AsyncSession): """Nested rollback preserves outer transaction.""" - async with get_transaction(db_session): + async with transaction(db_session): role1 = Role(name="preserved_role") db_session.add(role1) await db_session.flush() try: - async with get_transaction(db_session): + async with transaction(db_session): role2 = Role(name="rolled_back_role") db_session.add(role2) await db_session.flush() @@ -275,12 +441,13 @@ class TestLockMode: class TestLockTables: - """Tests for lock_tables context manager (PostgreSQL-specific).""" + """Tests for ``db.lock_tables`` (PostgreSQL-specific).""" @pytest.mark.anyio - async def test_lock_single_table(self, session_maker): + async def test_lock_single_table(self, engine, session_maker): """Lock a single table; changes inside are committed on context exit.""" - async with lock_tables(session_maker, [Role]) as session: + db = Database(engine=engine) + async with db.lock_tables([Role]) as session: role = Role(name="locked_role") session.add(role) @@ -289,9 +456,10 @@ class TestLockTables: assert result is not None @pytest.mark.anyio - async def test_lock_multiple_tables(self, session_maker): + async def test_lock_multiple_tables(self, engine, session_maker): """Lock multiple tables.""" - async with lock_tables(session_maker, [Role, User]) as session: + db = Database(engine=engine) + async with db.lock_tables([Role, User]) as session: role = Role(name="multi_lock_role") session.add(role) @@ -300,11 +468,10 @@ class TestLockTables: assert result is not None @pytest.mark.anyio - async def test_lock_with_custom_mode(self, session_maker): + async def test_lock_with_custom_mode(self, engine, session_maker): """Lock with custom lock mode.""" - async with lock_tables( - session_maker, [Role], mode=LockMode.EXCLUSIVE - ) as session: + db = Database(engine=engine) + async with db.lock_tables([Role], mode=LockMode.EXCLUSIVE) as session: role = Role(name="exclusive_lock_role") session.add(role) @@ -313,16 +480,15 @@ class TestLockTables: assert result is not None @pytest.mark.anyio - async def test_lock_rollback_on_exception(self, session_maker): + async def test_lock_rollback_on_exception(self, engine, session_maker): """Lock context rolls back on exception.""" - try: - async with lock_tables(session_maker, [Role]) as session: + db = Database(engine=engine) + with pytest.raises(ValueError): + async with db.lock_tables([Role]) as session: role = Role(name="lock_rollback_role") session.add(role) await session.flush() raise ValueError("Simulated error") - except ValueError: - pass async with session_maker() as verify: result = await RoleCrud.first(verify, [Role.name == "lock_rollback_role"]) @@ -367,6 +533,20 @@ class TestAdvisoryLock: assert a1 is True assert a2 is True + @pytest.mark.anyio + async def test_acquire_does_not_flush_pending(self, db_session: AsyncSession): + """Acquiring the lock must not autoflush the caller's pending ORM changes. + + Guards the SQLAlchemy 2.1 behavior where raw ``text()`` autoflushes too; + the helper wraps lock SQL in ``no_autoflush`` to preserve v4 semantics. + """ + role = Role(name="not_flushed_by_lock") + db_session.add(role) + + async with advisory_lock(db_session, 2001): + # The pending INSERT must still be unflushed inside the lock. + assert role in db_session.new + @pytest.mark.anyio async def test_tuple_key(self, db_session: AsyncSession): """(int, int) key variant acquires the lock.""" @@ -610,7 +790,7 @@ class TestM2MAdd: db_session.add_all([post, tag]) await db_session.flush() - async with get_transaction(db_session): + async with transaction(db_session): await m2m_add(db_session, post, Post.tags, tag) result = await db_session.execute( @@ -634,7 +814,7 @@ class TestM2MAdd: db_session.add_all([post, tag1, tag2, tag3]) await db_session.flush() - async with get_transaction(db_session): + async with transaction(db_session): await m2m_add(db_session, post, Post.tags, tag1, tag2, tag3) result = await db_session.execute( @@ -654,7 +834,7 @@ class TestM2MAdd: db_session.add(post) await db_session.flush() - async with get_transaction(db_session): + async with transaction(db_session): await m2m_add(db_session, post, Post.tags) # no related instances result = await db_session.execute( @@ -675,11 +855,11 @@ class TestM2MAdd: db_session.add_all([post, tag]) await db_session.flush() - async with get_transaction(db_session): + async with transaction(db_session): await m2m_add(db_session, post, Post.tags, tag) # Second call with ignore_conflicts=True must not raise - async with get_transaction(db_session): + async with transaction(db_session): await m2m_add(db_session, post, Post.tags, tag, ignore_conflicts=True) result = await db_session.execute( @@ -700,11 +880,11 @@ class TestM2MAdd: db_session.add_all([post, tag]) await db_session.flush() - async with get_transaction(db_session): + async with transaction(db_session): await m2m_add(db_session, post, Post.tags, tag) with pytest.raises(IntegrityError): - async with get_transaction(db_session): + async with transaction(db_session): await m2m_add(db_session, post, Post.tags, tag) @pytest.mark.anyio @@ -752,46 +932,36 @@ class TestDbErrors: """Tests for structured error handling in db utilities.""" @pytest.mark.anyio - async def test_pool_exhausted_on_get_db_raises_pool_exhausted_error(self): - """PoolExhaustedError is raised when the connection pool is exhausted on get_db.""" - engine = create_async_engine( - DATABASE_URL, pool_size=1, max_overflow=0, pool_timeout=0.1 - ) - session_factory = async_sessionmaker(engine, expire_on_commit=False) - get_db = create_db_dependency(session_factory) - + async def test_pool_exhausted_on_dependency_raises_pool_exhausted_error(self): + """PoolExhaustedError is raised when the pool is exhausted on the dependency.""" + db = Database(DATABASE_URL, pool_size=1, max_overflow=0, pool_timeout=0.1) try: - async with session_factory() as holder: - await holder.connection() # check out the single available connection + async with db.session(): # check out the single available connection with pytest.raises(PoolExhaustedError): - async for _ in get_db(): + async for _ in db(_make_request()): pass finally: - await engine.dispose() + await db.engine.dispose() @pytest.mark.anyio async def test_pool_exhausted_on_lock_tables_raises_pool_exhausted_error(self): - """PoolExhaustedError is raised when the connection pool is exhausted on lock_tables.""" - engine = create_async_engine( - DATABASE_URL, pool_size=1, max_overflow=0, pool_timeout=0.1 - ) - session_factory = async_sessionmaker(engine, expire_on_commit=False) - + """PoolExhaustedError is raised when the pool is exhausted on lock_tables.""" + db = Database(DATABASE_URL, pool_size=1, max_overflow=0, pool_timeout=0.1) try: - async with session_factory() as holder: - await holder.connection() # check out the single available connection + async with db.session(): # check out the single available connection with pytest.raises(PoolExhaustedError): - async with lock_tables(session_factory, [Role]) as _: + async with db.lock_tables([Role]): pass finally: - await engine.dispose() + await db.engine.dispose() @pytest.mark.anyio - async def test_lock_timeout_raises_lock_timeout_error(self, session_maker): + async def test_lock_timeout_raises_lock_timeout_error(self, engine, session_maker): """LockTimeoutError is raised when a table lock cannot be acquired within timeout.""" - async with lock_tables(session_maker, [Role]) as _: + db = Database(engine=engine) + async with db.lock_tables([Role]): with pytest.raises(LockTimeoutError): - async with lock_tables(session_maker, [Role], timeout="100ms") as _: + async with db.lock_tables([Role], timeout="100ms"): pass @@ -841,7 +1011,7 @@ class TestM2MRemove: session.add_all(tags) await session.flush() - async with get_transaction(session): + async with transaction(session): await m2m_add(session, post, Post.tags, *tags) return post, tags @@ -859,7 +1029,7 @@ class TestM2MRemove: db_session, "rm_author1", "rm1@test.com", "tag_rm_a", "tag_rm_b" ) - async with get_transaction(db_session): + async with transaction(db_session): await m2m_remove(db_session, post, Post.tags, tag1) remaining = await self._load_tags(db_session, post) @@ -873,7 +1043,7 @@ class TestM2MRemove: db_session, "rm_author2", "rm2@test.com", "tag_rm_c", "tag_rm_d", "tag_rm_e" ) - async with get_transaction(db_session): + async with transaction(db_session): await m2m_remove(db_session, post, Post.tags, tag1, tag3) remaining = await self._load_tags(db_session, post) @@ -887,7 +1057,7 @@ class TestM2MRemove: db_session, "rm_author3", "rm3@test.com", "tag_rm_f" ) - async with get_transaction(db_session): + async with transaction(db_session): await m2m_remove(db_session, post, Post.tags) remaining = await self._load_tags(db_session, post) @@ -904,7 +1074,7 @@ class TestM2MRemove: await db_session.flush() # tag2 was never associated — should not raise - async with get_transaction(db_session): + async with transaction(db_session): await m2m_remove(db_session, post, Post.tags, tag2) remaining = await self._load_tags(db_session, post) @@ -940,18 +1110,15 @@ class TestM2MRemove: session.add_all([owner, item1, item2]) await session.flush() - async with get_transaction(session): + async with transaction(session): await m2m_add(session, owner, _CompOwner.items, item1, item2) - async with get_transaction(session): + async with transaction(session): await m2m_remove(session, owner, _CompOwner.items, item1) await session.commit() async with session_factory() as verify: - from sqlalchemy import select - from sqlalchemy.orm import selectinload - result = await verify.execute( select(_CompOwner) .where(_CompOwner.id == owner.id) @@ -992,10 +1159,10 @@ class TestM2MSet: db_session.add_all([post, tag1, tag2, tag3]) await db_session.flush() - async with get_transaction(db_session): + async with transaction(db_session): await m2m_add(db_session, post, Post.tags, tag1, tag2) - async with get_transaction(db_session): + async with transaction(db_session): await m2m_set(db_session, post, Post.tags, tag3) remaining = await self._load_tags(db_session, post) @@ -1014,10 +1181,10 @@ class TestM2MSet: db_session.add_all([post, tag]) await db_session.flush() - async with get_transaction(db_session): + async with transaction(db_session): await m2m_add(db_session, post, Post.tags, tag) - async with get_transaction(db_session): + async with transaction(db_session): await m2m_set(db_session, post, Post.tags) remaining = await self._load_tags(db_session, post) @@ -1036,7 +1203,7 @@ class TestM2MSet: db_session.add_all([post, tag1, tag2]) await db_session.flush() - async with get_transaction(db_session): + async with transaction(db_session): await m2m_set(db_session, post, Post.tags, tag1, tag2) remaining = await self._load_tags(db_session, post) @@ -1055,3 +1222,278 @@ class TestM2MSet: with pytest.raises(TypeError, match="Many-to-Many"): await m2m_set(db_session, user, User.role, role) + + +STATE_ATTR = "test_db_session" + + +class _FakeSession: + """Records commit() calls into a shared event log.""" + + def __init__(self, events: list[str], *, in_txn: bool = True) -> None: + self.events = events + self._in_txn = in_txn + self.commits = 0 + + def in_transaction(self) -> bool: + return self._in_txn + + async def commit(self) -> None: + self.events.append("COMMIT") + self.commits += 1 + self._in_txn = False + + +async def _drive(app, scope, events: list[str]) -> list[str]: + """Run an ASGI app, appending the response messages it emits to *events* + (shared with the fake session so commit/response ordering is captured).""" + + async def receive(): # pragma: no cover - not exercised + return {"type": "http.disconnect"} + + async def send(message) -> None: + events.append(message["type"]) + + await app(scope, receive, send) + return events + + +class TestCommitOrdering: + """The commit must precede the forwarded response, and be skipped otherwise.""" + + @pytest.mark.anyio + async def test_commits_before_response_start(self): + events: list[str] = [] + session = _FakeSession(events) + + async def inner(scope, receive, send): + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"ok"}) + + app = _CommitOnResponseMiddleware(inner, state_attr=STATE_ATTR) + scope = {"type": "http", "state": {STATE_ATTR: session}} + + result = await _drive(app, scope, events) + + assert session.commits == 1 + assert result == ["COMMIT", "http.response.start", "http.response.body"] + + @pytest.mark.anyio + async def test_no_commit_when_not_in_transaction(self): + events: list[str] = [] + session = _FakeSession(events, in_txn=False) + + async def inner(scope, receive, send): + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b""}) + + app = _CommitOnResponseMiddleware(inner, state_attr=STATE_ATTR) + scope = {"type": "http", "state": {STATE_ATTR: session}} + + result = await _drive(app, scope, events) + + assert session.commits == 0 + assert result == ["http.response.start", "http.response.body"] + + @pytest.mark.anyio + async def test_no_session_is_noop(self): + events: list[str] = [] + + async def inner(scope, receive, send): + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b""}) + + app = _CommitOnResponseMiddleware(inner, state_attr=STATE_ATTR) + scope = {"type": "http", "state": {}} + + result = await _drive(app, scope, events) + + assert result == ["http.response.start", "http.response.body"] + + @pytest.mark.anyio + async def test_non_http_scope_passes_through(self): + called = False + + async def inner(scope, receive, send): + nonlocal called + called = True + + async def receive(): # pragma: no cover - not exercised + return {"type": "lifespan.startup"} + + async def send(message): # pragma: no cover - not exercised + return None + + app = _CommitOnResponseMiddleware(inner, state_attr=STATE_ATTR) + await app({"type": "lifespan"}, receive, send) + + assert called is True + + +class _ProbeMiddleware: + """Outer middleware that records, at response start, whether a row created + in the request is already visible to a *separate* session.""" + + def __init__(self, app, *, session_maker, name: str, result: dict) -> None: + self.app = app + self.session_maker = session_maker + self.name = name + self.result = result + + async def __call__(self, scope, receive, send): + async def send_wrapper(message): + if message["type"] == "http.response.start": + async with self.session_maker() as probe: + row = ( + await probe.execute(select(Role).where(Role.name == self.name)) + ).scalar_one_or_none() + self.result["visible_at_start"] = row is not None + await send(message) + + await self.app(scope, receive, send_wrapper) + + +def _build_app(db: Database) -> FastAPI: + """A FastAPI app wired with the Database dependency and commit middleware.""" + app = FastAPI() + + @app.post("/roles") + async def create_role( + body: RoleCreate, session: AsyncSession = Depends(db) + ) -> dict: + role = await RoleCrud.create(session, body) + return {"id": str(role.id), "name": role.name} + + @app.post("/roles-then-boom") + async def create_then_raise( + body: RoleCreate, session: AsyncSession = Depends(db) + ) -> dict: + await RoleCrud.create(session, body) + raise RuntimeError("boom after write") + + @app.post("/two-roles") + async def create_two_roles( + body: RoleCreate, session: AsyncSession = Depends(db) + ) -> dict: + # First write succeeds, second collides on the unique name and must + # take the whole request transaction down with it. + await RoleCrud.create(session, body) + await RoleCrud.create(session, body) + return {"ok": True} + + @app.post("/roles-self-commit") + async def create_then_self_commit( + body: RoleCreate, session: AsyncSession = Depends(db) + ) -> dict: + # Endpoint commits explicitly; the middleware must not double-commit or + # error — it finds no open transaction and no-ops. + role = await RoleCrud.create(session, body) + await session.commit() + return {"id": str(role.id), "name": role.name} + + @app.get("/roles-stream/{name}") + async def stream_role( + name: str, session: AsyncSession = Depends(db) + ) -> StreamingResponse: + # A write before the stream begins: the middleware commits it at + # response-start, before the generator runs. + await RoleCrud.create(session, RoleCreate(name=name)) + + async def gen(): + # Read-only DB use during the stream, via the request session. + row = ( + await session.execute(select(Role).where(Role.name == name)) + ).scalar_one() + yield f"data: {row.name}\n\n".encode() + + return StreamingResponse(gen(), media_type="text/event-stream") + + db.install(app) + return app + + +async def _row_exists(session_maker, name: str) -> bool: + async with session_maker() as session: + row = ( + await session.execute(select(Role).where(Role.name == name)) + ).scalar_one_or_none() + return row is not None + + +class TestCommitIntegration: + @pytest.mark.anyio + async def test_write_is_committed(self, engine, session_maker): + app = _build_app(Database(engine=engine)) + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as client: + resp = await client.post("/roles", json={"name": "committed_role"}) + + assert resp.status_code == 200 + assert await _row_exists(session_maker, "committed_role") + + @pytest.mark.anyio + async def test_visible_at_response_start(self, engine, session_maker): + """The write is visible to a separate session *before* the response is + sent — the read-after-write guarantee the middleware exists for.""" + app = _build_app(Database(engine=engine)) + result: dict = {} + app.add_middleware( + _ProbeMiddleware, + session_maker=session_maker, + name="probe_role", + result=result, + ) + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as client: + resp = await client.post("/roles", json={"name": "probe_role"}) + + assert resp.status_code == 200 + assert result.get("visible_at_start") is True + + @pytest.mark.anyio + async def test_error_rolls_back(self, engine, session_maker): + app = _build_app(Database(engine=engine)) + transport = ASGITransport(app=app, raise_app_exceptions=False) + async with AsyncClient(transport=transport, base_url="http://test") as client: + resp = await client.post("/roles-then-boom", json={"name": "ghost_role"}) + + assert resp.status_code == 500 + assert not await _row_exists(session_maker, "ghost_role") + + @pytest.mark.anyio + async def test_explicit_commit_in_endpoint(self, engine, session_maker): + """An endpoint that commits itself works: the middleware no-ops (no + double commit / error) and the write is persisted.""" + app = _build_app(Database(engine=engine)) + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as client: + resp = await client.post("/roles-self-commit", json={"name": "self_commit"}) + + assert resp.status_code == 200 + assert await _row_exists(session_maker, "self_commit") + + @pytest.mark.anyio + async def test_streaming_response_coexists(self, engine, session_maker): + """A read-only streaming endpoint works alongside the middleware: the + commit fires at stream start, the pre-stream write is committed, and the + generator can keep reading via the request session.""" + app = _build_app(Database(engine=engine)) + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as client: + resp = await client.get("/roles-stream/streamed_role") + + assert resp.status_code == 200 + assert "data: streamed_role" in resp.text + # The write made before the stream began is durably committed. + assert await _row_exists(session_maker, "streamed_role") + + @pytest.mark.anyio + async def test_multi_write_atomicity(self, engine, session_maker): + """When the 2nd write fails, the 1st must roll back too (one txn).""" + app = _build_app(Database(engine=engine)) + transport = ASGITransport(app=app, raise_app_exceptions=False) + async with AsyncClient(transport=transport, base_url="http://test") as client: + resp = await client.post("/two-roles", json={"name": "dup_role"}) + + assert resp.status_code >= 400 + assert not await _row_exists(session_maker, "dup_role") diff --git a/tests/test_example_pagination_search.py b/tests/test_example_pagination_search.py index c01d461..baee6d6 100644 --- a/tests/test_example_pagination_search.py +++ b/tests/test_example_pagination_search.py @@ -91,13 +91,19 @@ async def seed(session: AsyncSession): class TestAppSessionDep: @pytest.mark.anyio async def test_get_db_yields_async_session(self): - """get_db yields a real AsyncSession when called directly.""" - from docs_src.examples.pagination_search.db import get_db + """The Database dependency yields a real AsyncSession when called directly.""" + from starlette.requests import Request - gen = get_db() - session = await gen.__anext__() - assert isinstance(session, AsyncSession) - await gen.aclose() + from fastapi_toolsets.db import Database + + db = Database(DATABASE_URL) + try: + gen = db(Request({"type": "http", "headers": []})) + session = await gen.__anext__() + assert isinstance(session, AsyncSession) + await gen.aclose() + finally: + await db.engine.dispose() class TestOffsetPagination: diff --git a/tests/test_models.py b/tests/test_models.py index 8f98ecd..75222c2 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1506,8 +1506,8 @@ class TestListensFor: assert all(e["event"] == "change" for e in _listener_events) -class TestEventSessionWithGetTransaction: - """Verify callbacks fire correctly when using get_transaction / lock_tables.""" +class TestEventSessionWithTransaction: + """Verify callbacks fire correctly when using transaction / lock_tables.""" @pytest.fixture(autouse=True) def clear_events(self): @@ -1517,10 +1517,10 @@ class TestEventSessionWithGetTransaction: @pytest.mark.anyio async def test_callbacks_fire_after_outer_commit_not_savepoint(self, mixin_session): - """get_transaction creates a savepoint; callbacks fire only on outer commit.""" - from fastapi_toolsets.db import get_transaction + """transaction creates a savepoint; callbacks fire only on outer commit.""" + from fastapi_toolsets.db import transaction - async with get_transaction(mixin_session): + async with transaction(mixin_session): obj = WatchedModel(status="active", other="x") mixin_session.add(obj) @@ -1535,14 +1535,14 @@ class TestEventSessionWithGetTransaction: @pytest.mark.anyio async def test_nested_transactions_accumulate_events(self, mixin_session): - """Multiple get_transaction blocks accumulate events for a single commit.""" - from fastapi_toolsets.db import get_transaction + """Multiple transaction blocks accumulate events for a single commit.""" + from fastapi_toolsets.db import transaction - async with get_transaction(mixin_session): + async with transaction(mixin_session): obj1 = WatchedModel(status="first", other="x") mixin_session.add(obj1) - async with get_transaction(mixin_session): + async with transaction(mixin_session): obj2 = WatchedModel(status="second", other="y") mixin_session.add(obj2) @@ -1556,14 +1556,14 @@ class TestEventSessionWithGetTransaction: @pytest.mark.anyio async def test_savepoint_rollback_suppresses_events(self, mixin_session): """Objects from a rolled-back savepoint don't fire callbacks.""" - from fastapi_toolsets.db import get_transaction + from fastapi_toolsets.db import transaction survivor = WatchedModel(status="kept", other="x") mixin_session.add(survivor) await mixin_session.flush() try: - async with get_transaction(mixin_session): + async with transaction(mixin_session): doomed = WatchedModel(status="doomed", other="y") mixin_session.add(doomed) await mixin_session.flush() @@ -1590,9 +1590,9 @@ class TestEventSessionWithGetTransaction: assert len(creates) == 1 @pytest.mark.anyio - async def test_update_inside_get_transaction(self, mixin_session): - """UPDATE events fire with correct changes after get_transaction commit.""" - from fastapi_toolsets.db import get_transaction + async def test_update_inside_transaction(self, mixin_session): + """UPDATE events fire with correct changes after transaction commit.""" + from fastapi_toolsets.db import transaction obj = WatchedModel(status="initial", other="x") mixin_session.add(obj) @@ -1600,7 +1600,7 @@ class TestEventSessionWithGetTransaction: _test_events.clear() - async with get_transaction(mixin_session): + async with transaction(mixin_session): obj.status = "updated" await mixin_session.commit() @@ -1696,7 +1696,7 @@ class TestEventSessionWithNullableFields: class TestEventSessionWithFastAPIDependency: - """Verify EventSession works when session comes from create_db_dependency.""" + """Verify EventSession works when session comes from the Database dependency.""" @pytest.fixture(autouse=True) def clear_events(self): @@ -1706,31 +1706,24 @@ class TestEventSessionWithFastAPIDependency: @pytest.mark.anyio async def test_create_event_fires_via_dependency(self): - """CREATE callback fires when session is provided by create_db_dependency.""" + """CREATE callback fires when session is provided by the Database dependency.""" from fastapi import Depends, FastAPI from httpx import ASGITransport, AsyncClient - from sqlalchemy.ext.asyncio import ( - AsyncSession, - async_sessionmaker, - create_async_engine, - ) + from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine - from fastapi_toolsets.db import create_db_dependency + from fastapi_toolsets.db import Database from fastapi_toolsets.models import EventSession engine = create_async_engine(DATABASE_URL, echo=False) - session_factory = async_sessionmaker( - engine, expire_on_commit=False, class_=EventSession - ) async with engine.begin() as conn: await conn.run_sync(MixinBase.metadata.create_all) - get_db = create_db_dependency(session_factory) + db = Database(engine=engine, session_class=EventSession) app = FastAPI() @app.post("/watched") - async def create_watched(session: AsyncSession = Depends(get_db)): + async def create_watched(session: AsyncSession = Depends(db)): obj = WatchedModel(status="from-api", other="x") session.add(obj) return {"id": str(obj.id)} @@ -1753,40 +1746,33 @@ class TestEventSessionWithFastAPIDependency: @pytest.mark.anyio async def test_update_event_fires_via_dependency(self): - """UPDATE callback fires when session is provided by create_db_dependency.""" + """UPDATE callback fires when session is provided by the Database dependency.""" from fastapi import Depends, FastAPI from httpx import ASGITransport, AsyncClient - from sqlalchemy.ext.asyncio import ( - AsyncSession, - async_sessionmaker, - create_async_engine, - ) + from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine - from fastapi_toolsets.db import create_db_dependency + from fastapi_toolsets.db import Database from fastapi_toolsets.models import EventSession engine = create_async_engine(DATABASE_URL, echo=False) - session_factory = async_sessionmaker( - engine, expire_on_commit=False, class_=EventSession - ) async with engine.begin() as conn: await conn.run_sync(MixinBase.metadata.create_all) - get_db = create_db_dependency(session_factory) + db = Database(engine=engine, session_class=EventSession) app = FastAPI() # Pre-seed an object. - async with session_factory() as seed_session: + async with db.session() as seed_session: obj = WatchedModel(status="initial", other="x") seed_session.add(obj) - await seed_session.commit() + await seed_session.flush() obj_id = obj.id _test_events.clear() @app.put("/watched/{item_id}") - async def update_watched(item_id: str, session: AsyncSession = Depends(get_db)): + async def update_watched(item_id: str, session: AsyncSession = Depends(db)): from sqlalchemy import select stmt = select(WatchedModel).where(WatchedModel.id == item_id) diff --git a/tests/test_pytest.py b/tests/test_pytest.py index 6cd01e1..2063024 100644 --- a/tests/test_pytest.py +++ b/tests/test_pytest.py @@ -11,7 +11,7 @@ from sqlalchemy.engine import make_url from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship -from fastapi_toolsets.db import get_transaction +from fastapi_toolsets.db import transaction from fastapi_toolsets.fixtures import Context, FixtureRegistry, LoadStrategy from fastapi_toolsets.pytest import ( create_async_client, @@ -387,14 +387,14 @@ class TestCreateDbSession: assert session.autoflush is False @pytest.mark.anyio - async def test_get_transaction_commits_visible_to_separate_session(self): - """Data written via get_transaction() is committed and visible to other sessions.""" + async def test_transaction_commits_visible_to_separate_session(self): + """Data written via transaction() is committed and visible to other sessions.""" role_id = uuid.uuid4() async with create_db_session(DATABASE_URL, Base, drop_tables=False) as session: - # Simulate what _create_fixture_function does: insert via get_transaction + # Simulate what _create_fixture_function does: insert via transaction() # with no explicit commit afterward. - async with get_transaction(session): + async with transaction(session): role = Role(id=role_id, name="visible_to_other_session") session.add(role) @@ -409,9 +409,9 @@ class TestCreateDbSession: result = await other.execute(select(Role).where(Role.id == role_id)) fetched = result.scalar_one_or_none() assert fetched is not None, ( - "Fixture data inserted via get_transaction() must be committed " + "Fixture data inserted via transaction() must be committed " "and visible to a separate session. If create_db_session uses " - "create_db_context, auto-begin forces get_transaction() into " + "db.session(), auto-begin forces transaction() into " "savepoints instead of real commits." ) assert fetched.name == "visible_to_other_session"