From cd928688aff021f428bce42e932da9116fcaa44a Mon Sep 17 00:00:00 2001 From: d3vyce <44915747+d3vyce@users.noreply.github.com> Date: Wed, 3 Jun 2026 22:03:15 +0200 Subject: [PATCH] fix: pool exhaustion and lock timeout surface as 500 instead of structured 503 (#295) --- docs/module/db.md | 9 ++- docs/module/exceptions.md | 2 + docs/reference/exceptions.md | 6 ++ src/fastapi_toolsets/db.py | 40 +++++++++-- src/fastapi_toolsets/exceptions/__init__.py | 4 ++ src/fastapi_toolsets/exceptions/exceptions.py | 29 ++++++++ tests/test_db.py | 59 +++++++++++++++-- tests/test_exceptions.py | 66 +++++++++++++++++++ 8 files changed, 204 insertions(+), 11 deletions(-) diff --git a/docs/module/db.md b/docs/module/db.md index 5817268..230785c 100644 --- a/docs/module/db.md +++ b/docs/module/db.md @@ -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"): ... diff --git a/docs/module/exceptions.md b/docs/module/exceptions.md index f3b22fe..0501259 100644 --- a/docs/module/exceptions.md +++ b/docs/module/exceptions.md @@ -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 diff --git a/docs/reference/exceptions.md b/docs/reference/exceptions.md index 3ee24c8..b69c54b 100644 --- a/docs/reference/exceptions.md +++ b/docs/reference/exceptions.md @@ -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 diff --git a/src/fastapi_toolsets/db.py b/src/fastapi_toolsets/db.py index ae413e4..9ec30e1 100644 --- a/src/fastapi_toolsets/db.py +++ b/src/fastapi_toolsets/db.py @@ -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 diff --git a/src/fastapi_toolsets/exceptions/__init__.py b/src/fastapi_toolsets/exceptions/__init__.py index 2ee2d26..bca5d94 100644 --- a/src/fastapi_toolsets/exceptions/__init__.py +++ b/src/fastapi_toolsets/exceptions/__init__.py @@ -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", ] diff --git a/src/fastapi_toolsets/exceptions/exceptions.py b/src/fastapi_toolsets/exceptions/exceptions.py index 772aec3..6ed8015 100644 --- a/src/fastapi_toolsets/exceptions/exceptions.py +++ b/src/fastapi_toolsets/exceptions/exceptions.py @@ -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]]: diff --git a/tests/test_db.py b/tests/test_db.py index 01aeac6..b4bec40 100644 --- a/tests/test_db.py +++ b/tests/test_db.py @@ -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 diff --git a/tests/test_exceptions.py b/tests/test_exceptions.py index 700792a..2ae478e 100644 --- a/tests/test_exceptions.py +++ b/tests/test_exceptions.py @@ -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."""