From b1980085f2e464471cc2ab680edd71a57c11367a Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Wed, 4 Nov 2026 12:12:56 +0300 Subject: [PATCH] Use public transaction connections when available --- db_retry/transaction.py | 4 +++- tests/test_transaction.py | 36 ++++++++++++++++++++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/db_retry/transaction.py b/db_retry/transaction.py index e9e5ee0..28b6e7a 100644 --- a/db_retry/transaction.py +++ b/db_retry/transaction.py @@ -11,7 +11,9 @@ def _is_dead(session: sa_async.AsyncSession) -> bool: return False if not transaction.is_active: return True - # SQLAlchemy has no public accessor for the connections a transaction already holds. + connections: typing.Final = getattr(transaction, "connections", None) + if connections is not None: + return any(connection.invalidated for connection in connections) held: typing.Final = getattr(transaction, "_connections", {}) return any(connection.invalidated for connection, *_ in held.values()) diff --git a/tests/test_transaction.py b/tests/test_transaction.py index 405011a..cc39fe7 100644 --- a/tests/test_transaction.py +++ b/tests/test_transaction.py @@ -1,4 +1,5 @@ import contextlib +import types import typing import pytest @@ -120,3 +121,38 @@ async def test_entering_rolls_back_a_transaction_whose_flush_lost_the_connection assert not transaction.session.is_active async with transaction: assert await transaction.session.scalar(sqlalchemy.text("SELECT 1")) == 1 + + +@pytest.mark.parametrize("accessor", ["native", "public", "legacy"]) +@pytest.mark.parametrize("savepoint", [False, True]) +async def test_entering_recovers_an_invalidated_second_bind( + async_engine: sa_async.AsyncEngine, monkeypatch: pytest.MonkeyPatch, accessor: str, savepoint: bool +) -> None: + other_engine: typing.Final = sa_async.create_async_engine(async_engine.url) + try: + async with sa_async.AsyncSession(async_engine) as session: + await session.connection() + other_connection = await session.connection(bind_arguments={"bind": other_engine.sync_engine}) + await other_connection.invalidate() + root = session.sync_session.get_transaction() + assert root is not None + assert root.is_active + if savepoint: + await session.begin_nested() + with monkeypatch.context() as patch: + if accessor == "public": + public_transaction = types.SimpleNamespace( + is_active=True, connections=frozenset([other_connection.sync_connection]) + ) + patch.setattr(session.sync_session, "get_transaction", lambda: public_transaction) + elif accessor == "legacy": + patch.delattr(orm.SessionTransaction, "connections", raising=False) + async with Transaction(session=session): + assert ( + await session.scalar( + sqlalchemy.text("SELECT 1"), bind_arguments={"bind": other_engine.sync_engine} + ) + == 1 + ) + finally: + await other_engine.dispose()