fix: retry a connection SQLAlchemy invalidated

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.

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

db_retry/retriable.py

@@ -6,14 +6,16 @@
 
 
 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
     )
     return isinstance(candidate, RETRIABLE_ASYNCPG_ERRORS)
 
 
 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:

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.

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

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