@powersync/service-module-mongodb 0.21.2 → 0.22.0

package/src/replication/MongoSnapshotQuery.ts

@@ -15,12 +15,19 @@ export class ChunkedSnapshotQuery implements AsyncDisposable {
   private lastCursor: mongo.FindCursor | null = null;
   private collection: mongo.Collection;
   private batchSize: number;
+  private readonly filter: mongo.Filter<mongo.Document> | null;
 
-  public constructor(options: { collection: mongo.Collection; batchSize: number; key?: Uint8Array | null }) {
+  public constructor(options: {
+    collection: mongo.Collection;
+    batchSize: number;
+    key?: Uint8Array | null;
+    filter?: mongo.Filter<mongo.Document> | null;
+  }) {
     this.lastKey = options.key ? bson.deserialize(options.key, { useBigInt64: true })._id : null;
     this.lastCursor = null;
     this.collection = options.collection;
     this.batchSize = options.batchSize;
+    this.filter = options.filter ?? null;
   }
 
   async nextChunk(): Promise<
@@ -39,8 +46,14 @@ export class ChunkedSnapshotQuery implements AsyncDisposable {
       // any parsing as an operator.
       // Starting in MongoDB 5.0, this filter can use the _id index. Source:
       // https://www.mongodb.com/docs/manual/release-notes/5.0/#general-aggregation-improvements
-      const filter: mongo.Filter<mongo.Document> =
+      const continuationFilter: mongo.Filter<mongo.Document> =
         this.lastKey == null ? {} : { $expr: { $gt: ['$_id', { $literal: this.lastKey }] } };
+      const filter: mongo.Filter<mongo.Document> =
+        this.filter == null
+          ? continuationFilter
+          : this.lastKey == null
+            ? this.filter
+            : { $and: [this.filter, continuationFilter] };
       cursor = this.collection.find(filter, {
         readConcern: 'majority',
         limit: this.batchSize,
@@ -48,7 +61,10 @@ export class ChunkedSnapshotQuery implements AsyncDisposable {
         // See https://github.com/mongodb/node-mongodb-native/pull/4580
         batchSize: this.batchSize + 1,
         sort: { _id: 1 },
-        raw: true
+        raw: true,
+        // Replication predicates have stable binary semantics independent of a
+        // collection's default collation.
+        ...(this.filter == null ? {} : { collation: { locale: 'simple' } })
       });
       newCursor = true;
     }

package/src/replication/MongoSnapshotter.ts

@@ -16,10 +16,17 @@ import { PostImagesOption } from '../types/types.js';
 import { escapeRegExp } from '../utils.js';
 import { CheckpointImplementation } from './checkpoints/CheckpointImplementation.js';
 import { createCheckpointImplementation } from './checkpoints/create-checkpoint-implementation.js';
+import { estimateFilteredSnapshotCount } from './estimateFilteredSnapshotCount.js';
 import { MongoManager } from './MongoManager.js';
 import { getMongoRelation } from './MongoRelation.js';
+import {
+  DEFAULT_MONGO_REPLICATION_QUERY_PROVIDER,
+  MongoReplicationQueryProvider
+} from './MongoReplicationQueryProvider.js';
+import { openMongoReplicationStream } from './MongoReplicationStream.js';
+import type { MongoReplicationStreamBatch } from './MongoReplicationStreamTypes.js';
 import { ChunkedSnapshotQuery } from './MongoSnapshotQuery.js';
-import { ChangeStreamBatch, parseChangeDocument, rawChangeStream } from './RawChangeStream.js';
+import { ChangeStreamBatch } from './RawChangeStream.js';
 import { CHECKPOINTS_COLLECTION, detectDocumentDb } from './replication-utils.js';
 import { DirectSourceRowConverter, SourceRowConverter } from './SourceRowConverter.js';
 
@@ -34,6 +41,7 @@ export interface MongoSnapshotterOptions {
   checkpointStreamId: mongo.ObjectId;
   storageHooks?: storage.StorageHooks;
   snapshotHooks?: MongoSnapshotterHooks;
+  queryProvider?: MongoReplicationQueryProvider;
 }
 
 export interface MongoSnapshotterHooks {
@@ -59,6 +67,7 @@ export class MongoSnapshotter {
   private readonly defaultDb: mongo.Db;
   private readonly syncRules: HydratedSyncConfig;
   private readonly sourceRowConverter: SourceRowConverter;
+  private readonly queryProvider: MongoReplicationQueryProvider;
   private readonly maxAwaitTimeMS: number;
   private readonly snapshotChunkLength: number;
   private readonly abortSignal: AbortSignal;
@@ -90,6 +99,7 @@ export class MongoSnapshotter {
     this.checkpointStreamId = options.checkpointStreamId;
     this.storageHooks = options.storageHooks;
     this.snapshotHooks = options.snapshotHooks;
+    this.queryProvider = options.queryProvider ?? DEFAULT_MONGO_REPLICATION_QUERY_PROVIDER;
     this.changeStreamTimeout = Math.ceil(this.client.options.socketTimeoutMS * 0.9);
     this.syncRules = options.storage.getParsedSyncRules({
       defaultSchema: this.defaultDb.databaseName
@@ -116,7 +126,9 @@ export class MongoSnapshotter {
     return this.storage.storageConfig.softDeleteCurrentData;
   }
 
-  /** The active checkpoint strategy. Only valid after ensureDetected(). */
+  /**
+   * The active checkpoint strategy. Only valid after ensureDetected().
+   */
   private get checkpointImplementation(): CheckpointImplementation {
     if (this._checkpointImplementation == null) {
       throw new ServiceError(
@@ -150,6 +162,7 @@ export class MongoSnapshotter {
         );
       }
     }
+    await this.queryProvider.validateSource?.({ connectionManager: this.connections, isDocumentDb: this.isDocumentDb });
     this._checkpointImplementation = createCheckpointImplementation(this.isDocumentDb, {
       client: this.client,
       db: this.defaultDb,
@@ -244,9 +257,7 @@ export class MongoSnapshotter {
         totalEstimatedCount: count
       });
       this.queueTable(updated);
-      this.logger.info(
-        `To replicate: ${updated.qualifiedName}: ${updated.snapshotStatus?.replicatedCount}/~${updated.snapshotStatus?.totalEstimatedCount}`
-      );
+      this.logger.info(`To replicate: ${updated.qualifiedName} ${updated.formatSnapshotProgress()}`);
     }
   }
 
@@ -465,13 +476,20 @@ export class MongoSnapshotter {
     const bytesReplicatedMetric = this.metrics.getCounter(ReplicationMetric.DATA_REPLICATED_BYTES);
     const chunksReplicatedMetric = this.metrics.getCounter(ReplicationMetric.CHUNKS_REPLICATED);
 
-    const totalEstimatedCount = await this.estimatedCountNumber(table);
+    const filter = this.queryProvider.getSnapshotFilter(table);
+    // Reuse known filtered estimates to avoid repeating an expensive count and sample. Unknown totals
+    // (-1) can also represent uninitialized progress, so retry them. Unfiltered metadata counts are cheap
+    // and refreshed before copying rows.
+    const previousEstimate = table.snapshotStatus?.totalEstimatedCount ?? -1;
+    const totalEstimatedCount =
+      filter != null && previousEstimate >= 0 ? previousEstimate : await this.estimatedCountNumber(table, filter);
     let at = table.snapshotStatus?.replicatedCount ?? 0;
     const collection = this.client.db(table.schema).collection(table.name);
     await using query = new ChunkedSnapshotQuery({
       collection,
       key: table.snapshotStatus?.lastKey,
-      batchSize: this.snapshotChunkLength
+      batchSize: this.snapshotChunkLength,
+      filter
     });
     if (query.lastKey != null) {
       this.logger.info(
@@ -565,8 +583,20 @@ export class MongoSnapshotter {
     return result.tables;
   }
 
-  private async estimatedCountNumber(table: storage.SourceTable): Promise<number> {
-    return await this.client.db(table.schema).collection(table.name).estimatedDocumentCount();
+  private async estimatedCountNumber(
+    table: storage.SourceTable,
+    filter = this.queryProvider.getSnapshotFilter(table)
+  ): Promise<number> {
+    const collection = this.client.db(table.schema).collection(table.name);
+    if (filter != null) {
+      return await estimateFilteredSnapshotCount({
+        collection,
+        filter,
+        signal: this.abortSignal,
+        logger: this.logger
+      });
+    }
+    return await collection.estimatedDocumentCount();
   }
 
   private async getCollectionInfo(db: string, name: string): Promise<mongo.CollectionInfo | undefined> {
@@ -616,7 +646,7 @@ export class MongoSnapshotter {
     this.checkpointImplementation.seedPosition(null);
     const startStreamFromLsn = await this.checkpointImplementation.createFirstBarrier();
     const filters = this.getSourceNamespaceFilters();
-    const iter = this.rawChangeStreamBatches({
+    const iter = this.openChangeStream({
       lsn: startStreamFromLsn,
       maxAwaitTimeMS: 0,
       signal: this.abortSignal,
@@ -627,33 +657,48 @@ export class MongoSnapshotter {
     let eventsSeen = 0;
     let batchesSeen = 0;
 
-    for await (const { events } of iter) {
+    // Returns false once the timeout is reached. Otherwise creates another checkpoint when one is due.
+    const shouldKeepWaiting = async () => {
       if (performance.now() - startTime >= LSN_TIMEOUT_SECONDS * 1000) {
-        break;
+        return false;
       }
       if (performance.now() - lastCheckpointCreated >= LSN_CREATE_INTERVAL_SECONDS * 1000) {
         await this.checkpointImplementation.createBatchCheckpoint();
         lastCheckpointCreated = performance.now();
       }
-      batchesSeen += 1;
+      return true;
+    };
 
-      for (const rawChangeDocument of events) {
-        const changeDocument = parseChangeDocument(rawChangeDocument);
+    streamLoop: for await (const sourceBatch of iter) {
+      this.abortSignal.throwIfAborted();
+      for (const { event: changeDocument } of sourceBatch.events) {
+        this.abortSignal.throwIfAborted();
         const ns = 'ns' in changeDocument && 'coll' in changeDocument.ns ? changeDocument.ns : undefined;
 
         if (ns?.coll == CHECKPOINTS_COLLECTION && 'documentKey' in changeDocument) {
-          const kind = this.checkpointImplementation.event.observe(changeDocument);
-          if (kind != 'own-barrier') {
-            // Standalone events still feed the implementation's coordinate via
-            // event.observe above; we only resolve on our own barrier.
-            continue;
+          // Standalone events still feed the implementation's coordinate via
+          // event.observe; we only resolve on our own barrier.
+          if (this.checkpointImplementation.event.observe(changeDocument) == 'own-barrier') {
+            // Observing our own barrier has set the coordinate (the barrier event
+            // carries the sentinel counter directly), so the LSN is ready.
+            // This is checked before the timeout, so a barrier we already hold is never discarded.
+            return this.checkpointImplementation.event.lsn(changeDocument);
           }
-          // Observing our own barrier has set the coordinate (the barrier event
-          // carries the sentinel counter directly), so the LSN is ready.
-          return this.checkpointImplementation.event.lsn(changeDocument);
+        } else {
+          eventsSeen += 1;
         }
 
-        eventsSeen += 1;
+        if (!(await shouldKeepWaiting())) {
+          break streamLoop;
+        }
+      }
+      const resumeToken = sourceBatch.getResumeToken();
+      if (resumeToken != null) {
+        batchesSeen++;
+      }
+      // Empty batches have no events to trigger the per-event check.
+      if (!(await shouldKeepWaiting())) {
+        break;
       }
     }
 
@@ -668,13 +713,19 @@ export class MongoSnapshotter {
    * Given a snapshot LSN, validate that we can read from it, by opening a change stream.
    */
   private async validateSnapshotLsn(lsn: string) {
-    const stream = this.rawChangeStreamBatches({
-      lsn,
-      maxAwaitTimeMS: 0,
-      filters: this.getSourceNamespaceFilters()
-    });
-    for await (const _batch of stream) {
-      break;
+    const filters = this.getSourceNamespaceFilters();
+    const stream = this.openChangeStream({ lsn, maxAwaitTimeMS: 0, filters, signal: this.abortSignal });
+    for await (const batch of stream) {
+      this.abortSignal.throwIfAborted();
+      // Reading one complete event validates the position; closing the stream abandons batch progress.
+      for (const _change of batch.events) {
+        return;
+      }
+      // An empty or fully filtered batch can also validate the position through safe progress.
+      // If a split event is unfinished, progress is undefined: keep reading until it is complete.
+      if (batch.getResumeToken() != null) {
+        return;
+      }
     }
   }
 
@@ -722,73 +773,31 @@ export class MongoSnapshotter {
     return { $match: nsFilter, multipleDatabases };
   }
 
-  private rawChangeStreamBatches(options: {
+  private openChangeStream(options: {
     lsn: string | null;
     maxAwaitTimeMS?: number;
     batchSize?: number;
     filters: { $match: any; multipleDatabases: boolean };
     signal?: AbortSignal;
     tracer?: PerformanceTracer<'changestream'>;
-  }): AsyncIterableIterator<ChangeStreamBatch> {
-    const position = options.lsn ? this.checkpointImplementation.parseResumePosition(options.lsn) : null;
-    const startAfter = position?.startAfter ?? undefined;
-    const resumeAfter = position?.resumeAfter ?? undefined;
-
-    let fullDocument: 'required' | 'updateLookup';
-    if (this.isDocumentDb) {
-      // DocumentDB does not support changeStreamPreAndPostImages, so 'required' won't work.
-      fullDocument = 'updateLookup';
-    } else if (this.usePostImages) {
-      // 'read_only' or 'auto_configure'
-      // Configuration happens during snapshot, or when we see new
-      // collections.
-      fullDocument = 'required';
-    } else {
-      fullDocument = 'updateLookup';
-    }
-    const streamOptions: mongo.ChangeStreamOptions & mongo.Document = {
-      fullDocument
-    };
-    if (!this.isDocumentDb) {
-      // DocumentDB does not support showExpandedEvents.
-      streamOptions.showExpandedEvents = true;
-    }
-    const pipeline: mongo.Document[] = [{ $changeStream: streamOptions }, { $match: options.filters.$match }];
-    if (!this.isDocumentDb) {
-      // DocumentDB does not support $changeStreamSplitLargeEvent.
-      pipeline.push({ $changeStreamSplitLargeEvent: {} });
-    }
-
-    // Only one of these options can be supplied at a time.
-    if (resumeAfter) {
-      streamOptions.resumeAfter = resumeAfter;
-    } else if (startAfter != null) {
-      // Legacy: We don't persist lsns without resumeTokens anymore, but we do still handle the
-      // case if we have an old one. The sentinel implementation never produces a startAfter,
-      // and a fresh DocumentDB stream opens from "now" with neither option set.
-      streamOptions.startAtOperationTime = startAfter;
-    }
-
-    let watchDb: mongo.Db;
-    if (this.isDocumentDb || options.filters.multipleDatabases) {
-      // DocumentDB only supports cluster-level change streams.
-      // Requires readAnyDatabase@admin on Atlas.
-      watchDb = this.client.db('admin');
-      streamOptions.allChangesForCluster = true;
-    } else {
-      // Same general result, but requires less permissions than the above
-      watchDb = this.defaultDb;
-    }
-
-    const maxAwaitTimeMS = options.maxAwaitTimeMS ?? this.maxAwaitTimeMS;
-
-    return rawChangeStream(watchDb, pipeline, {
-      batchSize: options.batchSize ?? this.snapshotChunkLength,
-      maxAwaitTimeMS,
-      maxTimeMS: this.changeStreamTimeout,
-      signal: options.signal,
-      logger: this.logger,
-      tracer: options.tracer
+    onBatch?: (batch: ChangeStreamBatch) => Disposable | void;
+  }): AsyncIterableIterator<MongoReplicationStreamBatch> {
+    return openMongoReplicationStream({
+      db: this.defaultDb,
+      queryProvider: this.queryProvider,
+      namespaceFilter: options.filters,
+      isDocumentDb: this.isDocumentDb,
+      usePostImages: this.usePostImages,
+      position: options.lsn ? this.checkpointImplementation.parseResumePosition(options.lsn) : null,
+      onBatch: options.onBatch,
+      options: {
+        batchSize: options.batchSize ?? this.snapshotChunkLength,
+        maxAwaitTimeMS: options.maxAwaitTimeMS ?? this.maxAwaitTimeMS,
+        maxTimeMS: this.changeStreamTimeout,
+        signal: options.signal,
+        logger: this.logger,
+        tracer: options.tracer
+      }
     });
   }
 

package/src/replication/RawChangeStream.ts

@@ -364,7 +364,7 @@ type MapRawDocument<T> = T extends unknown ? { [K in keyof T]: Rawify<T[K]> } :
 export type ProjectedChangeStreamDocument =
   | Omit<mongo.ChangeStreamDropDocument, 'wallTime' | 'collectionUUID'>
   | Omit<mongo.ChangeStreamRenameDocument, 'wallTime'>
-  | Omit<mongo.ChangeStreamDeleteDocument<Buffer>, 'wallTime'>
+  | Omit<mongo.ChangeStreamDeleteDocument<Buffer>, 'wallTime' | 'collectionUUID'>
   | Omit<mongo.ChangeStreamInsertDocument<Buffer>, 'wallTime'>
   | Omit<mongo.ChangeStreamUpdateDocument<Buffer>, 'wallTime' | 'updateDescription'>
   | Omit<mongo.ChangeStreamReplaceDocument<Buffer>, 'wallTime'>;
@@ -432,7 +432,7 @@ export function mapChangeStreamError(e: unknown) {
   } else if (
     isMongoServerError(e) &&
     e.codeName == 'NoMatchingDocument' &&
-    e.errmsg?.includes('post-image was not found')
+    (e.errmsg?.includes('post-image was not found') || e.errmsg?.includes('pre-image was not found'))
   ) {
     throw new ChangeStreamInvalidatedError(e.errmsg, e);
   } else if (isMongoServerError(e) && e.hasErrorLabel('NonResumableChangeStreamError')) {

package/src/replication/estimateFilteredSnapshotCount.ts

@@ -0,0 +1,93 @@
+import type { mongo } from '@powersync/lib-service-mongodb';
+import type { Logger } from '@powersync/lib-services-framework';
+import { performance } from 'node:perf_hooks';
+
+const COUNT_TIMEOUT_MS = 10_000;
+/**
+ * Keep below 5% for MongoDB's random-cursor optimization on large collections.
+ * https://www.mongodb.com/docs/manual/reference/operator/aggregation/sample/
+ */
+const SAMPLE_FRACTION = 0.01;
+const MAX_SAMPLE_SIZE = 10_000;
+
+interface SampleCounts {
+  sampled: { count: number }[];
+  matches: { count: number }[];
+}
+
+/**
+ * Try an exact filtered count, then estimate selectivity from a bounded sample.
+ * Each attempt has a separate 10-second budget. Failures leave progress unknown (-1).
+ */
+export async function estimateFilteredSnapshotCount(options: {
+  collection: mongo.Collection;
+  filter: mongo.Filter<mongo.Document>;
+  signal: AbortSignal;
+  logger: Logger;
+}): Promise<number> {
+  const { collection, filter, signal, logger } = options;
+  signal.throwIfAborted();
+  try {
+    return await collection.countDocuments(filter, {
+      collation: { locale: 'simple' },
+      readConcern: 'majority',
+      timeoutMS: COUNT_TIMEOUT_MS,
+      signal
+    });
+  } catch (error) {
+    signal.throwIfAborted();
+    logger.debug(`Filtered snapshot count unavailable for ${collection.namespace}; trying a sample strategy`, error);
+  }
+
+  const sampleStartedAt = performance.now();
+  try {
+    const metadataOptions = { timeoutMS: COUNT_TIMEOUT_MS, signal };
+    const totalCount = await collection.estimatedDocumentCount(metadataOptions);
+    signal.throwIfAborted();
+    if (totalCount <= 0) {
+      // Metadata alone cannot establish that a filtered result is empty after the count failed.
+      return -1;
+    }
+    const remainingTimeMS = Math.floor(COUNT_TIMEOUT_MS - (performance.now() - sampleStartedAt));
+    if (remainingTimeMS <= 0) {
+      return -1;
+    }
+    // Keep $sample first and below 5% for MongoDB's random-cursor optimization on large collections.
+    // https://www.mongodb.com/docs/manual/reference/operator/aggregation/sample/
+    const sampleSize = Math.min(MAX_SAMPLE_SIZE, Math.max(1, Math.floor(totalCount * SAMPLE_FRACTION)));
+    const cursor = collection.aggregate<SampleCounts>(
+      [
+        { $sample: { size: sampleSize } },
+        {
+          $facet: {
+            sampled: [{ $count: 'count' }],
+            matches: [{ $match: filter }, { $count: 'count' }]
+          }
+        }
+      ],
+      {
+        collation: { locale: 'simple' },
+        readConcern: 'majority',
+        timeoutMS: remainingTimeMS,
+        signal,
+        allowDiskUse: false
+      }
+    );
+    try {
+      const sample = await cursor.next();
+      const sampledCount = sample?.sampled[0]?.count ?? 0;
+      const matchedCount = sample?.matches[0]?.count ?? 0;
+      if (sampledCount == 0 || matchedCount == 0) {
+        // A rare predicate can miss the whole sample; that does not establish an empty result.
+        return -1;
+      }
+      return Math.round((matchedCount / sampledCount) * totalCount);
+    } finally {
+      await cursor.close();
+    }
+  } catch (error) {
+    signal.throwIfAborted();
+    logger.debug(`Filtered snapshot sample unavailable for ${collection.namespace}; total remains unknown`, error);
+    return -1;
+  }
+}

package/src/replication/replication-index.ts

@@ -1,7 +1,14 @@
 export * from './ChangeStream.js';
 export * from './ChangeStreamReplicationJob.js';
 export * from './ChangeStreamReplicator.js';
+export * from './ConnectionManagerFactory.js';
 export * from './internal-mongodb-utils.js';
+export * from './MongoErrorRateLimiter.js';
 export * from './MongoRelation.js';
+export * from './MongoReplicationQueryProvider.js';
+export * from './MongoReplicationStream.js';
+export * from './MongoReplicationStreamTypes.js';
+export * from './MongoSnapshotter.js';
 export * from './RawChangeStream.js';
+export { CHECKPOINTS_COLLECTION } from './replication-utils.js';
 export * from './SourceRowConverter.js';