@rocicorp/zero 0.25.0 → 0.26.0-canary.0

package/out/zero-cache/src/services/view-syncer/pipeline-driver.js

@@ -1,510 +1,633 @@
-import { LogContext } from '@rocicorp/logger';
-import { assert, unreachable } from "../../../../shared/src/asserts.js";
+import { unreachable, assert } from "../../../../shared/src/asserts.js";
 import { deepEqual } from "../../../../shared/src/json.js";
 import { must } from "../../../../shared/src/must.js";
 import { buildPipeline } from "../../../../zql/src/builder/builder.js";
-import { Debug, runtimeDebugFlags, } from "../../../../zql/src/builder/debug-delegate.js";
+import { Debug, runtimeDebugFlags } from "../../../../zql/src/builder/debug-delegate.js";
 import { MeasurePushOperator } from "../../../../zql/src/query/measure-push-operator.js";
+import { createSQLiteCostModel } from "../../../../zqlite/src/sqlite-cost-model.js";
 import { TableSource } from "../../../../zqlite/src/table-source.js";
-import { reloadPermissionsIfChanged, } from "../../auth/load-permissions.js";
+import { reloadPermissionsIfChanged } from "../../auth/load-permissions.js";
 import { computeZqlSpecs, mustGetTableSpec } from "../../db/lite-tables.js";
-import { getOrCreateCounter, getOrCreateHistogram, } from "../../observability/metrics.js";
-import {} from "../../types/row-key.js";
+import { getOrCreateHistogram, getOrCreateCounter } from "../../observability/metrics.js";
+import "js-xxhash";
+import "../../../../shared/src/bigint-json.js";
+import { upstreamSchema } from "../../types/shards.js";
 import { getSubscriptionState } from "../replicator/schema/replication-state.js";
 import { checkClientSchema } from "./client-schema.js";
