fix: Transaction rolls back a dead transaction before joining it

README.md

@@ -45,7 +45,7 @@ pip install db-retry
 
 ### 1. Database operations with automatic retry
 
-Retry a coroutine on serialization failures and lost connections. Either error ends the whole Postgres transaction, so the retried coroutine has to own its transaction from start to finish: open it with `Transaction`, as below, or work on a fresh session it creates itself. A coroutine that runs on a session someone else began cannot be retried. Postgres has already discarded the statements that ran before it, and SQLAlchemy refuses the next attempt with `PendingRollbackError` instead of committing part of the work.
+Retry a coroutine on serialization failures and lost connections. Either error ends the whole Postgres transaction, so the retried coroutine has to own its transaction from start to finish: open it with `Transaction`, as below, or work on a fresh session it creates itself. A coroutine that runs on a session someone else began cannot be retried. Postgres has already discarded the statements that ran before it, and SQLAlchemy refuses the next attempt with `PendingRollbackError` instead of committing part of the work. One case is recovered for you: if the connection is lost on a query or flush that runs after the `Transaction` block, entering `Transaction` on the next attempt rolls back the dead transaction and starts on a fresh connection.
 
 ```python
 import asyncio
@@ -212,7 +212,7 @@ export DB_RETRY_RETRIES_NUMBER=5
 - `is_dsn_multihost(db_dsn)` - Checks if a DSN contains multiple hosts
 
 ### Transaction helper
-- `Transaction(session, isolation_level=None)` - Context manager for transaction handling; on exit, rolls back if still in a transaction, then closes the session
+- `Transaction(session, isolation_level=None)` - Context manager for transaction handling; on entry, rolls back a transaction that can no longer continue (a failed flush or a lost connection), otherwise joins the session's open transaction or begins one; on exit, rolls back if still in a transaction, then closes the session
 
 ## Requirements
 

db_retry/transaction.py

@@ -5,12 +5,26 @@
 from sqlalchemy.ext import asyncio as sa_async
 
 
+def _is_dead(session: sa_async.AsyncSession) -> bool:
+    transaction: typing.Final = session.sync_session.get_transaction()
+    if transaction is None:
+        return False
+    if not transaction.is_active:
+        return True
+    # SQLAlchemy has no public accessor for the connections a transaction already holds.
+    held: typing.Final = getattr(transaction, "_connections", {})
+    return any(connection.invalidated for connection, *_ in held.values())
+
+
 @dataclasses.dataclass(kw_only=True, frozen=True, slots=True)
 class Transaction:
     session: sa_async.AsyncSession
     isolation_level: IsolationLevel | None = None
 
     async def __aenter__(self) -> typing.Self:
+        if _is_dead(self.session):
+            await self.session.rollback()
+
         if self.isolation_level:
             await self.session.connection(execution_options={"isolation_level": self.isolation_level})
 

tests/conftest.py

@@ -1,7 +1,10 @@
+import asyncio
 import os
 import typing
 
+import asyncpg
 import pytest
+import sqlalchemy
 from sqlalchemy.ext import asyncio as sa_async
 
 
@@ -12,3 +15,36 @@ async def async_engine() -> typing.AsyncIterator[sa_async.AsyncEngine]:
         yield engine
     finally:
         await engine.dispose()
+
+
+@pytest.fixture
+def kill_idle_connection(async_engine: sa_async.AsyncEngine) -> typing.Callable[[], typing.Awaitable[None]]:
+    """Return a coroutine function that leaves the pool's idle connection with its backend terminated."""
+
+    async def _kill() -> 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)
+
+    return _kill
+
+
+@pytest.fixture
+async def lost_pooled_connection(kill_idle_connection: typing.Callable[[], typing.Awaitable[None]]) -> None:
+    """Leave the engine's pool holding one idle connection whose backend was terminated."""
+    await kill_idle_connection()

tests/test_retry.py

@@ -171,27 +171,10 @@ async def connect() -> None:
     assert attempts == expected_attempts
 
 
+@pytest.mark.usefixtures("lost_pooled_connection")
 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
 
@@ -209,6 +192,32 @@ async def select_one() -> int | None:
     assert attempts == 2  # noqa: PLR2004  # the closed connection, then a fresh one
 
 
+async def test_a_connection_lost_after_the_transaction_is_retried(
+    async_engine: sa_async.AsyncEngine,
+    kill_idle_connection: typing.Callable[[], typing.Awaitable[None]],
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    _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 read_after_the_block() -> int | None:
+            nonlocal attempts
+            attempts += 1
+            async with transaction:
+                await session.scalar(sqlalchemy.text("SELECT 1"))
+                await transaction.commit()
+            if attempts == 1:
+                await kill_idle_connection()
+            return await session.scalar(sqlalchemy.text("SELECT 2"))
+
+        assert await read_after_the_block() == 2  # noqa: PLR2004
+    assert attempts == 2  # noqa: PLR2004
+
+
 async def test_a_retriable_first_failure_waits_once_then_succeeds(monkeypatch: pytest.MonkeyPatch) -> None:
     waits: typing.Final = _record_backoff(monkeypatch)
     attempts = 0

tests/test_transaction.py

@@ -2,11 +2,23 @@
 import typing
 
 import pytest
+import sqlalchemy
+from sqlalchemy import orm
+from sqlalchemy.exc import DBAPIError
 from sqlalchemy.ext import asyncio as sa_async
 
 from db_retry import Transaction
 
 
+class _Base(orm.DeclarativeBase):
+    pass
+
+
+class _Item(_Base):
+    __tablename__ = "transaction_items"
+    id: orm.Mapped[int] = orm.mapped_column(primary_key=True)
+
+
 @pytest.fixture
 async def transaction(async_engine: sa_async.AsyncEngine) -> typing.AsyncIterator[Transaction]:
     async with sa_async.AsyncSession(async_engine, expire_on_commit=False, autoflush=False) as session:
@@ -53,3 +65,58 @@ async def test_transaction_with_exception(transaction: Transaction) -> None:
             msg: typing.Final = "some error"
             raise Exception(msg)  # noqa: TRY002
     assert not transaction.session.in_transaction()
+
+
+@pytest.mark.usefixtures("lost_pooled_connection")
+async def test_entering_rolls_back_a_transaction_whose_connection_was_lost(transaction: Transaction) -> None:
+    with pytest.raises(DBAPIError) as lost:
+        await transaction.session.scalar(sqlalchemy.text("SELECT 1"))
+    assert lost.value.connection_invalidated
+    async with transaction:
+        assert await transaction.session.scalar(sqlalchemy.text("SELECT 1")) == 1
+
+
+async def test_entering_joins_a_live_transaction(transaction: Transaction) -> None:
+    await transaction.session.execute(sqlalchemy.text("CREATE TEMP TABLE joined (x int)"))
+    await transaction.session.execute(sqlalchemy.text("INSERT INTO joined VALUES (1)"))
+    async with transaction:
+        assert await transaction.session.scalar(sqlalchemy.text("SELECT count(*) FROM joined")) == 1
+
+
+@pytest.mark.usefixtures("lost_pooled_connection")
+async def test_a_recovered_transaction_keeps_its_isolation_level(transaction_serializable: Transaction) -> None:
+    with pytest.raises(DBAPIError):
+        await transaction_serializable.session.scalar(sqlalchemy.text("SELECT 1"))
+    async with transaction_serializable:
+        level: typing.Final = await transaction_serializable.session.scalar(
+            sqlalchemy.text("SHOW transaction_isolation")
+        )
+        assert level == "serializable"
+
+
+async def test_entering_after_autobegin_keeps_the_isolation_level(transaction_serializable: Transaction) -> None:
+    transaction_serializable.session.add(_Item(id=1))
+    assert transaction_serializable.session.in_transaction()
+    async with transaction_serializable:
+        level: typing.Final = await transaction_serializable.session.scalar(
+            sqlalchemy.text("SHOW transaction_isolation")
+        )
+        assert level == "serializable"
+
+
+async def test_entering_a_session_with_only_mapper_binds(async_engine: sa_async.AsyncEngine) -> None:
+    async with sa_async.AsyncSession(binds={_Item: async_engine}, autoflush=False) as session:
+        session.add(_Item(id=1))
+        async with Transaction(session=session):
+            assert session.in_transaction()
+
+
+@pytest.mark.usefixtures("lost_pooled_connection")
+async def test_entering_rolls_back_a_transaction_whose_flush_lost_the_connection(transaction: Transaction) -> None:
+    transaction.session.add(_Item(id=1))
+    with pytest.raises(DBAPIError) as lost:
+        await transaction.session.flush()
+    assert lost.value.connection_invalidated
+    assert not transaction.session.is_active
+    async with transaction:
+        assert await transaction.session.scalar(sqlalchemy.text("SELECT 1")) == 1