@powersync/service-module-mongodb 0.9.1 → 0.10.1

package/src/replication/ChangeStreamReplicationJob.ts

@@ -1,5 +1,4 @@
-import { isMongoServerError } from '@powersync/lib-service-mongodb';
-import { container } from '@powersync/lib-services-framework';
+import { container, logger as defaultLogger } from '@powersync/lib-services-framework';
 import { replication } from '@powersync/service-core';
 
 import { ChangeStream, ChangeStreamInvalidatedError } from './ChangeStream.js';
@@ -11,18 +10,21 @@ export interface ChangeStreamReplicationJobOptions extends replication.AbstractR
 
 export class ChangeStreamReplicationJob extends replication.AbstractReplicationJob {
   private connectionFactory: ConnectionManagerFactory;
+  private lastStream: ChangeStream | null = null;
 
   constructor(options: ChangeStreamReplicationJobOptions) {
     super(options);
     this.connectionFactory = options.connectionFactory;
+    // We use a custom formatter to process the prefix
+    this.logger = defaultLogger.child({ prefix: `[powersync_${this.storage.group_id}] ` });
   }
 
   async cleanUp(): Promise<void> {
-    // TODO: Implement?
+    // Nothing needed here
   }
 
   async keepAlive() {
-    // TODO: Implement?
+    // Nothing needed here
   }
 
   private get slotName() {
@@ -72,14 +74,16 @@ export class ChangeStreamReplicationJob extends replication.AbstractReplicationJ
         abort_signal: this.abortController.signal,
         storage: this.options.storage,
         metrics: this.options.metrics,
-        connections: connectionManager
+        connections: connectionManager,
+        logger: this.logger
       });
