mirror of
https://github.com/d3vyce/fastapi-toolsets.git
synced 2026-08-04 23:54:09 +00:00
323 lines
11 KiB
Python
323 lines
11 KiB
Python
"""The ``Database`` facade: session lifecycle, dependency, middleware, transactions."""
|
|
|
|
from collections.abc import AsyncGenerator
|
|
from contextlib import AbstractAsyncContextManager, asynccontextmanager
|
|
from typing import Any
|
|
|
|
from pydantic import PostgresDsn
|
|
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. Accepts a plain string or a Pydantic
|
|
:class:`~pydantic.PostgresDsn`.
|
|
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``.
|
|
connect_args: DBAPI-level connection arguments forwarded to
|
|
:func:`create_async_engine` (URL mode only).
|
|
**engine_options: Extra keyword arguments forwarded to
|
|
:func:`create_async_engine` (URL mode only).
|
|
|
|
Raises:
|
|
TypeError: If neither or both of *url* and *engine* are given, or if
|
|
*connect_args*/*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 | PostgresDsn | None = None,
|
|
*,
|
|
engine: AsyncEngine | None = None,
|
|
session_class: type[AsyncSession] = AsyncSession,
|
|
expire_on_commit: bool = False,
|
|
autoflush: bool = True,
|
|
connect_args: dict[str, Any] | None = None,
|
|
**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 or connect_args is not None):
|
|
raise TypeError(
|
|
"connect_args/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
|
|
if connect_args is not None:
|
|
engine_options["connect_args"] = connect_args
|
|
# ``PostgresDsn`` (and other URL objects) are not str subclasses, so
|
|
# coerce to the string form SQLAlchemy expects.
|
|
self.engine = create_async_engine(str(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_), 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)
|