-import { ResetPipelinesSignal, Snapshotter, } from "./snapshotter.js";
-import { createSQLiteCostModel } from "../../../../zqlite/src/sqlite-cost-model.js";
-/**
- * Manages the state of IVM pipelines for a given ViewSyncer (i.e. client group).
- */
-export class PipelineDriver {
-    #tables = new Map();
-    // We probs need the original query hash
-    // so we can decide not to re-transform a custom query
-    // that is already hydrated.
-    #pipelines = new Map();
-    #lc;
-    #snapshotter;
-    #storage;
-    #shardID;
-    #logConfig;
-    #tableSpecs = new Map();
-    #costModels;
-    #streamer = null;
-    #replicaVersion = null;
-    #permissions = null;
-    #advanceTime = getOrCreateHistogram('sync', 'ivm.advance-time', {
-        description: 'Time to advance all queries for a given client group for in response to a single change.',
-        unit: 's',
-    });
-    #conflictRowsDeleted = getOrCreateCounter('sync', 'ivm.conflict-rows-deleted', 'Number of rows deleted because they conflicted with added row');
-    #inspectorDelegate;
-    constructor(lc, logConfig, snapshotter, shardID, storage, clientGroupID, inspectorDelegate, enablePlanner) {
-        this.#lc = lc.withContext('clientGroupID', clientGroupID);
-        this.#snapshotter = snapshotter;
-        this.#storage = storage;
-        this.#shardID = shardID;
-        this.#logConfig = logConfig;
-        this.#inspectorDelegate = inspectorDelegate;
-        this.#costModels = enablePlanner ? new WeakMap() : undefined;
+import { ResetPipelinesSignal } from "./snapshotter.js";
+const MIN_ADVANCEMENT_TIME_LIMIT_MS = 50;
+class PipelineDriver {
+  #tables = /* @__PURE__ */ new Map();
+  // We probs need the original query hash
+  // so we can decide not to re-transform a custom query
+  // that is already hydrated.
+  #pipelines = /* @__PURE__ */ new Map();
+  #lc;
+  #snapshotter;
+  #storage;
+  #shardID;
+  #logConfig;
+  #tableSpecs = /* @__PURE__ */ new Map();
+  #costModels;
+  #yieldThresholdMs;
+  #streamer = null;
+  #hydrateContext = null;
+  #advanceContext = null;
+  #replicaVersion = null;
+  #primaryKeys = null;
+  #permissions = null;
+  #advanceTime = getOrCreateHistogram("sync", "ivm.advance-time", {
+    description: "Time to advance all queries for a given client group for in response to a single change.",
+    unit: "s"
+  });
+  #conflictRowsDeleted = getOrCreateCounter(
+    "sync",
+    "ivm.conflict-rows-deleted",
+    "Number of rows deleted because they conflicted with added row"
+  );
+  #inspectorDelegate;
+  constructor(lc, logConfig, snapshotter, shardID, storage, clientGroupID, inspectorDelegate, yieldThresholdMs, enablePlanner) {
+    this.#lc = lc.withContext("clientGroupID", clientGroupID);
+    this.#snapshotter = snapshotter;
+    this.#storage = storage;
+    this.#shardID = shardID;
+    this.#logConfig = logConfig;
+    this.#inspectorDelegate = inspectorDelegate;
+    this.#costModels = enablePlanner ? /* @__PURE__ */ new WeakMap() : void 0;
+    this.#yieldThresholdMs = yieldThresholdMs;
+  }
+  /**
+   * Initializes the PipelineDriver to the current head of the database.
+   * Queries can then be added (i.e. hydrated) with {@link addQuery()}.
+   *
+   * Must only be called once.
+   */
+  init(clientSchema) {
+    assert(!this.#snapshotter.initialized(), "Already initialized");
+    this.#snapshotter.init();
+    this.#initAndResetCommon(clientSchema);
+  }
+  /**
+   * @returns Whether the PipelineDriver has been initialized.
+   */
+  initialized() {
+    return this.#snapshotter.initialized();
+  }
+  /**
+   * Clears the current pipelines and TableSources, returning the PipelineDriver
+   * to its initial state. This should be called in response to a schema change,
+   * as TableSources need to be recomputed.
+   */
+  reset(clientSchema) {
+    for (const { input } of this.#pipelines.values()) {
+      input.destroy();
     }
-    /**
-     * Initializes the PipelineDriver to the current head of the database.
-     * Queries can then be added (i.e. hydrated) with {@link addQuery()}.
-     *
-     * Must only be called once.
-     */
-    init(clientSchema) {
-        assert(!this.#snapshotter.initialized(), 'Already initialized');
-        const { db } = this.#snapshotter.init().current();
-        const fullTables = new Map();
-        computeZqlSpecs(this.#lc, db.db, this.#tableSpecs, fullTables);
-        if (clientSchema) {
-            checkClientSchema(this.#shardID, clientSchema, this.#tableSpecs, fullTables);
-        }
-        const { replicaVersion } = getSubscriptionState(db);
-        this.#replicaVersion = replicaVersion;
+    this.#pipelines.clear();
+    this.#tables.clear();
+    this.#initAndResetCommon(clientSchema);
+  }
+  #initAndResetCommon(clientSchema) {
+    const { db } = this.#snapshotter.current();
+    const fullTables = /* @__PURE__ */ new Map();
+    computeZqlSpecs(this.#lc, db.db, this.#tableSpecs, fullTables);
+    checkClientSchema(
+      this.#shardID,
+      clientSchema,
+      this.#tableSpecs,
+      fullTables
+    );
+    const primaryKeys = this.#primaryKeys ?? /* @__PURE__ */ new Map();
+    this.#primaryKeys = primaryKeys;
+    primaryKeys.clear();
+    for (const [table, spec] of this.#tableSpecs.entries()) {
+      if (table.startsWith(upstreamSchema(this.#shardID))) {
+        primaryKeys.set(table, spec.tableSpec.primaryKey);
+      }
     }
-    /**
-     * @returns Whether the PipelineDriver has been initialized.
-     */
-    initialized() {
-        return this.#snapshotter.initialized();
+    buildPrimaryKeys(clientSchema, primaryKeys);
+    const { replicaVersion } = getSubscriptionState(db);
+    this.#replicaVersion = replicaVersion;
+  }
+  /** @returns The replica version. The PipelineDriver must have been initialized. */
+  get replicaVersion() {
+    return must(this.#replicaVersion, "Not yet initialized");
+  }
+  /**
+   * Returns the current version of the database. This will reflect the
+   * latest version change when calling {@link advance()} once the
+   * iteration has begun.
+   */
+  currentVersion() {
+    assert(this.initialized(), "Not yet initialized");
+    return this.#snapshotter.current().version;
+  }
+  /**
+   * Returns the current supported schema version range of the database.  This
+   * will reflect changes to supported schema version range when calling
+   * {@link advance()} once the iteration has begun.
+   */
+  currentSchemaVersions() {
+    assert(this.initialized(), "Not yet initialized");
+    return this.#snapshotter.current().schemaVersions;
+  }
+  /**
+   * Returns the current upstream {app}.permissions, or `null` if none are defined.
+   */
+  currentPermissions() {
+    assert(this.initialized(), "Not yet initialized");
+    const res = reloadPermissionsIfChanged(
+      this.#lc,
+      this.#snapshotter.current().db,
+      this.#shardID.appID,
+      this.#permissions
+    );
+    if (res.changed) {
+      this.#permissions = res.permissions;
+      this.#lc.debug?.(
+        "Reloaded permissions",
+        JSON.stringify(this.#permissions)
+      );
     }
-    /** @returns The replica version. The PipelineDriver must have been initialized. */
-    get replicaVersion() {
-        return must(this.#replicaVersion, 'Not yet initialized');
+    return this.#permissions;
+  }
+  advanceWithoutDiff() {
+    const { db, version } = this.#snapshotter.advanceWithoutDiff().curr;
+    for (const table of this.#tables.values()) {
+      table.setDB(db.db);
     }
-    /**
-     * Returns the current version of the database. This will reflect the
-     * latest version change when calling {@link advance()} once the
-     * iteration has begun.
-     */
-    currentVersion() {
-        assert(this.initialized(), 'Not yet initialized');
-        return this.#snapshotter.current().version;
+    return version;
+  }
+  #ensureCostModelExistsIfEnabled(db) {
+    let existing = this.#costModels?.get(db);
+    if (existing) {
+      return existing;
     }
-    /**
-     * Returns the current supported schema version range of the database.  This
-     * will reflect changes to supported schema version range when calling
-     * {@link advance()} once the iteration has begun.
-     */
-    currentSchemaVersions() {
-        assert(this.initialized(), 'Not yet initialized');
-        return this.#snapshotter.current().schemaVersions;
+    if (this.#costModels) {
+      const costModel = createSQLiteCostModel(db, this.#tableSpecs);
+      this.#costModels.set(db, costModel);
+      return costModel;
     }
-    /**
-     * Returns the current upstream {app}.permissions, or `null` if none are defined.
-     */
-    currentPermissions() {
-        assert(this.initialized(), 'Not yet initialized');
-        const res = reloadPermissionsIfChanged(this.#lc, this.#snapshotter.current().db, this.#shardID.appID, this.#permissions);
-        if (res.changed) {
-            this.#permissions = res.permissions;
-            this.#lc.debug?.('Reloaded permissions', JSON.stringify(this.#permissions));
-        }
-        return this.#permissions;
+    return void 0;
+  }
+  /**
+   * Clears storage used for the pipelines. Call this when the
+   * PipelineDriver will no longer be used.
+   */
+  destroy() {
+    this.#storage.destroy();
+    this.#snapshotter.destroy();
+  }
+  /** @return The Set of query hashes for all added queries. */
+  addedQueries() {
+    const byOriginalHash = /* @__PURE__ */ new Map();
+    for (const pipeline of this.#pipelines.values()) {
+      const { originalHash, transformedAst, transformationHash } = pipeline;
+      if (!byOriginalHash.has(originalHash)) {
+        byOriginalHash.set(originalHash, []);
+      }
+      byOriginalHash.get(originalHash).push({
+        transformationHash,
+        transformedAst
+      });
     }
-    advanceWithoutDiff() {
-        const { db, version } = this.#snapshotter.advanceWithoutDiff().curr;
-        for (const table of this.#tables.values()) {
-            table.setDB(db.db);
-        }
-        return version;
+    return [new Set(this.#pipelines.keys()), byOriginalHash];
+  }
+  totalHydrationTimeMs() {
+    let total = 0;
+    for (const pipeline of this.#pipelines.values()) {
+      total += pipeline.hydrationTimeMs;
     }
-    #ensureCostModelExistsIfEnabled(db) {
-        let existing = this.#costModels?.get(db);
-        if (existing) {
-            return existing;
-        }
-        if (this.#costModels) {
-            const costModel = createSQLiteCostModel(db, this.#tableSpecs);
-            this.#costModels.set(db, costModel);
-            return costModel;
-        }
-        return undefined;
-    }
-    /**
-     * Clears the current pipelines and TableSources, returning the PipelineDriver
-     * to its initial state. This should be called in response to a schema change,
-     * as TableSources need to be recomputed.
-     */
-    reset(clientSchema) {
-        for (const { input } of this.#pipelines.values()) {
-            input.destroy();
-        }
-        this.#pipelines.clear();
-        this.#tables.clear();
-        const { db } = this.#snapshotter.current();
-        const fullTables = new Map();
-        computeZqlSpecs(this.#lc, db.db, this.#tableSpecs, fullTables);
-        if (clientSchema) {
-            checkClientSchema(this.#shardID, clientSchema, this.#tableSpecs, fullTables);
-        }
-        const { replicaVersion } = getSubscriptionState(db);
-        this.#replicaVersion = replicaVersion;
+    return total;
+  }
+  /**
+   * Adds a pipeline for the query. The method will hydrate the query using the
+   * driver's current snapshot of the database and return a stream of results.
+   * Henceforth, updates to the query will be returned when the driver is
+   * {@link advance}d. The query and its pipeline can be removed with
+   * {@link removeQuery()}.
+   *
+   * If a query with an identical hash has already been added, this method is a
+   * no-op and no RowChanges are generated.
+   *
+   * @param timer The caller-controlled {@link Timer} used to determine the
+   *        final hydration time. (The caller may pause and resume the timer
+   *        when yielding the thread for time-slicing).
+   * @return The rows from the initial hydration of the query.
+   */
+  *addQuery(transformationHash, queryID, query, timer) {
+    assert(this.initialized());
+    this.#inspectorDelegate.addQuery(transformationHash, queryID, query);
+    if (this.#pipelines.has(transformationHash)) {
+      this.#lc.info?.(`query ${transformationHash} already added`, query);
+      return;
     }
-    /**
-     * Clears storage used for the pipelines. Call this when the
-     * PipelineDriver will no longer be used.
-     */
-    destroy() {
-        this.#storage.destroy();
-        this.#snapshotter.destroy();
+    const debugDelegate = runtimeDebugFlags.trackRowsVended ? new Debug() : void 0;
+    const costModel = this.#ensureCostModelExistsIfEnabled(
+      this.#snapshotter.current().db.db
+    );
+    const input = buildPipeline(
+      query,
+      {
+        debug: debugDelegate,
+        enableNotExists: true,
+        // Server-side can handle NOT EXISTS
+        getSource: (name) => this.#getSource(name),
+        createStorage: () => this.#createStorage(),
+        decorateSourceInput: (input2, _queryID) => new MeasurePushOperator(
+          input2,
+          transformationHash,
+          this.#inspectorDelegate,
+          "query-update-server"
+        ),
+        decorateInput: (input2) => input2,
+        addEdge() {
+        },
+        decorateFilterInput: (input2) => input2
+      },
+      queryID,
+      costModel
+    );
+    const schema = input.getSchema();
+    input.setOutput({
+      push: (change) => {
+        const streamer = this.#streamer;
+        assert(streamer, "must #startAccumulating() before pushing changes");
+        streamer.accumulate(transformationHash, schema, [change]);
+        return [];
+      }
+    });
+    assert(this.#advanceContext === null);
+    this.#hydrateContext = {
+      timer
+    };
+    try {
+      yield* hydrateInternal(
+        input,
+        transformationHash,
+        must(this.#primaryKeys)
+      );
+    } finally {
+      this.#hydrateContext = null;
     }
-    /** @return The Set of query hashes for all added queries. */
-    addedQueries() {
-        const byOriginalHash = new Map();
-        for (const pipeline of this.#pipelines.values()) {
-            const { originalHash, transformedAst, transformationHash } = pipeline;
-            if (!byOriginalHash.has(originalHash)) {
-                byOriginalHash.set(originalHash, []);
-            }
-            byOriginalHash.get(originalHash).push({
-                transformationHash,
-                transformedAst,
-            });
+    const hydrationTimeMs = timer.totalElapsed();
+    if (runtimeDebugFlags.trackRowCountsVended) {
+      if (hydrationTimeMs > this.#logConfig.slowHydrateThreshold) {
+        let totalRowsConsidered = 0;
+        const lc = this.#lc.withContext("hash", transformationHash).withContext("hydrationTimeMs", hydrationTimeMs);
+        for (const tableName of this.#tables.keys()) {
+          const entries = Object.entries(
+            debugDelegate?.getVendedRowCounts()[tableName] ?? {}
+          );
+          totalRowsConsidered += entries.reduce(
+            (acc, entry) => acc + entry[1],
+            0
+          );
+          lc.info?.(tableName + " VENDED: ", entries);
         }
-        return [new Set(this.#pipelines.keys()), byOriginalHash];
+        lc.info?.(`Total rows considered: ${totalRowsConsidered}`);
+      }
     }
-    totalHydrationTimeMs() {
-        let total = 0;
-        for (const pipeline of this.#pipelines.values()) {
-            total += pipeline.hydrationTimeMs;
-        }
-        return total;
+    debugDelegate?.reset();
+    this.#pipelines.set(transformationHash, {
+      input,
+      hydrationTimeMs,
+      originalHash: queryID,
+      transformedAst: query,
+      transformationHash
+    });
+  }
+  /**
+   * Removes the pipeline for the query. This is a no-op if the query
+   * was not added.
+   */
+  removeQuery(hash) {
+    const pipeline = this.#pipelines.get(hash);
+    if (pipeline) {
+      this.#pipelines.delete(hash);
+      pipeline.input.destroy();
     }
-    /**
-     * Adds a pipeline for the query. The method will hydrate the query using the
-     * driver's current snapshot of the database and return a stream of results.
-     * Henceforth, updates to the query will be returned when the driver is
-     * {@link advance}d. The query and its pipeline can be removed with
-     * {@link removeQuery()}.
-     *
-     * If a query with an identical hash has already been added, this method is a
-     * no-op and no RowChanges are generated.
-     *
-     * @param timer The caller-controlled {@link Timer} used to determine the
-     *        final hydration time. (The caller may pause and resume the timer
-     *        when yielding the thread for time-slicing).
-     * @return The rows from the initial hydration of the query.
-     */
-    *addQuery(transformationHash, queryID, query, timer) {
-        assert(this.initialized());
-        this.#inspectorDelegate.addQuery(transformationHash, queryID, query);
-        if (this.#pipelines.has(transformationHash)) {
-            this.#lc.info?.(`query ${transformationHash} already added`, query);
-            return;
+  }
+  /**
+   * Returns the value of the row with the given primary key `pk`,
+   * or `undefined` if there is no such row. The pipeline must have been
+   * initialized.
+   */
+  getRow(table, pk) {
+    assert(this.initialized(), "Not yet initialized");
+    const source = must(this.#tables.get(table));
+    return source.getRow(pk);
+  }
+  /**
+   * Advances to the new head of the database.
+   *
+   * @param timer The caller-controlled {@link Timer} that will be used to
+   *        measure the progress of the advancement and abort with a
+   *        {@link ResetPipelinesSignal} if it is estimated to take longer
+   *        than a hydration.
+   * @return The resulting row changes for all added queries. Note that the
+   *         `changes` must be iterated over in their entirety in order to
+   *         advance the database snapshot.
+   */
+  advance(timer) {
+    assert(this.initialized());
+    const diff = this.#snapshotter.advance(this.#tableSpecs);
+    const { prev, curr, changes } = diff;
+    this.#lc.debug?.(
+      `advance ${prev.version} => ${curr.version}: ${changes} changes`
+    );
+    return {
+      version: curr.version,
+      numChanges: changes,
+      changes: this.#advance(diff, timer, changes)
+    };
+  }
+  *#advance(diff, timer, numChanges) {
+    assert(this.#hydrateContext === null);
+    this.#advanceContext = {
+      timer,
+      totalHydrationTimeMs: this.totalHydrationTimeMs(),
+      numChanges,
+      pos: 0
+    };
+    try {
+      for (const { table, prevValues, nextValue } of diff) {
+        if (this.#shouldAdvanceYieldMaybeAbortAdvance()) {
+          yield "yield";
         }
-        const debugDelegate = runtimeDebugFlags.trackRowsVended
-            ? new Debug()
-            : undefined;
-        const costModel = this.#ensureCostModelExistsIfEnabled(this.#snapshotter.current().db.db);
-        const input = buildPipeline(query, {
-            debug: debugDelegate,
-            enableNotExists: true, // Server-side can handle NOT EXISTS
-            getSource: name => this.#getSource(name),
-            createStorage: () => this.#createStorage(),
-            decorateSourceInput: (input, _queryID) => new MeasurePushOperator(input, transformationHash, this.#inspectorDelegate, 'query-update-server'),
-            decorateInput: input => input,
-            addEdge() { },
-            decorateFilterInput: input => input,
-        }, queryID, costModel);
-        const schema = input.getSchema();
-        input.setOutput({
-            push: change => {
-                const streamer = this.#streamer;
-                assert(streamer, 'must #startAccumulating() before pushing changes');
-                streamer.accumulate(transformationHash, schema, [change]);
-            },
-        });
-        yield* hydrate(input, transformationHash, this.#tableSpecs);
-        const hydrationTimeMs = timer.totalElapsed();
-        if (runtimeDebugFlags.trackRowCountsVended) {
-            if (hydrationTimeMs > this.#logConfig.slowHydrateThreshold) {
-                let totalRowsConsidered = 0;
-                const lc = this.#lc
-                    .withContext('hash', transformationHash)
-                    .withContext('hydrationTimeMs', hydrationTimeMs);
-                for (const tableName of this.#tables.keys()) {
-                    const entries = Object.entries(debugDelegate?.getVendedRowCounts()[tableName] ?? {});
-                    totalRowsConsidered += entries.reduce((acc, entry) => acc + entry[1], 0);
-                    lc.info?.(tableName + ' VENDED: ', entries);
-                }
-                lc.info?.(`Total rows considered: ${totalRowsConsidered}`);
+        const start = performance.now();
+        let type;
+        try {
+          const tableSource = this.#tables.get(table);
+          if (!tableSource) {
+            continue;
+          }
+          const primaryKey = mustGetPrimaryKey(this.#primaryKeys, table);
+          let editOldRow = void 0;
+          for (const prevValue of prevValues) {
+            if (nextValue && deepEqual(
+              getRowKey(primaryKey, prevValue),
+              getRowKey(primaryKey, nextValue)
+            )) {
+              editOldRow = prevValue;
+            } else {
+              if (nextValue) {
+                this.#conflictRowsDeleted.add(1);
+              }
+              yield* this.#push(tableSource, {
+                type: "remove",
+                row: prevValue
+              });
+            }
+          }
+          if (nextValue) {
+            if (editOldRow) {
+              yield* this.#push(tableSource, {
+                type: "edit",
+                row: nextValue,
+                oldRow: editOldRow
+              });
+            } else {
+              yield* this.#push(tableSource, {
+                type: "add",
+                row: nextValue
+              });
             }
+          }
+        } finally {
+          this.#advanceContext.pos++;
         }
-        debugDelegate?.reset();
-        // Note: This hydrationTime is a wall-clock overestimate, as it does
-        // not take time slicing into account. The view-syncer resets this
-        // to a more precise processing-time measurement with setHydrationTime().
-        this.#pipelines.set(transformationHash, {
-            input,
-            hydrationTimeMs,
-            originalHash: queryID,
-            transformedAst: query,
-            transformationHash,
+        const elapsed = performance.now() - start;
+        this.#advanceTime.record(elapsed / 1e3, {
+          table,
+          type
         });
+      }
+      const { curr } = diff;
+      for (const table of this.#tables.values()) {
+        table.setDB(curr.db.db);
+      }
+      this.#ensureCostModelExistsIfEnabled(curr.db.db);
+      this.#lc.debug?.(`Advanced to ${curr.version}`);
+    } finally {
+      this.#advanceContext = null;
     }
-    /**
-     * Removes the pipeline for the query. This is a no-op if the query
-     * was not added.
-     */
-    removeQuery(hash) {
-        const pipeline = this.#pipelines.get(hash);
-        if (pipeline) {
-            this.#pipelines.delete(hash);
-            pipeline.input.destroy();
-        }
+  }
+  /** Implements `BuilderDelegate.getSource()` */
+  #getSource(tableName) {
+    let source = this.#tables.get(tableName);
+    if (source) {
+      return source;
     }
-    /**
-     * Returns the value of the row with the given primary key `pk`,
-     * or `undefined` if there is no such row. The pipeline must have been
-     * initialized.
-     */
-    getRow(table, pk) {
-        assert(this.initialized(), 'Not yet initialized');
-        const source = must(this.#tables.get(table));
-        return source.getRow(pk);
+    const tableSpec = mustGetTableSpec(this.#tableSpecs, tableName);
+    const primaryKey = mustGetPrimaryKey(this.#primaryKeys, tableName);
+    const { db } = this.#snapshotter.current();
+    source = new TableSource(
+      this.#lc,
+      this.#logConfig,
+      db.db,
+      tableName,
+      tableSpec.zqlSpec,
+      primaryKey,
+      () => this.#shouldYield()
+    );
+    this.#tables.set(tableName, source);
+    this.#lc.debug?.(`created TableSource for ${tableName}`);
+    return source;
+  }
+  #shouldYield() {
+    if (this.#hydrateContext) {
+      return this.#hydrateContext.timer.elapsedLap() > this.#yieldThresholdMs;
     }
-    /**
-     * Advances to the new head of the database.
-     *
-     * @param timer The caller-controlled {@link Timer} that will be used to
-     *        measure the progress of the advancement and abort with a
-     *        {@link ResetPipelinesSignal} if it is estimated to take longer
-     *        than a hydration.
-     * @return The resulting row changes for all added queries. Note that the
-     *         `changes` must be iterated over in their entirety in order to
-     *         advance the database snapshot.
-     */
-    advance(timer) {
-        assert(this.initialized());
-        const diff = this.#snapshotter.advance(this.#tableSpecs);
-        const { prev, curr, changes } = diff;
-        this.#lc.debug?.(`${prev.version} => ${curr.version}: ${changes} changes`);
-        const totalHydrationTimeMs = this.totalHydrationTimeMs();
-        // Cancel the advancement processing if it takes longer than half the
-        // total hydration time to make it through half of the advancement.
-        // This serves as both a circuit breaker for very large transactions,
-        // as well as a bound on the amount of time the previous connection locks
-        // the inactive WAL file (as the lock prevents WAL2 from switching to the
-        // free WAL when the current one is over the size limit, which can make
-        // the WAL grow continuously and compound slowness).
-        //
-        // Note: 1/2 is a conservative estimate policy. A lower proportion would
-        // flag slowness sooner, at the expense of larger estimation error.
-        function checkProgress(pos) {
-            // Check every 10 changes
-            if (pos % 10 === 0) {
-                const elapsed = timer.totalElapsed();
-                if (elapsed > totalHydrationTimeMs / 2 && pos <= changes / 2) {
-                    throw new ResetPipelinesSignal(`advancement exceeded timeout at ${pos} of ${changes} changes (${elapsed} ms)`);
-                }
-            }
-        }
-        return {
-            version: curr.version,
-            numChanges: changes,
-            changes: this.#advance(diff, 
-            // Somewhat arbitrary: only check progress if there are at least 20
-            // changes (Note that the first check doesn't happen until 10 changes).
-            changes >= 20 ? checkProgress : () => { }),
-        };
+    if (this.#advanceContext) {
+      return this.#shouldAdvanceYieldMaybeAbortAdvance();
     }
-    *#advance(diff, onChange) {
-        let pos = 0;
-        for (const { table, prevValues, nextValue } of diff) {
-            const start = performance.now();
-            let type;
-            try {
-                const tableSpec = mustGetTableSpec(this.#tableSpecs, table).tableSpec;
-                const tableSource = this.#tables.get(table);
-                if (!tableSource) {
-                    // no pipelines read from this table, so no need to process the change
-                    continue;
-                }
-                let editOldRow = undefined;
-                for (const prevValue of prevValues) {
-                    // Rows are ultimately referred to by the union key (in #streamNodes())
-                    // so an update is represented as an `edit` if and only if the
-                    // unionKey-based row keys are the same in prevValue and nextValue.
-                    if (nextValue &&
-                        deepEqual(getRowKey(tableSpec.unionKey, prevValue), getRowKey(tableSpec.unionKey, nextValue))) {
-                        editOldRow = prevValue;
-                    }
-                    else {
-                        if (nextValue) {
-                            this.#conflictRowsDeleted.add(1);
-                        }
-                        yield* this.#push(tableSource, {
-                            type: 'remove',
-                            row: prevValue,
-                        });
-                    }
-                }
-                if (nextValue) {
-                    if (editOldRow) {
-                        yield* this.#push(tableSource, {
-                            type: 'edit',
-                            row: nextValue,
-                            oldRow: editOldRow,
-                        });
-                    }
-                    else {
-                        yield* this.#push(tableSource, {
-                            type: 'add',
-                            row: nextValue,
-                        });
-                    }
-                }
-            }
-            finally {
-                onChange(++pos);
-            }
-            const elapsed = performance.now() - start;
-            this.#advanceTime.record(elapsed / 1000, {
-                table,
-                type,
-            });
-        }
-        // Set the new snapshot on all TableSources.
-        const { curr } = diff;
-        for (const table of this.#tables.values()) {
-            table.setDB(curr.db.db);
-        }
-        this.#ensureCostModelExistsIfEnabled(curr.db.db);
-        this.#lc.debug?.(`Advanced to ${curr.version}`);
+    throw new Error("shouldYield called outside of hydration or advancement");
+  }
+  /**
+   * Cancel the advancement processing, by throwing a ResetPipelinesSignal, if
+   * it has taken longer than half the total hydration time to make it through
+   * half of the advancement, or if processing time exceeds total hydration
+   * time.  This serves as both a circuit breaker for very large transactions,
+   * as well as a bound on the amount of time the previous connection locks
+   * the inactive WAL file (as the lock prevents WAL2 from switching to the
+   * free WAL when the current one is over the size limit, which can make
+   * the WAL grow continuously and compound slowness).
+   * This is checked:
+   * 1. before starting to process each change in an advancement is processed
+   * 2. whenever a row is fetched from a TableSource during push processing
+   */
+  #shouldAdvanceYieldMaybeAbortAdvance() {
+    const {
+      pos,
+      numChanges,
+      timer: advanceTimer,
+      totalHydrationTimeMs
+    } = must(this.#advanceContext);
+    const elapsed = advanceTimer.totalElapsed();
+    if (elapsed > MIN_ADVANCEMENT_TIME_LIMIT_MS && (elapsed > totalHydrationTimeMs || elapsed > totalHydrationTimeMs / 2 && pos <= numChanges / 2)) {
+      throw new ResetPipelinesSignal(
+        `Advancement exceeded timeout at ${pos} of ${numChanges} changes after ${elapsed} ms. Advancement time limited based on total hydration time of ${totalHydrationTimeMs} ms.`
+      );
     }
-    /** Implements `BuilderDelegate.getSource()` */
-    #getSource(tableName) {
-        let source = this.#tables.get(tableName);
-        if (source) {
-            return source;
+    return advanceTimer.elapsedLap() > this.#yieldThresholdMs;
+  }
+  /** Implements `BuilderDelegate.createStorage()` */
+  #createStorage() {
+    return this.#storage.createStorage();
+  }
+  *#push(source, change) {
+    this.#startAccumulating();
+    try {
+      for (const val of source.genPush(change)) {
+        if (val === "yield") {
+          yield "yield";
         }
-        const tableSpec = mustGetTableSpec(this.#tableSpecs, tableName);
-        const { primaryKey } = tableSpec.tableSpec;
-        const { db } = this.#snapshotter.current();
-        source = new TableSource(this.#lc, this.#logConfig, db.db, tableName, tableSpec.zqlSpec, primaryKey);
-        this.#tables.set(tableName, source);
-        this.#lc.debug?.(`created TableSource for ${tableName}`);
-        return source;
-    }
-    /** Implements `BuilderDelegate.createStorage()` */
-    #createStorage() {
-        return this.#storage.createStorage();
-    }
-    *#push(source, change) {
-        this.#startAccumulating();
-        for (const _ of source.genPush(change)) {
-            yield* this.#stopAccumulating().stream();
-            this.#startAccumulating();
+        for (const changeOrYield of this.#stopAccumulating().stream()) {
+          yield changeOrYield;
         }
+        this.#startAccumulating();
+      }
+    } finally {
+      if (this.#streamer !== null) {
         this.#stopAccumulating();
+      }
     }
-    #startAccumulating() {
-        assert(this.#streamer === null);
-        this.#streamer = new Streamer(this.#tableSpecs);
-    }
-    #stopAccumulating() {
-        const streamer = this.#streamer;
-        assert(streamer);
-        this.#streamer = null;
-        return streamer;
-    }
+  }
+  #startAccumulating() {
+    assert(this.#streamer === null);
+    this.#streamer = new Streamer(must(this.#primaryKeys));
+  }
+  #stopAccumulating() {
+    const streamer = this.#streamer;
+    assert(streamer);
+    this.#streamer = null;
+    return streamer;
+  }
 }
 class Streamer {
-    #tableSpecs;
-    constructor(tableSpecs) {
-        this.#tableSpecs = tableSpecs;
+  #primaryKeys;
+  constructor(primaryKeys) {
+    this.#primaryKeys = primaryKeys;
+  }
+  #changes = [];
+  accumulate(hash, schema, changes) {
+    this.#changes.push([hash, schema, changes]);
+    return this;
+  }
+  *stream() {
+    for (const [hash, schema, changes] of this.#changes) {
+      yield* this.#streamChanges(hash, schema, changes);
     }
-    #changes = [];
-    accumulate(hash, schema, changes) {
-        this.#changes.push([hash, schema, changes]);
-        return this;
+  }
+  *#streamChanges(queryHash, schema, changes) {
+    if (schema.system === "permissions") {
+      return;
     }
-    *stream() {
-        for (const [hash, schema, changes] of this.#changes) {
-            yield* this.#streamChanges(hash, schema, changes);
+    for (const change of changes) {
+      if (change === "yield") {
+        yield change;
+        continue;
+      }
+      const { type } = change;
+      switch (type) {
+        case "add":
+        case "remove": {
+          yield* this.#streamNodes(queryHash, schema, type, () => [
+            change.node
+          ]);
+          break;
         }
-    }
-    *#streamChanges(queryHash, schema, changes) {
-        // We do not sync rows gathered by the permissions
-        // system to the client.
-        if (schema.system === 'permissions') {
-            return;
-        }
-        for (const change of changes) {
-            const { type } = change;
-            switch (type) {
-                case 'add':
-                case 'remove': {
-                    yield* this.#streamNodes(queryHash, schema, type, () => [
-                        change.node,
-                    ]);
-                    break;
-                }
-                case 'child': {
-                    const { child } = change;
-                    const childSchema = must(schema.relationships[child.relationshipName]);
-                    yield* this.#streamChanges(queryHash, childSchema, [child.change]);
-                    break;
-                }
-                case 'edit':
-                    yield* this.#streamNodes(queryHash, schema, type, () => [
-                        { row: change.node.row, relationships: {} },
-                    ]);
-                    break;
-                default:
-                    unreachable(type);
-            }
+        case "child": {
+          const { child } = change;
+          const childSchema = must(
+            schema.relationships[child.relationshipName]
+          );
+          yield* this.#streamChanges(queryHash, childSchema, [child.change]);
+          break;
         }
+        case "edit":
+          yield* this.#streamNodes(queryHash, schema, type, () => [
+            { row: change.node.row, relationships: {} }
+          ]);
+          break;
+        default:
+          unreachable();
+      }
     }
-    *#streamNodes(queryHash, schema, op, nodes) {
-        const { tableName: table, system } = schema;
-        // The primaryKey here is used for referencing rows in CVR and del-row
-        // patches sent in pokes. This is the "unionKey", i.e. the union of all
-        // columns in unique indexes. This allows clients to migrate from, e.g.
-        // pk1 to pk2, as del-patches will be keyed by [...pk1, ...pk2].
-        const primaryKey = must(this.#tableSpecs.get(table)).tableSpec.unionKey;
-        // We do not sync rows gathered by the permissions
-        // system to the client.
-        if (system === 'permissions') {
-            return;
-        }
-        for (const node of nodes()) {
-            const { relationships, row } = node;
-            const rowKey = getRowKey(primaryKey, row);
-            yield {
-                type: op,
-                queryHash,
-                table,
-                rowKey,
-                row: op === 'remove' ? undefined : row,
-            };
-            for (const [relationship, children] of Object.entries(relationships)) {
-                const childSchema = must(schema.relationships[relationship]);
-                yield* this.#streamNodes(queryHash, childSchema, op, children);
-            }
-        }
+  }
+  *#streamNodes(queryHash, schema, op, nodes) {
+    const { tableName: table, system } = schema;
+    const primaryKey = must(this.#primaryKeys.get(table));
+    if (system === "permissions") {
+      return;
+    }
+    for (const node of nodes()) {
+      if (node === "yield") {
+        yield node;
+        continue;
+      }
+      const { relationships, row } = node;
+      const rowKey = getRowKey(primaryKey, row);
+      yield {
+        type: op,
+        queryHash,
+        table,
+        rowKey,
+        row: op === "remove" ? void 0 : row
+      };
+      for (const [relationship, children] of Object.entries(relationships)) {
+        const childSchema = must(schema.relationships[relationship]);
+        yield* this.#streamNodes(queryHash, childSchema, op, children);
+      }
     }
+  }
 }
 function* toAdds(nodes) {
-    for (const node of nodes) {
-        yield { type: 'add', node };
+  for (const node of nodes) {
+    if (node === "yield") {
+      yield node;
+      continue;
     }
+    yield { type: "add", node };
+  }
 }
 function getRowKey(cols, row) {
-    return Object.fromEntries(cols.map(col => [col, must(row[col])]));
+  return Object.fromEntries(cols.map((col) => [col, must(row[col])]));
+}
+function* hydrate(input, hash, clientSchema) {
+  const res = input.fetch({});
+  const streamer = new Streamer(buildPrimaryKeys(clientSchema)).accumulate(
+    hash,
+    input.getSchema(),
+    toAdds(res)
+  );
+  yield* streamer.stream();
+}
+function* hydrateInternal(input, hash, primaryKeys) {
+  const res = input.fetch({});
+  const streamer = new Streamer(primaryKeys).accumulate(
+    hash,
+    input.getSchema(),
+    toAdds(res)
+  );
+  yield* streamer.stream();
+}
+function buildPrimaryKeys(clientSchema, primaryKeys = /* @__PURE__ */ new Map()) {
+  for (const [tableName, { primaryKey }] of Object.entries(clientSchema.tables)) {
+    primaryKeys.set(tableName, primaryKey);
+  }
+  return primaryKeys;
 }
-/**
- * Core hydration logic used by {@link PipelineDriver#addQuery}, extracted to a
- * function for reuse by bin-analyze so that bin-analyze's hydration logic
- * is as close as possible to zero-cache's real hydration logic.
- */
-export function* hydrate(input, hash, tableSpecs) {
-    const res = input.fetch({});
-    const streamer = new Streamer(tableSpecs).accumulate(hash, input.getSchema(), toAdds(res));
-    yield* streamer.stream();
+function mustGetPrimaryKey(primaryKeys, table) {
+  const pKeys = must(primaryKeys, "primaryKey map must be non-null");
+  return must(
+    pKeys.get(table),
+    `table '${table}' is not one of: ${[...pKeys.keys()].sort()}. Check the spelling and ensure that the table has a primary key.`
+  );
 }
-//# sourceMappingURL=pipeline-driver.js.map
+export {
+  PipelineDriver,
+  hydrate,
+  hydrateInternal
+};
+//# sourceMappingURL=pipeline-driver.js.map