@powersync/service-module-mongodb 0.0.0-dev-20250122110924 → 0.0.0-dev-20250227082606

package/src/replication/ChangeStream.ts

@@ -1,6 +1,7 @@
 import { mongo } from '@powersync/lib-service-mongodb';
 import {
   container,
+  DatabaseConnectionError,
   ErrorCode,
   logger,
   ReplicationAbortedError,
@@ -9,20 +10,13 @@ import {
 } from '@powersync/lib-services-framework';
 import { Metrics, SaveOperationTag, SourceEntityDescriptor, SourceTable, storage } from '@powersync/service-core';
 import { DatabaseInputRow, SqliteRow, SqlSyncRules, TablePattern } from '@powersync/service-sync-rules';
+import { MongoLSN } from '../common/MongoLSN.js';
 import { PostImagesOption } from '../types/types.js';
 import { escapeRegExp } from '../utils.js';
 import { MongoManager } from './MongoManager.js';
-import {
-  constructAfterRecord,
-  createCheckpoint,
-  getMongoLsn,
-  getMongoRelation,
-  mongoLsnToTimestamp
-} from './MongoRelation.js';
+import { constructAfterRecord, createCheckpoint, getCacheIdentifier, getMongoRelation } from './MongoRelation.js';
 import { CHECKPOINTS_COLLECTION } from './replication-utils.js';
 
-export const ZERO_LSN = '0000000000000000';
-
 export interface ChangeStreamOptions {
   connections: MongoManager;
   storage: storage.SyncRulesBucketStorage;
@@ -41,9 +35,9 @@ interface InitResult {
  *  * Some change stream documents do not have postImages.
  *  * startAfter/resumeToken is not valid anymore.
  */
-export class ChangeStreamInvalidatedError extends Error {
-  constructor(message: string) {
-    super(message);
+export class ChangeStreamInvalidatedError extends DatabaseConnectionError {
+  constructor(message: string, cause: any) {
+    super(ErrorCode.PSYNC_S1344, message, cause);
   }
 }
 
@@ -95,6 +89,10 @@ export class ChangeStream {
     return this.connections.options.postImages == PostImagesOption.AUTO_CONFIGURE;
   }
 
+  private get logPrefix() {
+    return `[powersync_${this.group_id}]`;
+  }
+
   /**
    * This resolves a pattern, persists the related metadata, and returns
    * the resulting SourceTables.
@@ -130,18 +128,13 @@ export class ChangeStream {
       .toArray();
 
     if (!tablePattern.isWildcard && collections.length == 0) {
-      logger.warn(`Collection ${schema}.${tablePattern.name} not found`);
+      logger.warn(`${this.logPrefix} Collection ${schema}.${tablePattern.name} not found`);
     }
 
     for (let collection of collections) {
       const table = await this.handleRelation(
         batch,
-        {
-          name: collection.name,
-          schema,
-          objectId: collection.name,
-          replicationColumns: [{ name: '_id' }]
-        } as SourceEntityDescriptor,
+        getMongoRelation({ db: schema, coll: collection.name }),
         // This is done as part of the initial setup - snapshot is handled elsewhere
         { snapshot: false, collectionInfo: collection }
       );
@@ -155,7 +148,7 @@ export class ChangeStream {
   async initSlot(): Promise<InitResult> {
     const status = await this.storage.getStatus();
     if (status.snapshot_done && status.checkpoint_lsn) {
-      logger.info(`Initial replication already done`);
+      logger.info(`${this.logPrefix} Initial replication already done`);
       return { needsInitialSync: false };
     }
 
@@ -164,7 +157,7 @@ export class ChangeStream {
 
   async estimatedCount(table: storage.SourceTable): Promise<string> {
     const db = this.client.db(table.schema);
-    const count = db.collection(table.table).estimatedDocumentCount();
+    const count = await db.collection(table.table).estimatedDocumentCount();
     return `~${count}`;
   }
 
@@ -207,7 +200,7 @@ export class ChangeStream {
     const session = await this.client.startSession();
     try {
       await this.storage.startBatch(
-        { zeroLSN: ZERO_LSN, defaultSchema: this.defaultDb.databaseName, storeCurrentData: false },
+        { zeroLSN: MongoLSN.ZERO.comparable, defaultSchema: this.defaultDb.databaseName, storeCurrentData: false },
         async (batch) => {
           // Start by resolving all tables.
           // This checks postImage configuration, and that should fail as
@@ -220,13 +213,13 @@ export class ChangeStream {
 
           for (let table of allSourceTables) {
             await this.snapshotTable(batch, table, session);
-            await batch.markSnapshotDone([table], ZERO_LSN);
+            await batch.markSnapshotDone([table], MongoLSN.ZERO.comparable);
 
             await touch();
           }
 
-          const lsn = getMongoLsn(snapshotTime);
-          logger.info(`Snapshot commit at ${snapshotTime.inspect()} / ${lsn}`);
+          const { comparable: lsn } = new MongoLSN({ timestamp: snapshotTime });
+          logger.info(`${this.logPrefix} Snapshot commit at ${snapshotTime.inspect()} / ${lsn}`);
           await batch.commit(lsn);
         }
       );
@@ -295,9 +288,10 @@ export class ChangeStream {
     table: storage.SourceTable,
     session?: mongo.ClientSession
   ) {
-    logger.info(`Replicating ${table.qualifiedName}`);
+    logger.info(`${this.logPrefix} Replicating ${table.qualifiedName}`);
     const estimatedCount = await this.estimatedCount(table);
     let at = 0;
+    let lastLogIndex = 0;
 
     const db = this.client.db(table.schema);
     const collection = db.collection(table.table);
@@ -310,8 +304,6 @@ export class ChangeStream {
         throw new ReplicationAbortedError(`Aborted initial replication`);
       }
 
-      at += 1;
-
       const record = constructAfterRecord(document);
 
       // This auto-flushes when the batch reaches its size limit
@@ -325,20 +317,26 @@ export class ChangeStream {
       });
 
       at += 1;
+      if (at - lastLogIndex >= 5000) {
+        logger.info(`${this.logPrefix}  Replicating ${table.qualifiedName} ${at}/${estimatedCount}`);
+        lastLogIndex = at;
+      }
       Metrics.getInstance().rows_replicated_total.add(1);
 
       await touch();
     }
 
     await batch.flush();
-    logger.info(`Replicated ${at} documents for ${table.qualifiedName}`);
+    logger.info(`${this.logPrefix} Replicated ${at} documents for ${table.qualifiedName}`);
   }
 
   private async getRelation(
     batch: storage.BucketStorageBatch,
-    descriptor: SourceEntityDescriptor
+    descriptor: SourceEntityDescriptor,
+    options: { snapshot: boolean }
   ): Promise<SourceTable> {
-    const existing = this.relation_cache.get(descriptor.objectId);
+    const cacheId = getCacheIdentifier(descriptor);
+    const existing = this.relation_cache.get(cacheId);
     if (existing != null) {
       return existing;
     }
@@ -347,7 +345,7 @@ export class ChangeStream {
     // missing values.
     const collection = await this.getCollectionInfo(descriptor.schema, descriptor.name);
 
-    return this.handleRelation(batch, descriptor, { snapshot: false, collectionInfo: collection });
+    return this.handleRelation(batch, descriptor, { snapshot: options.snapshot, collectionInfo: collection });
   }
 
   private async getCollectionInfo(db: string, name: string): Promise<mongo.CollectionInfo | undefined> {
@@ -378,7 +376,7 @@ export class ChangeStream {
         collMod: collectionInfo.name,
         changeStreamPreAndPostImages: { enabled: true }
       });
-      logger.info(`Enabled postImages on ${db}.${collectionInfo.name}`);
+      logger.info(`${this.logPrefix} Enabled postImages on ${db}.${collectionInfo.name}`);
     } else if (!enabled) {
       throw new ServiceError(ErrorCode.PSYNC_S1343, `postImages not enabled on ${db}.${collectionInfo.name}`);
     }
@@ -397,9 +395,6 @@ export class ChangeStream {
     }
 
     const snapshot = options.snapshot;
-    if (!descriptor.objectId && typeof descriptor.objectId != 'string') {
-      throw new ReplicationAssertionError('MongoDB replication - objectId expected');
-    }
     const result = await this.storage.resolveTable({
       group_id: this.group_id,
       connection_id: this.connection_id,
@@ -407,10 +402,16 @@ export class ChangeStream {
       entity_descriptor: descriptor,
       sync_rules: this.sync_rules
     });
-    this.relation_cache.set(descriptor.objectId, result.table);
+    this.relation_cache.set(getCacheIdentifier(descriptor), result.table);
 
-    // Drop conflicting tables. This includes for example renamed tables.
-    await batch.drop(result.dropTables);
+    // Drop conflicting collections.
+    // This is generally not expected for MongoDB source dbs, so we log an error.
+    if (result.dropTables.length > 0) {
+      logger.error(
+        `Conflicting collections found for ${JSON.stringify(descriptor)}. Dropping: ${result.dropTables.map((t) => t.id).join(', ')}`
+      );
+      await batch.drop(result.dropTables);
+    }
 
     // Snapshot if:
     // 1. Snapshot is requested (false for initial snapshot, since that process handles it elsewhere)
@@ -418,6 +419,7 @@ export class ChangeStream {
     // 3. The table is used in sync rules.
     const shouldSnapshot = snapshot && !result.table.snapshotComplete && result.table.syncAny;
     if (shouldSnapshot) {
+      logger.info(`${this.logPrefix} New collection: ${descriptor.schema}.${descriptor.name}`);
       // Truncate this table, in case a previous snapshot was interrupted.
       await batch.truncate([result.table]);
 
@@ -437,7 +439,7 @@ export class ChangeStream {
     change: mongo.ChangeStreamDocument
   ): Promise<storage.FlushedResult | null> {
     if (!table.syncAny) {
-      logger.debug(`Collection ${table.qualifiedName} not used in sync rules - skipping`);
+      logger.debug(`${this.logPrefix} Collection ${table.qualifiedName} not used in sync rules - skipping`);
       return null;
     }
 
@@ -513,7 +515,7 @@ export class ChangeStream {
         e.codeName == 'NoMatchingDocument' &&
         e.errmsg?.includes('post-image was not found')
       ) {
-        throw new ChangeStreamInvalidatedError(e.errmsg);
+        throw new ChangeStreamInvalidatedError(e.errmsg, e);
       }
       throw e;
     }
@@ -524,11 +526,14 @@ export class ChangeStream {
     await this.storage.autoActivate();
 
     await this.storage.startBatch(
-      { zeroLSN: ZERO_LSN, defaultSchema: this.defaultDb.databaseName, storeCurrentData: false },
+      { zeroLSN: MongoLSN.ZERO.comparable, defaultSchema: this.defaultDb.databaseName, storeCurrentData: false },
       async (batch) => {
-        const lastLsn = batch.lastCheckpointLsn;
-        const startAfter = mongoLsnToTimestamp(lastLsn) ?? undefined;
-        logger.info(`Resume streaming at ${startAfter?.inspect()} / ${lastLsn}`);
+        const { lastCheckpointLsn } = batch;
+        const lastLsn = lastCheckpointLsn ? MongoLSN.fromSerialized(lastCheckpointLsn) : null;
+        const startAfter = lastLsn?.timestamp;
+        const resumeAfter = lastLsn?.resumeToken;
+
+        logger.info(`${this.logPrefix} Resume streaming at ${startAfter?.inspect()} / ${lastLsn}`);
 
         const filters = this.getSourceNamespaceFilters();
 
@@ -551,12 +556,20 @@ export class ChangeStream {
         }
 
         const streamOptions: mongo.ChangeStreamOptions = {
-          startAtOperationTime: startAfter,
           showExpandedEvents: true,
-          useBigInt64: true,
           maxAwaitTimeMS: 200,
           fullDocument: fullDocument
         };
+
+        /**
+         * Only one of these options can be supplied at a time.
+         */
+        if (resumeAfter) {
+          streamOptions.resumeAfter = resumeAfter;
+        } else {
+          streamOptions.startAtOperationTime = startAfter;
+        }
+
         let stream: mongo.ChangeStream<mongo.Document>;
         if (filters.multipleDatabases) {
           // Requires readAnyDatabase@admin on Atlas
@@ -576,18 +589,24 @@ export class ChangeStream {
         });
 
         // Always start with a checkpoint.
-        // This helps us to clear erorrs when restarting, even if there is
+        // This helps us to clear errors when restarting, even if there is
         // no data to replicate.
         let waitForCheckpointLsn: string | null = await createCheckpoint(this.client, this.defaultDb);
 
         let splitDocument: mongo.ChangeStreamDocument | null = null;
 
+        let flexDbNameWorkaroundLogged = false;
+
         while (true) {
           if (this.abort_signal.aborted) {
             break;
           }
 
           const originalChangeDocument = await stream.tryNext();
+          // The stream was closed, we will only ever receive `null` from it
+          if (!originalChangeDocument && stream.closed) {
+            break;
+          }
 
           if (originalChangeDocument == null || this.abort_signal.aborted) {
             continue;
@@ -623,15 +642,61 @@ export class ChangeStream {
             throw new ReplicationAssertionError(`Incomplete splitEvent: ${JSON.stringify(splitDocument.splitEvent)}`);
           }
 
-          // console.log('event', changeDocument);
+          if (
+            !filters.multipleDatabases &&
+            'ns' in changeDocument &&
+            changeDocument.ns.db != this.defaultDb.databaseName &&
+            changeDocument.ns.db.endsWith(`_${this.defaultDb.databaseName}`)
+          ) {
+            // When all of the following conditions are met:
+            // 1. We're replicating from an Atlas Flex instance.
+            // 2. There were changestream events recorded while the PowerSync service is paused.
+            // 3. We're only replicating from a single database.
+            // Then we've observed an ns with for example {db: '67b83e86cd20730f1e766dde_ps'},
+            // instead of the expected {db: 'ps'}.
+            // We correct this.
+            changeDocument.ns.db = this.defaultDb.databaseName;
+
+            if (!flexDbNameWorkaroundLogged) {
+              flexDbNameWorkaroundLogged = true;
+              logger.warn(
+                `${this.logPrefix} Incorrect DB name in change stream: ${changeDocument.ns.db}. Changed to ${this.defaultDb.databaseName}.`
+              );
+            }
+          }
 
           if (
             (changeDocument.operationType == 'insert' ||
               changeDocument.operationType == 'update' ||
-              changeDocument.operationType == 'replace') &&
+              changeDocument.operationType == 'replace' ||
+              changeDocument.operationType == 'drop') &&
             changeDocument.ns.coll == CHECKPOINTS_COLLECTION
           ) {
-            const lsn = getMongoLsn(changeDocument.clusterTime!);
+            /**
+             * Dropping the database does not provide an `invalidate` event.
+             * We typically would receive `drop` events for the collection which we
+             * would process below.
+             *
+             * However we don't commit the LSN after collections are dropped.
+             * The prevents the `startAfter` or `resumeToken` from advancing past the drop events.
+             * The stream also closes after the drop events.
+             * This causes an infinite loop of processing the collection drop events.
+             *
+             * This check here invalidates the change stream if our `_checkpoints` collection
+             * is dropped. This allows for detecting when the DB is dropped.
+             */
+            if (changeDocument.operationType == 'drop') {
+              throw new ChangeStreamInvalidatedError(
+                'Internal collections have been dropped',
+                new Error('_checkpoints collection was dropped')
+              );
+            }
+
+            const { comparable: lsn } = new MongoLSN({
+              timestamp: changeDocument.clusterTime!,
+              resume_token: changeDocument._id
+            });
+
             if (waitForCheckpointLsn != null && lsn >= waitForCheckpointLsn) {
               waitForCheckpointLsn = null;
             }
@@ -646,28 +711,44 @@ export class ChangeStream {
               waitForCheckpointLsn = await createCheckpoint(this.client, this.defaultDb);
             }
             const rel = getMongoRelation(changeDocument.ns);
-            const table = await this.getRelation(batch, rel);
+            const table = await this.getRelation(batch, rel, {
+              // In most cases, we should not need to snapshot this. But if this is the first time we see the collection
+              // for whatever reason, then we do need to snapshot it.
+              // This may result in some duplicate operations when a collection is created for the first time after
+              // sync rules was deployed.
+              snapshot: true
+            });
             if (table.syncAny) {
               await this.writeChange(batch, table, changeDocument);
             }
           } else if (changeDocument.operationType == 'drop') {
             const rel = getMongoRelation(changeDocument.ns);
-            const table = await this.getRelation(batch, rel);
+            const table = await this.getRelation(batch, rel, {
+              // We're "dropping" this collection, so never snapshot it.
+              snapshot: false
+            });
             if (table.syncAny) {
               await batch.drop([table]);
-              this.relation_cache.delete(table.objectId);
+              this.relation_cache.delete(getCacheIdentifier(rel));
             }
           } else if (changeDocument.operationType == 'rename') {
             const relFrom = getMongoRelation(changeDocument.ns);
             const relTo = getMongoRelation(changeDocument.to);
-            const tableFrom = await this.getRelation(batch, relFrom);
+            const tableFrom = await this.getRelation(batch, relFrom, {
+              // We're "dropping" this collection, so never snapshot it.
+              snapshot: false
+            });
             if (tableFrom.syncAny) {
               await batch.drop([tableFrom]);
-              this.relation_cache.delete(tableFrom.objectId);
+              this.relation_cache.delete(getCacheIdentifier(relFrom));
             }
             // Here we do need to snapshot the new table
             const collection = await this.getCollectionInfo(relTo.schema, relTo.name);
-            await this.handleRelation(batch, relTo, { snapshot: true, collectionInfo: collection });
+            await this.handleRelation(batch, relTo, {
+              // This is a new (renamed) collection, so always snapshot it.
+              snapshot: true,
+              collectionInfo: collection
+            });
           }
         }
       }

package/src/replication/ChangeStreamReplicationJob.ts

@@ -1,4 +1,4 @@
-import { mongo } from '@powersync/lib-service-mongodb';
+import { isMongoServerError } from '@powersync/lib-service-mongodb';
 import { container } from '@powersync/lib-services-framework';
 import { replication } from '@powersync/service-core';
 
@@ -40,8 +40,8 @@ export class ChangeStreamReplicationJob extends replication.AbstractReplicationJ
       this.logger.error(`Replication failed`, e);
 
       if (e instanceof ChangeStreamInvalidatedError) {
-        // This stops replication on this slot, and creates a new slot
-        await this.options.storage.factory.slotRemoved(this.slotName);
+        // This stops replication and restarts with a new instance
+        await this.options.storage.factory.restartReplication(this.storage.group_id);
       }
     } finally {
       this.abortController.abort();
@@ -78,15 +78,15 @@ export class ChangeStreamReplicationJob extends replication.AbstractReplicationJ
       if (this.abortController.signal.aborted) {
         return;
       }
-      this.logger.error(`Replication error`, e);
+      this.logger.error(`${this.slotName} Replication error`, e);
       if (e.cause != null) {
         // Without this additional log, the cause may not be visible in the logs.
         this.logger.error(`cause`, e.cause);
       }
       if (e instanceof ChangeStreamInvalidatedError) {
         throw e;
-      } else if (e instanceof mongo.MongoError && e.hasErrorLabel('NonResumableChangeStreamError')) {
-        throw new ChangeStreamInvalidatedError(e.message);
+      } else if (isMongoServerError(e) && e.hasErrorLabel('NonResumableChangeStreamError')) {
+        throw new ChangeStreamInvalidatedError(e.message, e);
       } else {
         // Report the error if relevant, before retrying
         container.reporter.captureException(e, {

package/src/replication/MongoManager.ts

@@ -1,11 +1,9 @@
 import { mongo } from '@powersync/lib-service-mongodb';
 
 import { NormalizedMongoConnectionConfig } from '../types/types.js';
+import { BSON_DESERIALIZE_DATA_OPTIONS } from '@powersync/service-core';
 
 export class MongoManager {
-  /**
-   * Do not use this for any transactions.
-   */
   public readonly client: mongo.MongoClient;
   public readonly db: mongo.Db;
 
@@ -35,6 +33,9 @@ export class MongoManager {
 
       maxConnecting: 3,
       maxIdleTimeMS: 60_000,
+
+      ...BSON_DESERIALIZE_DATA_OPTIONS,
+
       ...overrides
     });
     this.db = this.client.db(options.database, {});

package/src/replication/MongoRelation.ts

@@ -3,31 +3,25 @@ import { storage } from '@powersync/service-core';
 import { JSONBig, JsonContainer } from '@powersync/service-jsonbig';
 import { SqliteRow, SqliteValue } from '@powersync/service-sync-rules';
 
-import { CHECKPOINTS_COLLECTION } from './replication-utils.js';
 import { ErrorCode, ServiceError } from '@powersync/lib-services-framework';
+import { MongoLSN } from '../common/MongoLSN.js';
+import { CHECKPOINTS_COLLECTION } from './replication-utils.js';
 
 export function getMongoRelation(source: mongo.ChangeStreamNameSpace): storage.SourceEntityDescriptor {
   return {
     name: source.coll,
     schema: source.db,
-    objectId: source.coll,
+    // Not relevant for MongoDB - we use db + coll name as the identifier
+    objectId: undefined,
     replicationColumns: [{ name: '_id' }]
   } satisfies storage.SourceEntityDescriptor;
 }
 
-export function getMongoLsn(timestamp: mongo.Timestamp) {
-  const a = timestamp.high.toString(16).padStart(8, '0');
-  const b = timestamp.low.toString(16).padStart(8, '0');
-  return a + b;
-}
-
-export function mongoLsnToTimestamp(lsn: string | null) {
-  if (lsn == null) {
-    return null;
-  }
-  const a = parseInt(lsn.substring(0, 8), 16);
-  const b = parseInt(lsn.substring(8, 16), 16);
-  return mongo.Timestamp.fromBits(b, a);
+/**
+ * For in-memory cache only.
+ */
+export function getCacheIdentifier(source: storage.SourceEntityDescriptor): string {
+  return `${source.schema}.${source.name}`;
 }
 
 export function constructAfterRecord(document: mongo.Document): SqliteRow {
@@ -174,7 +168,7 @@ export async function createCheckpoint(client: mongo.MongoClient, db: mongo.Db):
     );
     const time = session.operationTime!;
     // TODO: Use the above when we support custom write checkpoints
-    return getMongoLsn(time);
+    return new MongoLSN({ timestamp: time }).comparable;
   } finally {
     await session.endSession();
   }

package/test/src/change_stream.test.ts

@@ -239,6 +239,7 @@ bucket_definitions:
             - SELECT _id as id, description FROM "test_DATA"
       `);
 
+    await db.createCollection('test_DATA');
     await context.replicateSnapshot();
 
     context.startStreaming();
@@ -261,6 +262,7 @@ bucket_definitions:
           data:
             - SELECT _id as id, name, description FROM "test_data"
       `);
+    await db.createCollection('test_data');
 
     await context.replicateSnapshot();
     context.startStreaming();
@@ -371,6 +373,8 @@ bucket_definitions:
           - SELECT _id as id, name, other FROM "test_data"`);
     const { db } = context;
 
+    await db.createCollection('test_data');
+
     await context.replicateSnapshot();
 
     const collection = db.collection('test_data');
@@ -451,6 +455,8 @@ bucket_definitions:
 
     const data = await context.getBucketData('global[]');
     expect(data).toMatchObject([
+      // An extra op here, since this triggers a snapshot in addition to getting the event.
+      test_utils.putOp('test_data', { id: test_id!.toHexString(), description: 'test2' }),
       test_utils.putOp('test_data', { id: test_id!.toHexString(), description: 'test1' }),
       test_utils.putOp('test_data', { id: test_id!.toHexString(), description: 'test2' })
     ]);

package/test/src/change_stream_utils.ts

@@ -1,5 +1,5 @@
 import { mongo } from '@powersync/lib-service-mongodb';
-import { ActiveCheckpoint, BucketStorageFactory, OpId, SyncRulesBucketStorage } from '@powersync/service-core';
+import { BucketStorageFactory, OpId, ReplicationCheckpoint, SyncRulesBucketStorage } from '@powersync/service-core';
 import { test_utils } from '@powersync/service-core-tests';
 
 import { ChangeStream, ChangeStreamOptions } from '@module/replication/ChangeStream.js';
@@ -58,7 +58,7 @@ export class ChangeStreamTestContext {
   }
 
   async updateSyncRules(content: string) {
-    const syncRules = await this.factory.updateSyncRules({ content: content });
+    const syncRules = await this.factory.updateSyncRules({ content: content, validate: true });
     this.storage = this.factory.getInstance(syncRules);
     return this.storage!;
   }
@@ -85,7 +85,7 @@ export class ChangeStreamTestContext {
   }
 
   startStreaming() {
-    this.streamPromise = this.walStream.streamChanges();
+    return (this.streamPromise = this.walStream.streamChanges());
   }
 
   async getCheckpoint(options?: { timeout?: number }) {
@@ -138,7 +138,7 @@ export class ChangeStreamTestContext {
 export async function getClientCheckpoint(
   client: mongo.MongoClient,
   db: mongo.Db,
-  bucketStorage: BucketStorageFactory,
+  storageFactory: BucketStorageFactory,
   options?: { timeout?: number }
 ): Promise<OpId> {
   const start = Date.now();
@@ -147,14 +147,15 @@ export async function getClientCheckpoint(
   // Since we don't use LSNs anymore, the only way to get that is to wait.
 
   const timeout = options?.timeout ?? 50_000;
-  let lastCp: ActiveCheckpoint | null = null;
+  let lastCp: ReplicationCheckpoint | null = null;
 
   while (Date.now() - start < timeout) {
-    const cp = await bucketStorage.getActiveCheckpoint();
-    lastCp = cp;
-    if (!cp.hasSyncRules()) {
+    const storage = await storageFactory.getActiveStorage();
+    const cp = await storage?.getCheckpoint();
+    if (cp == null) {
       throw new Error('No sync rules available');
     }
+    lastCp = cp;
     if (cp.lsn && cp.lsn >= lsn) {
       return cp.checkpoint;
     }