From b7db5350abf8a5ffe52d38b4f405067ee05b7922 Mon Sep 17 00:00:00 2001 From: Jonas Daniels Date: Thu, 15 Oct 2026 11:36:51 -0700 Subject: [PATCH 1/9] perf: skip Decimal in js_number for safe integers and mid-range floats Integers below 2**53 print as str() does. Floats in [1e-4, 1e16) print as Python's repr does, minus a trailing ".0". js_keys returns the dict's own order when no key starts with a digit, and js_stringify reuses one encoder for strings instead of building one per json.dumps call. Co-Authored-By: Claude Opus 5.5 --- src/zero_server/shared/js.py | 27 ++++++++++++++++++++++----- tests/unit/test_shared.py | 9 +++++++++ 2 files changed, 31 insertions(+), 5 deletions(-) diff --git a/src/zero_server/shared/js.py b/src/zero_server/shared/js.py index 06dda13..072b374 100644 --- a/src/zero_server/shared/js.py +++ b/src/zero_server/shared/js.py @@ -22,8 +22,18 @@ def js_string(value: object) -> str: return str(value) +# Integers below 2**53 are exactly representable as a double, so JavaScript prints their digits as they are. +_MAX_EXACT = 2**53 + + def js_number(value: float) -> str: """`Number.prototype.toString()`: the shortest round-trip digits, laid out by ECMA-262's rules.""" + if type(value) is int and -_MAX_EXACT < value < _MAX_EXACT: + return str(value) + if type(value) is float and 1e-4 <= abs(value) < 1e16: + # Python's repr writes the same shortest digits, and in this range lays them out as JavaScript does. + text = repr(value) + return text[:-2] if text.endswith(".0") else text if value == 0: return "0" if value < 0: @@ -36,8 +46,6 @@ def js_number(value: float) -> str: n = k + exponent if k <= n <= 21: return digits + "0" * (n - k) - if 0 < n <= 21: - return digits[:n] + "." + digits[n:] if -6 < n <= 0: return "0." + "0" * -n + digits e = n - 1 @@ -91,6 +99,10 @@ def js_truthy(value: object) -> bool: return True +# `JSON.stringify`'s escaping of a string, without `json.dumps` building an encoder for every call. +_encode_string = json.JSONEncoder(ensure_ascii=False).encode + + def js_stringify(value: object) -> str: """`JSON.stringify(value)` for a JSON value: compact, numbers written as JavaScript writes them.""" match value: @@ -101,11 +113,11 @@ def js_stringify(value: object) -> str: case int() | float(): return js_number(value) case str(): - return json.dumps(value, ensure_ascii=False) + return _encode_string(value) case list() | tuple(): - return "[" + ",".join(js_stringify(item) for item in value) + "]" + return "[" + ",".join([js_stringify(item) for item in value]) + "]" case dict(): - members = (f"{json.dumps(k, ensure_ascii=False)}:{js_stringify(value[k])}" for k in js_keys(value)) + members = [f"{_encode_string(k)}:{js_stringify(value[k])}" for k in js_keys(value)] return "{" + ",".join(members) + "}" case _: raise TypeError(f"not a JSON value: {type(value).__name__}") @@ -117,5 +129,10 @@ def _is_array_index(key: str) -> bool: def js_keys(obj: dict[str, object]) -> list[str]: """`Object.keys(obj)`: array-index keys first in numeric order, then the rest as inserted.""" + for key in obj: + if type(key) is not str or key[:1].isdigit(): + break + else: + return list(obj) indexes = sorted((k for k in obj if _is_array_index(k)), key=int) return indexes + [k for k in obj if not _is_array_index(k)] diff --git a/tests/unit/test_shared.py b/tests/unit/test_shared.py index c333b1e..132358a 100644 --- a/tests/unit/test_shared.py +++ b/tests/unit/test_shared.py @@ -108,6 +108,15 @@ def test_error_details() -> None: (-1.25e-9, "-1.25e-9"), (0.0, "0"), (12.5, "12.5"), + (100.0, "100"), + (0.00001, "0.00001"), + (-0.0001, "-0.0001"), + (9999999999999998.0, "9999999999999998"), + (1e16, "10000000000000000"), + (-7, "-7"), + (2**53 - 1, "9007199254740991"), + # Past 2**53 an integer is a double first, as JavaScript would hold it. + (2**60 + 1, "1152921504606847000"), ], ) def test_js_number(value: float, expected: str) -> None: From 6a148b53fdd2fbc56bac29d0dd6f57114da8d477 Mon Sep 17 00:00:00 2001 From: Jonas Daniels Date: Thu, 15 Oct 2026 11:36:51 -0700 Subject: [PATCH 2/9] perf: cache type lookups, quoted identifiers and per-column casts Building a CRUD statement looked up and lowercased each column's type and rebuilt its cast for every value. The cast now comes from a cache keyed by column, split around the placeholder. Co-Authored-By: Claude Opus 5.5 --- src/zero_server/server/pg_types.py | 2 ++ src/zero_server/z2s/sql.py | 38 +++++++++++++++++++++++------- 2 files changed, 31 insertions(+), 9 deletions(-) diff --git a/src/zero_server/server/pg_types.py b/src/zero_server/server/pg_types.py index c7b9410..14dee7d 100644 --- a/src/zero_server/server/pg_types.py +++ b/src/zero_server/server/pg_types.py @@ -4,6 +4,7 @@ (905d08516ba0caaf314cd28e93a5a47e8acdd482). Apache-2.0, Copyright Rocicorp. """ +import functools from typing import Final from zero_server.schema import ValueType @@ -30,6 +31,7 @@ ) # fmt: skip +@functools.cache def _lookup(pg_type: str) -> str: """Lowercased, with any arguments (the `(32)` of `char(32)`) removed.""" lowered = pg_type.lower() diff --git a/src/zero_server/z2s/sql.py b/src/zero_server/z2s/sql.py index 915f048..27cd443 100644 --- a/src/zero_server/z2s/sql.py +++ b/src/zero_server/z2s/sql.py @@ -7,6 +7,7 @@ text, a timestamp epoch milliseconds, a JSON value its `JSON.stringify` output. """ +import functools from collections.abc import Sequence from typing import Any, Literal @@ -19,6 +20,7 @@ _LITERAL_PG_TYPE: dict[str, str] = {"boolean": "boolean", "number": "double precision", "string": "text"} +@functools.lru_cache(maxsize=4096) def ident(*names: str) -> str: """Quote an identifier as upstream's `sql.ident` does. A single dotted name is split into its parts.""" parts = names[0].split(".") if len(names) == 1 else list(names) @@ -38,6 +40,22 @@ def literal_type(value: object) -> LiteralType: return "object" +_PLACEHOLDER = "\x00" # never in a type name + + +@functools.cache +def _cast_around(column: ServerColumnSchema, comparison: bool) -> tuple[str, str]: + """The SQL before and after a value's placeholder that casts it to `column`'s type.""" + before, after = _cast(_PLACEHOLDER, column, comparison).split(_PLACEHOLDER) + return before, after + + +@functools.cache +def _binds_text(column: ServerColumnSchema) -> bool: + """Whether a scalar for `column` is bound as the text postgres.js sends, rather than as JSON.""" + return column.is_enum or is_string(column.type) + + def _cast(placeholder: str, column: ServerColumnSchema, comparison: bool) -> str: match column.type: case "timestamptz" | "timestamp with time zone": @@ -120,12 +138,12 @@ def __init__(self) -> None: def _bind(self, value: object, type_key: str, bound: str | None) -> int: key = (_same_value_key(value), type_key) - if key in self._seen: - return self._seen[key] - self._held.append(value) - self.values.append(bound) - self._seen[key] = len(self.values) - return len(self.values) + index = self._seen.get(key) + if index is None: + self._held.append(value) + self.values.append(bound) + index = self._seen[key] = len(self.values) + return index def column(self, column: ServerColumnSchema, value: Any, *, plural: bool = False, comparison: bool) -> str: """A value compared with or written to `column`, cast to the column's type.""" @@ -134,15 +152,17 @@ def column(self, column: ServerColumnSchema, value: Any, *, plural: bool = False bound = None elif plural: bound = js_stringify(value) - elif column.is_enum or is_string(column.type): + elif _binds_text(column): bound = value if isinstance(value, str) else _driver_text(value) else: bound = js_stringify(value) index = self._bind(value, column.type, bound) if value is None and plural: return f"${index}" - common = _cast("value" if plural else f"${index}", column, comparison) - return _plural(index, common) if plural else common + before, after = _cast_around(column, comparison) + if plural: + return _plural(index, f"{before}value{after}") + return f"{before}${index}{after}" def fail(self, error: Exception) -> str: """Record an error upstream raises while formatting, and stand in for the text it never wrote.""" From e0e9cb994f32e842fc18e12b16f68d7b8a4c383e Mon Sep 17 00:00:00 2001 From: Jonas Daniels Date: Thu, 15 Oct 2026 11:36:51 -0700 Subject: [PATCH 3/9] perf: write the mutate response without Pydantic's serializer A push answers with one MutationResponse per mutation, and each went through the wrap serializer three times. MutateSuccess, MutationResponse, MutationID and MutationOk now build their JSON directly, and fall back to model_dump for a float ID or a result with data. A unit test checks both paths against model_dump. Successful mutations share one MutationOk. Co-Authored-By: Claude Opus 5.5 --- src/zero_server/protocol/mutate_server.py | 8 ++++++++ src/zero_server/protocol/mutation.py | 6 ++++++ src/zero_server/protocol/mutation_id.py | 7 +++++++ src/zero_server/server/process_mutations.py | 3 ++- tests/unit/test_protocol.py | 17 +++++++++++++++++ 5 files changed, 40 insertions(+), 1 deletion(-) diff --git a/src/zero_server/protocol/mutate_server.py b/src/zero_server/protocol/mutate_server.py index e1a649a..9848a0c 100644 --- a/src/zero_server/protocol/mutate_server.py +++ b/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 diff --git a/src/zero_server/protocol/mutation.py b/src/zero_server/protocol/mutation.py index b3a66c2..e690ed9 100644 --- a/src/zero_server/protocol/mutation.py +++ b/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()} diff --git a/src/zero_server/protocol/mutation_id.py b/src/zero_server/protocol/mutation_id.py index 24bc9d8..f82a520 100644 --- a/src/zero_server/protocol/mutation_id.py +++ b/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() diff --git a/src/zero_server/server/process_mutations.py b/src/zero_server/server/process_mutations.py index c6ea7f0..9c35695 100644 --- a/src/zero_server/server/process_mutations.py +++ b/src/zero_server/server/process_mutations.py @@ -43,6 +43,7 @@ logger = logging.getLogger("zero_server.mutate") _cleanup_results_arg: TypeAdapter[CleanupResultsArg] = TypeAdapter(CleanupResultsArg) +_OK = MutationOk() # frozen, so every successful mutation can share it class TransactionProviderHooks(Protocol): @@ -363,7 +364,7 @@ async def run(tx: T, hooks: TransactionProviderHooks) -> MutationResponse: await callback(tx, mutation.name, mutation.args[0] if mutation.args else None) else: await hooks.write_mutation_result(_app_error_response(mutation, app_error)) - return MutationResponse(id=MutationID(clientID=mutation.client_id, id=mutation.id), result=MutationOk()) + return MutationResponse(id=MutationID(clientID=mutation.client_id, id=mutation.id), result=_OK) try: return await self._db.transaction( diff --git a/tests/unit/test_protocol.py b/tests/unit/test_protocol.py index 71936c7..838fcc4 100644 --- a/tests/unit/test_protocol.py +++ b/tests/unit/test_protocol.py @@ -1,3 +1,5 @@ +from typing import Any + import pytest from zero_server.protocol.application_error import ( @@ -16,6 +18,9 @@ map_condition, normalize_ast, ) +from zero_server.protocol.mutate_server import MutateSuccess +from zero_server.protocol.mutation import MutationAppError, MutationOk, MutationResponse +from zero_server.protocol.mutation_id import MutationID from zero_server.protocol.query_server import QuerySuccess from zero_server.shared.asserts import AssertionFailedError @@ -42,6 +47,18 @@ def test_user_id_is_reported_only_when_set() -> None: assert QuerySuccess(queries=[], user_id=None).to_json() == {"kind": "QueryResponse", "queries": [], "userID": None} +@pytest.mark.parametrize("user", [{}, {"userID": None}, {"userID": "u1"}]) +@pytest.mark.parametrize("mutation_id", [1, 2.5]) +@pytest.mark.parametrize( + "result", [MutationOk(), MutationOk(data={"a": [1]}), MutationAppError(message="no"), MutationAppError()] +) +def test_mutate_response_serializes_as_pydantic_does(user: dict[str, Any], mutation_id: float, result: Any) -> None: + response = MutationResponse(id=MutationID(id=mutation_id, clientID="c"), result=result) + success = MutateSuccess(mutations=[response], **user) + assert success.to_json() == success.model_dump(mode="json") + assert list(success.to_json()) == list(success.model_dump(mode="json")) + + def _eq(column: str, value: int) -> SimpleCondition: return SimpleCondition(op="=", left=ColumnReference(name=column), right=LiteralReference(value=value)) From 129fb1572b193da79b367f426edaf228313018a2 Mon Sep 17 00:00:00 2001 From: Jonas Daniels Date: Thu, 15 Oct 2026 11:36:51 -0700 Subject: [PATCH 4/9] perf: compute each wire model's optional keys once The before-validator walked every field on every parse to find the ones that may be absent but not null. It now reads them from a per-class cache. Co-Authored-By: Claude Opus 5.5 --- src/zero_server/protocol/wire.py | 23 +++++++++++++++++------ 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/src/zero_server/protocol/wire.py b/src/zero_server/protocol/wire.py index a562189..154839d 100644 --- a/src/zero_server/protocol/wire.py +++ b/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.""" From d1745fc4072c7cb9d92df16da7acceb0a2401567 Mon Sep 17 00:00:00 2001 From: Jonas Daniels Date: Thu, 15 Oct 2026 12:54:04 -0700 Subject: [PATCH 5/9] perf: send SQLAlchemy-on-asyncpg statements straight to asyncpg SQLAlchemy defers BEGIN to the first statement it runs itself, so that one still goes through exec_driver_sql. Once asyncpg reports the transaction open, statements go to the driver connection through AsyncpgTransaction, skipping SQLAlchemy's greenlet hop, prepared-statement cache and status parsing. A test checks a refused argument reads the same on either path. Co-Authored-By: Claude Opus 5.5 --- src/zero_server/server/adapters/sqlalchemy.py | 40 ++++++++++++------- tests/wire/test_mutate.py | 6 ++- 2 files changed, 31 insertions(+), 15 deletions(-) diff --git a/src/zero_server/server/adapters/sqlalchemy.py b/src/zero_server/server/adapters/sqlalchemy.py index 5bdfd50..dfd8cfa 100644 --- a/src/zero_server/server/adapters/sqlalchemy.py +++ b/src/zero_server/server/adapters/sqlalchemy.py @@ -17,21 +17,29 @@ 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. """ - 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 +54,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 - - 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") + from zero_server.server.adapters.asyncpg import AsyncpgTransaction + + 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: diff --git a/tests/wire/test_mutate.py b/tests/wire/test_mutate.py index 99037bf..37a4d37 100644 --- a/tests/wire/test_mutate.py +++ b/tests/wire/test_mutate.py @@ -323,10 +323,14 @@ async def callback(tx: Any, _hooks: Any) -> Any: wire_db.run(check()) -def test_asyncpg_reports_a_refused_argument_as_sqlalchemy_does(wire_db: WireDB) -> None: +@pytest.mark.parametrize("first", [True, False]) +def test_asyncpg_reports_a_refused_argument_as_sqlalchemy_does(wire_db: WireDB, first: bool) -> None: # A lone surrogate cannot be encoded, so asyncpg refuses it before Postgres sees the statement. + # SQLAlchemy runs a transaction's first statement itself and hands the rest to asyncpg. async def message(adapter: str) -> str: async def callback(tx: Any, _hooks: Any) -> Any: + if not first: + await tx.db_transaction.query("SELECT 1") await tx.db_transaction.query("SELECT $1::text", ["a\ud800"]) with pytest.raises(PostgresError) as error: From 5625720b22116ea4f76078c8b8cd1e0bd0604180 Mon Sep 17 00:00:00 2001 From: Jonas Daniels Date: Thu, 15 Oct 2026 12:25:32 -0700 Subject: [PATCH 6/9] perf: reuse each mutation's ID and drop two frames from the mutate path The push already builds a MutationID per mutation for its failure response; the success and error responses now reuse it instead of validating a new one. The application-error wrapper and the lastMutationID check run inline, so every resumption of a suspended statement passes through two fewer coroutine frames. CPU per put mutation on native Postgres, above the adapter's own transaction: 17 to 11-14 us. Co-Authored-By: Claude Opus 5.5 --- src/zero_server/server/process_mutations.py | 74 ++++++++++----------- 1 file changed, 37 insertions(+), 37 deletions(-) diff --git a/src/zero_server/server/process_mutations.py b/src/zero_server/server/process_mutations.py index 9c35695..ef4b09e 100644 --- a/src/zero_server/server/process_mutations.py +++ b/src/zero_server/server/process_mutations.py @@ -131,15 +131,6 @@ def __init__(self, phase: _DatabaseTransactionPhase, cause: BaseException) -> No _PASSES_THROUGH = (DatabaseTransactionError, OutOfOrderMutation, MutationAlreadyProcessedError, ApplicationError) -async def _application_error_wrapper[R](fn: Callable[[], Awaitable[R]]) -> R: - try: - return await fn() - except _PASSES_THROUGH: - raise - except Exception as error: - raise wrap_with_application_error(error) from error - - def _refuse_constant(name: str) -> None: raise ValueError(f"{name} is not valid JSON") @@ -230,7 +221,7 @@ async def handle_mutate[T]( try: transactor = _Transactor(db_provider, push_body, params) - for m in push_body.mutations: + for m, mutation_id in zip(push_body.mutations, mutation_ids, strict=True): if m.type == "custom" and m.name == CLEANUP_RESULTS_MUTATION_NAME: try: await _process_cleanup_results_mutation(db_provider, m, params) @@ -242,7 +233,7 @@ async def handle_mutate[T]( check(m.type == "custom", "Expected custom mutation") assert isinstance(m, CustomMutation) - response = await _process_one(transactor, handler, m) + response = await _process_one(transactor, handler, m, mutation_id) responses.append(response) processed_count += 1 @@ -269,28 +260,34 @@ async def handle_mutate[T]( async def _process_one[T]( - transactor: "_Transactor[T]", handler: MutateRequestHandler[T], m: CustomMutation + transactor: "_Transactor[T]", handler: MutateRequestHandler[T], m: CustomMutation, mutation_id: MutationID ) -> MutationResponse: phase: _MutationPhase = "preTransaction" async def transact(callback: TransactFnCallback[T]) -> MutationResponse: nonlocal phase phase = "transactionPending" - result = await transactor.transact(m, callback) + result = await transactor.transact(m, mutation_id, callback) phase = "postCommit" return result try: - return await _application_error_wrapper(lambda: handler(transact, m)) + # Upstream's `applicationErrorWrapper`, inline. + try: + return await handler(transact, m) + except _PASSES_THROUGH: + raise + except Exception as error: + raise wrap_with_application_error(error) from error except ApplicationError as error: if phase == "preTransaction": - await transactor.persist_pre_transaction_failure(m, error) + await transactor.persist_pre_transaction_failure(m, mutation_id, error) elif phase == "postCommit": # pragma: no branch - `transact` never raises an ApplicationError logger.error( "Post-commit mutation handler failed for mutation %s for client %s", m.id, m.client_id, exc_info=error ) logger.warning("Application error processing mutation %s for client %s", m.id, m.client_id, exc_info=error) - return _app_error_response(m, error) + return _app_error_response(mutation_id, error) class _Transactor[T]: @@ -299,11 +296,13 @@ def __init__(self, db_provider: Database[T], push: PushBody, params: MutateParam self._push = push self._params = params - async def transact(self, mutation: CustomMutation, callback: TransactFnCallback[T]) -> MutationResponse: + async def transact( + self, mutation: CustomMutation, mutation_id: MutationID, callback: TransactFnCallback[T] + ) -> MutationResponse: app_error: ApplicationError | None = None while True: try: - result = await self._transact_impl(mutation, callback, app_error) + result = await self._transact_impl(mutation, mutation_id, callback, app_error) if app_error is not None: logger.warning( "Mutation %s for client %s was retried after an error", @@ -311,7 +310,7 @@ async def transact(self, mutation: CustomMutation, callback: TransactFnCallback[ mutation.client_id, exc_info=app_error, ) - return _app_error_response(mutation, app_error) + return _app_error_response(mutation_id, app_error) return result except OutOfOrderMutation as error: logger.error("%s", error) @@ -319,7 +318,7 @@ async def transact(self, mutation: CustomMutation, callback: TransactFnCallback[ except MutationAlreadyProcessedError as error: logger.warning("%s", error) return MutationResponse( - id=MutationID(clientID=mutation.client_id, id=mutation.id), + id=mutation_id, result=MutationZeroError(error="alreadyProcessed", details=str(error)), ) except Exception as error: @@ -343,13 +342,19 @@ async def transact(self, mutation: CustomMutation, callback: TransactFnCallback[ exc_info=app_error, ) - async def persist_pre_transaction_failure(self, mutation: CustomMutation, app_error: ApplicationError) -> None: + async def persist_pre_transaction_failure( + self, mutation: CustomMutation, mutation_id: MutationID, app_error: ApplicationError + ) -> None: """The handler raised before calling `transact`. The ID still advances and the failure is recorded.""" - await self._transact_impl(mutation, None, app_error) + await self._transact_impl(mutation, mutation_id, None, app_error) async def _transact_impl( - self, mutation: CustomMutation, callback: TransactFnCallback[T] | None, app_error: ApplicationError | None + self, + mutation: CustomMutation, + mutation_id: MutationID, + callback: TransactFnCallback[T] | None, + app_error: ApplicationError | None, ) -> MutationResponse: """Bump the ID, then run `callback`, or with `app_error` set, record that error in its place.""" phase: _DatabaseTransactionPhase = "open" @@ -357,14 +362,18 @@ async def _transact_impl( async def run(tx: T, hooks: TransactionProviderHooks) -> MutationResponse: nonlocal phase phase = "execute" - await self._check_and_increment_last_mutation_id(hooks, mutation.client_id, mutation.id) + last_mutation_id = await hooks.update_client_mutation_id() + if mutation.id < last_mutation_id: + raise MutationAlreadyProcessedError(mutation.client_id, mutation.id, last_mutation_id) + if mutation.id > last_mutation_id: + raise OutOfOrderMutation(mutation.client_id, mutation.id, last_mutation_id) if app_error is None: assert callback is not None logger.debug("Executing mutator '%s' (id=%s)", mutation.name, mutation.id) await callback(tx, mutation.name, mutation.args[0] if mutation.args else None) else: - await hooks.write_mutation_result(_app_error_response(mutation, app_error)) - return MutationResponse(id=MutationID(clientID=mutation.client_id, id=mutation.id), result=_OK) + await hooks.write_mutation_result(_app_error_response(mutation_id, app_error)) + return MutationResponse(id=mutation_id, result=_OK) try: return await self._db.transaction( @@ -381,19 +390,10 @@ async def run(tx: T, hooks: TransactionProviderHooks) -> MutationResponse: except Exception as error: raise DatabaseTransactionError(phase, error) from error - async def _check_and_increment_last_mutation_id( - self, hooks: TransactionProviderHooks, client_id: str, received: float - ) -> None: - last_mutation_id = await hooks.update_client_mutation_id() - if received < last_mutation_id: - raise MutationAlreadyProcessedError(client_id, received, last_mutation_id) - if received > last_mutation_id: - raise OutOfOrderMutation(client_id, received, last_mutation_id) - -def _app_error_response(m: CustomMutation, error: ApplicationError) -> MutationResponse: +def _app_error_response(mutation_id: MutationID, error: ApplicationError) -> MutationResponse: return MutationResponse( - id=MutationID(clientID=m.client_id, id=m.id), + id=mutation_id, result=MutationAppError(message=error.message, details=error.details if js_truthy(error.details) else None), ) From d0499f3607ca8b614a28250058b9032a7a5ff4a4 Mon Sep 17 00:00:00 2001 From: Jonas Daniels Date: Thu, 15 Oct 2026 12:29:57 -0700 Subject: [PATCH 7/9] perf: trim the per-column work of a CRUD insert One cache lookup per bound value gives both how it binds and its cast. An insert reads Object.keys once, and the column and mutator checks format their message only when they fail. A 15-column insert's SQL builds in 15.5 us, from 18. Co-Authored-By: Claude Opus 5.5 --- src/zero_server/server/crud.py | 16 +++++++++------- src/zero_server/server/push_processor.py | 5 +++-- src/zero_server/z2s/sql.py | 17 ++++++----------- 3 files changed, 18 insertions(+), 20 deletions(-) diff --git a/src/zero_server/server/crud.py b/src/zero_server/server/crud.py index 68f9f9b..6094813 100644 --- a/src/zero_server/server/crud.py +++ b/src/zero_server/server/crud.py @@ -15,7 +15,7 @@ from zero_server.schema import Schema, TableSchema from zero_server.server.server_schema import ServerSchema, ServerTableSchema from zero_server.server.transaction import DBTransaction -from zero_server.shared.asserts import check +from zero_server.shared.asserts import AssertionFailedError from zero_server.shared.js import js_keys from zero_server.z2s.sql import Params, ident @@ -55,8 +55,8 @@ def _server_name(table: TableSchema) -> str: def _server_column(table: TableSchema, column: str) -> str: schema = table.columns.get(column) - check(schema, f"Column {column} was not found in the Zero schema for the table {table.name}") - assert schema is not None + if schema is None: # `check`, without formatting the message for every column that passes + raise AssertionFailedError(f"Column {column} was not found in the Zero schema for the table {table.name}") return schema.server_name if schema.server_name is not None else column @@ -95,11 +95,13 @@ async def _insert(self, value: dict[str, Any], *, upsert: bool) -> None: pk = [_server_column(self._table, k) for k in self._table.primary_key] params = Params() values = ", ".join( - params.column(self._server_table[server], value[k], comparison=False) - for k, server in zip(keys, targeted, strict=True) + [ + params.column(self._server_table[server], value[k], comparison=False) + for k, server in zip(keys, targeted, strict=True) + ] ) action = "DO NOTHING" - updated = self._non_pk(value) + updated = [k for k in keys if k not in self._table.primary_key] if upsert and updated: sets = ", ".join( f"{ident(server)} = {params.column(self._server_table[server], value[k], comparison=False)}" @@ -107,7 +109,7 @@ async def _insert(self, value: dict[str, Any], *, upsert: bool) -> None: ) action = f"DO UPDATE SET {sets}" sql = ( - f"INSERT INTO {ident(_server_name(self._table))} ({','.join(ident(c) for c in targeted)}) " + f"INSERT INTO {ident(_server_name(self._table))} ({','.join([ident(c) for c in targeted])}) " f"VALUES ({values}) ON CONFLICT ({', '.join(ident(c) for c in pk)}) {action}" ) await self._tx.query(sql, params.values) diff --git a/src/zero_server/server/push_processor.py b/src/zero_server/server/push_processor.py index 0ee3b0b..307698c 100644 --- a/src/zero_server/server/push_processor.py +++ b/src/zero_server/server/push_processor.py @@ -13,7 +13,7 @@ from zero_server.mutators import Mutator from zero_server.protocol.mutation import CustomMutation, MutationResponse from zero_server.server.process_mutations import Database, MutateResponse, TransactFn, handle_mutate -from zero_server.shared.asserts import check +from zero_server.shared.asserts import AssertionFailedError, check from zero_server.shared.json import JSONValue from zero_server.shared.registry import Registry, value_at_path @@ -42,7 +42,8 @@ async def dispatch(tx: T, name: str, args: JSONValue | None) -> None: async def _dispatch(self, mutators: Mapping[str, Any], tx: T, key: str, args: JSONValue | None) -> None: mutator = value_at_path(mutators, key, SEPARATOR) - check(callable(mutator) and not isinstance(mutator, Mapping), f"could not find mutator {key}") + if not callable(mutator) or isinstance(mutator, Mapping): # `check`, formatting only on failure + raise AssertionFailedError(f"could not find mutator {key}") if isinstance(mutator, Mutator): await mutator.fn(args=args, ctx=self._context, tx=tx) else: diff --git a/src/zero_server/z2s/sql.py b/src/zero_server/z2s/sql.py index 27cd443..4ec1ce7 100644 --- a/src/zero_server/z2s/sql.py +++ b/src/zero_server/z2s/sql.py @@ -44,16 +44,11 @@ def literal_type(value: object) -> LiteralType: @functools.cache -def _cast_around(column: ServerColumnSchema, comparison: bool) -> tuple[str, str]: - """The SQL before and after a value's placeholder that casts it to `column`'s type.""" +def _binding(column: ServerColumnSchema, comparison: bool) -> tuple[bool, str, str]: + """How a value for `column` is bound: whether a scalar goes as the text postgres.js sends rather than + as JSON, and the SQL before and after its placeholder that casts it to the column's type.""" before, after = _cast(_PLACEHOLDER, column, comparison).split(_PLACEHOLDER) - return before, after - - -@functools.cache -def _binds_text(column: ServerColumnSchema) -> bool: - """Whether a scalar for `column` is bound as the text postgres.js sends, rather than as JSON.""" - return column.is_enum or is_string(column.type) + return column.is_enum or is_string(column.type), before, after def _cast(placeholder: str, column: ServerColumnSchema, comparison: bool) -> str: @@ -148,18 +143,18 @@ def _bind(self, value: object, type_key: str, bound: str | None) -> int: def column(self, column: ServerColumnSchema, value: Any, *, plural: bool = False, comparison: bool) -> str: """A value compared with or written to `column`, cast to the column's type.""" plural = plural or column.is_array + as_text, before, after = _binding(column, comparison) if value is None: bound = None elif plural: bound = js_stringify(value) - elif _binds_text(column): + elif as_text: bound = value if isinstance(value, str) else _driver_text(value) else: bound = js_stringify(value) index = self._bind(value, column.type, bound) if value is None and plural: return f"${index}" - before, after = _cast_around(column, comparison) if plural: return _plural(index, f"{before}value{after}") return f"{before}${index}{after}" From 6c3cf24c88ae88e952d72516c9a6553d3c186b5e Mon Sep 17 00:00:00 2001 From: Jonas Daniels Date: Thu, 15 Oct 2026 12:55:39 -0700 Subject: [PATCH 8/9] perf: run a push's mutations on one pooled connection handle_mutate_request (and so PushProcessor) now runs its mutation loop inside a session of the database, when the database opens one. ZQLDatabase does, through its DBConnection, and each adapter's session takes a connection from the pool at the first transaction and returns it once, when the push ends. asyncpg's pool resets it then: one reset per push instead of one per mutation, and no shielded release task per mutation. A transaction that fails returns its connection, so the next one starts on a fresh one, as it would without a session. Direct users of a database's transaction() are unchanged: each still takes and returns its own connection. A Database or DBConnection without session() runs as before. Co-Authored-By: Claude Opus 5.5 --- src/zero_server/server/adapters/_session.py | 68 ++++++++++++++++ src/zero_server/server/adapters/asyncpg.py | 17 +++- src/zero_server/server/adapters/psycopg.py | 19 ++++- src/zero_server/server/adapters/sqlalchemy.py | 17 ++++ src/zero_server/server/process_mutations.py | 56 ++++++++----- src/zero_server/server/zql_database.py | 52 +++++++++++- tests/wire/mutators.py | 12 ++- tests/wire/test_adapters.py | 79 +++++++++++++++++++ 8 files changed, 292 insertions(+), 28 deletions(-) create mode 100644 src/zero_server/server/adapters/_session.py create mode 100644 tests/wire/test_adapters.py diff --git a/src/zero_server/server/adapters/_session.py b/src/zero_server/server/adapters/_session.py new file mode 100644 index 0000000..8c127e7 --- /dev/null +++ b/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() diff --git a/src/zero_server/server/adapters/asyncpg.py b/src/zero_server/server/adapters/asyncpg.py index d93dfa4..d41c4a1 100644 --- a/src/zero_server/server/adapters/asyncpg.py +++ b/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]): diff --git a/src/zero_server/server/adapters/psycopg.py b/src/zero_server/server/adapters/psycopg.py index aaee444..01a3b99 100644 --- a/src/zero_server/server/adapters/psycopg.py +++ b/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]]): diff --git a/src/zero_server/server/adapters/sqlalchemy.py b/src/zero_server/server/adapters/sqlalchemy.py index dfd8cfa..1a45ccc 100644 --- a/src/zero_server/server/adapters/sqlalchemy.py +++ b/src/zero_server/server/adapters/sqlalchemy.py @@ -4,12 +4,14 @@ """ 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 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 @@ -78,6 +80,21 @@ 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(): + return await fn(await sqlalchemy_transaction(connection)) + class SQLAlchemyDatabase(ZQLDatabase[AsyncConnection]): """A `ZQLDatabase` on `engine`, whose URL names the asyncpg or the psycopg driver diff --git a/src/zero_server/server/process_mutations.py b/src/zero_server/server/process_mutations.py index ef4b09e..5d2900a 100644 --- a/src/zero_server/server/process_mutations.py +++ b/src/zero_server/server/process_mutations.py @@ -14,6 +14,7 @@ import json import logging from collections.abc import Awaitable, Callable, Mapping +from contextlib import AbstractAsyncContextManager, nullcontext from dataclasses import dataclass from typing import Literal, Protocol @@ -69,6 +70,9 @@ class Database[T](Protocol): `transaction` opens one transaction, calls `callback` inside it with the app's transaction object and the bookkeeping hooks, and commits when `callback` returns. + + A database may also have a `session()`: an async context manager yielding a `Database` whose + transactions all run on one connection. A push's mutations run in one session, `ZQLDatabase`'s. """ async def transaction[R]( @@ -78,6 +82,19 @@ async def transaction[R]( ) -> R: ... +class _SessionOpener[T](Protocol): + def session(self) -> AbstractAsyncContextManager[Database[T]]: ... + + +def _session[T](db_provider: Database[T]) -> AbstractAsyncContextManager[Database[T]]: + """The database a push's mutations run on: a session of `db_provider` if it opens them, so the push + takes one connection from the pool for all its transactions, or else `db_provider` itself.""" + if hasattr(db_provider, "session"): + opener: _SessionOpener[T] = db_provider # ty: ignore[invalid-assignment] + return opener.session() + return nullcontext(db_provider) + + type TransactFnCallback[T] = Callable[[T, str, JSONValue | None], Awaitable[None]] type TransactFn[T] = Callable[[TransactFnCallback[T]], Awaitable[MutationResponse]] type MutateRequestHandler[T] = Callable[[TransactFn[T], CustomMutation], Awaitable[MutationResponse]] @@ -219,27 +236,28 @@ async def handle_mutate[T]( processed_count = 0 try: - transactor = _Transactor(db_provider, push_body, params) - - for m, mutation_id in zip(push_body.mutations, mutation_ids, strict=True): - if m.type == "custom" and m.name == CLEANUP_RESULTS_MUTATION_NAME: - try: - await _process_cleanup_results_mutation(db_provider, m, params) - except Exception as error: - # A failed cleanup costs disk, not correctness, so it does not fail the push. - logger.warning("Failed to process cleanup mutation for client %s", m.client_id, exc_info=error) + async with _session(db_provider) as db: + transactor = _Transactor(db, push_body, params) + + for m, mutation_id in zip(push_body.mutations, mutation_ids, strict=True): + if m.type == "custom" and m.name == CLEANUP_RESULTS_MUTATION_NAME: + try: + await _process_cleanup_results_mutation(db, m, params) + except Exception as error: + # A failed cleanup costs disk, not correctness, so it does not fail the push. + logger.warning("Failed to process cleanup mutation for client %s", m.client_id, exc_info=error) + processed_count += 1 + continue + + check(m.type == "custom", "Expected custom mutation") + assert isinstance(m, CustomMutation) + response = await _process_one(transactor, handler, m, mutation_id) + responses.append(response) processed_count += 1 - continue - - check(m.type == "custom", "Expected custom mutation") - assert isinstance(m, CustomMutation) - response = await _process_one(transactor, handler, m, mutation_id) - responses.append(response) - processed_count += 1 - if report_user_id: - return MutateSuccess(mutations=responses, userID=user_id) - return MutateSuccess(mutations=responses) + if report_user_id: + return MutateSuccess(mutations=responses, userID=user_id) + return MutateSuccess(mutations=responses) except Exception as error: logger.error("Failed to process push request", exc_info=error) details = get_error_details(error) diff --git a/src/zero_server/server/zql_database.py b/src/zero_server/server/zql_database.py index 7dee893..4b74bcd 100644 --- a/src/zero_server/server/zql_database.py +++ b/src/zero_server/server/zql_database.py @@ -8,7 +8,8 @@ supported driver. """ -from collections.abc import Awaitable, Callable +from collections.abc import AsyncGenerator, Awaitable, Callable +from contextlib import AbstractAsyncContextManager, asynccontextmanager from typing import Any, Protocol, overload from zero_server.protocol.mutation import CleanupResultsArg, CleanupResultsBulkArg, MutationResponse @@ -34,11 +35,19 @@ def __init__(self, message: str) -> None: class DBConnection[W](Protocol): - """Opens one transaction, runs `fn` in it, and commits when `fn` returns or rolls back when it raises.""" + """Opens one transaction, runs `fn` in it, and commits when `fn` returns or rolls back when it raises. + + A connection may also have a `session()`: an async context manager yielding a `DBConnection` whose + transactions all run on one connection from the pool, in turn, as a push's mutations do. + """ async def transaction[R](self, fn: Callable[[DBTransaction[W]], Awaitable[R]]) -> R: ... +class _SessionOpener[W](Protocol): + def session(self) -> AbstractAsyncContextManager[DBConnection[W]]: ... + + class _Hooks: def __init__(self, tx: DBTransaction[Any], transaction_input: TransactionProviderInput) -> None: self._tx = tx @@ -103,6 +112,28 @@ async def transaction[R]( self, callback: Callable[[ServerTransaction[W], TransactionProviderHooks], Awaitable[R]], transaction_input: TransactionProviderInput | None = None, + ) -> R: + return await self._transaction(self.connection, callback, transaction_input) + + @asynccontextmanager + async def session(self) -> "AsyncGenerator[_Session[W] | ZQLDatabase[W]]": + """This database, with every transaction on one connection until the block ends. + + `handle_mutate_request` runs a push's mutations in one, so the push takes a connection from the + pool once rather than once per mutation. A connection without sessions runs as before. + """ + if not hasattr(self.connection, "session"): + yield self + return + opener: _SessionOpener[W] = self.connection # ty: ignore[invalid-assignment] + async with opener.session() as connection: + yield _Session(self, connection) + + async def _transaction[R]( + self, + connection: DBConnection[W], + callback: Callable[[ServerTransaction[W], TransactionProviderHooks], Awaitable[R]], + transaction_input: TransactionProviderInput | None, ) -> R: transaction_input = transaction_input or TransactionProviderInput( upstream_schema="", client_group_id="", client_id="", mutation_id=0 @@ -116,7 +147,7 @@ async def run(tx: DBTransaction[W]) -> R: ) return await callback(server_tx, _Hooks(tx, transaction_input)) - return await self.connection.transaction(run) + return await connection.transaction(run) @overload async def run[R](self, query: OneQuery[R, Any]) -> R | None: ... @@ -129,3 +160,18 @@ async def run(tx: ServerTransaction[W], _hooks: TransactionProviderHooks) -> Any return await tx.run(query) return await self.transaction(run) + + +class _Session[W]: + """A `ZQLDatabase`'s transactions, on the one connection of a session of its `DBConnection`.""" + + def __init__(self, database: ZQLDatabase[W], connection: DBConnection[W]) -> None: + self._database = database + self._connection = connection + + async def transaction[R]( + self, + callback: Callable[[ServerTransaction[W], TransactionProviderHooks], Awaitable[R]], + transaction_input: TransactionProviderInput | None = None, + ) -> R: + return await self._database._transaction(self._connection, callback, transaction_input) diff --git a/tests/wire/mutators.py b/tests/wire/mutators.py index 96e1c85..166eb5a 100644 --- a/tests/wire/mutators.py +++ b/tests/wire/mutators.py @@ -1,6 +1,7 @@ """The Python twin of `harness/mutate.ts`: the same mutators, handler and failing provider.""" -from collections.abc import Awaitable, Callable +from collections.abc import AsyncGenerator, Awaitable, Callable +from contextlib import asynccontextmanager from typing import Any from sqlalchemy.ext.asyncio import AsyncConnection @@ -120,7 +121,7 @@ async def run(tx: Tx, name: str, mutator_args: Any) -> None: class FailingForCfail: """Fails to open any transaction for client `cfail`, as the TypeScript twin does.""" - def __init__(self, db: ZQLDatabase[Any]) -> None: + def __init__(self, db: Database[Tx]) -> None: self.db = db async def transaction( @@ -132,6 +133,13 @@ async def transaction( raise Exception("cannot open") return await self.db.transaction(callback, transaction_input) + @asynccontextmanager + async def session(self) -> AsyncGenerator["FailingForCfail"]: + # The Python database runs a push in a session; this forwards it, failing as before. + assert isinstance(self.db, ZQLDatabase) + async with self.db.session() as db: + yield FailingForCfail(db) + def provider(db: ZQLDatabase[Any]) -> Database[Tx]: return FailingForCfail(db) diff --git a/tests/wire/test_adapters.py b/tests/wire/test_adapters.py new file mode 100644 index 0000000..7da3613 --- /dev/null +++ b/tests/wire/test_adapters.py @@ -0,0 +1,79 @@ +"""Sessions: a push's transactions on one connection from the adapter's pool.""" + +import logging +from typing import Any + +import asyncpg +import pytest + +from tests.wire.conftest import ADAPTERS, WireDB +from zero_server.server.adapters.asyncpg import AsyncpgDatabase +from zero_server.server.zql_database import ZQLDatabase + + +async def _connection(tx: Any, _hooks: Any) -> Any: + return tx.db_transaction.wrapped_transaction + + +async def _fail(_tx: Any, _hooks: Any) -> Any: + raise RuntimeError("no") + + +def _asyncpg_url(wire_db: WireDB) -> str: + return wire_db.py_engine.url.set(drivername="postgresql").render_as_string(hide_password=False) + + +@pytest.mark.parametrize("adapter", ADAPTERS) +def test_a_session_runs_its_transactions_on_one_connection(wire_db: WireDB, adapter: str) -> None: + async def check() -> None: + database = wire_db.py[adapter](None) + async with database.session() as db: + first = await db.transaction(_connection) + assert await db.transaction(_connection) is first + + wire_db.run(check()) + + +def test_a_session_takes_and_returns_its_connection_once(wire_db: WireDB) -> None: + async def check() -> None: + pool = await asyncpg.create_pool(_asyncpg_url(wire_db), min_size=2, max_size=2) + async with AsyncpgDatabase(pool).session() as db: + assert pool.get_idle_size() == 2 # nothing taken before the first transaction + first = await db.transaction(_connection) + assert pool.get_idle_size() == 1 + with pytest.raises(RuntimeError): + await db.transaction(_fail) + # A failed transaction hands its connection back; the next one takes a connection afresh. + assert pool.get_idle_size() == 2 + assert await db.transaction(_connection) is not first + assert pool.get_idle_size() == 2 + await pool.close() + + wire_db.run(check()) + + +def test_a_session_logs_a_connection_the_pool_could_not_take_back( + wire_db: WireDB, caplog: pytest.LogCaptureFixture +) -> None: + async def refuse(_connection: Any) -> None: + raise RuntimeError("cannot reset") + + async def check() -> None: + pool = await asyncpg.create_pool(_asyncpg_url(wire_db), min_size=1, max_size=1, reset=refuse) + async with AsyncpgDatabase(pool).session() as db: + await db.transaction(_connection) + await pool.close() + + with caplog.at_level(logging.WARNING, logger="zero_server.adapters"): + wire_db.run(check()) + assert "Failed to return a connection to the pool" in caplog.text + + +async def test_a_connection_without_sessions_runs_each_transaction_on_its_own() -> None: + class Connection: + async def transaction(self, fn: Any) -> Any: + return await fn(None) + + database = ZQLDatabase(Connection()) + async with database.session() as db: + assert db is database From 8a3eeff3a8720f202d14d1f2a2c424e3eeeed4bb Mon Sep 17 00:00:00 2001 From: Jonas Daniels Date: Thu, 15 Oct 2026 13:08:41 -0700 Subject: [PATCH 9/9] perf: send SQLAlchemy's BEGIN up front in a push session SQLAlchemy's asyncpg dialect defers BEGIN to the first statement it runs itself, so each mutation's first statement took SQLAlchemy's execution path. In a push session the adapter now calls the dialect's private _start_transaction right after begin(), and every statement goes straight to asyncpg: a push of 100 puts went from 1.35x to 1.13x of TypeScript. If a release drops that method, the first statement begins the transaction as before; a test covers that fallback. The README and handle_mutate_request now say what sharing one connection per push means: mutators in a push share session state, and a handler that takes a second connection from the same pool can deadlock when the pool is smaller than the number of concurrent pushes. Co-Authored-By: Claude Opus 5.5 --- README.md | 2 +- src/zero_server/server/adapters/sqlalchemy.py | 18 +++++++++++++++++- src/zero_server/server/process_mutations.py | 4 ++++ tests/wire/test_adapters.py | 7 +++++++ 4 files changed, 29 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 6716926..98994bf 100644 --- a/README.md +++ b/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. diff --git a/src/zero_server/server/adapters/sqlalchemy.py b/src/zero_server/server/adapters/sqlalchemy.py index 1a45ccc..fc71982 100644 --- a/src/zero_server/server/adapters/sqlalchemy.py +++ b/src/zero_server/server/adapters/sqlalchemy.py @@ -9,6 +9,7 @@ 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 @@ -23,7 +24,8 @@ class SQLAlchemyTransaction: 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. + 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__( @@ -93,9 +95,23 @@ 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 (`postgresql+asyncpg://` or `postgresql+psycopg://`).""" diff --git a/src/zero_server/server/process_mutations.py b/src/zero_server/server/process_mutations.py index 5d2900a..a9146c4 100644 --- a/src/zero_server/server/process_mutations.py +++ b/src/zero_server/server/process_mutations.py @@ -189,6 +189,10 @@ async def handle_mutate_request[T]( A caller that needs push-level fields such as `timestamp` parses once with `PushBody.from_json` and passes the result. If that parse fails, pass the raw body instead and return the parse failure this produces. + + With a database that opens sessions, as `ZQLDatabase` does, the push's mutations share one pooled + connection: they see each other's session state, and a handler that takes another connection from + the same pool can deadlock when the pool is smaller than the number of concurrent pushes. """ return await handle_mutate(db_provider, handler, query, body, user_id, report_user_id=True) diff --git a/tests/wire/test_adapters.py b/tests/wire/test_adapters.py index 7da3613..bc61921 100644 --- a/tests/wire/test_adapters.py +++ b/tests/wire/test_adapters.py @@ -1,6 +1,7 @@ """Sessions: a push's transactions on one connection from the adapter's pool.""" import logging +from types import SimpleNamespace from typing import Any import asyncpg @@ -8,6 +9,7 @@ from tests.wire.conftest import ADAPTERS, WireDB from zero_server.server.adapters.asyncpg import AsyncpgDatabase +from zero_server.server.adapters.sqlalchemy import _send_begin from zero_server.server.zql_database import ZQLDatabase @@ -77,3 +79,8 @@ async def transaction(self, fn: Any) -> Any: database = ZQLDatabase(Connection()) async with database.session() as db: assert db is database + + +async def test_sqlalchemy_begins_on_its_first_statement_without_start_transaction() -> None: + # A SQLAlchemy without the dialect's private `_start_transaction` sends nothing early. + await _send_begin(SimpleNamespace(dbapi_connection=object())) # ty: ignore[invalid-argument-type]