perf: create a direct chat with one INSERT ... ON CONFLICT DO UPDATE

AGENTS.md

@@ -135,12 +135,10 @@ back. `app`/`client` fixtures build the real app and run it through
 container_fixture="request_container")` (`tests/use_cases/conftest.py`)
 exposes every repository/use case provider as a same-named pytest fixture —
 the template predates this and hand-assembles dependencies instead. The
-race-simulation pattern used to test the concurrent-retry path without a second
-real connection is the `_Racing*Repository` classes in
-`tests/use_cases/test_create_chat.py`. `tests/use_cases/test_create_message.py`
-uses a second real connection instead: it holds an uncommitted insert and polls
-`pg_stat_activity` until the use case waits on its lock. The invariant each one
-pins is in the docstring on the test that uses it.
+concurrent-create tests in `tests/use_cases/test_create_chat.py` and
+`tests/use_cases/test_create_message.py` use a second real connection: it holds
+an uncommitted insert, and the test polls `pg_stat_activity` until the use case
+waits on its lock. The invariant each one pins is in its docstring.
 
 **Migrations**: `migrations/env.py` reads the shared `METADATA` and rewrites
 the DSN driver from `postgresql+asyncpg` → `postgresql` (Alembic uses sync

app/repositories/chats_repository.py

@@ -2,9 +2,11 @@
 from collections.abc import Sequence
 
 import sqlalchemy as sa
-from advanced_alchemy.repository import SQLAlchemyAsyncRepository
+from advanced_alchemy.exceptions import wrap_sqlalchemy_exception
+from advanced_alchemy.repository import DEFAULT_ERROR_MESSAGE_TEMPLATES, SQLAlchemyAsyncRepository
 from advanced_alchemy.service import SQLAlchemyAsyncRepositoryService
 from sqlalchemy import orm
+from sqlalchemy.dialects import postgresql
 
 from app.database import tables
 
@@ -19,13 +21,26 @@ async def fetch_with_members(self, chat_id: int) -> tables.ChatsTable:
         return await self.get_one(
             tables.ChatsTable.id == chat_id,
             load=[orm.selectinload(tables.ChatsTable.members)],
+            # A chat already in the identity map keeps its noload members without this.
+            execution_options={"populate_existing": True},
         )
 
-    async def fetch_direct_by_key(self, direct_key: str) -> tables.ChatsTable | None:
-        return await self.get_one_or_none(
-            tables.ChatsTable.direct_key == direct_key,
-            load=[orm.selectinload(tables.ChatsTable.members)],
+    async def get_or_create_direct(self, *, direct_key: str, created_by_id: int) -> tuple[tables.ChatsTable, bool]:
+        statement = postgresql.insert(tables.ChatsTable).values(
+            chat_type=tables.ChatType.DIRECT, created_by_id=created_by_id, direct_key=direct_key
+        )
+        statement = statement.on_conflict_do_update(
+            index_elements=[tables.ChatsTable.direct_key],
+            set_={"direct_key": statement.excluded.direct_key},
         )
+        # xmax is 0 only on a row version this statement inserted.
+        created: typing.Final = sa.literal_column("xmax = 0", sa.Boolean)
+        with wrap_sqlalchemy_exception(error_messages=DEFAULT_ERROR_MESSAGE_TEMPLATES, dialect_name="postgresql"):
+            result: typing.Final = await self.repository.session.execute(
+                statement.returning(tables.ChatsTable, created)
+            )
+        chat, was_created = result.one()
+        return chat, was_created
 
     async def list_for_user(self, user_id: int) -> Sequence[tables.ChatsTable]:
         unread_count: typing.Final = (

app/use_cases/create_chat.py

@@ -1,7 +1,6 @@
 import dataclasses
 import typing
 
-from advanced_alchemy.exceptions import DuplicateKeyError
 from db_retry import Transaction, postgres_retry
 
 from app.actor import Actor
@@ -32,40 +31,19 @@ class CreateChatUseCase:
     @postgres_retry
     async def __call__(self, *, actor: Actor, data: CreateChatRequest) -> tuple[tables.ChatsTable, bool]:
         member_ids: typing.Final = {actor.id, *data.member_ids}
-        direct_key: typing.Final = _direct_key(member_ids) if data.chat_type is tables.ChatType.DIRECT else None
-
-        chat: tables.ChatsTable | None = None
         async with self.transaction:
-            if direct_key is not None:
-                existing = await self.chats_repository.fetch_direct_by_key(direct_key)
-                if existing is not None:
-                    await self.transaction.commit()
-                    return existing, False
-
-            try:
-                chat = await self.chats_repository.create(
-                    tables.ChatsTable(
-                        chat_type=data.chat_type,
-                        title=data.title if data.chat_type is tables.ChatType.GROUP else None,
-                        created_by_id=actor.id,
-                        direct_key=direct_key,
-                    )
+            if data.chat_type is tables.ChatType.DIRECT:
+                chat, created = await self.chats_repository.get_or_create_direct(
+                    direct_key=_direct_key(member_ids), created_by_id=actor.id
                 )
-            except DuplicateKeyError:
-                # Group chats have no unique constraint here, so a collision is not recoverable.
-                if direct_key is None:
-                    raise
-                await self.transaction.rollback()
             else:
+                chat = await self.chats_repository.create(
+                    tables.ChatsTable(chat_type=data.chat_type, title=data.title, created_by_id=actor.id)
+                )
+                created = True
+            if created:
                 for user_id in sorted(member_ids):
                     await self.chat_members_repository.create(tables.ChatMembersTable(chat_id=chat.id, user_id=user_id))
-                await self.transaction.commit()
-
-        if chat is None:
-            existing = await self.chats_repository.fetch_direct_by_key(direct_key)  # ty: ignore[invalid-argument-type]
-            if existing is None:
-                msg = "Direct chat creation raced but the resulting row could not be found"
-                raise RuntimeError(msg)
-            return existing, False
-
-        return await self.chats_repository.fetch_with_members(chat.id), True
+            chat_with_members: typing.Final = await self.chats_repository.fetch_with_members(chat.id)
+            await self.transaction.commit()
+        return chat_with_members, created

docs/adr/0005-upsert-via-duplicate-key-recovery.md

@@ -1,3 +1,7 @@
+---
+superseded_by: 0013-direct-chat-create-is-one-upsert
+---
+
 # Upsert by recovering from DuplicateKeyError
 
 `CreateChatUseCase` reads first to see whether the direct chat already exists

docs/adr/0009-coverage-exclusions-structural-only.md

@@ -5,10 +5,9 @@ The suite runs at `--cov-fail-under=100`, and exemptions are structural only:
 all, each carrying an inline reason in `pyproject.toml` where the path does not
 explain itself, plus `exclude_also` for `if typing.TYPE_CHECKING:` blocks. No
 `# pragma: no cover` appears in `app/`, `tests/` or `migrations/`. "Awkward to
