fastapi-toolsets 4.1.3__tar.gz → 5.0.0b1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (52) hide show
  1. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/PKG-INFO +1 -1
  2. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/pyproject.toml +1 -1
  3. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/__init__.py +6 -3
  4. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/crud/factory.py +5 -5
  5. fastapi_toolsets-5.0.0b1/src/fastapi_toolsets/db/__init__.py +18 -0
  6. fastapi_toolsets-5.0.0b1/src/fastapi_toolsets/db/core.py +315 -0
  7. fastapi_toolsets-5.0.0b1/src/fastapi_toolsets/db/locks.py +185 -0
  8. fastapi_toolsets-5.0.0b1/src/fastapi_toolsets/db/m2m.py +170 -0
  9. fastapi_toolsets-5.0.0b1/src/fastapi_toolsets/db/testing.py +69 -0
  10. fastapi_toolsets-5.0.0b1/src/fastapi_toolsets/db/watch.py +90 -0
  11. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/fixtures/utils.py +2 -2
  12. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/pytest/plugin.py +2 -2
  13. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/pytest/utils.py +1 -1
  14. fastapi_toolsets-4.1.3/src/fastapi_toolsets/db.py +0 -591
  15. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/LICENSE +0 -0
  16. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/README.md +0 -0
  17. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/_imports.py +0 -0
  18. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/cli/__init__.py +0 -0
  19. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/cli/app.py +0 -0
  20. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/cli/commands/__init__.py +0 -0
  21. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/cli/commands/fixtures.py +0 -0
  22. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/cli/config.py +0 -0
  23. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/cli/pyproject.py +0 -0
  24. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/cli/utils.py +0 -0
  25. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/crud/__init__.py +0 -0
  26. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/crud/search.py +0 -0
  27. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/dependencies.py +0 -0
  28. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/exceptions/__init__.py +0 -0
  29. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/exceptions/exceptions.py +0 -0
  30. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/exceptions/handler.py +0 -0
  31. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/fixtures/__init__.py +0 -0
  32. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/fixtures/enum.py +0 -0
  33. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/fixtures/registry.py +0 -0
  34. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/logger.py +0 -0
  35. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/metrics/__init__.py +0 -0
  36. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/metrics/handler.py +0 -0
  37. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/metrics/registry.py +0 -0
  38. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/models/__init__.py +0 -0
  39. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/models/columns.py +0 -0
  40. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/models/watched.py +0 -0
  41. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/py.typed +0 -0
  42. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/pytest/__init__.py +0 -0
  43. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/schemas.py +0 -0
  44. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/security/__init__.py +0 -0
  45. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/security/abc.py +0 -0
  46. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/security/oauth.py +0 -0
  47. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/security/sources/__init__.py +0 -0
  48. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/security/sources/bearer.py +0 -0
  49. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/security/sources/cookie.py +0 -0
  50. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/security/sources/header.py +0 -0
  51. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/security/sources/multi.py +0 -0
  52. {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/types.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: fastapi-toolsets
3
- Version: 4.1.3
3
+ Version: 5.0.0b1
4
4
  Summary: Production-ready utilities for FastAPI applications
5
5
  Keywords: fastapi,sqlalchemy,postgresql
6
6
  Author: d3vyce
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "fastapi-toolsets"
3
- version = "4.1.3"
3
+ version = "5.0.0b1"
4
4
  description = "Production-ready utilities for FastAPI applications"
5
5
  readme = "README.md"
6
6
  license = "MIT"
@@ -7,18 +7,21 @@ Example usage:
7
7
  from fastapi import FastAPI, Depends
8
8
  from fastapi_toolsets.exceptions import init_exceptions_handlers
9
9
  from fastapi_toolsets.crud import CrudFactory
10
- from fastapi_toolsets.db import create_db_dependency
10
+ from fastapi_toolsets.db import Database
11
11
  from fastapi_toolsets.schemas import Response
12
12
 
13
+ db = Database("postgresql+asyncpg://postgres:postgres@localhost/app")
14
+
13
15
  app = FastAPI()
16
+ db.install(app)
14
17
  init_exceptions_handlers(app)
15
18
 
16
19
  UserCrud = CrudFactory(User)
17
20
 
18
21
  @app.get("/users/{user_id}", response_model=Response[dict])
19
- async def get_user(user_id: int, session = Depends(get_db)):
22
+ async def get_user(user_id: int, session = Depends(db)):
20
23
  user = await UserCrud.get(session, [User.id == user_id])
21
24
  return Response(data={"user": user.username}, message="Success")
22
25
  """
23
26
 
24
- __version__ = "4.1.3"
27
+ __version__ = "5.0.0b1"
@@ -22,7 +22,7 @@ from sqlalchemy.orm import DeclarativeBase, QueryableAttribute, selectinload
22
22
  from sqlalchemy.sql.base import ExecutableOption
23
23
  from sqlalchemy.sql.roles import WhereHavingRole
24
24
 
25
- from ..db import get_transaction
25
+ from ..db import transaction
26
26
  from ..exceptions import InvalidOrderFieldError, NotFoundError
27
27
  from ..schemas import (
28
28
  CursorPaginatedResponse,
@@ -716,7 +716,7 @@ class AsyncCrud(Generic[ModelType]):
716
716
  Returns:
717
717
  Created model instance, or ``Response[schema]`` when ``schema`` is given.
718
718
  """
719
- async with get_transaction(session):
719
+ async with transaction(session):
720
720
  m2m_exclude = cls._m2m_schema_fields()
721
721
  data = (
722
722
  obj.model_dump(exclude=m2m_exclude) if m2m_exclude else obj.model_dump()
@@ -1067,7 +1067,7 @@ class AsyncCrud(Generic[ModelType]):
1067
1067
  Raises:
1068
1068
  NotFoundError: If no record found
1069
1069
  """
1070
- async with get_transaction(session):
1070
+ async with transaction(session):
1071
1071
  m2m_exclude = cls._m2m_schema_fields()
1072
1072
 
1073
1073
  # Eagerly load M2M relationships that will be updated so that
@@ -1127,7 +1127,7 @@ class AsyncCrud(Generic[ModelType]):
1127
1127
  Returns:
1128
1128
  Model instance
1129
1129
  """
1130
- async with get_transaction(session):
1130
+ async with transaction(session):
1131
1131
  values = obj.model_dump(exclude_unset=True)
1132
1132
  q = insert(cls.model).values(**values)
1133
1133
  if set_:
@@ -1189,7 +1189,7 @@ class AsyncCrud(Generic[ModelType]):
1189
1189
  Returns:
1190
1190
  ``None``, or ``Response[None]`` when ``return_response=True``.
1191
1191
  """
1192
- async with get_transaction(session):
1192
+ async with transaction(session):
1193
1193
  result = await session.execute(select(cls.model).where(and_(*filters)))
1194
1194
  objects = result.scalars().all()
1195
1195
  for obj in objects:
@@ -0,0 +1,18 @@
1
+ """Database package: the ``Database`` facade plus PostgreSQL power-tools."""
2
+
3
+ from .core import Database, transaction
4
+ from .locks import LockMode, advisory_lock, lock_tables
5
+ from .m2m import m2m_add, m2m_remove, m2m_set
6
+ from .watch import wait_for_row_change
7
+
8
+ __all__ = [
9
+ "Database",
10
+ "LockMode",
11
+ "advisory_lock",
12
+ "lock_tables",
13
+ "m2m_add",
14
+ "m2m_remove",
15
+ "m2m_set",
16
+ "transaction",
17
+ "wait_for_row_change",
18
+ ]
@@ -0,0 +1,315 @@
1
+ """The ``Database`` facade: session lifecycle, dependency, middleware, transactions."""
2
+
3
+ from collections.abc import AsyncGenerator
4
+ from contextlib import AbstractAsyncContextManager, asynccontextmanager
5
+ from typing import Any
6
+
7
+ from sqlalchemy import exc as sa_exc
8
+ from sqlalchemy.ext.asyncio import (
9
+ AsyncEngine,
10
+ AsyncSession,
11
+ async_sessionmaker,
12
+ create_async_engine,
13
+ )
14
+ from sqlalchemy.orm import DeclarativeBase
15
+ from starlette.requests import Request
16
+ from starlette.types import ASGIApp, Message, Receive, Scope, Send
17
+
18
+ from ..exceptions import PoolExhaustedError
19
+ from .locks import LockMode, lock_tables
20
+
21
+
22
+ @asynccontextmanager
23
+ async def transaction(
24
+ session: AsyncSession,
25
+ ) -> AsyncGenerator[AsyncSession, None]:
26
+ """Run a block inside a savepoint-aware transaction.
27
+
28
+ If *session* is already in a transaction, a nested transaction (savepoint)
29
+ is opened so the block can roll back independently. Otherwise a top-level
30
+ transaction is started. Commits on clean exit, rolls back on exception.
31
+
32
+ Args:
33
+ session: AsyncSession instance.
34
+
35
+ Yields:
36
+ The session within the transaction context.
37
+
38
+ Example:
39
+ ```python
40
+ from fastapi_toolsets.db import transaction
41
+
42
+ async with transaction(session):
43
+ session.add(model)
44
+ ```
45
+ """
46
+ if session.in_transaction():
47
+ async with session.begin_nested():
48
+ yield session
49
+ else:
50
+ async with session.begin():
51
+ yield session
52
+
53
+
54
+ class _CommitOnResponseMiddleware:
55
+ """Commit the request's DB session before the response is sent."""
56
+
57
+ def __init__(self, app: ASGIApp, *, state_attr: str) -> None:
58
+ self.app = app
59
+ self.state_attr = state_attr
60
+
61
+ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
62
+ if scope["type"] != "http":
63
+ await self.app(scope, receive, send)
64
+ return
65
+
66
+ async def send_wrapper(message: Message) -> None:
67
+ if message["type"] == "http.response.start":
68
+ # ``scope["state"]`` is the same dict ``request.state`` writes
69
+ # to, so this is the session stashed by the dependency.
70
+ state = scope.get("state")
71
+ session = state.get(self.state_attr) if state else None
72
+ if session is not None and session.in_transaction():
73
+ await session.commit()
74
+ await send(message)
75
+
76
+ await self.app(scope, receive, send_wrapper)
77
+
78
+
79
+ class Database:
80
+ """One object that owns the engine, sessions, dependency, and middleware.
81
+
82
+ Provide exactly one of *url* (the facade builds and disposes the engine) or
83
+ *engine* (an engine you own, e.g. for Alembic or ``event.listen``, left
84
+ untouched).
85
+
86
+ Args:
87
+ url: Database connection URL (e.g. ``"postgresql+asyncpg://..."``).
88
+ engine: An existing :class:`AsyncEngine` to reuse instead of *url*.
89
+ session_class: Session class for the sessionmaker (e.g. ``EventSession``).
90
+ expire_on_commit: Expire attributes after commit. Defaults to ``False``.
91
+ autoflush: Autoflush the session before queries. Defaults to ``True``.
92
+ **engine_options: Extra keyword arguments forwarded to
93
+ :func:`create_async_engine` (URL mode only, e.g. ``pool_size``,
94
+ ``echo``, ``connect_args``).
95
+
96
+ Raises:
97
+ TypeError: If neither or both of *url* and *engine* are given, or if
98
+ *engine_options* are passed together with *engine*.
99
+
100
+ Example:
101
+ ```python
102
+ from fastapi import Depends, FastAPI
103
+ from fastapi_toolsets.db import Database
104
+
105
+ db = Database("postgresql+asyncpg://postgres:postgres@localhost/app")
106
+
107
+ app = FastAPI()
108
+ db.install(app)
109
+
110
+ @app.get("/users/{user_id}")
111
+ async def get_user(user_id: int, session=Depends(db)):
112
+ return await UserCrud.get(session, [User.id == user_id])
113
+ ```
114
+ """
115
+
116
+ def __init__(
117
+ self,
118
+ url: str | None = None,
119
+ *,
120
+ engine: AsyncEngine | None = None,
121
+ session_class: type[AsyncSession] = AsyncSession,
122
+ expire_on_commit: bool = False,
123
+ autoflush: bool = True,
124
+ **engine_options: Any,
125
+ ) -> None:
126
+ if (url is None) == (engine is None):
127
+ raise TypeError(
128
+ "Database requires exactly one of 'url' or 'engine' "
129
+ "(got both or neither)."
130
+ )
131
+ if engine is not None and engine_options:
132
+ raise TypeError(
133
+ "engine_options are only valid in URL mode; configure the "
134
+ "engine you pass via 'engine=' yourself."
135
+ )
136
+
137
+ if engine is not None:
138
+ self._owns_engine = False
139
+ self.engine: AsyncEngine = engine
140
+ else:
141
+ assert url is not None # guaranteed by the XOR check above
142
+ self._owns_engine = True
143
+ self.engine = create_async_engine(url, **engine_options)
144
+ self._sessionmaker: async_sessionmaker[AsyncSession] = async_sessionmaker(
145
+ self.engine,
146
+ class_=session_class,
147
+ expire_on_commit=expire_on_commit,
148
+ autoflush=autoflush,
149
+ )
150
+ # Private, per-instance state attribute; cannot collide with another
151
+ # Database or be mismatched against the middleware.
152
+ self._state_attr = f"_ft_db_session_{id(self):x}"
153
+ self._middleware_installed = False
154
+ self._disposed = False
155
+
156
+ async def _dispose(self) -> None:
157
+ """Dispose the engine once, only if we own it (idempotent)."""
158
+ if self._owns_engine and not self._disposed:
159
+ self._disposed = True
160
+ await self.engine.dispose()
161
+
162
+ @asynccontextmanager
163
+ async def lifespan(self, app: Any) -> AsyncGenerator[None, None]:
164
+ """Dispose the engine on shutdown; use as ``FastAPI(lifespan=db.lifespan)``.
165
+
166
+ Args:
167
+ app: The ASGI application (unused; required by the lifespan protocol).
168
+
169
+ Yields:
170
+ Control to the application for its lifetime.
171
+
172
+ Example:
173
+ ```python
174
+ app = FastAPI(lifespan=db.lifespan)
175
+ ```
176
+ """
177
+ try:
178
+ yield
179
+ finally:
180
+ await self._dispose()
181
+
182
+ def install(self, app: Any) -> None:
183
+ """Wire the commit middleware and engine disposal onto *app*.
184
+
185
+ Args:
186
+ app: The FastAPI/Starlette application to wire.
187
+
188
+ Example:
189
+ ```python
190
+ @asynccontextmanager
191
+ async def lifespan(app):
192
+ ... # your startup
193
+ yield
194
+ ... # your shutdown
195
+
196
+ app = FastAPI(lifespan=lifespan)
197
+ db.install(app)
198
+ ```
199
+ """
200
+ app.add_middleware(_CommitOnResponseMiddleware, state_attr=self._state_attr)
201
+ self._middleware_installed = True
202
+
203
+ inner_lifespan = app.router.lifespan_context
204
+
205
+ @asynccontextmanager
206
+ async def _composed(app_: Any) -> AsyncGenerator[None, None]:
207
+ async with self.lifespan(app_):
208
+ async with inner_lifespan(app_):
209
+ yield
210
+
211
+ app.router.lifespan_context = _composed
212
+
213
+ @asynccontextmanager
214
+ async def _open(self) -> AsyncGenerator[AsyncSession, None]:
215
+ """Open a session and eagerly acquire a connection (fail-fast on pool)."""
216
+ async with self._sessionmaker() as session:
217
+ try:
218
+ await session.connection()
219
+ except sa_exc.TimeoutError as e:
220
+ raise PoolExhaustedError() from e
221
+ yield session
222
+
223
+ async def __call__(self, request: Request) -> AsyncGenerator[AsyncSession, None]:
224
+ """FastAPI dependency: yield a session and commit once at the right time.
225
+
226
+ Args:
227
+ request: The incoming request (injected by FastAPI).
228
+
229
+ Yields:
230
+ An AsyncSession for the duration of the request.
231
+
232
+ Example:
233
+ ```python
234
+ @app.get("/users/{user_id}")
235
+ async def get_user(user_id: int, session=Depends(db)):
236
+ return await UserCrud.get(session, [User.id == user_id])
237
+ ```
238
+ """
239
+ async with self._open() as session:
240
+ setattr(request.state, self._state_attr, session)
241
+ yield session
242
+ if not self._middleware_installed and session.in_transaction():
243
+ await session.commit()
244
+
245
+ @asynccontextmanager
246
+ async def session(self) -> AsyncGenerator[AsyncSession, None]:
247
+ """Open a session outside request handlers (background tasks, CLI, tests).
248
+
249
+ Commits on clean exit, rolls back on exception.
250
+
251
+ Yields:
252
+ An AsyncSession ready for database operations.
253
+
254
+ Example:
255
+ ```python
256
+ async with db.session() as session:
257
+ user = await UserCrud.get(session, [User.id == 1])
258
+ ```
259
+ """
260
+ async with self._open() as session:
261
+ yield session
262
+ if session.in_transaction():
263
+ await session.commit()
264
+
265
+ @asynccontextmanager
266
+ async def begin(self) -> AsyncGenerator[AsyncSession, None]:
267
+ """Open a session already inside a transaction (sugar for the common case).
268
+
269
+ Equivalent to ``session()`` + :func:`transaction`. Commits on clean exit,
270
+ rolls back on exception.
271
+
272
+ Yields:
273
+ An AsyncSession open within a transaction.
274
+
275
+ Example:
276
+ ```python
277
+ async with db.begin() as session:
278
+ session.add(User(name="ada"))
279
+ ```
280
+ """
281
+ async with self.session() as session, transaction(session):
282
+ yield session
283
+
284
+ def lock_tables(
285
+ self,
286
+ tables: list[type[DeclarativeBase]],
287
+ *,
288
+ mode: LockMode = LockMode.SHARE_UPDATE_EXCLUSIVE,
289
+ timeout: str = "5s",
290
+ ) -> AbstractAsyncContextManager[AsyncSession]:
291
+ """Lock PostgreSQL tables for the duration of a dedicated transaction.
292
+
293
+ Opens its own session from the facade's sessionmaker, changes are
294
+ committed when the context exits.
295
+
296
+ Args:
297
+ tables: List of SQLAlchemy model classes to lock.
298
+ mode: Lock mode (default: ``SHARE UPDATE EXCLUSIVE``).
299
+ timeout: Lock timeout (default: ``"5s"``).
300
+
301
+ Yields:
302
+ The dedicated session, open within the locked transaction.
303
+
304
+ Raises:
305
+ LockTimeoutError: If the lock cannot be acquired within *timeout*.
306
+ PoolExhaustedError: If the connection pool is exhausted.
307
+
308
+ Example:
309
+ ```python
310
+ async with db.lock_tables([User, Account]) as session:
311
+ user = await UserCrud.get(session, [User.id == 1])
312
+ user.balance += 100
313
+ ```
314
+ """
315
+ return lock_tables(self._sessionmaker, tables, mode=mode, timeout=timeout)
@@ -0,0 +1,185 @@
1
+ """PostgreSQL locking helpers: table locks and advisory locks."""
2
+
3
+ from collections.abc import AsyncGenerator
4
+ from contextlib import AbstractAsyncContextManager, asynccontextmanager
5
+ from enum import Enum
6
+ from typing import TypeVar
7
+
8
+ import asyncpg
9
+ from sqlalchemy import exc as sa_exc
10
+ from sqlalchemy import text
11
+ from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
12
+ from sqlalchemy.orm import DeclarativeBase
13
+
14
+ from ..exceptions import LockTimeoutError, PoolExhaustedError
15
+
16
+ _SessionT = TypeVar("_SessionT", bound=AsyncSession)
17
+
18
+
19
+ def _is_lock_not_available(e: sa_exc.DBAPIError) -> bool:
20
+ return e.orig is not None and isinstance(
21
+ e.orig.__cause__, asyncpg.exceptions.LockNotAvailableError
22
+ )
23
+
24
+
25
+ class LockMode(str, Enum):
26
+ """PostgreSQL table lock modes.
27
+
28
+ See: https://www.postgresql.org/docs/current/explicit-locking.html
29
+ """
30
+
31
+ ACCESS_SHARE = "ACCESS SHARE"
32
+ ROW_SHARE = "ROW SHARE"
33
+ ROW_EXCLUSIVE = "ROW EXCLUSIVE"
34
+ SHARE_UPDATE_EXCLUSIVE = "SHARE UPDATE EXCLUSIVE"
35
+ SHARE = "SHARE"
36
+ SHARE_ROW_EXCLUSIVE = "SHARE ROW EXCLUSIVE"
37
+ EXCLUSIVE = "EXCLUSIVE"
38
+ ACCESS_EXCLUSIVE = "ACCESS EXCLUSIVE"
39
+
40
+
41
+ def lock_tables(
42
+ session_maker: async_sessionmaker[_SessionT],
43
+ tables: list[type[DeclarativeBase]],
44
+ *,
45
+ mode: LockMode = LockMode.SHARE_UPDATE_EXCLUSIVE,
46
+ timeout: str = "5s",
47
+ ) -> AbstractAsyncContextManager[_SessionT]:
48
+ """Lock PostgreSQL tables for the duration of a transaction.
49
+
50
+ Prefer the method on a :class:`Database` instance; use this
51
+ directly only when you manage your own session factory.
52
+
53
+ Args:
54
+ session_maker: Async session factory used to create the dedicated
55
+ session.
56
+ tables: List of SQLAlchemy model classes to lock.
57
+ mode: Lock mode (default: SHARE UPDATE EXCLUSIVE).
58
+ timeout: Lock timeout (default: "5s").
59
+
60
+ Yields:
61
+ The dedicated session, open within the locked transaction.
62
+
63
+ Raises:
64
+ LockTimeoutError: If the lock cannot be acquired within *timeout*.
65
+ PoolExhaustedError: If the connection pool is exhausted.
66
+
67
+ Example:
68
+ ```python
69
+ from fastapi_toolsets.db import lock_tables
70
+
71
+ async with lock_tables(session_maker, [User, Account]) as session:
72
+ user = await UserCrud.get(session, [User.id == 1])
73
+ user.balance += 100
74
+ ```
75
+ """
76
+ table_names = ",".join(table.__tablename__ for table in tables)
77
+
78
+ @asynccontextmanager
79
+ async def _lock() -> AsyncGenerator[_SessionT, None]:
80
+ async with session_maker() as session:
81
+ try:
82
+ await session.execute(text(f"SET LOCAL lock_timeout='{timeout}'"))
83
+ await session.execute(text(f"LOCK {table_names} IN {mode.value} MODE"))
84
+ yield session
85
+ await session.commit()
86
+ except sa_exc.TimeoutError as e:
87
+ await session.rollback()
88
+ raise PoolExhaustedError(
89
+ f"Connection pool exhausted while locking '{table_names}'. "
90
+ ) from e
91
+ except sa_exc.DBAPIError as e:
92
+ await session.rollback()
93
+ if _is_lock_not_available(e):
94
+ raise LockTimeoutError(
95
+ f"Lock on '{table_names}' could not be acquired within {timeout}."
96
+ ) from e
97
+ raise # pragma: no cover
98
+ except BaseException:
99
+ await session.rollback()
100
+ raise
101
+
102
+ return _lock()
103
+
104
+
105
+ @asynccontextmanager
106
+ async def advisory_lock(
107
+ session: AsyncSession,
108
+ key: int | tuple[int, int],
109
+ *,
110
+ shared: bool = False,
111
+ nowait: bool = False,
112
+ timeout: str | None = None,
113
+ ) -> AsyncGenerator[bool, None]:
114
+ """Acquire a PostgreSQL session-level advisory lock.
115
+
116
+ Args:
117
+ session: AsyncSession instance.
118
+ key: Lock key, either a single ``int`` (bigint) or a ``(int, int)`` pair for namespacing.
119
+ shared: Acquire a shared lock (multiple holders allowed). Default is exclusive.
120
+ nowait: Return ``False`` immediately if the lock is unavailable instead of waiting.
121
+ timeout: Maximum wait time (e.g. ``"5s"``, ``"500ms"``). Raises ``DBAPIError``
122
+ if exceeded. Ignored when *nowait* is ``True``.
123
+
124
+ Yields:
125
+ ``True`` if the lock was acquired, ``False`` if *nowait* is ``True`` and the lock
126
+ is already held.
127
+
128
+ Raises:
129
+ LockTimeoutError: If *timeout* is set and the lock cannot be acquired in time.
130
+
131
+ Example:
132
+ ```python
133
+ from fastapi_toolsets.db import advisory_lock
134
+
135
+ async with advisory_lock(session, 42):
136
+ ...
137
+
138
+ async with advisory_lock(session, 42, nowait=True) as acquired:
139
+ if not acquired:
140
+ raise HTTPException(409, "Resource is locked")
141
+
142
+ async with advisory_lock(session, 42, timeout="5s"):
143
+ ...
144
+
145
+ async with advisory_lock(session, (1, user_id), shared=True):
146
+ ...
147
+ ```
148
+ """
149
+ suffix = "_shared" if shared else ""
150
+ acquire_fn = f"{'pg_try_advisory_lock' if nowait else 'pg_advisory_lock'}{suffix}"
151
+ release_fn = f"pg_advisory_unlock{suffix}"
152
+
153
+ if isinstance(key, tuple):
154
+ k1, k2 = key
155
+ args = "CAST(:k1 AS integer), CAST(:k2 AS integer)"
156
+ params: dict[str, int] = {"k1": k1, "k2": k2}
157
+ else:
158
+ args = ":k"
159
+ params = {"k": key}
160
+
161
+ acquire_sql = text(f"SELECT {acquire_fn}({args})")
162
+ release_sql = text(f"SELECT {release_fn}({args})")
163
+
164
+ # Lock management runs raw SQL on the caller's session. Guard it with
165
+ # ``no_autoflush`` so acquiring or releasing the lock never flushes the
166
+ # caller's pending ORM changes; SQLAlchemy 2.1 autoflushes on raw
167
+ # ``text()`` too, where 2.0 did not.
168
+ try:
169
+ with session.no_autoflush:
170
+ if timeout is not None and not nowait:
171
+ await session.execute(text(f"SET LOCAL lock_timeout='{timeout}'"))
172
+ result = await session.execute(acquire_sql, params)
173
+ except sa_exc.DBAPIError as e:
174
+ if _is_lock_not_available(e):
175
+ raise LockTimeoutError(
176
+ f"Advisory lock {key!r} could not be acquired within {timeout}."
177
+ ) from e
178
+ raise # pragma: no cover
179
+ acquired = result.scalar() if nowait else True
180
+ try:
181
+ yield acquired
182
+ finally:
183
+ if acquired:
184
+ with session.no_autoflush:
185
+ await session.execute(release_sql, params)