From 048b0a27baf566548a0177bb702ee7aabebad8b8 Mon Sep 17 00:00:00 2001 From: Jonas Daniels Date: Wed, 28 Oct 2026 13:59:40 -0700 Subject: [PATCH 1/2] perf: intern built queries so rebuilding one hands back the same objects Co-Authored-By: Claude Opus 5.5 --- docs/compat.md | 2 +- src/zero_server/protocol/ast.py | 5 +- src/zero_server/query/builder.py | 234 +++++++++++++++++++++++++++---- tests/unit/test_ergonomics.py | 8 ++ tests/wire/test_query.py | 32 +++++ 5 files changed, 249 insertions(+), 32 deletions(-) diff --git a/docs/compat.md b/docs/compat.md index 38fc7d1..fce72c4 100644 --- a/docs/compat.md +++ b/docs/compat.md @@ -70,7 +70,7 @@ These are the cases where the proof holds Python to something narrower than "ide - **Interned queries.** Upstream's builder hands back the query it built before when a step repeats one, keyed so that some values JavaScript's types rule out share a key: a `flip` of `'yes'` with `true`, any `orderBy` direction but `'asc'` with `'desc'`. In one process, such a value can come back as whichever - was built first. Python builds what it is given. The harness gives each program a fresh schema, which + was built first. Python interns as well, but by exact keys, so it builds what it is given. The harness gives each program a fresh schema, which upstream interns by, so that no comparison depends on the programs before it. - **A mutator returning `null`.** Upstream records a mutator's return value as result data unless it is `undefined`, so a TypeScript mutator can return `null` as data. Python has only `None`, which the port diff --git a/src/zero_server/protocol/ast.py b/src/zero_server/protocol/ast.py index 53badee..a339b12 100644 --- a/src/zero_server/protocol/ast.py +++ b/src/zero_server/protocol/ast.py @@ -11,7 +11,7 @@ import functools from collections.abc import Callable, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, ClassVar, Literal, Protocol from zero_server.protocol.wire import dump_json, json_model, parse_wire @@ -192,6 +192,9 @@ class CorrelatedSubquery: hidden: bool | None = None subquery: AST system: System | None = None + # The builder query `subquery` is the AST of, kept alive with it: the builder interns a query by + # the identity of its subqueries' ASTs, so they must outlive the query keyed by them. + source: Any = field(default=None, repr=False, compare=False) def to_json(self) -> dict[str, Any]: out: dict[str, Any] = {"correlation": self.correlation.to_json()} diff --git a/src/zero_server/query/builder.py b/src/zero_server/query/builder.py index 579e8ec..dfe5e02 100644 --- a/src/zero_server/query/builder.py +++ b/src/zero_server/query/builder.py @@ -6,12 +6,18 @@ The errors match upstream's message for message, because a query endpoint reports them to the client. A query's AST is always normalized, as upstream's is: each step keeps it so, sorting a new `related` -into place and normalizing the conditions of a `where`. Upstream also interns queries, so that -rebuilding the same chain hands back the same objects; that only saves work, and is not ported. +into place and normalizing the conditions of a `where`. + +Queries are interned, as upstream's are, so that rebuilding the same chain hands back the same objects +rather than building them again. Each query holds the query it was derived from, and weakly the ones +derived from it, keyed by the step and its arguments. Python's keys are exact where upstream's are not: +upstream folds a `flip` of `'yes'` into `true`, and Python keeps them apart (see docs/compat.md). """ import math import re +import weakref +from collections import deque from collections.abc import Callable, Mapping from dataclasses import dataclass, field from typing import Any, Generic @@ -28,6 +34,7 @@ CorrelatedSubqueryCondition, Correlation, Direction, + LiteralReference, SimpleCondition, System, insert_related, @@ -47,6 +54,12 @@ class Format: relationships: dict[str, "Format"] = field(default_factory=dict) +# The queries built last, each holding the chain it came from. V8 keeps a query nothing holds until +# its next collection, long enough for upstream to find it when a request builds it again; CPython +# frees it at once. Holding the last few thousand gives a server that builds the same queries request +# after request the same reuse, in bounded memory. +_RECENT: deque[Any] = deque(maxlen=4096) + _MANY = Format() _ONE = Format(singular=True) _object_new = object.__new__ @@ -67,7 +80,18 @@ class UnknownColumnError(Exception): class Query(Generic[R, C]): """An immutable query on one table. Every method returns a new query.""" - __slots__ = ("_ast", "_format", "_junction", "_schema", "_strict", "_system", "_table") + __slots__ = ( + "__weakref__", + "_ast", + "_children", + "_format", + "_junction", + "_parent", + "_schema", + "_strict", + "_system", + "_table", + ) def __init__( self, @@ -87,6 +111,8 @@ def __init__( self._junction = current_junction self._format = format or _MANY self._strict = strict + self._parent: Query[Any, Any] | None = None + self._children: dict[Any, weakref.ref[Any]] | None = None def __repr__(self) -> str: # python-only: for debugging return f"{type(self).__name__}({self._ast.to_json()!r})" @@ -124,6 +150,35 @@ def _copy[Q: Query[Any, Any]](self, cls: type[Q], ast: AST, format: Format) -> Q query._junction = self._junction query._format = format query._strict = self._strict + query._parent = None + query._children = None + return query + + def _derive[T](self, key: Any, build: Callable[[], T]) -> T: + """The query `build` makes from this one, or the one it made before for the same `key`. + + `key` names the step and every argument it was given, exactly. + """ + found = self._find(key) + return self._keep(key, build()) if found is None else found + + def _find(self, key: Any) -> Any: + """The query derived from this one for `key` before, while it lives, or `None`.""" + children = self._children + if children is None: + return None + ref = children.get(key) + return None if ref is None else ref() + + def _keep[T](self, key: Any, query: T) -> T: + """Remember `query` as the one derived from this query for `key`.""" + children = self._children + if children is None: + children = self._children = {} + query._parent = self # ty: ignore[unresolved-attribute] + # Dropped once the query is: a parent keeps no query alive that nothing else holds. + children[key] = weakref.ref(query, lambda _, children=children, key=key: children.pop(key, None)) + _RECENT.append(query) return query def _with[Q: Query[Any, Any]](self: Q, field: str, value: Any, format: Format | None = None) -> Q: @@ -134,33 +189,55 @@ def _with[Q: Query[Any, Any]](self: Q, field: str, value: Any, format: Format | def _with_relationship[Q: Query[Any, Any]](self: Q, name: str, sub: AnyQuery, entry: CorrelatedSubquery) -> Q: format = Format(self._format.singular, {**self._format.relationships, name: sub._format}) - return self._with("related", insert_related(self._ast.related, entry), format) + query = self._with("related", insert_related(self._ast.related, entry), format) + return self._keep(("related", name, id(sub._ast)), query) def one(self) -> "OneQuery[R, C]": """At most one row: `tx.run` returns it, or `None`.""" - ast = self._ast._copy() - ast.limit = 1 - return self._copy(OneQuery, ast, Format(True, self._format.relationships)) + + def build() -> OneQuery[R, C]: + ast = self._ast._copy() + ast.limit = 1 + return self._copy(OneQuery, ast, Format(True, self._format.relationships)) + + return self._derive(("one",), build) def where[Q: Query[Any, Any]]( self: Q, field_or_factory: C | Callable[[ExpressionBuilder], Condition], *args: Any ) -> Q: """`where(field, value)`, `where(field, op, value)`, or `where(lambda eb: ...)`.""" cond: Condition + key: Any = None # Anything but a function is a field, as upstream has it: a number fails once the column is looked up. if isinstance(field_or_factory, str) or not callable(field_or_factory): if self._strict: # python-only: strict mode self._check_column(field_or_factory) if len(args) == 1: - cond = SimpleCondition(op="=", left=ColumnReference(name=field_or_factory), right=_position(args[0])) + op, value = "=", args[0] else: check(len(args) >= 2, "Invalid condition. Too few arguments.") - cond = cmp(field_or_factory, args[0], args[1]) + op, value = args[0], args[1] + # A column compared with a scalar, by far the usual condition, is looked up before it is built. + kind = type(value) + if kind in _SCALAR_TYPES and type(field_or_factory) is str and type(op) is str: + key = ("where", field_or_factory, op, kind, value) + found = self._find(key) + if found is not None: + return found + if len(args) == 1: + cond = SimpleCondition(op="=", left=ColumnReference(name=field_or_factory), right=_position(value)) + else: + cond = cmp(field_or_factory, op, value) else: cond = field_or_factory(self.expression_builder()) - if self._ast.where is not None: - cond = and_(self._ast.where, cond) - return self._with("where", normalize_condition(simplify_condition(cond))) + if key is None: + key = ("where", _condition_key(cond)) + found = self._find(key) + if found is not None: + return found + existing = self._ast.where + merged = cond if existing is None else and_(existing, cond) + return self._keep(key, self._with("where", normalize_condition(simplify_condition(merged)))) def where_exists[Q: Query[Any, Any]]( self: Q, @@ -188,11 +265,14 @@ def related[Q: Query[Any, Any]](self: Q, relationship: str, cb: SubqueryFn | Non related = self._relationship(relationship) if len(related) == 1: (hop,) = related - q = self._new( - hop.dest_schema, - AST(table=hop.dest_schema, alias=relationship), - None, - _ONE if hop.cardinality == "one" else _MANY, + q = self._derive( + ("related_base", relationship), + lambda: self._new( + hop.dest_schema, + AST(table=hop.dest_schema, alias=relationship), + None, + _ONE if hop.cardinality == "one" else _MANY, + ), ) sub = _as_query(cb(q) if cb else q) check(len(hop.source_field) >= 1, "The source of a relationship must specify at last 1 field") @@ -201,26 +281,37 @@ def related[Q: Query[Any, Any]](self: Q, relationship: str, cb: SubqueryFn | Non len(hop.source_field) == len(hop.dest_field), "The source and destination of a relationship must have the same number of fields", ) + found = self._find(("related", relationship, id(sub._ast))) + if found is not None: + return found entry = CorrelatedSubquery( system=self._system, correlation=Correlation(parent_field=hop.source_field, child_field=hop.dest_field), subquery=sub._ast, + source=sub, ) return self._with_relationship(relationship, sub, entry) first, second = related - q = self._new( - second.dest_schema, - AST(table=second.dest_schema, alias=relationship), - relationship, - _ONE if second.cardinality == "one" else _MANY, + q = self._derive( + ("related_base", relationship), + lambda: self._new( + second.dest_schema, + AST(table=second.dest_schema, alias=relationship), + relationship, + _ONE if second.cardinality == "one" else _MANY, + ), ) sub = _as_query(cb(q) if cb else q) _check_keys(first, second) + found = self._find(("related", relationship, id(sub._ast))) + if found is not None: + return found entry = CorrelatedSubquery( system=self._system, correlation=Correlation(parent_field=first.source_field, child_field=first.dest_field), hidden=True, + source=sub, subquery=AST( table=first.dest_schema, alias=relationship, @@ -229,6 +320,7 @@ def related[Q: Query[Any, Any]](self: Q, relationship: str, cb: SubqueryFn | Non system=self._system, correlation=Correlation(parent_field=second.source_field, child_field=second.dest_field), subquery=sub._ast, + source=sub, ), ), ), @@ -241,7 +333,9 @@ def start[Q: Query[Any, Any]](self: Q, row: Mapping[str, Any], *, inclusive: boo for column in row: self._check_column(column) row = dict(row) - return self._with("start", Bound(row=row, exclusive=not inclusive)) + return self._derive( + ("start", _tag(inclusive), _tag(row)), lambda: self._with("start", Bound(row=row, exclusive=not inclusive)) + ) def limit[Q: Query[Any, Any]](self: Q, limit: int | float) -> Q: if not (type(limit) is int and 0 <= limit < 2**31): # The usual limit passes both checks. @@ -255,7 +349,7 @@ def limit[Q: Query[Any, Any]](self: Q, limit: int | float) -> Q: "Limit is not supported in junction relationships yet. Junction relationship being limited: " + self._junction ) - return self._with("limit", limit) + return self._derive(("limit", _tag(limit)), lambda: self._with("limit", limit)) def order_by[Q: Query[Any, Any]](self: Q, field: C, direction: Direction) -> Q: if self._strict: # python-only: strict mode @@ -266,7 +360,10 @@ def order_by[Q: Query[Any, Any]](self: Q, field: C, direction: Direction) -> Q: + self._junction ) order_by = self._ast.order_by - return self._with("order_by", (*order_by, (field, direction)) if order_by else ((field, direction),)) + return self._derive( + ("order_by", _tag(field), _tag(direction)), + lambda: self._with("order_by", (*order_by, (field, direction)) if order_by else ((field, direction),)), + ) def _exists( self, @@ -281,13 +378,19 @@ def _exists( (hop,) = related check(len(hop.source_field) >= 1, "Invalid relationship") check(len(hop.dest_field) >= 1, "Invalid relationship") - q = self._new(hop.dest_schema, AST(table=hop.dest_schema, alias=f"{SUBQ_PREFIX}{relationship}"), None) + q = self._derive( + ("exists_base", relationship), + lambda: self._new( + hop.dest_schema, AST(table=hop.dest_schema, alias=f"{SUBQ_PREFIX}{relationship}"), None + ), + ) sub = _as_query(cb(q) if cb else q) return CorrelatedSubqueryCondition( related=CorrelatedSubquery( system=self._system, correlation=Correlation(parent_field=hop.source_field, child_field=hop.dest_field), subquery=sub._ast, + source=sub, ), op="EXISTS", flip=flip, @@ -297,12 +400,17 @@ def _exists( first, second = related _check_keys(first, second) alias = f"{SUBQ_PREFIX}zhidden_{relationship}" - q = self._new(second.dest_schema, AST(table=second.dest_schema, alias=alias), relationship) + q = self._derive( + ("exists_base", relationship), + lambda: self._new(second.dest_schema, AST(table=second.dest_schema, alias=alias), relationship), + ) to_dest = _as_query(cb(q) if cb else q) return CorrelatedSubqueryCondition( related=CorrelatedSubquery( system=self._system, correlation=Correlation(parent_field=first.source_field, child_field=first.dest_field), + # The junction AST is built here, new every time; the query it leads to stands for it. + source=to_dest, subquery=AST( table=first.dest_schema, alias=f"{SUBQ_PREFIX}{relationship}", @@ -312,6 +420,7 @@ def _exists( system=self._system, correlation=Correlation(parent_field=second.source_field, child_field=second.dest_field), subquery=to_dest._ast, + source=to_dest, ), op="EXISTS", flip=flip, @@ -331,6 +440,66 @@ class OneQuery(Query[R, C]): """A query `one()` narrowed to a single row. Its methods keep it singular.""" +_SCALAR_TYPES = (str, int, bool, type(None)) + + +def _tag(value: object) -> Any: + """`value` as an exact, hashable key. + + Each part carries its type, so `1`, `1.0`, `True` and `"1"` stay apart, and a float is spelled + exactly, so `0.0` and `-0.0` do too. A dict keeps its order, which its JSON has. + """ + kind = type(value) + if kind in _SCALAR_TYPES: + return kind, value + if kind is float: + return kind, value.hex() # ty: ignore[unresolved-attribute] + if kind is list or kind is tuple: + return kind, tuple([_tag(v) for v in value]) # ty: ignore[not-iterable] + if kind is dict: + return kind, tuple([(_tag(k), _tag(v)) for k, v in value.items()]) # ty: ignore[unresolved-attribute] + # The query keyed by it holds the value in its AST, so the identity is not reused while the key is. + return object, id(value) # python-only: a value of no JSON type, which only Python code passes + + +def _position_key(position: object) -> Any: + kind = type(position) + if kind is ColumnReference: + return kind, _tag(position.name) # ty: ignore[unresolved-attribute] + if kind is LiteralReference: + return kind, _tag(position.value) # ty: ignore[unresolved-attribute] + return kind, _tag(position.anchor), _tag(position.field) # ty: ignore[unresolved-attribute] + + +def _condition_key(condition: Condition) -> Any: + """`condition` as an exact key. A subquery is named by the identity of its AST, which the condition + keeps alive. For a two-hop `exists`, whose junction AST is new every time, that is the AST of the + query it leads to (`CorrelatedSubquery.source`).""" + kind = type(condition) + if kind is SimpleCondition: + return ( + kind, + _position_key(condition.left), # ty: ignore[unresolved-attribute] + _tag(condition.op), # ty: ignore[unresolved-attribute] + _position_key(condition.right), # ty: ignore[unresolved-attribute] + ) + if kind is CorrelatedSubqueryCondition: + related: CorrelatedSubquery = condition.related # ty: ignore[unresolved-attribute] + return ( + kind, + _tag(related.correlation.parent_field), + _tag(related.correlation.child_field), + _tag(related.hidden), + _tag(related.system), + # A condition the builder did not make has no source; its own subquery stands for itself. + id(getattr(related.source, "_ast", related.subquery)), + _tag(condition.op), # ty: ignore[unresolved-attribute] + _tag(condition.flip), # ty: ignore[unresolved-attribute] + _tag(condition.scalar), # ty: ignore[unresolved-attribute] + ) + return kind, tuple([_condition_key(c) for c in condition.conditions]) # ty: ignore[unresolved-attribute] + + def _check_keys(first: Connection, second: Connection) -> None: for key in (first.source_field, first.dest_field, second.source_field, second.dest_field): check(len(key) >= 1, "Invalid relationship") @@ -385,14 +554,19 @@ class Builder: def __init__(self, schema: Schema, *, strict: bool = False) -> None: self._schema = schema self._strict = strict + # Every query this builder makes descends from one of these, so they hold the interned queries. + self._roots: dict[str, Query[Any, Any]] = {} def __getattr__(self, table: str) -> Query[Any, Any]: return self.table(table) def table(self, name: str) -> Query[Any, Any]: - if name not in self._schema.tables: # python-only: upstream's builder has `undefined` for it instead - raise AttributeError(name) - return Query(self._schema, name, AST(table=name), strict=self._strict) + root = self._roots.get(name) + if root is None: + if name not in self._schema.tables: # python-only: upstream's builder has `undefined` for it instead + raise AttributeError(name) + root = self._roots[name] = Query(self._schema, name, AST(table=name), strict=self._strict) + return root def create_builder(schema: Schema, *, strict: bool = False) -> Builder: diff --git a/tests/unit/test_ergonomics.py b/tests/unit/test_ergonomics.py index 33c6335..e7f4051 100644 --- a/tests/unit/test_ergonomics.py +++ b/tests/unit/test_ergonomics.py @@ -227,3 +227,11 @@ def test_routes_only_for_what_is_given() -> None: assert [r.path for r in zero_routes(schema=load_schema(SCHEMA), authenticate=lambda _r: None)] == [] with pytest.raises(ValueError, match="needs a database"): zero_routes(schema=load_schema(SCHEMA), mutators=define_mutators({}), authenticate=lambda _r: None) + + +def test_builder_reuses_a_query_built_from_the_same_object() -> None: + """A value of no JSON type, which only Python code passes, is keyed by its identity.""" + zql = create_builder(load_schema(json.dumps(SCHEMA))) + value, other = object(), object() + assert zql.issue.where("meta", value) is zql.issue.where("meta", value) + assert zql.issue.where("meta", value) is not zql.issue.where("meta", other) diff --git a/tests/wire/test_query.py b/tests/wire/test_query.py index 8dad0c7..20209e4 100644 --- a/tests/wire/test_query.py +++ b/tests/wire/test_query.py @@ -203,6 +203,38 @@ def test_async_handlers_match_upstream(harness: Harness, schema: Schema) -> None assert _comparable(actual.to_json()) == _comparable(expected) +# Every kind of step, for the builder to hand back what it built before when a batch builds it again. +REBUILT = [ + _q( + "issue", + ["where", "open", True], + ["where", "priority", ">", 2.5], + ["whereFn", ["or", [eq, ["cmpLit", "a", "=", {"$param": ["authData", "sub"]}], exists]]], + ["whereExists", "labels", {"ops": [["where", "name", "!=", "zzz"]]}, {"flip": True}], + ["whereRaw", {"type": "and", "conditions": [_raw_eq]}], + ["related", "owner", None], + ["related", "labels", {"ops": [["where", "name", "bug"]]}], + ["related", "comments", {"ops": [["orderBy", "id", "desc"], ["limit", 5]]}], + ["orderBy", "priority", "desc"], + ["start", {"priority": 1.5, "id": "i1", "meta": [1, {"a": None}]}, True], + ["limit", 20], + ), + _q("issue", ["where", "open", True], ["where", "priority", ">", 2.5], ["one"]), + _q("issue", ["where", "meta", [1, "x"]], ["where", "title", "IN", ["a", "b"]]), +] + + +def test_rebuilt_queries_match_upstream(harness: Harness, schema: Schema) -> None: + """The same programs, three times over in one batch, built by one builder.""" + body = [ + "transform", + [{"id": str(n), "name": "prog", "args": [program]} for n, program in enumerate(REBUILT * 3)], + ] + expected = harness.call("query", body=body, userID="u") + actual = asyncio.run(handle_query_request(handler=make_handler(schema), schema=schema, body=body, user_id="u")) + assert _comparable(actual.to_json()) == _comparable(expected) + + def test_each_thrown_shape_matches_upstream(harness: Harness, schema: Schema) -> None: names = [ *(f"throw.{shape}" for shape in SHAPES), From 688312630321743f4dc449a7116d7fafe177460f Mon Sep 17 00:00:00 2001 From: Jonas Daniels Date: Wed, 28 Oct 2026 14:02:37 -0700 Subject: [PATCH 2/2] fix: key a two-hop exists by its junction condition, which carries scalar Co-Authored-By: Claude Opus 5.5 --- src/zero_server/query/builder.py | 15 +++++++++++---- tests/wire/test_query.py | 3 +++ 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/src/zero_server/query/builder.py b/src/zero_server/query/builder.py index dfe5e02..89e610d 100644 --- a/src/zero_server/query/builder.py +++ b/src/zero_server/query/builder.py @@ -473,8 +473,7 @@ def _position_key(position: object) -> Any: def _condition_key(condition: Condition) -> Any: """`condition` as an exact key. A subquery is named by the identity of its AST, which the condition - keeps alive. For a two-hop `exists`, whose junction AST is new every time, that is the AST of the - query it leads to (`CorrelatedSubquery.source`).""" + keeps alive (`CorrelatedSubquery.source`).""" kind = type(condition) if kind is SimpleCondition: return ( @@ -491,8 +490,7 @@ def _condition_key(condition: Condition) -> Any: _tag(related.correlation.child_field), _tag(related.hidden), _tag(related.system), - # A condition the builder did not make has no source; its own subquery stands for itself. - id(getattr(related.source, "_ast", related.subquery)), + _subquery_key(related), _tag(condition.op), # ty: ignore[unresolved-attribute] _tag(condition.flip), # ty: ignore[unresolved-attribute] _tag(condition.scalar), # ty: ignore[unresolved-attribute] @@ -500,6 +498,15 @@ def _condition_key(condition: Condition) -> Any: return kind, tuple([_condition_key(c) for c in condition.conditions]) # ty: ignore[unresolved-attribute] +def _subquery_key(related: CorrelatedSubquery) -> Any: + subquery = related.subquery + # A condition the builder did not make has no source, and its subquery stands for itself. + source = getattr(related.source, "_ast", subquery) + # A two-hop `exists` builds its junction AST new every time, so it is keyed by its condition, + # which names the query it leads to. + return id(subquery) if source is subquery else _condition_key(subquery.where) # ty: ignore[invalid-argument-type] + + def _check_keys(first: Connection, second: Connection) -> None: for key in (first.source_field, first.dest_field, second.source_field, second.dest_field): check(len(key) >= 1, "Invalid relationship") diff --git a/tests/wire/test_query.py b/tests/wire/test_query.py index 20209e4..ba4bdec 100644 --- a/tests/wire/test_query.py +++ b/tests/wire/test_query.py @@ -221,6 +221,9 @@ def test_async_handlers_match_upstream(harness: Harness, schema: Schema) -> None ), _q("issue", ["where", "open", True], ["where", "priority", ">", 2.5], ["one"]), _q("issue", ["where", "meta", [1, "x"]], ["where", "title", "IN", ["a", "b"]]), + # A two-hop `exists` puts `scalar` only on its junction's condition. + _q("issue", ["whereExists", "labels", None, {"scalar": True}]), + _q("issue", ["whereExists", "labels", None, {"scalar": False}]), ]