diff --git a/CONTEXT.md b/CONTEXT.md index 4e8806b..cb6f03b 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -13,9 +13,9 @@ package uses it. **Retriable error**: A PostgreSQL failure that can be re-run unchanged, because it says nothing about the request being -wrong: a serialization failure (`40001`) or a lost connection (class `08`). -`RETRIABLE_ASYNCPG_ERRORS` is the whole of it — an error outside that tuple is not retriable, -however temporary it looks. +wrong: a serialization failure (`40001`) or a lost connection (class `08`, or a `DBAPIError` +SQLAlchemy marked `connection_invalidated`). `RETRIABLE_ASYNCPG_ERRORS` and that flag are the whole +of it — an error outside them is not retriable, however temporary it looks. _Avoid_: transient error — it describes how long the failure lasts, not whether repeating the work is safe, and the two come apart: `StatementCompletionUnknownError` (`40003`) is as short-lived as `40001` and must never be retried. diff --git a/README.md b/README.md index 05a4c1e..59661e7 100644 --- a/README.md +++ b/README.md @@ -22,7 +22,7 @@ Retry helpers for PostgreSQL and SQLAlchemy, plus multi-host connection and tran ## Features -- `postgres_retry`: retries a coroutine on asyncpg serialization and connection errors. +- `postgres_retry`: retries a coroutine on serialization failures and lost connections. - `build_connection_factory`: multi-host asyncpg connections with failover on timeout. - `build_db_dsn` and `is_dsn_multihost`: DSN parsing and manipulation. - `Transaction`: transaction context manager that rolls back and closes the session on exit. @@ -45,14 +45,14 @@ pip install db-retry ### 1. Database operations with automatic retry -Retry a coroutine on asyncpg serialization or connection errors: +Retry a coroutine on serialization failures and lost connections. The retried coroutine must open its own transaction with `Transaction`, so that a failed attempt is rolled back before the next one starts. If it reuses a session that is still inside the failed transaction, the next attempt raises `PendingRollbackError` and the retry gives up. ```python import asyncio import sqlalchemy as sa from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column -from db_retry import postgres_retry +from db_retry import Transaction, postgres_retry class Base(DeclarativeBase): @@ -69,8 +69,9 @@ class User(Base): # Total attempts come from DB_RETRY_RETRIES_NUMBER (default 3: the first call plus 2 retries) @postgres_retry -async def get_user_by_email(session: AsyncSession, email: str) -> User: - return await session.scalar(sa.select(User).where(User.email == email)) +async def get_user_by_email(session: AsyncSession, email: str) -> User | None: + async with Transaction(session=session): + return await session.scalar(sa.select(User).where(User.email == email)) async def main(): @@ -195,6 +196,8 @@ export DB_RETRY_RETRIES_NUMBER=5 ### Retry decorator - `@postgres_retry` - Decorator for async functions that should retry on database errors (total attempts, first call plus retries, from `DB_RETRY_RETRIES_NUMBER`) - `@postgres_retry(retries=N)` - Override the total attempts (first call plus retries) per callsite +- Retries a serialization failure (`40001`), a connection error (class `08`), and any `DBAPIError` that SQLAlchemy marks with `connection_invalidated`. The last one covers asyncpg's `InterfaceError: connection is closed` on a pooled connection that died while idle, which is what `pool_pre_ping=True` otherwise catches +- The retried coroutine must open its own `Transaction`; a retry on a session left in a failed transaction raises `PendingRollbackError` - Retries also fire when the retriable `asyncpg` error is wrapped by [`advanced-alchemy`](https://github.com/litestar-org/advanced-alchemy)'s `wrap_sqlalchemy_exception()` (i.e. surfaced as `RepositoryError` / `IntegrityError`); the handler walks the `__cause__` / `__context__` chain. ### Connection utilities diff --git a/db_retry/retriable.py b/db_retry/retriable.py index 176c6bf..5fdcf24 100644 --- a/db_retry/retriable.py +++ b/db_retry/retriable.py @@ -6,6 +6,8 @@ def _is_retriable_link(exception: BaseException) -> bool: + if isinstance(exception, DBAPIError) and exception.connection_invalidated: + return True candidate = ( exception.orig.__cause__ if isinstance(exception, DBAPIError) and exception.orig is not None else exception ) @@ -13,7 +15,7 @@ def _is_retriable_link(exception: BaseException) -> bool: def is_retriable(exception: BaseException) -> bool: - """Walk __cause__/__context__; True if any link is a retriable asyncpg error, raw or wrapped in a DBAPIError.""" + """Walk __cause__/__context__; True if any link is a retriable asyncpg error or a disconnect SQLAlchemy detected.""" current: BaseException | None = exception seen: set[int] = set() while current is not None and id(current) not in seen: diff --git a/docs/adr/0002-retriable-taxonomy-is-asyncpg-classes.md b/docs/adr/0002-retriable-taxonomy-is-asyncpg-classes.md index aa3fcb3..b97c430 100644 --- a/docs/adr/0002-retriable-taxonomy-is-asyncpg-classes.md +++ b/docs/adr/0002-retriable-taxonomy-is-asyncpg-classes.md @@ -11,3 +11,12 @@ silently, losing subclass coverage as it did; the rationale it was meant to hold the executed `INVARIANT:` docstring in `tests/test_retriable.py` that pins the boundary against `StatementCompletionUnknownError` (`40003`). Only a second axis the hierarchy does not encode, a per-error retry budget or a different backoff per code, is beyond what the tuple can carry. + +Amended by #68: SQLAlchemy's disconnect flag is that second axis. A pooled connection that died +while idle fails its next use with asyncpg's `InterfaceError: connection is closed`, which has no +SQLSTATE and shares its class with deterministic misuse such as "another operation is in progress" +or a wrong argument count, so no class in the hierarchy names it. SQLAlchemy's asyncpg dialect +already decides it, from `is_closed()` on the driver connection, and exposes the answer as +`DBAPIError.connection_invalidated`; `_is_retriable_link` accepts that flag before the +`isinstance`. A raw asyncpg error carries no such flag and stays class-only, because there the +closed case could be told apart only by matching message text. diff --git a/tests/test_retriable.py b/tests/test_retriable.py index 8dd6a1a..7b6c8c5 100644 --- a/tests/test_retriable.py +++ b/tests/test_retriable.py @@ -41,6 +41,11 @@ def test_a_statement_of_unknown_outcome_is_never_retriable() -> None: _make_dbapi_error(asyncpg.ConnectionDoesNotExistError()), True, id="connection_does_not_exist_08003" ), pytest.param(_make_dbapi_error(asyncpg.PostgresError()), False, id="non_retriable_postgres_error"), + pytest.param( + _make_dbapi_error(asyncpg.InterfaceError("cannot perform operation: another operation is in progress")), + False, + id="interface_error_on_a_live_connection", + ), pytest.param(ValueError("not a db error"), False, id="bare_non_dbapi_exception"), pytest.param(asyncpg.SerializationError(), True, id="raw_serialization_error_40001"), pytest.param(asyncpg.PostgresConnectionError(), True, id="raw_postgres_connection_error_08000"), @@ -52,6 +57,12 @@ def test_is_retriable(exception: BaseException, expected: bool) -> None: assert is_retriable(exception) == expected +def test_a_connection_sqlalchemy_invalidated_is_retriable() -> None: + orig = Exception("connection is closed") + orig.__cause__ = asyncpg.InterfaceError("cannot perform operation: connection is closed") + assert is_retriable(DBAPIError("SELECT 1", None, orig, connection_invalidated=True)) is True + + def test_is_retriable_rewrapped_cause() -> None: class RepositoryError(Exception): pass diff --git a/tests/test_retry.py b/tests/test_retry.py index e95d6f0..d9b2b90 100644 --- a/tests/test_retry.py +++ b/tests/test_retry.py @@ -7,6 +7,7 @@ from sqlalchemy.exc import DBAPIError from sqlalchemy.ext import asyncio as sa_async +from db_retry import Transaction from db_retry.retry import postgres_retry @@ -169,3 +170,41 @@ async def connect() -> None: finally: await engine.dispose() assert attempts == expected_attempts + + +async def test_a_connection_closed_while_pooled_is_retried_through_a_transaction( + async_engine: sa_async.AsyncEngine, monkeypatch: pytest.MonkeyPatch +) -> None: + async with async_engine.connect() as connection: + pooled: typing.Final = (await connection.get_raw_connection()).driver_connection + assert pooled is not None + pid: typing.Final = await connection.scalar(sqlalchemy.text("SELECT pg_backend_pid()")) + closed: typing.Final = asyncio.Event() + + def _on_close(_: asyncpg.Connection) -> None: + closed.set() + + pooled.add_termination_listener(_on_close) + killer: typing.Final = await asyncpg.connect( + async_engine.url.set(drivername="postgresql").render_as_string(hide_password=False) + ) + try: + await killer.execute("SELECT pg_terminate_backend($1)", pid) + finally: + await killer.close() + await asyncio.wait_for(closed.wait(), timeout=5) + _record_backoff(monkeypatch) + attempts = 0 + + async with sa_async.AsyncSession(async_engine) as session: + transaction: typing.Final = Transaction(session=session) + + @postgres_retry(retries=2) + async def select_one() -> int | None: + nonlocal attempts + attempts += 1 + async with transaction: + return await session.scalar(sqlalchemy.text("SELECT 1")) + + assert await select_one() == 1 + assert attempts == 2 # noqa: PLR2004 # the closed connection, then a fresh one