perf: cut avoidable SQL round trips from every endpoint (#63)

* test: pin SQL statement counts per endpoint and a BEGIN-free revocation check

* perf: cut avoidable SQL round trips from every endpoint

Author: lesnik512

AGENTS.md

@@ -109,6 +109,10 @@ a session, a transaction or a repository directly. Provider scopes:
 - `Database.database_session` — request-scoped, finalizer closes the session.
 - `Database.transaction` — request-scoped `db_retry.Transaction`, the object
   every use case wraps its work in.
+- `Database.autocommit_transaction` — the same, opened at `AUTOCOMMIT` so the
+  token revocation check sends no BEGIN or COMMIT. `rolled_back_db_session`
+  overrides `Database.autocommit_isolation_level` to `None`, because switching
+  the shared test connection would roll back its outer transaction.
 - `Repositories.*` — request-scoped, `auto_commit=False` (see above).
 - `UseCases.*` — request-scoped, one class per operation.
 

app/database/resources.py

@@ -1,6 +1,7 @@
 import asyncio
 import typing
 
+from sqlalchemy.engine.interfaces import IsolationLevel
 from sqlalchemy.ext import asyncio as sa
 
 from app.settings import settings
@@ -31,6 +32,10 @@ def create_session(engine: sa.AsyncEngine | sa.AsyncConnection) -> sa.AsyncSessi
     )
 
 
+def autocommit_isolation_level() -> IsolationLevel | None:
+    return "AUTOCOMMIT"
+
+
 async def close_session(session: sa.AsyncSession) -> None:
     task: typing.Final = asyncio.create_task(session.close())
     await asyncio.shield(task)

app/ioc.py

@@ -39,6 +39,15 @@ class Database(Group):
         creator=Transaction,
         kwargs={"session": database_session},
     )
+    autocommit_isolation_level = providers.Factory(
+        creator=database_resources.autocommit_isolation_level, bound_type=None
+    )
+    autocommit_transaction = providers.Factory(
+        scope=Scope.REQUEST,
+        creator=Transaction,
+        bound_type=None,
+        kwargs={"session": database_session, "isolation_level": autocommit_isolation_level},
+    )
 
 
 class Repositories(Group, scope=Scope.REQUEST):
@@ -77,7 +86,9 @@ class UseCases(Group, scope=Scope.REQUEST):
     fetch_user_use_case = providers.Factory(creator=FetchUserUseCase)
     mark_read_use_case = providers.Factory(creator=MarkReadUseCase)
     revoke_token_use_case = providers.Factory(creator=RevokeTokenUseCase)
-    check_token_revoked_use_case = providers.Factory(creator=CheckTokenRevokedUseCase)
+    check_token_revoked_use_case = providers.Factory(
+        creator=CheckTokenRevokedUseCase, kwargs={"transaction": Database.autocommit_transaction}
+    )
 
 
 ALL_GROUPS: typing.Final[list[type[Group]]] = [Database, Repositories, UseCases]

app/repositories/chat_members_repository.py

@@ -1,8 +1,11 @@
 import typing
+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 app.database import tables
 
@@ -16,8 +19,29 @@ class BaseRepository(SQLAlchemyAsyncRepository[tables.ChatMembersTable]):
     async def is_member(self, chat_id: int, user_id: int) -> bool:
         return await self.exists(chat_id=chat_id, user_id=user_id)
 