-reach" is not a warrant: where a branch looked untestable the answer was a
-repository subclass that raises the condition the database would raise, which is
-how the `DuplicateKeyError` recovery path in `CreateChatUseCase` and its
-defensive `is None` guard became executed code. Pragmas on unreachable-in-production guards were argued
+reach" is not a warrant: where a branch looks untestable, the answer is to
+produce the database condition for real, as the concurrent create tests do with
+a second connection holding an uncommitted insert. Pragmas on unreachable-in-production guards were argued
 for twice and refused, because an excluded branch is one nobody notices when it
 stops being unreachable, and the coverage number then asserts something untrue
 about what the tests exercise; tests written only to move the number were

docs/adr/0013-direct-chat-create-is-one-upsert.md

@@ -0,0 +1,15 @@
+# A direct chat is created with one upsert
+
+`ChatsRepository.get_or_create_direct` uses the statement shape from
+[0012](0012-message-send-is-one-upsert.md): `INSERT ... ON CONFLICT (direct_key)
+DO UPDATE SET direct_key = EXCLUDED.direct_key RETURNING ..., xmax = 0`. A
+concurrent request for the same pair waits on the row lock and gets the
+winner's chat. Members are inserted only when the chat is new, and
+`fetch_with_members` loads the result inside the same transaction. This replaces
+the `DuplicateKeyError` recovery in
+[0005](0005-upsert-via-duplicate-key-recovery.md), whose re-read ran outside the
+retried transaction. The cost lands on replay, the common path: opening an
+existing DM now takes three statements (upsert, SELECT, selectin) where the
+pre-check took two. Keeping replay at two meant inserting or selecting members
+by hand and attaching them to the viewonly relationship, which was rejected as
+more machinery than one statement is worth.