+      this.lastStream = stream;
       await stream.replicate();
     } catch (e) {
       if (this.abortController.signal.aborted) {
         return;
       }
-      this.logger.error(`${this.slotName} Replication error`, e);
+      this.logger.error(`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);
@@ -98,4 +102,8 @@ export class ChangeStreamReplicationJob extends replication.AbstractReplicationJ
       await connectionManager.end();
     }
   }
+
+  async getReplicationLagMillis(): Promise<number | undefined> {
+    return this.lastStream?.getReplicationLagMillis();
+  }
 }

package/src/replication/ChangeStreamReplicator.ts

@@ -3,6 +3,8 @@ import { ChangeStreamReplicationJob } from './ChangeStreamReplicationJob.js';
 import { ConnectionManagerFactory } from './ConnectionManagerFactory.js';
 import { MongoErrorRateLimiter } from './MongoErrorRateLimiter.js';
 import { MongoModule } from '../module/MongoModule.js';
+import { MongoLSN } from '../common/MongoLSN.js';
+import { timestampToDate } from './replication-utils.js';
 
 export interface ChangeStreamReplicatorOptions extends replication.AbstractReplicatorOptions {
   connectionFactory: ConnectionManagerFactory;
@@ -39,4 +41,25 @@ export class ChangeStreamReplicator extends replication.AbstractReplicator<Chang
   async testConnection() {
     return await MongoModule.testConnection(this.connectionFactory.dbConnectionConfig);
   }
+
+  async getReplicationLagMillis(): Promise<number | undefined> {
+    const lag = await super.getReplicationLagMillis();
+    if (lag != null) {
+      return lag;
+    }
+
+    // Booting or in an error loop. Check last active replication status.
+    // This includes sync rules in an ERROR state.
+    const content = await this.storage.getActiveSyncRulesContent();
+    if (content == null) {
+      return undefined;
+    }
+    // Measure the lag from the last resume token's time
+    const lsn = content.last_checkpoint_lsn;
+    if (lsn == null) {
+      return undefined;
+    }
+    const { timestamp } = MongoLSN.fromSerialized(lsn);
+    return Date.now() - timestampToDate(timestamp).getTime();
+  }
 }

package/src/replication/MongoRelation.ts

@@ -20,7 +20,10 @@ export function getMongoRelation(source: mongo.ChangeStreamNameSpace): storage.S
 /**
  * For in-memory cache only.
  */
-export function getCacheIdentifier(source: storage.SourceEntityDescriptor): string {
+export function getCacheIdentifier(source: storage.SourceEntityDescriptor | storage.SourceTable): string {
+  if (source instanceof storage.SourceTable) {
+    return `${source.schema}.${source.table}`;
+  }
   return `${source.schema}.${source.name}`;
 }
 

package/src/replication/MongoSnapshotQuery.ts

@@ -0,0 +1,59 @@
+import { mongo } from '@powersync/lib-service-mongodb';
+import { ReplicationAssertionError } from '@powersync/lib-services-framework';
+import { bson } from '@powersync/service-core';
+
+/**
+ * Performs a collection snapshot query, chunking by ranges of _id.
+ *
+ * This may miss some rows if they are modified during the snapshot query.
+ * In that case, the change stream replication will pick up those rows afterwards.
+ */
+export class ChunkedSnapshotQuery implements AsyncDisposable {
+  lastKey: any = null;
+  private lastCursor: mongo.FindCursor | null = null;
+  private collection: mongo.Collection;
+  private batchSize: number;
+
+  public constructor(options: { collection: mongo.Collection; batchSize: number; key?: Uint8Array | null }) {
+    this.lastKey = options.key ? bson.deserialize(options.key, { useBigInt64: true })._id : null;
+    this.lastCursor = null;
+    this.collection = options.collection;
+    this.batchSize = options.batchSize;
+  }
+
+  async nextChunk(): Promise<{ docs: mongo.Document[]; lastKey: Uint8Array } | { docs: []; lastKey: null }> {
+    let cursor = this.lastCursor;
+    let newCursor = false;
+    if (cursor == null || cursor.closed) {
+      const filter: mongo.Filter<mongo.Document> = this.lastKey == null ? {} : { _id: { $gt: this.lastKey as any } };
+      cursor = this.collection.find(filter, {
+        batchSize: this.batchSize,
+        readConcern: 'majority',
+        limit: this.batchSize,
+        sort: { _id: 1 }
+      });
+      newCursor = true;
+    }
+    const hasNext = await cursor.hasNext();
+    if (!hasNext) {
+      this.lastCursor = null;
+      if (newCursor) {
+        return { docs: [], lastKey: null };
+      } else {
+        return this.nextChunk();
+      }
+    }
+    const docBatch = cursor.readBufferedDocuments();
+    this.lastCursor = cursor;
+    if (docBatch.length == 0) {
+      throw new ReplicationAssertionError(`MongoDB snapshot query returned an empty batch, but hasNext() was true.`);
+    }
+    const lastKey = docBatch[docBatch.length - 1]._id;
+    this.lastKey = lastKey;
+    return { docs: docBatch, lastKey: bson.serialize({ _id: lastKey }) };
+  }
+
+  async [Symbol.asyncDispose](): Promise<void> {
+    await this.lastCursor?.close();
+  }
+}

package/src/replication/replication-utils.ts

@@ -1,6 +1,7 @@
 import { ErrorCode, ServiceError } from '@powersync/lib-services-framework';
 import { MongoManager } from './MongoManager.js';
 import { PostImagesOption } from '../types/types.js';
+import * as bson from 'bson';
 
 export const CHECKPOINTS_COLLECTION = '_powersync_checkpoints';
 
@@ -86,3 +87,7 @@ export async function checkSourceConfiguration(connectionManager: MongoManager):
       .toArray();
   }
 }
+
+export function timestampToDate(timestamp: bson.Timestamp) {
+  return new Date(timestamp.getHighBitsUnsigned() * 1000);
+}

package/test/src/change_stream.test.ts

@@ -8,8 +8,7 @@ import { test_utils } from '@powersync/service-core-tests';
 
 import { PostImagesOption } from '@module/types/types.js';
 import { ChangeStreamTestContext } from './change_stream_utils.js';
-import { env } from './env.js';
-import { INITIALIZED_MONGO_STORAGE_FACTORY, INITIALIZED_POSTGRES_STORAGE_FACTORY } from './util.js';
+import { describeWithStorage } from './util.js';
 
 const BASIC_SYNC_RULES = `
 bucket_definitions:
@@ -18,12 +17,8 @@ bucket_definitions:
       - SELECT _id as id, description FROM "test_data"
 `;
 
-describe.skipIf(!env.TEST_MONGO_STORAGE)('change stream - mongodb', { timeout: 20_000 }, function () {
-  defineChangeStreamTests(INITIALIZED_MONGO_STORAGE_FACTORY);
-});
-
-describe.skipIf(!env.TEST_POSTGRES_STORAGE)('change stream - postgres', { timeout: 20_000 }, function () {
-  defineChangeStreamTests(INITIALIZED_POSTGRES_STORAGE_FACTORY);
+describe('change stream', () => {
+  describeWithStorage({ timeout: 20_000 }, defineChangeStreamTests);
 });
 
 function defineChangeStreamTests(factory: storage.TestStorageFactory) {
@@ -97,7 +92,9 @@ bucket_definitions:
   });
 
   test('updateLookup - no fullDocument available', async () => {
-    await using context = await ChangeStreamTestContext.open(factory, { postImages: PostImagesOption.OFF });
+    await using context = await ChangeStreamTestContext.open(factory, {
+      mongoOptions: { postImages: PostImagesOption.OFF }
+    });
     const { db, client } = context;
     await context.updateSyncRules(`
 bucket_definitions:
@@ -141,7 +138,9 @@ bucket_definitions:
   test('postImages - autoConfigure', async () => {
     // Similar to the above test, but with postImages enabled.
     // This resolves the consistency issue.
-    await using context = await ChangeStreamTestContext.open(factory, { postImages: PostImagesOption.AUTO_CONFIGURE });
+    await using context = await ChangeStreamTestContext.open(factory, {
+      mongoOptions: { postImages: PostImagesOption.AUTO_CONFIGURE }
+    });
     const { db, client } = context;
     await context.updateSyncRules(`
 bucket_definitions:
@@ -187,7 +186,9 @@ bucket_definitions:
   test('postImages - on', async () => {
     // Similar to postImages - autoConfigure, but does not auto-configure.
     // changeStreamPreAndPostImages must be manually configured.
-    await using context = await ChangeStreamTestContext.open(factory, { postImages: PostImagesOption.READ_ONLY });
+    await using context = await ChangeStreamTestContext.open(factory, {
+      mongoOptions: { postImages: PostImagesOption.READ_ONLY }
+    });
     const { db, client } = context;
     await context.updateSyncRules(`
 bucket_definitions:
@@ -432,7 +433,9 @@ bucket_definitions:
   });
 
   test('postImages - new collection with postImages enabled', async () => {
-    await using context = await ChangeStreamTestContext.open(factory, { postImages: PostImagesOption.AUTO_CONFIGURE });
+    await using context = await ChangeStreamTestContext.open(factory, {
+      mongoOptions: { postImages: PostImagesOption.AUTO_CONFIGURE }
+    });
     const { db } = context;
     await context.updateSyncRules(`
 bucket_definitions:
@@ -463,7 +466,9 @@ bucket_definitions:
   });
 
   test('postImages - new collection with postImages disabled', async () => {
-    await using context = await ChangeStreamTestContext.open(factory, { postImages: PostImagesOption.AUTO_CONFIGURE });
+    await using context = await ChangeStreamTestContext.open(factory, {
+      mongoOptions: { postImages: PostImagesOption.AUTO_CONFIGURE }
+    });
     const { db } = context;
     await context.updateSyncRules(`
 bucket_definitions:

package/test/src/change_stream_utils.ts

@@ -4,9 +4,11 @@ import {
   createCoreReplicationMetrics,
   initializeCoreReplicationMetrics,
   InternalOpId,
+  OplogEntry,
   ProtocolOpId,
   ReplicationCheckpoint,
-  SyncRulesBucketStorage
+  SyncRulesBucketStorage,
+  TestStorageOptions
 } from '@powersync/service-core';
 import { METRICS_HELPER, test_utils } from '@powersync/service-core-tests';
 
@@ -29,17 +31,27 @@ export class ChangeStreamTestContext {
    *
    * This configures all the context, and tears it down afterwards.
    */
-  static async open(factory: () => Promise<BucketStorageFactory>, options?: Partial<NormalizedMongoConnectionConfig>) {
-    const f = await factory();
-    const connectionManager = new MongoManager({ ...TEST_CONNECTION_OPTIONS, ...options });
+  static async open(
+    factory: (options: TestStorageOptions) => Promise<BucketStorageFactory>,
+    options?: {
+      doNotClear?: boolean;
+      mongoOptions?: Partial<NormalizedMongoConnectionConfig>;
+      streamOptions?: Partial<ChangeStreamOptions>;
+    }
+  ) {
+    const f = await factory({ doNotClear: options?.doNotClear });
+    const connectionManager = new MongoManager({ ...TEST_CONNECTION_OPTIONS, ...options?.mongoOptions });
 
-    await clearTestDb(connectionManager.db);
-    return new ChangeStreamTestContext(f, connectionManager);
+    if (!options?.doNotClear) {
+      await clearTestDb(connectionManager.db);
+    }
+    return new ChangeStreamTestContext(f, connectionManager, options?.streamOptions);
   }
 
   constructor(
     public factory: BucketStorageFactory,
-    public connectionManager: MongoManager
+    public connectionManager: MongoManager,
+    private streamOptions?: Partial<ChangeStreamOptions>
   ) {
     createCoreReplicationMetrics(METRICS_HELPER.metricsEngine);
     initializeCoreReplicationMetrics(METRICS_HELPER.metricsEngine);
@@ -74,6 +86,16 @@ export class ChangeStreamTestContext {
     return this.storage!;
   }
 
+  async loadNextSyncRules() {
+    const syncRules = await this.factory.getNextSyncRulesContent();
+    if (syncRules == null) {
+      throw new Error(`Next sync rules not available`);
+    }
+
+    this.storage = this.factory.getInstance(syncRules);
+    return this.storage!;
+  }
+
   get walStream() {
     if (this.storage == null) {
       throw new Error('updateSyncRules() first');
@@ -88,7 +110,8 @@ export class ChangeStreamTestContext {
       abort_signal: this.abortController.signal,
       // Specifically reduce this from the default for tests on MongoDB <= 6.0, otherwise it can take
       // a long time to abort the stream.
-      maxAwaitTimeMS: 200
+      maxAwaitTimeMS: this.streamOptions?.maxAwaitTimeMS ?? 200,
+      snapshotChunkLength: this.streamOptions?.snapshotChunkLength
     };
     this._walStream = new ChangeStream(options);
     return this._walStream!;
@@ -122,23 +145,25 @@ export class ChangeStreamTestContext {
     return test_utils.fromAsync(this.storage!.getBucketDataBatch(checkpoint, map));
   }
 
-  async getBucketData(
-    bucket: string,
-    start?: ProtocolOpId | InternalOpId | undefined,
-    options?: { timeout?: number; limit?: number; chunkLimitBytes?: number }
-  ) {
+  async getBucketData(bucket: string, start?: ProtocolOpId | InternalOpId | undefined, options?: { timeout?: number }) {
     start ??= 0n;
     if (typeof start == 'string') {
       start = BigInt(start);
     }
-    let checkpoint = await this.getCheckpoint(options);
+    const checkpoint = await this.getCheckpoint(options);
     const map = new Map<string, InternalOpId>([[bucket, start]]);
-    const batch = this.storage!.getBucketDataBatch(checkpoint, map, {
-      limit: options?.limit,
-      chunkLimitBytes: options?.chunkLimitBytes
-    });
-    const batches = await test_utils.fromAsync(batch);
-    return batches[0]?.chunkData.data ?? [];
+    let data: OplogEntry[] = [];
+    while (true) {
+      const batch = this.storage!.getBucketDataBatch(checkpoint, map);
+
+      const batches = await test_utils.fromAsync(batch);
+      data = data.concat(batches[0]?.chunkData.data ?? []);
+      if (batches.length == 0 || !batches[0]!.chunkData.has_more) {
+        break;
+      }
+      map.set(bucket, BigInt(batches[0]!.chunkData.next_after));
+    }
+    return data;
   }
 
   async getChecksums(buckets: string[], options?: { timeout?: number }) {

package/test/src/chunked_snapshot.test.ts

@@ -0,0 +1,153 @@
+import { mongo } from '@powersync/lib-service-mongodb';
+import { reduceBucket, TestStorageFactory } from '@powersync/service-core';
+import { METRICS_HELPER } from '@powersync/service-core-tests';
+import { JSONBig } from '@powersync/service-jsonbig';
+import { SqliteJsonValue } from '@powersync/service-sync-rules';
+import * as timers from 'timers/promises';
+import { describe, expect, test } from 'vitest';
+import { ChangeStreamTestContext } from './change_stream_utils.js';
+import { describeWithStorage } from './util.js';
+
+describe('chunked snapshots', () => {
+  describeWithStorage({ timeout: 120_000 }, defineBatchTests);
+});
+
+function defineBatchTests(factory: TestStorageFactory) {
+  // This is not as sensitive to the id type as postgres, but we still test a couple of cases
+  test('chunked snapshot (int32)', async () => {
+    await testChunkedSnapshot({
+      generateId(i) {
+        return i;
+      },
+      idToSqlite(id: number) {
+        return BigInt(id);
+      }
+    });
+  });
+
+  test('chunked snapshot (Timestamp)', async () => {
+    await testChunkedSnapshot({
+      generateId(i) {
+        return mongo.Timestamp.fromBits(Math.floor(i / 1000), i % 1000);
+      },
+      idToSqlite(id: mongo.Timestamp) {
+        return id.toBigInt();
+      }
+    });
+  });
+
+  test('chunked snapshot (compound)', async () => {
+    await testChunkedSnapshot({
+      generateId(i) {
+        return { a: Math.floor(i / 100), b: i % 100 };
+      },
+      idToSqlite(id: any) {
+        return JSON.stringify(id);
+      }
+    });
+  });
+
+  test('chunked snapshot (float)', async () => {
+    await testChunkedSnapshot({
+      generateId(i) {
+        // Floating-point operations are not exact, but it should be consistent at least
+        return i / Math.PI;
+      },
+      idToSqlite(id: any) {
+        return id;
+      }
+    });
+  });
+
+  async function testChunkedSnapshot(options: {
+    generateId: (i: number) => any;
+    idToSqlite?: (id: any) => SqliteJsonValue;
+  }) {
+    // This is not quite as much of an edge cases as with Postgres. We do still test that
+    // updates applied while replicating are applied correctly.
+    const idToSqlite = options.idToSqlite ?? ((n) => n);
+    const idToString = (id: any) => String(idToSqlite(id));
+
+    await using context = await ChangeStreamTestContext.open(factory, {
+      // We need to use a smaller chunk size here, so that we can run a query in between chunks
+      streamOptions: { snapshotChunkLength: 100 }
+    });
+
+    await context.updateSyncRules(`bucket_definitions:
+  global:
+    data:
+      - SELECT _id as id, description FROM test_data`);
+    const { db } = context;
+
+    let batch = db.collection('test_data').initializeUnorderedBulkOp();
+
+    // 1. Start with 2k rows...
+    for (let i = 1; i <= 2000; i++) {
+      batch.insert({ _id: options.generateId(i), description: 'foo' });
+    }
+    await batch.execute();
+
+    // 2. Replicate one batch of rows
+    // Our "stopping point" here is not quite deterministic.
+    const p = context.replicateSnapshot();
+
+    const stopAfter = 100;
+    const startRowCount = (await METRICS_HELPER.getMetricValueForTests('powersync_rows_replicated_total')) ?? 0;
+
+    while (true) {
+      const count =
+        ((await METRICS_HELPER.getMetricValueForTests('powersync_rows_replicated_total')) ?? 0) - startRowCount;
+
+      if (count >= stopAfter) {
+        break;
+      }
+      await timers.setTimeout(1);
+    }
+
+    // 3. Update some records
+    const idA = options.generateId(2000);
+    const idB = options.generateId(1);
+    await db.collection('test_data').updateOne({ _id: idA }, { $set: { description: 'bar' } });
+    await db.collection('test_data').updateOne({ _id: idB }, { $set: { description: 'baz' } });
+
+    // 4. Delete
+    const idC = options.generateId(1999);
+    await db.collection('test_data').deleteOne({ _id: idC });
+
+    // 5. Insert
+    const idD = options.generateId(2001);
+    await db.collection('test_data').insertOne({ _id: idD, description: 'new' });
+
+    // 4. Replicate the rest of the table.
+    await p;
+
+    context.startStreaming();
+
+    const data = await context.getBucketData('global[]');
+    const reduced = reduceBucket(data);
+
+    expect(reduced.find((row) => row.object_id == idToString(idA))?.data).toEqual(
+      JSONBig.stringify({
+        id: idToSqlite(idA),
+        description: 'bar'
+      })
+    );
+
+    expect(reduced.find((row) => row.object_id == idToString(idB))?.data).toEqual(
+      JSONBig.stringify({
+        id: idToSqlite(idB),
+        description: 'baz'
+      })
+    );
+
+    expect(reduced.find((row) => row.object_id == idToString(idC))).toBeUndefined();
+
+    expect(reduced.find((row) => row.object_id == idToString(idD))?.data).toEqual(
+      JSONBig.stringify({
+        id: idToSqlite(idD),
+        description: 'new'
+      })
+    );
+    expect(reduced.length).toEqual(2001);
+  }
+}

package/test/src/resume.test.ts

@@ -1,96 +1,14 @@
-import { MongoLSN, ZERO_LSN } from '@module/common/MongoLSN.js';
-
+import { ChangeStreamInvalidatedError } from '@module/replication/ChangeStream.js';
 import { MongoManager } from '@module/replication/MongoManager.js';
 import { normalizeConnectionConfig } from '@module/types/types.js';
-import { isMongoServerError, mongo } from '@powersync/lib-service-mongodb';
 import { BucketStorageFactory, TestStorageOptions } from '@powersync/service-core';
-import { describe, expect, test, vi } from 'vitest';
+import { describe, expect, test } from 'vitest';
 import { ChangeStreamTestContext } from './change_stream_utils.js';
 import { env } from './env.js';
-import { INITIALIZED_MONGO_STORAGE_FACTORY, INITIALIZED_POSTGRES_STORAGE_FACTORY } from './util.js';
-import { ChangeStreamInvalidatedError } from '@module/replication/ChangeStream.js';
-
-describe('mongo lsn', () => {
-  test('LSN with resume tokens should be comparable', () => {
-    // Values without a resume token should be comparable
-    expect(
-      new MongoLSN({
-        timestamp: mongo.Timestamp.fromNumber(1)
-      }).comparable <
-        new MongoLSN({
-          timestamp: mongo.Timestamp.fromNumber(10)
-        }).comparable
-    ).true;
-
-    // Values with resume tokens should correctly compare
-    expect(
-      new MongoLSN({
-        timestamp: mongo.Timestamp.fromNumber(1),
-        resume_token: { _data: 'resume1' }
-      }).comparable <
-        new MongoLSN({
-          timestamp: mongo.Timestamp.fromNumber(10),
-          resume_token: { _data: 'resume2' }
-        }).comparable
-    ).true;
-
-    // The resume token should not affect comparison
-    expect(
-      new MongoLSN({
-        timestamp: mongo.Timestamp.fromNumber(1),
-        resume_token: { _data: '2' }
-      }).comparable <
-        new MongoLSN({
-          timestamp: mongo.Timestamp.fromNumber(10),
-          resume_token: { _data: '1' }
-        }).comparable
-    ).true;
-
-    // Resume token should not be required for comparison
-    expect(
-      new MongoLSN({
-        timestamp: mongo.Timestamp.fromNumber(10),
-        resume_token: { _data: '2' }
-      }).comparable > // Switching the order to test this case
-        new MongoLSN({
-          timestamp: mongo.Timestamp.fromNumber(9)
-        }).comparable
-    ).true;
+import { describeWithStorage } from './util.js';
 
-    // Comparison should be backwards compatible with old LSNs
-    expect(
-      new MongoLSN({
-        timestamp: mongo.Timestamp.fromNumber(10),
-        resume_token: { _data: '2' }
-      }).comparable > ZERO_LSN
-    ).true;
-    expect(
-      new MongoLSN({
-        timestamp: mongo.Timestamp.fromNumber(10),
-        resume_token: { _data: '2' }
-      }).comparable >
-        new MongoLSN({
-          timestamp: mongo.Timestamp.fromNumber(1)
-        }).comparable.split('|')[0] // Simulate an old LSN
-    ).true;
-    expect(
-      new MongoLSN({
-        timestamp: mongo.Timestamp.fromNumber(1),
-        resume_token: { _data: '2' }
-      }).comparable <
-        new MongoLSN({
-          timestamp: mongo.Timestamp.fromNumber(10)
-        }).comparable.split('|')[0] // Simulate an old LSN
-    ).true;
-  });
-});
-
-describe.skipIf(!env.TEST_MONGO_STORAGE)('MongoDB resume - mongo storage', () => {
-  defineResumeTest(INITIALIZED_MONGO_STORAGE_FACTORY);
-});
-
-describe.skipIf(!env.TEST_POSTGRES_STORAGE)('MongoDB resume - postgres storage', () => {
-  defineResumeTest(INITIALIZED_POSTGRES_STORAGE_FACTORY);
+describe('mongodb resuming replication', () => {
+  describeWithStorage({}, defineResumeTest);
 });
 
 function defineResumeTest(factoryGenerator: (options?: TestStorageOptions) => Promise<BucketStorageFactory>) {
@@ -112,13 +30,7 @@ function defineResumeTest(factoryGenerator: (options?: TestStorageOptions) => Pr
     await collection.insertOne({ description: 'test1', num: 1152921504606846976n });
 
     // Wait for the item above to be replicated. The commit should store a resume token.
-    await vi.waitFor(
-      async () => {
-        const checkpoint = await context.storage?.getCheckpoint();
-        expect(MongoLSN.fromSerialized(checkpoint!.lsn!).resumeToken).exist;
-      },
-      { timeout: 5000 }
-    );
+    await context.getCheckpoint();
 
     // Done with this context for now
     await context.dispose();
@@ -145,6 +57,7 @@ function defineResumeTest(factoryGenerator: (options?: TestStorageOptions) => Pr
     const activeContent = await factory.getActiveSyncRulesContent();
     context2.storage = factory.getInstance(activeContent!);
 
+    // If this test times out, it likely didn't throw the expected error here.
     const error = await context2.startStreaming().catch((ex) => ex);
     // The ChangeStreamReplicationJob will detect this and throw a ChangeStreamInvalidatedError
     expect(error).toBeInstanceOf(ChangeStreamInvalidatedError);