-    async def fetch_member(self, chat_id: int, user_id: int) -> tables.ChatMembersTable | None:
-        return await self.get_one_or_none(chat_id=chat_id, user_id=user_id)
+    async def add_members(self, chat: tables.ChatsTable, user_ids: Sequence[int]) -> None:
+        with wrap_sqlalchemy_exception(error_messages=DEFAULT_ERROR_MESSAGE_TEMPLATES, dialect_name="postgresql"):
+            members: typing.Final = await self.repository.session.scalars(
+                sa.insert(tables.ChatMembersTable).returning(tables.ChatMembersTable),
+                [{"chat_id": chat.id, "user_id": user_id} for user_id in user_ids],
+            )
+        orm.attributes.set_committed_value(chat, "members", members.all())
+
+    async def fetch_member_and_message_exists(
+        self, chat_id: int, user_id: int, message_id: int
+    ) -> tuple[tables.ChatMembersTable, bool] | None:
+        message_exists: typing.Final = (
+            sa.exists()
+            .where(tables.MessagesTable.id == message_id, tables.MessagesTable.chat_id == chat_id)
+            .label("message_exists")
+        )
+        result: typing.Final = await self.repository.session.execute(
+            sa.select(tables.ChatMembersTable, message_exists).where(
+                tables.ChatMembersTable.chat_id == chat_id, tables.ChatMembersTable.user_id == user_id
+            )
+        )
+        row: typing.Final = result.one_or_none()
+        return None if row is None else (row[0], row[1])
 
     async def mark_read(self, member_id: int, requested_message_id: int) -> tables.ChatMembersTable:
         """Advance last_read_message_id to GREATEST(current, requested), atomically.

app/repositories/chats_repository.py

@@ -1,3 +1,4 @@
+import datetime as dt
 import typing
 from collections.abc import Sequence
 
@@ -11,18 +12,45 @@
 from app.database import tables
 
 
+_WITH_MEMBERS: typing.Final[dict[str, typing.Any]] = {
+    "load": [orm.joinedload(tables.ChatsTable.members)],
+    "uniquify": True,
+    # A chat already in the identity map keeps its noload members without this.
+    "execution_options": {"populate_existing": True},
+}
+
+
 class ChatsRepository(SQLAlchemyAsyncRepositoryService[tables.ChatsTable]):
     class BaseRepository(SQLAlchemyAsyncRepository[tables.ChatsTable]):
         model_type = tables.ChatsTable
 
     repository_type = BaseRepository
 
     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},
+        return await self.get_one(tables.ChatsTable.id == chat_id, **_WITH_MEMBERS)
+
+    async def find_with_members(self, chat_id: int) -> tables.ChatsTable | None:
+        return await self.get_one_or_none(tables.ChatsTable.id == chat_id, **_WITH_MEMBERS)
+
+    async def set_last_message(self, chat_id: int, message_id: int) -> None:
+        await self.repository.session.execute(
+            sa.update(tables.ChatsTable)
+            .where(tables.ChatsTable.id == chat_id)
+            .values(last_message_id=message_id, updated_at=dt.datetime.now(tz=dt.UTC))
+        )
+
+    async def repoint_last_message(self, chat_id: int, deleted_message_id: int) -> None:
+        newest_active: typing.Final = (
+            sa.select(tables.MessagesTable.id)
+            .where(tables.MessagesTable.chat_id == chat_id, tables.MessagesTable.deleted_at.is_(None))
+            .order_by(tables.MessagesTable.id.desc())
+            .limit(1)
+            .scalar_subquery()
+        )
+        await self.repository.session.execute(
+            sa.update(tables.ChatsTable)
+            .where(tables.ChatsTable.id == chat_id, tables.ChatsTable.last_message_id == deleted_message_id)
+            .values(last_message_id=newest_active, updated_at=dt.datetime.now(tz=dt.UTC))
         )
 
     async def get_or_create_direct(self, *, direct_key: str, created_by_id: int) -> tuple[tables.ChatsTable, bool]:
@@ -59,7 +87,7 @@ async def list_for_user(self, user_id: int) -> Sequence[tables.ChatsTable]:
             sa.select(tables.ChatsTable)
             .options(
                 orm.with_expression(tables.ChatsTable.unread_count, unread_count),
-                orm.selectinload(tables.ChatsTable.last_message),
+                orm.joinedload(tables.ChatsTable.last_message),
             )
             .join(tables.ChatMembersTable, tables.ChatMembersTable.chat_id == tables.ChatsTable.id)
             .where(tables.ChatMembersTable.user_id == user_id)

app/repositories/messages_repository.py

@@ -1,9 +1,10 @@
+import datetime as dt
 import typing
 import uuid
 from collections.abc import Sequence
 
 import sqlalchemy as sa
-from advanced_alchemy.exceptions import wrap_sqlalchemy_exception
+from advanced_alchemy.exceptions import NotFoundError, wrap_sqlalchemy_exception
 from advanced_alchemy.filters import LimitOffset
 from advanced_alchemy.repository import DEFAULT_ERROR_MESSAGE_TEMPLATES, SQLAlchemyAsyncRepository
 from advanced_alchemy.service import SQLAlchemyAsyncRepositoryService
@@ -37,18 +38,37 @@ async def get_or_create(
         message, was_created = result.one()
         return message, was_created
 
-    async def fetch_latest_active(self, chat_id: int) -> tables.MessagesTable | None:
-        """Return the newest non-deleted message in a chat, or None if none remains.
+    async def fetch_with_membership(self, message_id: int, user_id: int) -> tuple[tables.MessagesTable, bool]:
+        is_member: typing.Final = (
+            sa.exists()
+            .where(
+                tables.ChatMembersTable.chat_id == tables.MessagesTable.chat_id,
+                tables.ChatMembersTable.user_id == user_id,
+            )
+            .label("is_member")
+        )
+        result: typing.Final = await self.repository.session.execute(
+            sa.select(tables.MessagesTable, is_member).where(tables.MessagesTable.id == message_id)
+        )
+        row: typing.Final = result.one_or_none()
+        if row is None:
+            msg = "No item found when one was expected"
+            raise NotFoundError(msg)
+        return row[0], row[1]
 
