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.
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/PKG-INFO +1 -1
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/pyproject.toml +1 -1
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/__init__.py +6 -3
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/crud/factory.py +5 -5
- fastapi_toolsets-5.0.0b1/src/fastapi_toolsets/db/__init__.py +18 -0
- fastapi_toolsets-5.0.0b1/src/fastapi_toolsets/db/core.py +315 -0
- fastapi_toolsets-5.0.0b1/src/fastapi_toolsets/db/locks.py +185 -0
- fastapi_toolsets-5.0.0b1/src/fastapi_toolsets/db/m2m.py +170 -0
- fastapi_toolsets-5.0.0b1/src/fastapi_toolsets/db/testing.py +69 -0
- fastapi_toolsets-5.0.0b1/src/fastapi_toolsets/db/watch.py +90 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/fixtures/utils.py +2 -2
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/pytest/plugin.py +2 -2
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/pytest/utils.py +1 -1
- fastapi_toolsets-4.1.3/src/fastapi_toolsets/db.py +0 -591
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/LICENSE +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/README.md +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/_imports.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/cli/__init__.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/cli/app.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/cli/commands/__init__.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/cli/commands/fixtures.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/cli/config.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/cli/pyproject.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/cli/utils.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/crud/__init__.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/crud/search.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/dependencies.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/exceptions/__init__.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/exceptions/exceptions.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/exceptions/handler.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/fixtures/__init__.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/fixtures/enum.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/fixtures/registry.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/logger.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/metrics/__init__.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/metrics/handler.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/metrics/registry.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/models/__init__.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/models/columns.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/models/watched.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/py.typed +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/pytest/__init__.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/schemas.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/security/__init__.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/security/abc.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/security/oauth.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/security/sources/__init__.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/security/sources/bearer.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/security/sources/cookie.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/security/sources/header.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/security/sources/multi.py +0 -0
- {fastapi_toolsets-4.1.3 → fastapi_toolsets-5.0.0b1}/src/fastapi_toolsets/types.py +0 -0
|
@@ -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
|
|
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(
|
|
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__ = "
|
|
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
|
|
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
|
|
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
|
|
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
|
|
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
|
|
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)
|