fix: open a Transaction in the read use cases so a retry can recover

AGENTS.md

@@ -30,6 +30,11 @@ template breaks the transaction model or the DI wiring.
    `Transaction.__aexit__`'s unconditional rollback-on-open-transaction
    behavior (returning a loaded ORM object from inside an uncommitted `async
    with self.transaction:` block detaches it).
+   Read-only use cases open a `Transaction` too and commit before returning.
+   `@postgres_retry` cannot roll back a session it does not see, so a retry
+   on a session still inside the failed transaction ends in
+   `PendingRollbackError`; the `Transaction` exit is what rolls it back.
+   `tests/use_cases/test_connection_loss.py` pins this.
 
 ## Commands
 
@@ -103,7 +108,7 @@ a session, a transaction or a repository directly. Provider scopes:
   the engine.
 - `Database.database_session` — request-scoped, finalizer closes the session.
 - `Database.transaction` — request-scoped `db_retry.Transaction`, the object
-  every write-side use case wraps its commit in.
+  every use case wraps its work in.
 - `Repositories.*` — request-scoped, `auto_commit=False` (see above).
 - `UseCases.*` — request-scoped, one class per operation.
 

app/use_cases/authenticate_user.py

@@ -1,7 +1,7 @@
 import dataclasses
 import typing
 
-from db_retry import postgres_retry
+from db_retry import Transaction, postgres_retry
 
 from app import security
 from app.database import tables
@@ -10,11 +10,14 @@
 
 @dataclasses.dataclass(kw_only=True, frozen=True, slots=True)
 class AuthenticateUserUseCase:
+    transaction: Transaction
     users_repository: UsersRepository
 
     @postgres_retry
     async def __call__(self, *, username: str, password: str) -> tables.UsersTable | None:
-        user: typing.Final = await self.users_repository.get_one_or_none(username=username)
+        async with self.transaction:
+            user: typing.Final = await self.users_repository.get_one_or_none(username=username)
+            await self.transaction.commit()
         if user is None:
             # Discarded, but skipping the argon2 work would turn login timing into a username oracle.
             security.hash_password(password)

app/use_cases/check_token_revoked.py

@@ -1,14 +1,19 @@
 import dataclasses
+import typing
 
-from db_retry import postgres_retry
+from db_retry import Transaction, postgres_retry
 
 from app.repositories.revoked_tokens_repository import RevokedTokensRepository
 
 
 @dataclasses.dataclass(kw_only=True, frozen=True, slots=True)
 class CheckTokenRevokedUseCase:
+    transaction: Transaction
     revoked_tokens_repository: RevokedTokensRepository
 
     @postgres_retry
     async def __call__(self, *, jti: str) -> bool:
-        return await self.revoked_tokens_repository.is_revoked(jti)
+        async with self.transaction:
+            revoked: typing.Final = await self.revoked_tokens_repository.is_revoked(jti)
+            await self.transaction.commit()
+        return revoked

app/use_cases/fetch_chat.py

@@ -1,6 +1,7 @@
 import dataclasses
+import typing
 
-from db_retry import postgres_retry
+from db_retry import Transaction, postgres_retry
 
 from app.actor import Actor
 from app.database import tables
@@ -11,12 +12,16 @@
 
 @dataclasses.dataclass(kw_only=True, frozen=True, slots=True)
 class FetchChatUseCase:
+    transaction: Transaction
     chats_repository: ChatsRepository
     chat_members_repository: ChatMembersRepository
 
     @postgres_retry
     async def __call__(self, *, actor: Actor, chat_id: int) -> tables.ChatsTable:
-        if not await self.chat_members_repository.is_member(chat_id, actor.id):
-            msg = "Not a member of this chat"
-            raise PermissionDeniedError(msg)
-        return await self.chats_repository.fetch_with_members(chat_id)
+        async with self.transaction:
+            if not await self.chat_members_repository.is_member(chat_id, actor.id):
+                msg = "Not a member of this chat"
+                raise PermissionDeniedError(msg)
+            chat: typing.Final = await self.chats_repository.fetch_with_members(chat_id)
+            await self.transaction.commit()
+        return chat

app/use_cases/fetch_chats.py

@@ -1,7 +1,8 @@
 import dataclasses
+import typing
 from collections.abc import Sequence
 
-from db_retry import postgres_retry
+from db_retry import Transaction, postgres_retry
 
 from app.actor import Actor
 from app.database import tables
@@ -10,8 +11,12 @@
 
 @dataclasses.dataclass(kw_only=True, frozen=True, slots=True)
 class FetchChatsUseCase:
+    transaction: Transaction
     chats_repository: ChatsRepository
 
     @postgres_retry
     async def __call__(self, *, actor: Actor) -> Sequence[tables.ChatsTable]:
-        return await self.chats_repository.list_for_user(actor.id)
+        async with self.transaction:
+            chats: typing.Final = await self.chats_repository.list_for_user(actor.id)
+            await self.transaction.commit()
+        return chats

app/use_cases/fetch_messages.py

@@ -2,7 +2,7 @@
 import typing
 from collections.abc import Sequence
 
-from db_retry import postgres_retry
+from db_retry import Transaction, postgres_retry
 
 from app.actor import Actor
 from app.database import tables
@@ -16,6 +16,7 @@
 
 @dataclasses.dataclass(kw_only=True, frozen=True, slots=True)
 class FetchMessagesUseCase:
+    transaction: Transaction
     chat_members_repository: ChatMembersRepository
     messages_repository: MessagesRepository
 
@@ -29,15 +30,18 @@ async def __call__(
         after_id: int | None = None,
         limit: int = 50,
     ) -> Sequence[tables.MessagesTable]:
-        if not await self.chat_members_repository.is_member(chat_id, actor.id):
-            msg = "Not a member of this chat"
-            raise PermissionDeniedError(msg)
-        if before_id is not None and after_id is not None:
-            msg = "before_id and after_id are mutually exclusive"
-            raise ValidationError(msg)
-        if limit < 1:
-            msg = "limit must be at least 1"
-            raise ValidationError(msg)
-        return await self.messages_repository.list_page(
-            chat_id, before_id=before_id, after_id=after_id, limit=min(limit, MAX_PAGE_SIZE)
-        )
+        async with self.transaction:
+            if not await self.chat_members_repository.is_member(chat_id, actor.id):
+                msg = "Not a member of this chat"
+                raise PermissionDeniedError(msg)
+            if before_id is not None and after_id is not None:
+                msg = "before_id and after_id are mutually exclusive"
+                raise ValidationError(msg)
+            if limit < 1:
+                msg = "limit must be at least 1"
+                raise ValidationError(msg)
+            messages: typing.Final = await self.messages_repository.list_page(
+                chat_id, before_id=before_id, after_id=after_id, limit=min(limit, MAX_PAGE_SIZE)
+            )
+            await self.transaction.commit()
+        return messages

app/use_cases/fetch_user.py

@@ -1,6 +1,7 @@
 import dataclasses
+import typing
 
-from db_retry import postgres_retry
+from db_retry import Transaction, postgres_retry
 
 from app.actor import Actor
 from app.database import tables
@@ -9,8 +10,12 @@
 
 @dataclasses.dataclass(kw_only=True, frozen=True, slots=True)
 class FetchUserUseCase:
+    transaction: Transaction
     users_repository: UsersRepository
 
     @postgres_retry
     async def __call__(self, *, actor: Actor) -> tables.UsersTable | None:
-        return await self.users_repository.get_one_or_none(id=actor.id)
+        async with self.transaction:
+            user: typing.Final = await self.users_repository.get_one_or_none(id=actor.id)
+            await self.transaction.commit()
+        return user

tests/use_cases/test_connection_loss.py

@@ -0,0 +1,87 @@
+import asyncio
+import typing
+
+import advanced_alchemy.exceptions
+import asyncpg
+import modern_di
+import pytest
+import sqlalchemy as sa
+from sqlalchemy.exc import InterfaceError
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app import ioc
+from app.actor import Actor
+from app.database.resources import create_database_engine
+from app.exceptions import PermissionDeniedError
+
+
+pytestmark = pytest.mark.xfail(
+    raises=(advanced_alchemy.exceptions.IntegrityError, InterfaceError),
+    reason="the released db-retry does not retry a closed connection; modern-python/db-retry#69 does",
+    strict=True,
+)
+
+UNKNOWN_ACTOR: typing.Final = Actor(id=-1)
+
+
+@pytest.fixture
+async def request_container(di_container: modern_di.Container) -> typing.AsyncIterator[modern_di.Container]:
+    engine: typing.Final = create_database_engine()
+    di_container.override(ioc.Database.database_engine, engine)
+    try:
+        async with di_container.build_child_container(scope=modern_di.Scope.REQUEST) as container:
+            yield container
+    finally:
+        di_container.reset_override(ioc.Database.database_engine)
+        await engine.dispose()
+
+
+@pytest.fixture
+async def lost_connection(request_container: modern_di.Container) -> None:
+    """Leave the request session holding a pooled connection whose backend was terminated."""
+    session: typing.Final = request_container.resolve_provider(ioc.Database.database_session)
+    assert isinstance(session, AsyncSession)
+    connection: typing.Final = await session.connection()
+    pid: typing.Final = await connection.scalar(sa.text("SELECT pg_backend_pid()"))
+    driver_connection: typing.Final = (await connection.get_raw_connection()).driver_connection
+    assert driver_connection is not None
+    closed: typing.Final = asyncio.Event()
+    driver_connection.add_termination_listener(lambda _: closed.set())
+    killer: typing.Final = await asyncpg.connect(
+        connection.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)
+
+
+@pytest.mark.usefixtures("lost_connection")
+@pytest.mark.parametrize(
+    ("provider", "kwargs", "expected"),
+    [
+        (ioc.UseCases.authenticate_user_use_case, {"username": "nobody", "password": "hunter2hunter2"}, None),
+        (ioc.UseCases.check_token_revoked_use_case, {"jti": "unknown"}, False),
+        (ioc.UseCases.fetch_chats_use_case, {"actor": UNKNOWN_ACTOR}, []),
+        (ioc.UseCases.fetch_user_use_case, {"actor": UNKNOWN_ACTOR}, None),
+    ],
+)
+async def test_read_use_case_recovers_from_a_lost_connection(
+    request_container: modern_di.Container,
+    provider: modern_di.providers.Factory[typing.Any],
+    kwargs: dict[str, typing.Any],
+    expected: object,
+) -> None:
+    use_case: typing.Final = request_container.resolve_provider(provider)
+    assert await use_case(**kwargs) == expected
+
+
+@pytest.mark.usefixtures("lost_connection")
+@pytest.mark.parametrize("provider", [ioc.UseCases.fetch_chat_use_case, ioc.UseCases.fetch_messages_use_case])
+async def test_membership_gated_use_case_recovers_from_a_lost_connection(
+    request_container: modern_di.Container, provider: modern_di.providers.Factory[typing.Any]
+) -> None:
+    use_case: typing.Final = request_container.resolve_provider(provider)
+    with pytest.raises(PermissionDeniedError):
+        await use_case(actor=UNKNOWN_ACTOR, chat_id=-1)