-        Used to repoint ChatsTable.last_message_id after a delete removes the current pointer.
-        """
-        messages = await self.get_many(
-            tables.MessagesTable.chat_id == chat_id,
-            tables.MessagesTable.deleted_at.is_(None),
-            LimitOffset(limit=1, offset=0),
-            order_by=[("id", True)],
+    async def edit_text(self, message_id: int, text: str, edited_at: dt.datetime) -> tables.MessagesTable:
+        result: typing.Final = await self.repository.session.execute(
+            sa.update(tables.MessagesTable)
+            .where(tables.MessagesTable.id == message_id)
+            .values(text=text, edited_at=edited_at)
+            .returning(tables.MessagesTable)
+        )
+        return result.scalar_one()
+
+    async def soft_delete(self, message_id: int, deleted_at: dt.datetime) -> None:
+        await self.repository.session.execute(
+            sa.update(tables.MessagesTable).where(tables.MessagesTable.id == message_id).values(deleted_at=deleted_at)
         )
-        return messages[0] if messages else None
 
     async def list_page(
         self,

app/repositories/revoked_tokens_repository.py

@@ -1,5 +1,6 @@
 import datetime as dt
 
+import sqlalchemy as sa
 from advanced_alchemy.repository import SQLAlchemyAsyncRepository
 from advanced_alchemy.service import SQLAlchemyAsyncRepositoryService
 from sqlalchemy.dialects import postgresql
@@ -25,4 +26,6 @@ async def revoke(self, jti: str, expires_at: dt.datetime) -> None:
         )
 
     async def prune_expired(self, now: dt.datetime) -> None:
-        await self.delete_where(tables.RevokedTokensTable.expires_at < now)
+        await self.repository.session.execute(
+            sa.delete(tables.RevokedTokensTable).where(tables.RevokedTokensTable.expires_at < now)
+        )

app/use_cases/create_chat.py

@@ -38,12 +38,13 @@ async def __call__(self, *, actor: Actor, data: CreateChatRequest) -> tuple[tabl
                 )
             else:
                 chat = await self.chats_repository.create(
-                    tables.ChatsTable(chat_type=data.chat_type, title=data.title, created_by_id=actor.id)
+                    tables.ChatsTable(chat_type=data.chat_type, title=data.title, created_by_id=actor.id),
+                    auto_refresh=False,
                 )
                 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))
-            chat_with_members: typing.Final = await self.chats_repository.fetch_with_members(chat.id)
+                await self.chat_members_repository.add_members(chat, sorted(member_ids))
+            else:
+                chat = await self.chats_repository.fetch_with_members(chat.id)
             await self.transaction.commit()
-        return chat_with_members, created
+        return chat, created

app/use_cases/create_message.py

@@ -31,10 +31,6 @@ async def __call__(
                 chat_id=chat_id, user_id=actor.id, idempotency_key=data.idempotency_key, text=data.text
             )
             if created:
-                await self.chats_repository.update(
-                    tables.ChatsTable(id=chat_id, last_message_id=message.id),
-                    item_id=chat_id,
-                    attribute_names=["last_message_id"],
-                )
+                await self.chats_repository.set_last_message(chat_id, message.id)
             await self.transaction.commit()
         return message, created

app/use_cases/delete_message.py

@@ -4,8 +4,6 @@
 from db_retry import Transaction, postgres_retry
 
 from app.actor import Actor
-from app.database import tables
-from app.repositories.chat_members_repository import ChatMembersRepository
 from app.repositories.chats_repository import ChatsRepository
 from app.repositories.messages_repository import MessagesRepository
 from app.use_cases.message_authorization import fetch_message_for_author
@@ -15,31 +13,20 @@
 class DeleteMessageUseCase:
     transaction: Transaction
     messages_repository: MessagesRepository
-    chat_members_repository: ChatMembersRepository
     chats_repository: ChatsRepository
 
     @postgres_retry
     async def __call__(self, *, actor: Actor, message_id: int) -> None:
         async with self.transaction:
             message = await fetch_message_for_author(
                 messages_repository=self.messages_repository,
-                chat_members_repository=self.chat_members_repository,
                 actor=actor,
                 message_id=message_id,
                 action="delete",
             )
             if message.deleted_at is not None:
                 # DELETE is idempotent; a second delete is not an error, unlike an edit.
                 return
-            message.deleted_at = dt.datetime.now(tz=dt.UTC)
-            await self.messages_repository.update(message, item_id=message_id)
-
-            chat = await self.chats_repository.get_one(id=message.chat_id)
-            if chat.last_message_id == message_id:
-                newest = await self.messages_repository.fetch_latest_active(message.chat_id)
-                await self.chats_repository.update(
-                    tables.ChatsTable(id=message.chat_id, last_message_id=newest.id if newest is not None else None),
-                    item_id=message.chat_id,
-                    attribute_names=["last_message_id"],
-                )
+            await self.messages_repository.soft_delete(message_id, dt.datetime.now(tz=dt.UTC))
+            await self.chats_repository.repoint_last_message(message.chat_id, message_id)
             await self.transaction.commit()