perf: faster mutate path, one pooled connection per push

README.md

@@ -76,7 +76,7 @@ app = Starlette(
 
 Point `ZERO_QUERY_URL` and `ZERO_MUTATE_URL` at `/query` and `/mutate`. With FastAPI, `app.router.routes.extend(zero_routes(...))`. With another framework, call `handle_query_request` and `handle_mutate_request` yourself; `zero_server/integrations/starlette.py` is short enough to copy. The types to annotate those handlers (`QueryRequestHandler`, `TransactFn`, `MutateResponse`, `MutationResponse`, `ServerPushFailed` and the rest) import from `zero_server`, and `server_ast(query, schema)` returns the AST zero-cache receives for a query.
 
-- **Databases.** `SQLAlchemyDatabase` (SQLAlchemy's asyncio engine on asyncpg or psycopg; the `sqlalchemy` extra installs neither driver), `AsyncpgDatabase` (an asyncpg pool) and `PsycopgDatabase` (a psycopg 3 async pool), under `zero_server.server.adapters`. Another driver plugs into `ZQLDatabase` as a `DBConnection` whose `query` takes `$1` placeholders.
+- **Databases.** `SQLAlchemyDatabase` (SQLAlchemy's asyncio engine on asyncpg or psycopg; the `sqlalchemy` extra installs neither driver), `AsyncpgDatabase` (an asyncpg pool) and `PsycopgDatabase` (a psycopg 3 async pool), under `zero_server.server.adapters`. A push's mutations run on one connection from the pool, taken at its first transaction and returned when the push ends. So mutators in one push share session state (a `SET` without `LOCAL`, session advisory locks, temp tables), and a handler that takes a second connection from the same pool while its push holds one can deadlock once the pool has fewer connections than there are concurrent pushes. Another driver plugs into `ZQLDatabase` as a `DBConnection` whose `query` takes `$1` placeholders.
 - **Validation.** `@define_query(Model)` and `@define_mutator(Model)` parse args with a Pydantic model or `TypeAdapter` first; a failure reaches the client as upstream's validation error.
 - **Strict building.** `create_builder(schema, strict=True)` refuses an unknown column when the query is built. Upstream only refuses it when the endpoint maps the query, failing the whole batch, so this is for your tests.
 

src/zero_server/protocol/mutate_server.py

@@ -11,6 +11,7 @@
 from zero_server.protocol.error import PushFailedBody
 from zero_server.protocol.mutation import MutationResponse
 from zero_server.protocol.wire import WireModel
+from zero_server.shared.json import JSONValue
 
 
 class MutateSuccess(WireModel):
@@ -22,6 +23,13 @@ class MutateSuccess(WireModel):
     user_id: StrictStr | None = Field(default=None, alias="userID")
     mutations: list[MutationResponse]
 
+    def to_json(self) -> dict[str, JSONValue]:
+        out: dict[str, JSONValue] = {"kind": self.kind}
+        if self.user_id is not None or "user_id" in self.model_fields_set:
+            out["userID"] = self.user_id
+        out["mutations"] = [m.to_json() for m in self.mutations]
+        return out
+
 
 # The wire body. Handlers see `zero_server.MutateResponse`, what `handle_mutate_request` returns.
 type MutateResponse = MutateSuccess | PushFailedBody

src/zero_server/protocol/mutation.py

@@ -106,6 +106,9 @@ class MutationOk(WireModel):
     # Whatever the mutator returned, if anything.
     data: JSONValue | None = None
 
+    def to_json(self) -> dict[str, JSONValue]:
+        return {} if self.data is None else super().to_json()
+
 
 class MutationAppError(WireModel):
     error: Literal["app"] = "app"
@@ -125,3 +128,6 @@ class MutationZeroError(WireModel):
 class MutationResponse(WireModel):
     id: MutationID
     result: MutationResult
+
+    def to_json(self) -> dict[str, JSONValue]:
+        return {"id": self.id.to_json(), "result": self.result.to_json()}

src/zero_server/protocol/mutation_id.py

@@ -5,8 +5,15 @@
 from pydantic import Field, StrictStr
 
 from zero_server.protocol.wire import Number, WireModel
+from zero_server.shared.json import JSONValue
 
 
 class MutationID(WireModel):
     id: Number
     client_id: StrictStr = Field(alias="clientID")
+
+    def to_json(self) -> dict[str, JSONValue]:
+        # A push's response holds one per mutation; an integer ID needs nothing of Pydantic's serializer.
+        if type(self.id) is int:
+            return {"id": self.id, "clientID": self.client_id}
+        return super().to_json()

src/zero_server/protocol/wire.py

@@ -9,6 +9,7 @@
 refused unless the field is `nullable`.
 """
 
+import functools
 from typing import Annotated, Any, ClassVar, Self
 
 from pydantic import (
@@ -64,12 +65,9 @@ class WireModel(BaseModel):
     @classmethod
     def _refuse_null_for_optional(cls, data: Any, info: ValidationInfo) -> Any:
         if info.context == WIRE and isinstance(data, dict):
-            for name, field in cls.model_fields.items():
-                if name in cls.nullable or field.is_required():
-                    continue
-                for key in {name, field.alias or name}:
-                    if key in data and data[key] is None:
-                        raise ValueError(f"{key} may be absent but not null")
+            for key in _optional_keys(cls):
+                if key in data and data[key] is None:
+                    raise ValueError(f"{key} may be absent but not null")
         return data
 
     @model_serializer(mode="wrap")
@@ -95,6 +93,19 @@ def from_json(cls, value: object) -> Self:
         return cls.model_validate(value, context=WIRE)
 
 
+@functools.cache
+def _optional_keys(cls: type[WireModel]) -> tuple[str, ...]:
+    """The keys, by name or alias, of the fields that may be absent but not null."""
+    return tuple(
+        dict.fromkeys(
+            key
+            for name, field in cls.model_fields.items()
+            if not (name in cls.nullable or field.is_required())
+            for key in (name, field.alias or name)
+        )
+    )
+
+
 type Number = StrictInt | StrictFloat
 """A JSON number. `bool` is not one, though Python says otherwise."""
 

src/zero_server/server/adapters/_session.py

@@ -0,0 +1,68 @@
+"""A session: one pooled connection that a push's transactions run on, one after another."""
+
+import logging
+from collections.abc import AsyncGenerator, Awaitable, Callable
+from contextlib import asynccontextmanager
+from typing import Protocol
+
+from zero_server.server.transaction import DBTransaction
+
+logger = logging.getLogger("zero_server.adapters")
+
+
+class RunTransaction[C, W](Protocol):
+    """Runs `fn` in one transaction on `connection`, as the adapter's own `transaction` does."""
+
+    def __call__[R](self, connection: C, fn: Callable[[DBTransaction[W]], Awaitable[R]], /) -> Awaitable[R]: ...
+
+
+class SessionConnection[C, W]:
+    """A `DBConnection` whose transactions all run on one connection from a pool.
+
+    The connection is taken from the pool for the first transaction. After a transaction fails it goes
+    back, as the connection may be what failed, and the next transaction takes a fresh one.
+    """
+
+    def __init__(
+        self,
+        acquire: Callable[[], Awaitable[C]],
+        release: Callable[[C], Awaitable[object]],
+        run: RunTransaction[C, W],
+    ) -> None:
+        self._acquire = acquire
+        self._release = release
+        self._run = run
+        self._connection: C | None = None
+
+    async def transaction[R](self, fn: Callable[[DBTransaction[W]], Awaitable[R]]) -> R:
+        if self._connection is None:
+            self._connection = await self._acquire()
+        try:
+            return await self._run(self._connection, fn)
+        except BaseException:
+            await self.release()
+            raise
+
+    async def release(self) -> None:
+        connection, self._connection = self._connection, None
+        if connection is None:
+            return
+        try:
+            await self._release(connection)
+        except Exception as error:
+            # The push's mutations have committed; the pool failing to take the connection back
+            # must not report them failed.
+            logger.warning("Failed to return a connection to the pool", exc_info=error)
+
+
+@asynccontextmanager
+async def session[C, W](
+    acquire: Callable[[], Awaitable[C]],
+    release: Callable[[C], Awaitable[object]],
+    run: RunTransaction[C, W],
+) -> AsyncGenerator[SessionConnection[C, W]]:
+    connection = SessionConnection(acquire, release, run)
+    try:
+        yield connection
+    finally:
+        await connection.release()

src/zero_server/server/adapters/asyncpg.py

@@ -1,11 +1,13 @@
 """asyncpg, directly. Needs the `zero-server[asyncpg]` extra."""
 
 from collections.abc import Awaitable, Callable, Sequence
+from contextlib import AbstractAsyncContextManager
 from typing import Any
 
 import asyncpg
 
 from zero_server.schema import Schema
+from zero_server.server.adapters._session import SessionConnection, session
 from zero_server.server.transaction import DBTransaction, Row
 from zero_server.server.zql_database import PostgresError, ZQLDatabase
 
@@ -33,8 +35,19 @@ def __init__(self, pool: asyncpg.Pool) -> None:
         self.pool = pool
 
     async def transaction[R](self, fn: Callable[[DBTransaction[asyncpg.Connection]], Awaitable[R]]) -> R:
-        async with self.pool.acquire() as connection, connection.transaction():
-            return await fn(AsyncpgTransaction(connection))
+        async with self.pool.acquire() as connection:
+            return await _transaction(connection, fn)
+
+    def session(self) -> AbstractAsyncContextManager[SessionConnection[asyncpg.Connection, asyncpg.Connection]]:
+        """One connection for a push's transactions, released (and so reset by the pool) once, at the end."""
+        return session(self.pool.acquire, self.pool.release, _transaction)
+
+
+async def _transaction[R](
+    connection: asyncpg.Connection, fn: Callable[[DBTransaction[asyncpg.Connection]], Awaitable[R]]
+) -> R:
+    async with connection.transaction():
+        return await fn(AsyncpgTransaction(connection))
 
 
 class AsyncpgDatabase(ZQLDatabase[asyncpg.Connection]):

src/zero_server/server/adapters/psycopg.py

@@ -5,6 +5,7 @@
 """
 
 from collections.abc import Awaitable, Callable, Sequence
+from contextlib import AbstractAsyncContextManager
 from typing import TYPE_CHECKING, Any
 
 import psycopg
@@ -18,6 +19,7 @@
     from psycopg_pool import AsyncConnectionPool
 
 from zero_server.schema import Schema
+from zero_server.server.adapters._session import SessionConnection, session
 from zero_server.server.transaction import DBTransaction, Row
 from zero_server.server.zql_database import PostgresError, ZQLDatabase
 
@@ -61,8 +63,21 @@ def __init__(self, pool: "AsyncConnectionPool") -> None:
         self.pool = pool
 
     async def transaction[R](self, fn: Callable[[DBTransaction[psycopg.AsyncConnection[Any]]], Awaitable[R]]) -> R:
-        async with self.pool.connection() as connection, connection.transaction():
-            return await fn(PsycopgTransaction(connection))
+        async with self.pool.connection() as connection:
+            return await _transaction(connection, fn)
+
+    def session(
+        self,
+    ) -> AbstractAsyncContextManager[SessionConnection[psycopg.AsyncConnection[Any], psycopg.AsyncConnection[Any]]]:
+        """One connection for a push's transactions, returned to the pool once, at the end."""
+        return session(self.pool.getconn, self.pool.putconn, _transaction)
+
+
+async def _transaction[R](
+    connection: psycopg.AsyncConnection[Any], fn: Callable[[DBTransaction[psycopg.AsyncConnection[Any]]], Awaitable[R]]
+) -> R:
+    async with connection.transaction():
+        return await fn(PsycopgTransaction(connection))
 
 
 class PsycopgDatabase(ZQLDatabase[psycopg.AsyncConnection[Any]]):

src/zero_server/server/adapters/sqlalchemy.py

@@ -4,34 +4,46 @@
 """
 
 from collections.abc import Awaitable, Callable, Sequence
+from contextlib import AbstractAsyncContextManager
 from typing import Any
 
 from sqlalchemy.exc import DBAPIError
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
+from sqlalchemy.pool import PoolProxiedConnection
 
 from zero_server.schema import Schema
+from zero_server.server.adapters._session import SessionConnection, session
 from zero_server.server.transaction import DBTransaction, Row
 from zero_server.server.zql_database import PostgresError, ZQLDatabase
 
 
 class SQLAlchemyTransaction:
     """One transaction of an `AsyncConnection`.
 
-    On asyncpg a statement goes through `exec_driver_sql`, which hands the `$1` placeholders to asyncpg
-    as they are. psycopg's paramstyle is `%s`, so on psycopg it goes through psycopg's raw cursor on the
-    same connection, inside the transaction SQLAlchemy opened.
+    psycopg's paramstyle is `%s`, so on psycopg a statement goes through psycopg's raw cursor on the
+    same connection, inside the transaction SQLAlchemy opened. On asyncpg, SQLAlchemy defers `BEGIN` to
+    the first statement it runs itself, so until the transaction is open a statement goes through
+    `exec_driver_sql`, which hands the `$1` placeholders to asyncpg as they are. After that, statements
+    go straight to the asyncpg connection, as the asyncpg adapter sends them. In a push's session,
+    `BEGIN` goes out before the callback runs, so every statement does.
     """
 
-    def __init__(self, connection: AsyncConnection, driver_transaction: DBTransaction[Any] | None = None) -> None:
+    def __init__(
+        self,
+        connection: AsyncConnection,
+        driver_transaction: DBTransaction[Any] | None = None,
+        driver_in_transaction: Callable[[], bool] = lambda: True,
+    ) -> None:
         self._connection = connection
         self._driver_transaction = driver_transaction
+        self._driver_in_transaction = driver_in_transaction
 
     @property
     def wrapped_transaction(self) -> AsyncConnection:
         return self._connection
 
     async def query(self, sql: str, params: Sequence[Any] = ()) -> list[Row]:
-        if self._driver_transaction is not None:
+        if self._driver_transaction is not None and self._driver_in_transaction():
             return await self._driver_transaction.query(sql, params)
         try:
             result = await self._connection.exec_driver_sql(sql, tuple(params))
@@ -46,16 +58,20 @@ async def query(self, sql: str, params: Sequence[Any] = ()) -> list[Row]:
 async def sqlalchemy_transaction(connection: AsyncConnection) -> SQLAlchemyTransaction:
     """Wrap a connection that is already in a transaction, such as one an `AsyncSession` holds."""
     driver = connection.dialect.driver
+    if driver not in ("asyncpg", "psycopg"):
+        raise ValueError(f"SQLAlchemy driver {driver!r} is not supported; use asyncpg or psycopg")
+    driver_connection = (await connection.get_raw_connection()).driver_connection
+    # `None` only once SQLAlchemy has invalidated the connection, which it never does mid-transaction.
+    assert driver_connection is not None
     if driver == "asyncpg":
-        return SQLAlchemyTransaction(connection)
-    if driver == "psycopg":
-        from zero_server.server.adapters.psycopg import PsycopgTransaction
+        from zero_server.server.adapters.asyncpg import AsyncpgTransaction
 
-        driver_connection = (await connection.get_raw_connection()).driver_connection
-        # `None` only once SQLAlchemy has invalidated the connection, which it never does mid-transaction.
-        assert driver_connection is not None
-        return SQLAlchemyTransaction(connection, PsycopgTransaction(driver_connection))
-    raise ValueError(f"SQLAlchemy driver {driver!r} is not supported; use asyncpg or psycopg")
+        return SQLAlchemyTransaction(
+            connection, AsyncpgTransaction(driver_connection), driver_connection.is_in_transaction
+        )
+    from zero_server.server.adapters.psycopg import PsycopgTransaction
+
+    return SQLAlchemyTransaction(connection, PsycopgTransaction(driver_connection))
 
 
 class SQLAlchemyConnection:
@@ -66,6 +82,35 @@ async def transaction[R](self, fn: Callable[[DBTransaction[AsyncConnection]], Aw
         async with self.engine.begin() as connection:
             return await fn(await sqlalchemy_transaction(connection))
 
+    def session(self) -> AbstractAsyncContextManager[SessionConnection[AsyncConnection, AsyncConnection]]:
+        """One connection for a push's transactions, checked back into the pool once, at the end."""
+        return session(self.engine.connect, _close, _transaction)
+
+
+async def _close(connection: AsyncConnection) -> None:
+    await connection.close()
+
+
+async def _transaction[R](
+    connection: AsyncConnection, fn: Callable[[DBTransaction[AsyncConnection]], Awaitable[R]]
+) -> R:
+    async with connection.begin():
+        if connection.dialect.driver == "asyncpg":
+            await _send_begin(await connection.get_raw_connection())
+        return await fn(await sqlalchemy_transaction(connection))
+
+
+async def _send_begin(pooled: PoolProxiedConnection) -> None:
+    """Have SQLAlchemy's asyncpg dialect send `BEGIN` now, so every statement can go straight to asyncpg.
+
+    The dialect defers `BEGIN` to the first statement it runs itself and has no public way to send it
+    sooner; that first statement through SQLAlchemy cost a push of 100 puts about 20%.
+    """
+    # Private: the asyncpg dialect's `_start_transaction`. A public way to begin, or its removal, retires this.
+    start = getattr(pooled.dbapi_connection, "_start_transaction", None)
+    if start is not None:
+        await start()
+
 
 class SQLAlchemyDatabase(ZQLDatabase[AsyncConnection]):
     """A `ZQLDatabase` on `engine`, whose URL names the asyncpg or the psycopg driver