mirror of
https://github.com/d3vyce/fastapi-toolsets.git
synced 2026-08-04 15:44:09 +00:00
fix: pool exhaustion and lock timeout surface as 500 instead of structured 503 (#295)
This commit is contained in:
+8
-1
@@ -69,6 +69,13 @@ async with lock_tables(session_maker=session_maker, tables=[User], mode=LockMode
|
||||
|
||||
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:
|
||||
|
||||
```python
|
||||
async with lock_tables(session_maker, [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.
|
||||
@@ -85,7 +92,7 @@ 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 DBAPIError 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"):
|
||||
...
|
||||
|
||||
|
||||
@@ -39,6 +39,8 @@ It also patches `app.openapi()` to replace the default Pydantic 422 schema with
|
||||
| [`NoSearchableFieldsError`](../reference/exceptions.md#fastapi_toolsets.exceptions.exceptions.NoSearchableFieldsError) | 400 | No Searchable Fields |
|
||||
| [`InvalidFacetFilterError`](../reference/exceptions.md#fastapi_toolsets.exceptions.exceptions.InvalidFacetFilterError) | 400 | Invalid Facet Filter |
|
||||
| [`InvalidOrderFieldError`](../reference/exceptions.md#fastapi_toolsets.exceptions.exceptions.InvalidOrderFieldError) | 422 | Invalid Order Field |
|
||||
| [`PoolExhaustedError`](../reference/exceptions.md#fastapi_toolsets.exceptions.exceptions.PoolExhaustedError) | 503 | Service Unavailable |
|
||||
| [`LockTimeoutError`](../reference/exceptions.md#fastapi_toolsets.exceptions.exceptions.LockTimeoutError) | 503 | Service Unavailable |
|
||||
|
||||
### Per-instance overrides
|
||||
|
||||
|
||||
@@ -15,6 +15,8 @@ from fastapi_toolsets.exceptions import (
|
||||
InvalidSearchColumnError,
|
||||
InvalidFacetFilterError,
|
||||
InvalidOrderFieldError,
|
||||
PoolExhaustedError,
|
||||
LockTimeoutError,
|
||||
generate_error_responses,
|
||||
init_exceptions_handlers,
|
||||
)
|
||||
@@ -38,6 +40,10 @@ from fastapi_toolsets.exceptions import (
|
||||
|
||||
## ::: fastapi_toolsets.exceptions.exceptions.InvalidOrderFieldError
|
||||
|
||||
## ::: fastapi_toolsets.exceptions.exceptions.PoolExhaustedError
|
||||
|
||||
## ::: fastapi_toolsets.exceptions.exceptions.LockTimeoutError
|
||||
|
||||
## ::: fastapi_toolsets.exceptions.exceptions.generate_error_responses
|
||||
|
||||
## ::: fastapi_toolsets.exceptions.handler.init_exceptions_handlers
|
||||
|
||||
@@ -6,13 +6,22 @@ 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 NotFoundError
|
||||
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",
|
||||
@@ -65,7 +74,10 @@ def create_db_dependency(
|
||||
|
||||
async def get_db() -> AsyncGenerator[_SessionT, None]:
|
||||
async with session_maker() as session:
|
||||
await session.connection()
|
||||
try:
|
||||
await session.connection()
|
||||
except sa_exc.TimeoutError as e:
|
||||
raise PoolExhaustedError() from e
|
||||
yield session
|
||||
if session.in_transaction():
|
||||
await session.commit()
|
||||
@@ -198,6 +210,18 @@ def lock_tables(
|
||||
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
|
||||
@@ -229,8 +253,7 @@ async def advisory_lock(
|
||||
is already held.
|
||||
|
||||
Raises:
|
||||
sqlalchemy.exc.DBAPIError: If *timeout* is set and the lock cannot be acquired
|
||||
in time.
|
||||
LockTimeoutError: If *timeout* is set and the lock cannot be acquired in time.
|
||||
|
||||
Example:
|
||||
```python
|
||||
@@ -268,7 +291,14 @@ async def advisory_lock(
|
||||
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)
|
||||
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
|
||||
|
||||
@@ -8,8 +8,10 @@ from .exceptions import (
|
||||
InvalidFacetFilterError,
|
||||
InvalidOrderFieldError,
|
||||
InvalidSearchColumnError,
|
||||
LockTimeoutError,
|
||||
NoSearchableFieldsError,
|
||||
NotFoundError,
|
||||
PoolExhaustedError,
|
||||
UnauthorizedError,
|
||||
UnsupportedFacetTypeError,
|
||||
generate_error_responses,
|
||||
@@ -26,8 +28,10 @@ __all__ = [
|
||||
"InvalidFacetFilterError",
|
||||
"InvalidOrderFieldError",
|
||||
"InvalidSearchColumnError",
|
||||
"LockTimeoutError",
|
||||
"NoSearchableFieldsError",
|
||||
"NotFoundError",
|
||||
"PoolExhaustedError",
|
||||
"UnauthorizedError",
|
||||
"UnsupportedFacetTypeError",
|
||||
]
|
||||
|
||||
@@ -223,6 +223,35 @@ class InvalidOrderFieldError(ApiException):
|
||||
)
|
||||
|
||||
|
||||
class PoolExhaustedError(ApiException):
|
||||
"""HTTP 503 - Database connection pool is exhausted."""
|
||||
|
||||
api_error = ApiError(
|
||||
code=503,
|
||||
msg="Service Unavailable",
|
||||
desc=(
|
||||
"The database connection pool is exhausted. "
|
||||
"Too many concurrent requests are holding connections. "
|
||||
"Retry shortly or contact support if the issue persists."
|
||||
),
|
||||
err_code="DB-503-POOL",
|
||||
)
|
||||
|
||||
|
||||
class LockTimeoutError(ApiException):
|
||||
"""HTTP 503 - A database lock could not be acquired within the timeout."""
|
||||
|
||||
api_error = ApiError(
|
||||
code=503,
|
||||
msg="Service Unavailable",
|
||||
desc=(
|
||||
"A database lock could not be acquired within the allowed timeout. "
|
||||
"The resource is under heavy contention. Retry shortly."
|
||||
),
|
||||
err_code="DB-503-LOCK",
|
||||
)
|
||||
|
||||
|
||||
def generate_error_responses(
|
||||
*errors: type[ApiException],
|
||||
) -> dict[int | str, dict[str, Any]]:
|
||||
|
||||
+54
-5
@@ -39,7 +39,11 @@ from fastapi_toolsets.db import (
|
||||
m2m_set,
|
||||
wait_for_row_change,
|
||||
)
|
||||
from fastapi_toolsets.exceptions import NotFoundError
|
||||
from fastapi_toolsets.exceptions import (
|
||||
LockTimeoutError,
|
||||
NotFoundError,
|
||||
PoolExhaustedError,
|
||||
)
|
||||
from fastapi_toolsets.pytest import create_db_session
|
||||
|
||||
from .conftest import DATABASE_URL, Base, Post, Role, RoleCrud, Tag, User, UserCrud
|
||||
@@ -399,15 +403,13 @@ class TestAdvisoryLock:
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_timeout_raises_when_contended(self, session_maker):
|
||||
"""timeout= raises when the lock cannot be acquired within the interval."""
|
||||
from sqlalchemy.exc import DBAPIError
|
||||
|
||||
"""timeout= raises LockTimeoutError when the lock cannot be acquired."""
|
||||
async with session_maker() as holder:
|
||||
async with holder.begin():
|
||||
async with advisory_lock(holder, 1006):
|
||||
async with session_maker() as contender:
|
||||
async with contender.begin():
|
||||
with pytest.raises(DBAPIError):
|
||||
with pytest.raises(LockTimeoutError):
|
||||
async with advisory_lock(
|
||||
contender, 1006, timeout="10ms"
|
||||
):
|
||||
@@ -746,6 +748,53 @@ class TestM2MAdd:
|
||||
assert loaded.tags[0].name == "locked_tag"
|
||||
|
||||
|
||||
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)
|
||||
|
||||
try:
|
||||
async with session_factory() as holder:
|
||||
await holder.connection() # check out the single available connection
|
||||
with pytest.raises(PoolExhaustedError):
|
||||
async for _ in get_db():
|
||||
pass
|
||||
finally:
|
||||
await 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)
|
||||
|
||||
try:
|
||||
async with session_factory() as holder:
|
||||
await holder.connection() # check out the single available connection
|
||||
with pytest.raises(PoolExhaustedError):
|
||||
async with lock_tables(session_factory, [Role]) as _:
|
||||
pass
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_lock_timeout_raises_lock_timeout_error(self, session_maker):
|
||||
"""LockTimeoutError is raised when a table lock cannot be acquired within timeout."""
|
||||
async with lock_tables(session_maker, [Role]) as _:
|
||||
with pytest.raises(LockTimeoutError):
|
||||
async with lock_tables(session_maker, [Role], timeout="100ms") as _:
|
||||
pass
|
||||
|
||||
|
||||
class _LocalBase(DeclarativeBase):
|
||||
pass
|
||||
|
||||
|
||||
@@ -10,7 +10,9 @@ from fastapi_toolsets.exceptions import (
|
||||
ConflictError,
|
||||
ForbiddenError,
|
||||
InvalidOrderFieldError,
|
||||
LockTimeoutError,
|
||||
NotFoundError,
|
||||
PoolExhaustedError,
|
||||
UnauthorizedError,
|
||||
generate_error_responses,
|
||||
init_exceptions_handlers,
|
||||
@@ -216,6 +218,70 @@ class TestApiExceptionGuard:
|
||||
assert err.api_error.code == 404
|
||||
|
||||
|
||||
class TestDbExceptions:
|
||||
"""Tests for database-related exception classes."""
|
||||
|
||||
def test_pool_exhausted_error_attributes(self):
|
||||
"""PoolExhaustedError has 503 status and DB-503-POOL error code."""
|
||||
error = PoolExhaustedError()
|
||||
assert error.api_error.code == 503
|
||||
assert error.api_error.err_code == "DB-503-POOL"
|
||||
assert error.api_error.msg == "Service Unavailable"
|
||||
|
||||
def test_pool_exhausted_error_with_detail(self):
|
||||
"""PoolExhaustedError accepts a detail string that overrides msg."""
|
||||
error = PoolExhaustedError("pool full")
|
||||
assert error.api_error.msg == "pool full"
|
||||
assert PoolExhaustedError.api_error.msg == "Service Unavailable"
|
||||
|
||||
def test_lock_timeout_error_attributes(self):
|
||||
"""LockTimeoutError has 503 status and DB-503-LOCK error code."""
|
||||
error = LockTimeoutError()
|
||||
assert error.api_error.code == 503
|
||||
assert error.api_error.err_code == "DB-503-LOCK"
|
||||
assert error.api_error.msg == "Service Unavailable"
|
||||
|
||||
def test_lock_timeout_error_with_detail(self):
|
||||
"""LockTimeoutError accepts a detail string that overrides msg."""
|
||||
error = LockTimeoutError("contended")
|
||||
assert error.api_error.msg == "contended"
|
||||
assert LockTimeoutError.api_error.msg == "Service Unavailable"
|
||||
|
||||
def test_pool_exhausted_handled_as_503(self):
|
||||
"""init_exceptions_handlers turns PoolExhaustedError into a 503 response."""
|
||||
from fastapi import FastAPI
|
||||
from fastapi_toolsets.exceptions import init_exceptions_handlers
|
||||
|
||||
app = FastAPI()
|
||||
init_exceptions_handlers(app)
|
||||
|
||||
@app.get("/db")
|
||||
async def endpoint():
|
||||
raise PoolExhaustedError()
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.get("/db")
|
||||
assert response.status_code == 503
|
||||
assert response.json()["error_code"] == "DB-503-POOL"
|
||||
|
||||
def test_lock_timeout_handled_as_503(self):
|
||||
"""init_exceptions_handlers turns LockTimeoutError into a 503 response."""
|
||||
from fastapi import FastAPI
|
||||
from fastapi_toolsets.exceptions import init_exceptions_handlers
|
||||
|
||||
app = FastAPI()
|
||||
init_exceptions_handlers(app)
|
||||
|
||||
@app.get("/lock")
|
||||
async def endpoint():
|
||||
raise LockTimeoutError()
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.get("/lock")
|
||||
assert response.status_code == 503
|
||||
assert response.json()["error_code"] == "DB-503-LOCK"
|
||||
|
||||
|
||||
class TestBuiltInExceptions:
|
||||
"""Tests for built-in exception classes."""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user