tests/use_cases/test_connection_loss.py

@@ -145,7 +145,7 @@ async def test_membership_gated_use_case_recovers_from_a_connection_lost_while_p
 
 
 @pytest.mark.usefixtures("lost_pooled_connection")
-async def test_direct_chat_lookup_recovers_from_a_connection_lost_while_pooled(
+async def test_direct_chat_upsert_recovers_from_a_connection_lost_while_pooled(
     request_container: modern_di.Container,
 ) -> None:
     use_case: typing.Final = request_container.resolve_provider(ioc.UseCases.create_chat_use_case)

tests/use_cases/test_create_chat.py

@@ -1,62 +1,22 @@
+import asyncio
+import typing
+
 import pytest
-from advanced_alchemy.exceptions import DuplicateKeyError
+import sqlalchemy as sa
+from db_retry import Transaction
+from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession
 
+from app import security
 from app.actor import Actor
 from app.database import tables
+from app.database.resources import create_database_engine, create_session
 from app.exceptions import ValidationError
+from app.repositories.chat_members_repository import ChatMembersRepository
 from app.repositories.chats_repository import ChatsRepository
 from app.schemas import api as schemas
 from app.use_cases.create_chat import CreateChatUseCase
-
-
-class _RacingChatsRepository(ChatsRepository):
-    """Simulates losing a create-direct-chat race.
-
-    The pre-check misses (as if the winner's row weren't committed/visible yet), the insert
-    then collides with the winner's now-committed row (DuplicateKeyError), and the recovery
-    re-read must find it.
-    """
-
-    _missed_precheck: bool = False
-
-    async def fetch_direct_by_key(self, direct_key: str) -> tables.ChatsTable | None:
-        if not self._missed_precheck:
-            self._missed_precheck = True
-            return None
-        return await super().fetch_direct_by_key(direct_key)
-
-    async def create(self, *_args: object, **_kwargs: object) -> tables.ChatsTable:
-        msg = "simulated race: another request already created this direct chat"
-        raise DuplicateKeyError(msg)
-
-
-class _NeverFoundChatsRepository(ChatsRepository):
-    """Simulates a race whose recovery re-read can never find the winner's row.
-
-    Both `fetch_direct_by_key` calls (the pre-check and the post-rollback recovery re-read)
-    return `None`, and `create()` always raises `DuplicateKeyError` - a state the unique
-    constraint on `direct_key` should make unreachable in production, exercised here only to
-    prove `CreateChatUseCase` raises `RuntimeError` rather than returning `None` silently.
-    """
-
-    async def fetch_direct_by_key(self, direct_key: str) -> tables.ChatsTable | None:  # noqa: ARG002
-        return None
-
-    async def create(self, *_args: object, **_kwargs: object) -> tables.ChatsTable:
-        msg = "simulated race: another request already created this direct chat"
-        raise DuplicateKeyError(msg)
-
-
-class _AlwaysDuplicateChatsRepository(ChatsRepository):
-    """Stub whose create() always raises DuplicateKeyError.
-
-    Simulates an unexpected unique-constraint violation at the seam the real repository
-    would raise it from.
-    """
-
-    async def create(self, *_args: object, **_kwargs: object) -> tables.ChatsTable:
-        msg = "simulated duplicate key"
-        raise DuplicateKeyError(msg)
+from tests.api.helpers import count_statements
+from tests.factories import UserFactory
 
 
 async def test_direct_chat_is_created_with_both_members(
@@ -86,65 +46,6 @@ async def test_direct_chat_is_idempotent_for_the_same_pair(
     assert {member.user_id for member in second.members} == {alice.id, bob.id}
 
 
-async def test_direct_chat_creation_recovers_from_a_concurrent_duplicate_key(
-    create_chat_use_case: CreateChatUseCase, alice: Actor, bob: Actor
-) -> None:
-    """INVARIANT: a direct-chat insert losing the uq_chats_direct_key race recovers the winner's row.
-
-    Broken by returning from inside the `async with self.transaction:` block, or
-    by re-reading before the rollback: Transaction.__aexit__ rolls back and closes
-    the session on an open transaction, expiring every loaded attribute, so the
-    recovery read must happen after the block exits.
-    """
-    winner, winner_created = await create_chat_use_case(
-        actor=alice, data=schemas.CreateChatRequest(chat_type=tables.ChatType.DIRECT, member_ids=[bob.id])
-    )
-    assert winner_created is True
-    # Captured now: the racer's recovery rollback() expires `winner` on the shared session.
-    winner_id = winner.id
-
-    racer = CreateChatUseCase(
-        transaction=create_chat_use_case.transaction,
-        chats_repository=_RacingChatsRepository(
-            session=create_chat_use_case.chats_repository.repository.session, auto_commit=False
-        ),
-        chat_members_repository=create_chat_use_case.chat_members_repository,
-    )
-    loser, loser_created = await racer(
-        actor=bob, data=schemas.CreateChatRequest(chat_type=tables.ChatType.DIRECT, member_ids=[alice.id])
-    )
-    assert loser_created is False
-    assert loser.id == winner_id
-
-
-async def test_direct_chat_recovery_raises_if_the_winners_row_is_unreadable(
-    create_chat_use_case: CreateChatUseCase, alice: Actor, bob: Actor
-) -> None:
-    broken = CreateChatUseCase(
-        transaction=create_chat_use_case.transaction,
-        chats_repository=_NeverFoundChatsRepository(
-            session=create_chat_use_case.chats_repository.repository.session, auto_commit=False
-        ),
-        chat_members_repository=create_chat_use_case.chat_members_repository,
-    )
-    with pytest.raises(RuntimeError, match="could not be found"):
-        await broken(actor=alice, data=schemas.CreateChatRequest(chat_type=tables.ChatType.DIRECT, member_ids=[bob.id]))
-
-
-async def test_group_chat_reraises_an_unexpected_duplicate_key(
-    create_chat_use_case: CreateChatUseCase, alice: Actor, bob: Actor
-) -> None:
-    broken = CreateChatUseCase(
-        transaction=create_chat_use_case.transaction,
-        chats_repository=_AlwaysDuplicateChatsRepository(
-            session=create_chat_use_case.chats_repository.repository.session, auto_commit=False
-        ),
-        chat_members_repository=create_chat_use_case.chat_members_repository,
-    )
-    with pytest.raises(DuplicateKeyError):
-        await broken(actor=alice, data=schemas.CreateChatRequest(chat_type=tables.ChatType.GROUP, member_ids=[bob.id]))
-
-
 async def test_direct_chat_rejects_more_than_two_members(
     create_chat_use_case: CreateChatUseCase,
     alice: Actor,
@@ -170,3 +71,90 @@ async def test_group_chat_includes_the_creator(
     assert created is True
     assert chat.direct_key is None
     assert {member.user_id for member in chat.members} == {alice.id, bob.id, carol.id}
+
+
+async def test_direct_chat_replay_starts_with_the_upsert(
+    create_chat_use_case: CreateChatUseCase, alice: Actor, bob: Actor, db_session: AsyncSession
+) -> None:
+    await create_chat_use_case(
+        actor=alice, data=schemas.CreateChatRequest(chat_type=tables.ChatType.DIRECT, member_ids=[bob.id])
+    )
+    with count_statements(db_session) as statements:
+        await create_chat_use_case(
+            actor=bob, data=schemas.CreateChatRequest(chat_type=tables.ChatType.DIRECT, member_ids=[alice.id])
+        )
+
+    round_trips: typing.Final = [statement for statement in statements if "SAVEPOINT" not in statement]
+    assert "ON CONFLICT" in round_trips[0], round_trips
+
+
+async def _wait_until_blocked_on_a_lock(engine: AsyncEngine, pid: int) -> None:
+    wait_event_type = None
+    while wait_event_type != "Lock":
+        await asyncio.sleep(0.01)
+        async with engine.connect() as connection:
+            wait_event_type = await connection.scalar(
+                sa.text("SELECT wait_event_type FROM pg_stat_activity WHERE pid = :pid"), {"pid": pid}
+            )
+
+
+async def test_direct_chat_returns_the_chat_a_concurrent_request_created_for_the_same_pair() -> None:
+    """INVARIANT: a direct-chat create racing an uncommitted insert of the same pair waits and returns that chat.
+
+    It reports created=False and inserts no members of its own. Broken by switching to ON CONFLICT
+    DO NOTHING, whose RETURNING yields no row for the request that loses the race.
+    """
+    engine: typing.Final = create_database_engine()
+    setup_session: typing.Final = create_session(engine)
+    users: typing.Final = [
+        UserFactory.build(
+            username=f"racer-{index}-{id(setup_session)}",
+            password_hash=security.hash_password("hunter2hunter2"),
+            display_name="Racer",
+        )
+        for index in range(2)
+    ]
+    setup_session.add_all(users)
+    await setup_session.commit()
+    await setup_session.close()
+    first, second = (user.id for user in users)
+    direct_key: typing.Final = tables.build_direct_key(first, second)
+    try:
+        async with engine.connect() as concurrent_request:
+            chat_id: typing.Final = await concurrent_request.scalar(
+                sa.insert(tables.ChatsTable)
+                .values(chat_type=tables.ChatType.DIRECT, created_by_id=first, direct_key=direct_key)
+                .returning(tables.ChatsTable.id)
+            )
+            await concurrent_request.execute(
+                sa.insert(tables.ChatMembersTable),
+                [{"chat_id": chat_id, "user_id": first}, {"chat_id": chat_id, "user_id": second}],
+            )
+            session: typing.Final = create_session(engine)
+            pid: typing.Final = await (await session.connection()).scalar(sa.text("SELECT pg_backend_pid()"))
+            use_case: typing.Final = CreateChatUseCase(
+                transaction=Transaction(session=session),
+                chats_repository=ChatsRepository(session=session, auto_commit=False),
+                chat_members_repository=ChatMembersRepository(session=session, auto_commit=False),
+            )
+            pending: typing.Final = asyncio.create_task(
+                use_case(
+                    actor=Actor(id=second),
+                    data=schemas.CreateChatRequest(chat_type=tables.ChatType.DIRECT, member_ids=[first]),
+                )
+            )
+            await asyncio.wait_for(_wait_until_blocked_on_a_lock(engine, pid), timeout=5)
+            await concurrent_request.commit()
+            chat, created = await pending
+            await session.close()
+        assert chat.id == chat_id
+        assert created is False
+        assert sorted(member.user_id for member in chat.members) == sorted([first, second])
+    finally:
+        async with engine.begin() as connection:
+            await connection.execute(
+                sa.delete(tables.ChatMembersTable).where(tables.ChatMembersTable.user_id.in_([first, second]))
+            )
+            await connection.execute(sa.delete(tables.ChatsTable).where(tables.ChatsTable.direct_key == direct_key))
+            await connection.execute(sa.delete(tables.UsersTable).where(tables.UsersTable.id.in_([first, second])))
+        await engine.dispose()