diff --git a/README.md b/README.md index d7b566c..1451e51 100644 --- a/README.md +++ b/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 diff --git a/db_retry/transaction.py b/db_retry/transaction.py index efcaa89..e9e5ee0 100644 --- a/db_retry/transaction.py +++ b/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}) diff --git a/tests/conftest.py b/tests/conftest.py index fbf341d..7188a32 100644 --- a/tests/conftest.py +++ b/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() diff --git a/tests/test_retry.py b/tests/test_retry.py index 1e7beda..e9e7713 100644 --- a/tests/test_retry.py +++ b/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 diff --git a/tests/test_transaction.py b/tests/test_transaction.py index 5d6068d..405011a 100644 --- a/tests/test_transaction.py +++ b/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