@powersync/service-module-mongodb-storage 0.0.0-dev-20260602073133 → 0.0.0-dev-20260909133214

package/src/storage/implementation/v1/MongoBucketBatchV1.ts

@@ -1,5 +1,7 @@
+import { mongo } from '@powersync/lib-service-mongodb';
 import { ReplicationAssertionError } from '@powersync/lib-services-framework';
-import { ColumnDescriptor, SourceTable, storage } from '@powersync/service-core';
+import { ColumnDescriptor, InternalOpId, SourceTable, storage } from '@powersync/service-core';
+import { HydratedSyncConfig } from '@powersync/service-sync-rules';
 import * as bson from 'bson';
 import { mongoTableId } from '../../../utils/util.js';
 import { calculateCheckpointState } from '../CheckpointState.js';
@@ -12,19 +14,32 @@ import { VersionedPowerSyncMongoV1 } from './VersionedPowerSyncMongoV1.js';
 import { SourceTableDocumentV1, SyncRuleDocumentV1 } from './models.js';
 
 export class MongoBucketBatchV1 extends MongoBucketBatch {
-  declare public readonly db: VersionedPowerSyncMongoV1;
+  declare readonly db: VersionedPowerSyncMongoV1;
 
   private readonly store: SourceRecordStore;
   private needsActivation = true;
   private lastWaitingLogThrottled = 0;
 
+  /**
+   * The last persisted op that is not yet covered by a checkpoint, if any.
+   *
+   * Seeded from the persisted keepalive_op, so that ops persisted before a restart can still be
+   * included in the next checkpoint.
+   */
+  private persisted_op: InternalOpId | null = null;
+
   constructor(options: MongoBucketBatchOptions) {
     super(options);
-    this.store = new SourceRecordStoreV1(this.db, this.group_id);
+    this.persisted_op = options.keepaliveOp ?? null;
+    this.store = new SourceRecordStoreV1(this.db, this.replicationStreamId);
+  }
+
+  protected override recordPersistedOp(lastOp: InternalOpId): void {
+    this.persisted_op = lastOp;
   }
 
   protected createPersistedBatch(writtenSize: number): PersistedBatch {
-    return new PersistedBatchV1(this.db, this.group_id, this.mapping, writtenSize, {
+    return new PersistedBatchV1(this.db, this.replicationStreamId, this.mapping, writtenSize, {
       logger: this.logger
     });
   }
@@ -32,15 +47,45 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
   protected get sourceRecordStore(): SourceRecordStore {
     return this.store;
   }
-
   protected async cleanupDroppedSourceTables(_tables: SourceTable[]) {
     // No-op for V1: source records live in a shared collection.
   }
 
+  protected override async batchCreateCustomWriteCheckpoints(
+    session: mongo.ClientSession,
+    opId: InternalOpId
+  ): Promise<void> {
+    await this.db.custom_write_checkpoints.bulkWrite(
+      this.write_checkpoint_batch.map((checkpointOptions) => {
+        const set: Record<string, unknown> = {
+          checkpoint: checkpointOptions.checkpoint,
+          sync_rules_id: checkpointOptions.sync_rules_id,
+          op_id: opId
+        };
+        if (checkpointOptions.checkpoint_requested_at != null) {
+          set.checkpoint_requested_at = checkpointOptions.checkpoint_requested_at;
+        }
+
+        return {
+          updateOne: {
+            filter: { user_id: checkpointOptions.user_id, sync_rules_id: checkpointOptions.sync_rules_id },
+            update: {
+              $set: set,
+              ...(checkpointOptions.checkpoint_requested_at == null ? { $unset: { checkpoint_requested_at: 1 } } : {})
+            },
+            upsert: true
+          }
+        };
+      }),
+      { session }
+    );
+  }
+
   async resolveTables(options: storage.ResolveTablesOptions): Promise<storage.ResolveTablesResult> {
-    const syncRules = options.syncRules ?? this.sync_rules;
+    const syncRules = options.parsedSyncConfig?.hydratedSyncConfig ?? this.sync_rules;
+    const reconcile = options.reconcileSourceTables ?? storage.defaultSourceTableReconciler;
     const { connection_id, source } = options;
-    const { schema, name, objectId, replicaIdColumns, connectionTag } = source;
+    const { schema, name, objectId, replicaIdColumns, connectionTag, sendsCompleteRows } = source;
 
     const normalizedReplicaIdColumns = replicaIdColumns.map((column) => ({
       name: column.name,
@@ -50,23 +95,41 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
 
     let result: storage.ResolveTablesResult | null = null;
     await this.db.client.withSession(async (session) => {
-      const col = this.db.commonSourceTables(this.group_id);
-      const filter: any = {
-        group_id: this.group_id,
-        connection_id,
-        schema_name: schema,
-        table_name: name,
-        replica_id_columns2: normalizedReplicaIdColumns
-      };
+      const col = this.db.sourceTablesV1(this.replicationStreamId);
+
+      // Find records that overlap by name or relation id.
+      const orClauses: Record<string, unknown>[] = [{ schema_name: schema, table_name: name }];
       if (objectId != null) {
-        filter.relation_id = objectId;
+        orClauses.push({ relation_id: objectId });
+      }
+      const candidateDocs = await col
+        .find({ group_id: this.replicationStreamId, connection_id, $or: orClauses }, { session })
+        .toArray();
+
+      const candidateTables = candidateDocs.map((doc) =>
+        this.hydrateSourceTable(doc, connectionTag, syncRules, sendsCompleteRows)
+      );
+      const candidates = candidateTables.map((table) => table.clone());
+      const resolution = await reconcile({ source, candidates });
+      storage.validateSourceTableCandidateResolution(candidates, resolution);
+
+      for (const { id, sourceMetadata } of storage.diffSourceTableUpdates(candidateTables, resolution)) {
+        await col.updateOne({ _id: mongoTableId(id) }, { $set: { source_metadata: sourceMetadata } }, { session });
       }
 
-      let doc = await col.findOne(filter, { session });
-      if (doc == null) {
-        doc = {
+      const materializedResolution = storage.materializeSourceTableResolution(candidateTables, resolution);
+      const compatibleTables = materializedResolution.compatibleTables;
+
+      // V1 keeps one record, preferring one that has already been snapshotted.
+      const reuse = compatibleTables.find((table) => table.snapshotComplete) ?? compatibleTables[0] ?? null;
+
+      let sourceTable: storage.SourceTable;
+      if (reuse != null) {
+        sourceTable = reuse;
+      } else {
+        const doc: SourceTableDocumentV1 = {
           _id: options.idGenerator ? (options.idGenerator() as bson.ObjectId) : new bson.ObjectId(),
-          group_id: this.group_id,
+          group_id: this.replicationStreamId,
           connection_id,
           relation_id: objectId,
           schema_name: schema,
@@ -74,68 +137,18 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
           replica_id_columns: null,
           replica_id_columns2: normalizedReplicaIdColumns,
           snapshot_done: false,
-          snapshot_status: undefined
+          snapshot_status: undefined,
+          source_metadata: resolution.newTableValues.sourceMetadata
         };
         await col.insertOne(doc, { session });
+        sourceTable = this.hydrateSourceTable(doc, connectionTag, syncRules, sendsCompleteRows);
       }
 
-      const sourceTable = new storage.SourceTable({
-        id: doc._id,
-        ref: source,
-        objectId,
-        replicaIdColumns,
-        snapshotComplete: doc.snapshot_done ?? true,
-        ...syncRules.getMatchingSources(source)
-      });
-      sourceTable.syncEvent = syncRules.tableTriggersEvent(source);
-      sourceTable.syncData = sourceTable.bucketDataSources.length > 0;
-      sourceTable.syncParameters = sourceTable.parameterLookupSources.length > 0;
-      sourceTable.snapshotStatus =
-        doc.snapshot_status == null
-          ? undefined
-          : {
-              lastKey: doc.snapshot_status.last_key?.buffer ?? null,
-              totalEstimatedCount: doc.snapshot_status.total_estimated_count,
-              replicatedCount: doc.snapshot_status.replicated_count
-            };
-
-      const truncateFilter = [{ schema_name: schema, table_name: name }] as any[];
-      if (objectId != null) {
-        truncateFilter.push({ relation_id: objectId });
-      }
-      const truncate = await col
-        .find(
-          {
-            group_id: this.group_id,
-            connection_id,
-            _id: { $ne: doc._id },
-            $or: truncateFilter
-          },
-          { session }
-        )
-        .toArray();
-      const dropTables = truncate.map((dropDoc) => {
-        const ref = {
-          connectionTag,
-          schema: dropDoc.schema_name,
-          name: dropDoc.table_name
-        };
-        const table = new storage.SourceTable({
-          id: dropDoc._id,
-          ref,
-          objectId: dropDoc.relation_id,
-          replicaIdColumns:
-            dropDoc.replica_id_columns2?.map(
-              (c) => ({ name: c.name, typeId: c.type_oid, type: c.type }) satisfies ColumnDescriptor
-            ) ?? [],
-          snapshotComplete: dropDoc.snapshot_done ?? true,
-          ...syncRules.getMatchingSources(ref)
-        });
-        table.syncEvent = syncRules.tableTriggersEvent(ref);
-        table.syncData = table.bucketDataSources.length > 0;
-        table.syncParameters = table.parameterLookupSources.length > 0;
-        return table;
-      });
+      const dropTables = [
+        ...materializedResolution.incompatibleTables,
+        // V1 only keeps one SourceTable per physical table.
+        ...compatibleTables.filter((table) => !storage.sourceTableIdEquals(table.id, sourceTable.id))
+      ];
 
       result = { tables: [sourceTable], dropTables };
     });
@@ -143,38 +156,33 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
     return result!;
   }
 
-  async getSourceTableStatus(table: storage.SourceTable): Promise<storage.SourceTable | null> {
-    const doc = (await this.db.commonSourceTables(this.group_id).findOne(
-      {
-        group_id: this.group_id,
-        _id: mongoTableId(table.id)
-      },
-      { session: this.session }
-    )) as SourceTableDocumentV1 | null;
-    if (doc == null) {
-      return null;
-    }
-
-    const ref = {
-      connectionTag: table.ref.connectionTag,
-      schema: doc.schema_name,
-      name: doc.table_name
-    };
-    const sourceTable = new storage.SourceTable({
+  /**
+   * Build a SourceTable from its persisted v1 record.
+   */
+  private hydrateSourceTable(
+    doc: SourceTableDocumentV1,
+    connectionTag: string,
+    syncRules: HydratedSyncConfig,
+    sendsCompleteRows: boolean | undefined
+  ): storage.SourceTable {
+    const ref = { connectionTag, schema: doc.schema_name, name: doc.table_name };
+    const table = new storage.SourceTable({
       id: doc._id,
       ref,
       objectId: doc.relation_id,
       replicaIdColumns:
         doc.replica_id_columns2?.map(
-          (c) => ({ name: c.name, typeId: c.type_oid, type: c.type }) satisfies ColumnDescriptor
+          (column) => ({ name: column.name, typeId: column.type_oid, type: column.type }) satisfies ColumnDescriptor
         ) ?? [],
       snapshotComplete: doc.snapshot_done ?? true,
-      ...this.sync_rules.getMatchingSources(ref)
+      sourceMetadata: doc.source_metadata ?? null,
+      ...syncRules.getMatchingSources(ref)
     });
-    sourceTable.syncEvent = this.sync_rules.tableTriggersEvent(ref);
-    sourceTable.syncData = sourceTable.bucketDataSources.length > 0;
-    sourceTable.syncParameters = sourceTable.parameterLookupSources.length > 0;
-    sourceTable.snapshotStatus =
+    table.syncEvent = syncRules.tableTriggersEvent(ref);
+    table.syncData = table.bucketDataSources.length > 0;
+    table.syncParameters = table.parameterLookupSources.length > 0;
+    table.storeCurrentData = sendsCompleteRows !== true;
+    table.snapshotStatus =
       doc.snapshot_status == null
         ? undefined
         : {
@@ -182,7 +190,23 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
             totalEstimatedCount: doc.snapshot_status.total_estimated_count,
             replicatedCount: doc.snapshot_status.replicated_count
           };
-    return sourceTable;
+    return table;
+  }
+
+  async getSourceTableStatus(table: storage.SourceTable): Promise<storage.SourceTable | null> {
+    const doc = await this.db.sourceTablesV1(this.replicationStreamId).findOne(
+      {
+        group_id: this.replicationStreamId,
+        _id: mongoTableId(table.id)
+      },
+      { session: this.session }
+    );
+    if (doc == null) {
+      return null;
+    }
+
+    // Use the default storeCurrentData value when reading outside resolution.
+    return this.hydrateSourceTable(doc, table.ref.connectionTag, this.sync_rules, undefined);
   }
 
   async commit(lsn: string, options?: storage.BucketBatchCommitOptions): Promise<storage.CheckpointResult> {
@@ -190,6 +214,8 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
 
     await this.flush(options);
 
+    using _ = this.tracer.span('storage', 'commit');
+
     const now = new Date();
 
     await this.db.write_checkpoints.updateMany(
@@ -242,7 +268,7 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
     };
 
     const preUpdateDocument = (await this.db.sync_rules.findOneAndUpdate(
-      { _id: this.group_id },
+      { _id: this.replicationStreamId },
       [
         {
           $set: {
@@ -295,7 +321,7 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
 
     if (preUpdateDocument == null) {
       throw new ReplicationAssertionError(
-        'Failed to update checkpoint - no matching sync_rules document for _id: ' + this.group_id
+        'Failed to update checkpoint - no matching sync_rules document for _id: ' + this.replicationStreamId
       );
     }
 
@@ -329,7 +355,6 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
       await this.autoActivate(lsn);
       await this.db.notifyCheckpoint();
       this.persisted_op = null;
-      this.last_checkpoint_lsn = lsn;
       if (checkpointState.newLastCheckpoint != null) {
         await this.sourceRecordStore.postCommitCleanup(checkpointState.newLastCheckpoint, this.logger);
       }
@@ -346,23 +371,30 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
   }
 
   async setResumeLsn(lsn: string): Promise<void> {
+    using _ = this.tracer.span('storage', 'set_resume_lsn');
     await this.db.sync_rules.updateOne(
       {
-        _id: this.group_id
+        _id: this.replicationStreamId
       },
       {
         $set: {
           snapshot_lsn: lsn
         }
       },
-      { session: this.session }
+      {
+        session: this.session,
+        // Losing occasional resume LSN is fine. That may mean reprocessing
+        // some source changes in some edge cases, which is not an issue since
+        // changes are processed in an idempotent way.
+        writeConcern: { w: 1 }
+      }
     );
   }
 
   async markAllSnapshotDone(no_checkpoint_before_lsn: string): Promise<void> {
     await this.db.sync_rules.updateOne(
       {
-        _id: this.group_id
+        _id: this.replicationStreamId
       },
       {
         $set: {
@@ -380,17 +412,24 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
   async markSnapshotDone(no_checkpoint_before_lsn: string, options?: { throwOnConflict?: boolean }): Promise<void> {
     await this.withTransaction(async () => {
       // Protect against race conditions
-      const count = await this.db.commonSourceTables(this.group_id).countDocuments(
-        {
-          group_id: this.group_id,
-          snapshot_done: false
-        },
-        { session: this.session }
-      );
-      if (count > 0) {
+      const blockingTables = await this.db
+        .sourceTablesV1(this.replicationStreamId)
+        .find(
+          {
+            group_id: this.replicationStreamId,
+            snapshot_done: false
+          },
+          {
+            session: this.session,
+            projection: { schema_name: 1, table_name: 1 }
+          }
+        )
+        .toArray();
+
+      if (blockingTables.length > 0) {
         if (options?.throwOnConflict ?? true) {
           throw new ReplicationAssertionError(
-            `Cannot mark snapshot done while ${count} source table${count == 1 ? '' : 's'} still require snapshotting`
+            `Cannot mark snapshot done while source tables still require snapshotting. Tables: ${blockingTables.map((t) => `${t.schema_name}.${t.table_name}`).join(', ')}`
           );
         } else {
           return;
@@ -404,7 +443,7 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
   async markTableSnapshotRequired(_table: storage.SourceTable): Promise<void> {
     await this.db.sync_rules.updateOne(
       {
-        _id: this.group_id
+        _id: this.replicationStreamId
       },
       {
         $set: {
@@ -423,7 +462,7 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
     const ids = tables.map((table) => mongoTableId(table.id));
 
     await this.withTransaction(async () => {
-      await this.db.commonSourceTables(this.group_id).updateMany(
+      await this.db.sourceTablesV1(this.replicationStreamId).updateMany(
         { _id: { $in: ids } },
         {
           $set: {
@@ -439,7 +478,7 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
       if (no_checkpoint_before_lsn != null) {
         await this.db.sync_rules.updateOne(
           {
-            _id: this.group_id
+            _id: this.replicationStreamId
           },
           {
             $set: {
@@ -467,12 +506,20 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
 
     const session = this.session;
     let activated = false;
+    let needsFutureActivationCheck = true;
     await session.withTransaction(async () => {
-      const doc = (await this.db.sync_rules.findOne({ _id: this.group_id }, { session })) as SyncRuleDocumentV1;
+      // Reset on transaction retries.
+      activated = false;
+      needsFutureActivationCheck = true;
+
+      const doc = (await this.db.sync_rules.findOne(
+        { _id: this.replicationStreamId },
+        { session }
+      )) as SyncRuleDocumentV1;
       if (doc && doc.state == storage.SyncRuleState.PROCESSING && doc.snapshot_done && doc.last_checkpoint != null) {
         await this.db.sync_rules.updateOne(
           {
-            _id: this.group_id
+            _id: this.replicationStreamId
           },
           {
             $set: {
@@ -484,7 +531,7 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
 
         await this.db.sync_rules.updateMany(
           {
-            _id: { $ne: this.group_id },
+            _id: { $ne: this.replicationStreamId },
             state: { $in: [storage.SyncRuleState.ACTIVE, storage.SyncRuleState.ERRORED] }
           },
           {
@@ -496,13 +543,15 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
         );
         activated = true;
       } else if (doc?.state != storage.SyncRuleState.PROCESSING) {
-        this.needsActivation = false;
+        needsFutureActivationCheck = false;
       }
     });
     if (activated) {
       this.logger.info(`Activated new replication stream at ${lsn}`);
       await this.db.notifyCheckpoint();
       this.needsActivation = false;
+    } else if (!needsFutureActivationCheck) {
+      this.needsActivation = false;
     }
   }
 }

package/src/storage/implementation/v1/MongoChecksumsV1.ts

@@ -1,20 +1,31 @@
+import * as lib_mongo from '@powersync/lib-service-mongodb';
 import {
+  addPartialChecksums,
   bson,
   BucketChecksum,
   FetchPartialBucketChecksum,
   InternalOpId,
-  PartialChecksumMap
+  isPartialChecksum,
+  PartialChecksumMap,
+  PartialOrFullChecksum
 } from '@powersync/service-core';
-import { FetchPartialBucketChecksumByBucket, MongoChecksums } from '../MongoChecksums.js';
+import {
+  checksumFromAggregate,
+  DEFAULT_OPERATION_BATCH_LIMIT,
+  FetchPartialBucketChecksumByBucket,
+  MongoChecksums,
+  MongoChecksumSessionContext
+} from '../MongoChecksums.js';
 import { VersionedPowerSyncMongoV1 } from './VersionedPowerSyncMongoV1.js';
-
 export class MongoChecksumsV1 extends MongoChecksums {
   declare protected readonly db: VersionedPowerSyncMongoV1;
 
   async computePartialChecksumsDirectByBucket(
-    batch: FetchPartialBucketChecksumByBucket[]
+    batch: FetchPartialBucketChecksumByBucket[],
+    context: MongoChecksumSessionContext
   ): Promise<PartialChecksumMap> {
-    return this.computePartialChecksumsForCollection(batch, this.db.bucketDataV1, (request) => ({
+    const collection = this.db.bucketDataV1;
+    const createFilter = (request: FetchPartialBucketChecksumByBucket) => ({
       _id: {
         $gt: {
           g: this.group_id,
@@ -27,11 +38,134 @@ export class MongoChecksumsV1 extends MongoChecksums {
           o: request.end
         }
       }
-    }));
+    });
+
+    const batchLimit = this.options?.operationBatchLimit ?? DEFAULT_OPERATION_BATCH_LIMIT;
+
+    // Map requests by bucket. We adjust this as we get partial results.
+    let requests = new Map<string, FetchPartialBucketChecksumByBucket>();
+    for (let request of batch) {
+      requests.set(request.bucket, request);
+    }
+
+    const partialChecksums = new Map<string, PartialOrFullChecksum>();
+
+    while (requests.size > 0) {
+      const filters = Array.from(requests.values(), createFilter);
+
+      // Historically, checksum may be stored as 'int' or 'double'.
+      // More recently, this should be a 'long'.
+      // $toLong ensures that we always sum it as a long, avoiding inaccuracies in the calculations.
+      const checksumLong = this.storageConfig.longChecksums ? '$checksum' : { $toLong: '$checksum' };
+
+      // Aggregate over a max of `batchLimit` operations at a time.
+      // Let's say we have 3 buckets (A, B, C), each with 10 operations, and our batch limit is 12.
+      // Then we'll do three batches:
+      // 1. Query: A[1-end], B[1-end], C[1-end]
+      //    Returns: A[1-10], B[1-2]
+      // 2. Query: B[3-end], C[1-end]
+      //    Returns: B[3-10], C[1-4]
+      // 3. Query: C[5-end]
+      //    Returns: C[5-10]
+      const aggregate = await collection
+        .aggregate(
+          [
+            {
+              $match: {
+                $or: filters
+              }
+            },
+            // sort and limit _before_ grouping
+            { $sort: { _id: 1 } },
+            { $limit: batchLimit },
+            {
+              $group: {
+                _id: '$_id.b',
+                checksum_total: { $sum: checksumLong },
+                count: { $sum: 1 },
+                has_clear_op: {
+                  $max: {
+                    $cond: [{ $eq: ['$op', 'CLEAR'] }, 1, 0]
+                  }
+                },
+                last_op: { $max: '$_id.o' }
+              }
+            },
+            // Sort the aggregated results (100 max, so should be fast).
+            // This is important to identify which buckets we have partial data for.
+            { $sort: { _id: 1 } }
+          ],
+          {
+            ...context.readOptions,
+            maxTimeMS: lib_mongo.MONGO_CHECKSUM_TIMEOUT_MS
+          }
+        )
+        .toArray()
+        .catch((e) => {
+          throw lib_mongo.mapQueryError(e, 'while reading checksums');
+        });
+
+      let batchCount = 0;
+      let limitReached = false;
+      for (let doc of aggregate) {
+        const bucket = doc._id;
+        const checksum = checksumFromAggregate(doc);
+
+        const existing = partialChecksums.get(bucket);
+        if (existing != null) {
+          partialChecksums.set(bucket, addPartialChecksums(bucket, existing, checksum));
+        } else {
+          partialChecksums.set(bucket, checksum);
+        }
+
+        batchCount += doc.count;
+        if (batchCount == batchLimit) {
+          // Limit reached. Request more in the next batch.
+          // Note that this only affects the _last_ bucket in a batch.
+          limitReached = true;
+          const req = requests.get(bucket);
+          requests.set(bucket, {
+            ...req!,
+            start: doc.last_op
+          });
+        } else {
+          // All done for this bucket
+          requests.delete(bucket);
+        }
+      }
+      if (!limitReached) {
+        break;
+      }
+    }
+
+    return new Map<string, PartialOrFullChecksum>(
+      batch.map((request) => {
+        const bucket = request.bucket;
+        let partialChecksum = partialChecksums.get(bucket);
+        // Could be null if we got no data
+        if (partialChecksum == null) {
+          partialChecksum = {
+            bucket,
+            partialCount: 0,
+            partialChecksum: 0
+          };
+        }
+        if (request.start == null && isPartialChecksum(partialChecksum)) {
+          partialChecksum = {
+            bucket,
+            count: partialChecksum.partialCount,
+            checksum: partialChecksum.partialChecksum
+          };
+        }
+
+        return [bucket, partialChecksum];
+      })
+    );
   }
 
   protected async fetchPreStates(
-    batch: FetchPartialBucketChecksum[]
+    batch: FetchPartialBucketChecksum[],
+    context: MongoChecksumSessionContext
   ): Promise<Map<string, { opId: InternalOpId; checksum: BucketChecksum }>> {
     const preFilters = batch
       .filter((request) => request.start == null)
@@ -49,9 +183,14 @@ export class MongoChecksumsV1 extends MongoChecksums {
     }
 
     const states = await this.db.bucketStateV1
-      .find({
-        $or: preFilters
-      })
+      .find(
+        {
+          $or: preFilters
+        },
+        {
+          ...context.readOptions
+        }
+      )
       .toArray();
 
     for (const state of states) {
@@ -69,7 +208,10 @@ export class MongoChecksumsV1 extends MongoChecksums {
     return preStates;
   }
 
-  protected async computePartialChecksumsInternal(batch: FetchPartialBucketChecksum[]): Promise<PartialChecksumMap> {
-    return this.computePartialChecksumsDirectByBucket(batch);
+  protected async computePartialChecksumsInternal(
+    batch: FetchPartialBucketChecksum[],
+    context: MongoChecksumSessionContext
+  ): Promise<PartialChecksumMap> {
+    return this.computePartialChecksumsDirectByBucket(batch, context);
   }
 }