@powersync/service-module-mongodb 0.15.4 → 0.16.0

package/src/replication/ChangeStream.ts

@@ -1,8 +1,7 @@
-import { isMongoNetworkTimeoutError, isMongoServerError, mongo } from '@powersync/lib-service-mongodb';
+import { mongo } from '@powersync/lib-service-mongodb';
 import {
   container,
   DatabaseConnectionError,
-  logger as defaultLogger,
   ErrorCode,
   Logger,
   ReplicationAbortedError,
@@ -11,6 +10,7 @@ import {
 } from '@powersync/lib-services-framework';
 import {
   MetricsEngine,
+  PerformanceTracer,
   RelationCache,
   ReplicationLagTracker,
   SaveOperationTag,
@@ -18,29 +18,23 @@ import {
   SourceTable,
   storage
 } from '@powersync/service-core';
-import {
-  DatabaseInputRow,
-  HydratedSyncRules,
-  SqliteInputRow,
-  SqliteRow,
-  TablePattern
-} from '@powersync/service-sync-rules';
+import { HydratedSyncRules, TablePattern } from '@powersync/service-sync-rules';
 import { ReplicationMetric } from '@powersync/service-types';
+import { performance } from 'node:perf_hooks';
 import { MongoLSN } from '../common/MongoLSN.js';
 import { PostImagesOption } from '../types/types.js';
 import { escapeRegExp } from '../utils.js';
-import { trackChangeStreamBsonBytes } from './internal-mongodb-utils.js';
 import { MongoManager } from './MongoManager.js';
-import {
-  constructAfterRecord,
-  createCheckpoint,
-  getCacheIdentifier,
-  getMongoRelation,
-  STANDALONE_CHECKPOINT_ID
-} from './MongoRelation.js';
+import { createCheckpoint, getCacheIdentifier, getMongoRelation, STANDALONE_CHECKPOINT_ID } from './MongoRelation.js';
 import { ChunkedSnapshotQuery } from './MongoSnapshotQuery.js';
+import {
+  ChangeStreamBatch,
+  parseChangeDocument,
+  ProjectedChangeStreamDocument,
+  rawChangeStream
+} from './RawChangeStream.js';
 import { CHECKPOINTS_COLLECTION, timestampToDate } from './replication-utils.js';
-
+import { DirectSourceRowConverter, SourceRowConverter } from './SourceRowConverter.js';
 export interface ChangeStreamOptions {
   connections: MongoManager;
   storage: storage.SyncRulesBucketStorage;
@@ -110,6 +104,8 @@ export class ChangeStream {
 
   private changeStreamTimeout: number;
 
+  private readonly sourceRowConverter: SourceRowConverter;
+
   constructor(options: ChangeStreamOptions) {
     this.storage = options.storage;
     this.metrics = options.metrics;
@@ -122,6 +118,8 @@ export class ChangeStream {
     this.sync_rules = options.storage.getParsedSyncRules({
       defaultSchema: this.defaultDb.databaseName
     });
+    this.sourceRowConverter = new DirectSourceRowConverter(this.sync_rules.compatibility);
+
     // The change stream aggregation command should timeout before the socket times out,
     // so we use 90% of the socket timeout value.
     this.changeStreamTimeout = Math.ceil(this.client.options.socketTimeoutMS * 0.9);
@@ -135,7 +133,7 @@ export class ChangeStream {
       { once: true }
     );
 
-    this.logger = options.logger ?? defaultLogger;
+    this.logger = options.logger ?? this.storage.logger;
   }
 
   get stopped() {
@@ -259,48 +257,53 @@ export class ChangeStream {
 
     // Create a checkpoint, and open a change stream using startAtOperationTime with the checkpoint's operationTime.
     const firstCheckpointLsn = await createCheckpoint(this.client, this.defaultDb, this.checkpointStreamId);
-    await using streamManager = this.openChangeStream({ lsn: firstCheckpointLsn, maxAwaitTimeMs: 0 });
 
-    const { stream } = streamManager;
     const startTime = performance.now();
     let lastCheckpointCreated = performance.now();
     let eventsSeen = 0;
+    let batchesSeen = 0;
 
-    while (performance.now() - startTime < LSN_TIMEOUT_SECONDS * 1000) {
+    const filters = this.getSourceNamespaceFilters();
+    const iter = this.rawChangeStreamBatches({
+      lsn: firstCheckpointLsn,
+      maxAwaitTimeMS: 0,
+      signal: this.abort_signal,
+      filters
+    });
+    for await (let { events } of iter) {
+      if (performance.now() - startTime >= LSN_TIMEOUT_SECONDS * 1000) {
+        break;
+      }
       if (performance.now() - lastCheckpointCreated >= LSN_CREATE_INTERVAL_SECONDS * 1000) {
         await createCheckpoint(this.client, this.defaultDb, this.checkpointStreamId);
         lastCheckpointCreated = performance.now();
       }
+      batchesSeen += 1;
 
-      // tryNext() doesn't block, while next() / hasNext() does block until there is data on the stream
-      const changeDocument = await stream.tryNext().catch((e) => {
-        throw mapChangeStreamError(e);
-      });
-      if (changeDocument == null) {
-        continue;
-      }
-
-      const ns = 'ns' in changeDocument && 'coll' in changeDocument.ns ? changeDocument.ns : undefined;
+      for (let rawChangeDocument of events) {
+        const changeDocument = parseChangeDocument(rawChangeDocument);
+        const ns = 'ns' in changeDocument && 'coll' in changeDocument.ns ? changeDocument.ns : undefined;
 
-      if (ns?.coll == CHECKPOINTS_COLLECTION && 'documentKey' in changeDocument) {
-        const checkpointId = changeDocument.documentKey._id as string | mongo.ObjectId;
-        if (!this.checkpointStreamId.equals(checkpointId)) {
-          continue;
+        if (ns?.coll == CHECKPOINTS_COLLECTION && 'documentKey' in changeDocument) {
+          const checkpointId = changeDocument.documentKey._id as string | mongo.ObjectId;
+          if (!this.checkpointStreamId.equals(checkpointId)) {
+            continue;
+          }
+          const { comparable: lsn } = new MongoLSN({
+            timestamp: changeDocument.clusterTime!,
+            resume_token: changeDocument._id
+          });
+          return lsn;
         }
-        const { comparable: lsn } = new MongoLSN({
-          timestamp: changeDocument.clusterTime!,
-          resume_token: changeDocument._id
-        });
-        return lsn;
-      }
 
-      eventsSeen += 1;
+        eventsSeen += 1;
+      }
     }
 
     // Could happen if there is a very large replication lag?
     throw new ServiceError(
       ErrorCode.PSYNC_S1301,
-      `Timeout after while waiting for checkpoint document for ${LSN_TIMEOUT_SECONDS}s. Streamed events = ${eventsSeen}`
+      `Timeout after while waiting for checkpoint document for ${LSN_TIMEOUT_SECONDS}s. Streamed events = ${eventsSeen}, batches = ${batchesSeen}`
     );
   }
 
@@ -308,21 +311,24 @@ export class ChangeStream {
    * Given a snapshot LSN, validate that we can read from it, by opening a change stream.
    */
   private async validateSnapshotLsn(lsn: string) {
-    await using streamManager = this.openChangeStream({ lsn: lsn, maxAwaitTimeMs: 0 });
-    const { stream } = streamManager;
-    try {
-      // tryNext() doesn't block, while next() / hasNext() does block until there is data on the stream
-      await stream.tryNext();
-    } catch (e) {
-      // Note: A timeout here is not handled as a ChangeStreamInvalidatedError, even though
-      // we possibly cannot recover from it.
-      throw mapChangeStreamError(e);
+    const filters = this.getSourceNamespaceFilters();
+    const stream = this.rawChangeStreamBatches({
+      lsn: lsn,
+      // maxAwaitTimeMS should never actually be used here
+      maxAwaitTimeMS: 0,
+      filters
+    });
+    for await (let _batch of stream) {
+      // We got a response from the aggregate command, so consider the LSN valid.
+      // Close the stream immediately.
+      break;
     }
   }
 
   async initialReplication(snapshotLsn: string | null) {
     const sourceTables = this.sync_rules.getSourceTables();
     await this.client.connect();
+    const tracer = new PerformanceTracer('MongoDB initial replication');
 
     const flushResult = await this.storage.startBatch(
       {
@@ -330,7 +336,8 @@ export class ChangeStream {
         zeroLSN: MongoLSN.ZERO.comparable,
         defaultSchema: this.defaultDb.databaseName,
         storeCurrentData: false,
-        skipExistingRows: true
+        skipExistingRows: true,
+        tracer
       },
       async (batch) => {
         if (snapshotLsn == null) {
@@ -465,12 +472,6 @@ export class ChangeStream {
     return { $match: nsFilter, multipleDatabases };
   }
 
-  static *getQueryData(results: Iterable<DatabaseInputRow>): Generator<SqliteInputRow> {
-    for (let row of results) {
-      yield constructAfterRecord(row);
-    }
-  }
-
   private async snapshotTable(batch: storage.BucketStorageBatch, table: storage.SourceTable) {
     const rowsReplicatedMetric = this.metrics.getCounter(ReplicationMetric.ROWS_REPLICATED);
     const bytesReplicatedMetric = this.metrics.getCounter(ReplicationMetric.DATA_REPLICATED_BYTES);
@@ -510,8 +511,8 @@ export class ChangeStream {
 
       // Pre-fetch next batch, so that we can read and write concurrently
       nextChunkPromise = query.nextChunk();
-      for (let document of docBatch) {
-        const record = this.constructAfterRecord(document);
+      for (let buffer of docBatch) {
+        const { row: record, replicaId: replicaId } = this.rawToSqliteRow(buffer);
 
         // This auto-flushes when the batch reaches its size limit
         await batch.save({
@@ -520,7 +521,7 @@ export class ChangeStream {
           before: undefined,
           beforeReplicaId: undefined,
           after: record,
-          afterReplicaId: document._id
+          afterReplicaId: replicaId
         });
       }
 
@@ -632,7 +633,7 @@ export class ChangeStream {
     // Snapshot if:
     // 1. Snapshot is requested (false for initial snapshot, since that process handles it elsewhere)
     // 2. Snapshot is not already done, AND:
-    // 3. The table is used in sync rules.
+    // 3. The table is used in sync config.
     const shouldSnapshot = snapshot && !result.table.snapshotComplete && result.table.syncAny;
     if (shouldSnapshot) {
       this.logger.info(`New collection: ${descriptor.schema}.${descriptor.name}`);
@@ -649,30 +650,28 @@ export class ChangeStream {
     return result.table;
   }
 
-  private constructAfterRecord(document: mongo.Document): SqliteRow {
-    const inputRow = constructAfterRecord(document);
-    return this.sync_rules.applyRowContext<never>(inputRow);
-  }
-
   async writeChange(
     batch: storage.BucketStorageBatch,
     table: storage.SourceTable,
-    change: mongo.ChangeStreamDocument
+    change: ProjectedChangeStreamDocument
   ): Promise<storage.FlushedResult | null> {
     if (!table.syncAny) {
-      this.logger.debug(`Collection ${table.qualifiedName} not used in sync rules - skipping`);
+      this.logger.debug(`Collection ${table.qualifiedName} not used in sync config - skipping`);
       return null;
     }
 
     this.metrics.getCounter(ReplicationMetric.ROWS_REPLICATED).add(1);
     if (change.operationType == 'insert') {
-      const baseRecord = this.constructAfterRecord(change.fullDocument);
+      const { row: baseRecord, replicaId: _replicaId } = this.rawToSqliteRow(change.fullDocument);
       return await batch.save({
         tag: SaveOperationTag.INSERT,
         sourceTable: table,
         before: undefined,
         beforeReplicaId: undefined,
         after: baseRecord,
+        // Same as _replicaId
+        // We specifically need to use the source _id, not the converted one in baseRecord,
+        // to preserve _id uniqueness properties.
         afterReplicaId: change.documentKey._id
       });
     } else if (change.operationType == 'update' || change.operationType == 'replace') {
@@ -685,14 +684,14 @@ export class ChangeStream {
           beforeReplicaId: change.documentKey._id
         });
       }
-      const after = this.constructAfterRecord(change.fullDocument!);
+      const { row: after, replicaId: _replicaId } = this.rawToSqliteRow(change.fullDocument!);
       return await batch.save({
         tag: SaveOperationTag.UPDATE,
         sourceTable: table,
         before: undefined,
         beforeReplicaId: undefined,
         after: after,
-        afterReplicaId: change.documentKey._id
+        afterReplicaId: change.documentKey._id // Same as _replicaId
       });
     } else if (change.operationType == 'delete') {
       return await batch.save({
@@ -728,7 +727,7 @@ export class ChangeStream {
       }
       const { lastOpId } = await this.initialReplication(result.snapshotLsn);
       if (lastOpId != null) {
-        // Populate the cache _after_ initial replication, but _before_ we switch to this sync rules.
+        // Populate the cache _after_ initial replication, but _before_ we switch to this replication stream.
         await this.storage.populatePersistentChecksumCache({
           signal: this.abort_signal,
           // No checkpoint yet, but we do have the opId.
@@ -753,19 +752,19 @@ export class ChangeStream {
     }
   }
 
-  private openChangeStream(options: { lsn: string | null; maxAwaitTimeMs?: number }) {
+  private rawChangeStreamBatches(options: {
+    lsn: string | null;
+    maxAwaitTimeMS?: number;
+    batchSize?: number;
+    filters: { $match: any; multipleDatabases: boolean };
+    signal?: AbortSignal;
+    tracer?: PerformanceTracer<'changestream'>;
+  }): AsyncIterableIterator<ChangeStreamBatch> {
     const lastLsn = options.lsn ? MongoLSN.fromSerialized(options.lsn) : null;
     const startAfter = lastLsn?.timestamp;
     const resumeAfter = lastLsn?.resumeToken;
 
-    const filters = this.getSourceNamespaceFilters();
-
-    const pipeline: mongo.Document[] = [
-      {
-        $match: filters.$match
-      },
-      { $changeStreamSplitLargeEvent: {} }
-    ];
+    const filters = options.filters;
 
     let fullDocument: 'required' | 'updateLookup';
 
@@ -777,12 +776,19 @@ export class ChangeStream {
     } else {
       fullDocument = 'updateLookup';
     }
-    const streamOptions: mongo.ChangeStreamOptions = {
+    const streamOptions: mongo.ChangeStreamOptions & mongo.Document = {
       showExpandedEvents: true,
-      maxAwaitTimeMS: options.maxAwaitTimeMs ?? this.maxAwaitTimeMS,
-      fullDocument: fullDocument,
-      maxTimeMS: this.changeStreamTimeout
+      fullDocument: fullDocument
     };
+    const pipeline: mongo.Document[] = [
+      {
+        $changeStream: streamOptions
+      },
+      {
+        $match: filters.$match
+      },
+      { $changeStreamSplitLargeEvent: {} }
+    ];
 
     /**
      * Only one of these options can be supplied at a time.
@@ -796,45 +802,27 @@ export class ChangeStream {
       streamOptions.startAtOperationTime = startAfter;
     }
 
-    let stream: mongo.ChangeStream<mongo.Document>;
+    let watchDb: mongo.Db;
     if (filters.multipleDatabases) {
-      // Requires readAnyDatabase@admin on Atlas
-      stream = this.client.watch(pipeline, streamOptions);
+      watchDb = this.client.db('admin');
+      streamOptions.allChangesForCluster = true;
     } else {
-      // Same general result, but requires less permissions than the above
-      stream = this.defaultDb.watch(pipeline, streamOptions);
+      watchDb = this.defaultDb;
     }
 
-    this.abort_signal.addEventListener('abort', () => {
-      stream.close();
-    });
+    return rawChangeStream(watchDb, pipeline, {
+      batchSize: options.batchSize ?? this.snapshotChunkLength,
+      maxAwaitTimeMS: options.maxAwaitTimeMS ?? this.maxAwaitTimeMS,
+      maxTimeMS: this.changeStreamTimeout,
 
-    return {
-      stream,
-      filters,
-      [Symbol.asyncDispose]: async () => {
-        return stream.close();
-      }
-    };
+      signal: options.signal,
+      logger: this.logger,
+      tracer: options.tracer
+    });
   }
 
-  private getBufferedChangeCount(stream: mongo.ChangeStream<mongo.Document>): number {
-    // The driver keeps fetched change stream documents on the underlying cursor, but does
-    // not expose that through the public ChangeStream API. We use this to detect backlog
-    // building up before we have processed the corresponding source changes locally.
-    // If the driver API changes, we'll have a hard error here.
-    // We specifically want to avoid a silent performance regression if the driver behavior changes.
-    const cursor = (
-      stream as mongo.ChangeStream<mongo.Document> & {
-        cursor: mongo.AbstractCursor<mongo.ChangeStreamDocument<mongo.Document>>;
-      }
-    ).cursor;
-    if (cursor == null || typeof cursor.bufferedCount != 'function') {
-      throw new ReplicationAssertionError(
-        'MongoDB ChangeStream no longer exposes an internal cursor with bufferedCount'
-      );
-    }
-    return cursor.bufferedCount();
+  private rawToSqliteRow(row: Buffer) {
+    return this.sourceRowConverter.rawToSqliteRow(row);
   }
 
   async streamChangesInternal() {
@@ -842,13 +830,15 @@ export class ChangeStream {
     const bytesReplicatedMetric = this.metrics.getCounter(ReplicationMetric.DATA_REPLICATED_BYTES);
     const chunksReplicatedMetric = this.metrics.getCounter(ReplicationMetric.CHUNKS_REPLICATED);
 
+    const tracer = new PerformanceTracer('MongoDB streaming replication');
     await this.storage.startBatch(
       {
         logger: this.logger,
         zeroLSN: MongoLSN.ZERO.comparable,
         defaultSchema: this.defaultDb.databaseName,
         // We get a complete postimage for every change, so we don't need to store the current data.
-        storeCurrentData: false
+        storeCurrentData: false,
+        tracer
       },
       async (batch) => {
         const { resumeFromLsn } = batch;
@@ -857,6 +847,7 @@ export class ChangeStream {
         }
         const lastLsn = MongoLSN.fromSerialized(resumeFromLsn);
         const startAfter = lastLsn?.timestamp;
+        let outerSpan = tracer.span('batch');
 
         // It is normal for this to be a minute or two old when there is a low volume
         // of ChangeStream events.
@@ -864,16 +855,13 @@ export class ChangeStream {
 
         this.logger.info(`Resume streaming at ${startAfter?.inspect()} / ${lastLsn}  | Token age: ${tokenAgeSeconds}s`);
 
-        await using streamManager = this.openChangeStream({ lsn: resumeFromLsn });
-        const { stream, filters } = streamManager;
-        if (this.abort_signal.aborted) {
-          await stream.close();
-          return;
-        }
-        trackChangeStreamBsonBytes(stream, (bytes) => {
-          bytesReplicatedMetric.add(bytes);
-          // Each of these represent a single response message from MongoDB.
-          chunksReplicatedMetric.add(1);
+        const filters = this.getSourceNamespaceFilters();
+        // This is closed when the for loop below returns/breaks/throws
+        const batchStream = this.rawChangeStreamBatches({
+          lsn: resumeFromLsn,
+          filters,
+          signal: this.abort_signal,
+          tracer
         });
 
         // Always start with a checkpoint.
@@ -885,44 +873,30 @@ export class ChangeStream {
           this.checkpointStreamId
         );
 
-        let splitDocument: mongo.ChangeStreamDocument | null = null;
+        let splitDocument: ProjectedChangeStreamDocument | null = null;
 
         let flexDbNameWorkaroundLogged = false;
-        let changesSinceLastCheckpoint = 0;
 
         let lastEmptyResume = performance.now();
         let lastTxnKey: string | null = null;
 
-        while (true) {
-          if (this.abort_signal.aborted) {
-            break;
-          }
-
-          const originalChangeDocument = await stream.tryNext().catch((e) => {
-            throw mapChangeStreamError(e);
-          });
-          // The stream was closed, we will only ever receive `null` from it
-          if (!originalChangeDocument && stream.closed) {
-            break;
-          }
+        for await (let eventBatch of batchStream) {
+          const { events, resumeToken } = eventBatch;
+          using batchSpan = tracer.span('processing');
 
+          bytesReplicatedMetric.add(eventBatch.byteSize);
+          chunksReplicatedMetric.add(1);
           if (this.abort_signal.aborted) {
             break;
           }
-
-          if (originalChangeDocument == null) {
-            // We get a new null document after `maxAwaitTimeMS` if there were no other events.
-            // In this case, stream.resumeToken is the resume token associated with the last response.
-            // stream.resumeToken is not updated if stream.tryNext() returns data, while stream.next()
-            // does update it.
-            // From observed behavior, the actual resumeToken changes around once every 10 seconds.
+          this.touch();
+          if (events.length == 0) {
+            // No changes in this batch, but we still want to keep the connection alive.
+            // We do this by persisting a keepalive checkpoint.
             // If we don't update it on empty events, we do keep consistency, but resuming the stream
             // with old tokens may cause connection timeouts.
-            // We throttle this further by only persisting a keepalive once a minute.
-            // We add an additional check for waitForCheckpointLsn == null, to make sure we're not
-            // doing a keepalive in the middle of a transaction.
             if (waitForCheckpointLsn == null && performance.now() - lastEmptyResume > 60_000) {
-              const { comparable: lsn, timestamp } = MongoLSN.fromResumeToken(stream.resumeToken);
+              const { comparable: lsn, timestamp } = MongoLSN.fromResumeToken(resumeToken);
               await batch.keepalive(lsn);
               this.touch();
               lastEmptyResume = performance.now();
@@ -933,226 +907,251 @@ export class ChangeStream {
               );
               this.replicationLag.markStarted();
             }
-            continue;
-          }
 
-          this.touch();
-
-          if (startAfter != null && originalChangeDocument.clusterTime?.lte(startAfter)) {
-            continue;
+            // If we have no changes, we can just persist the keepalive.
+            // This is throttled to once per minute.
+            if (performance.now() - lastEmptyResume < 60_000) {
+              continue;
+            }
           }
 
-          let changeDocument = originalChangeDocument;
-          if (originalChangeDocument?.splitEvent != null) {
-            // Handle split events from $changeStreamSplitLargeEvent.
-            // This is only relevant for very large update operations.
-            const splitEvent = originalChangeDocument?.splitEvent;
+          this.touch();
 
-            if (splitDocument == null) {
-              splitDocument = originalChangeDocument;
-            } else {
-              splitDocument = Object.assign(splitDocument, originalChangeDocument);
+          for (let eventIndex = 0; eventIndex < events.length; eventIndex++) {
+            const rawChangeDocument = events[eventIndex];
+            const originalChangeDocument = parseChangeDocument(rawChangeDocument);
+            if (this.abort_signal.aborted) {
+              break;
             }
 
-            if (splitEvent.fragment == splitEvent.of) {
-              // Got all fragments
-              changeDocument = splitDocument;
-              splitDocument = null;
-            } else {
-              // Wait for more fragments
+            if (startAfter != null && originalChangeDocument.clusterTime?.lte(startAfter)) {
               continue;
             }
-          } else if (splitDocument != null) {
-            // We were waiting for fragments, but got a different event
-            throw new ReplicationAssertionError(`Incomplete splitEvent: ${JSON.stringify(splitDocument.splitEvent)}`);
-          }
 
-          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;
-              this.logger.warn(
-                `Incorrect DB name in change stream: ${changeDocument.ns.db}. Changed to ${this.defaultDb.databaseName}.`
-              );
-            }
-          }
+            let changeDocument = originalChangeDocument;
+            if (originalChangeDocument?.splitEvent != null) {
+              // Handle split events from $changeStreamSplitLargeEvent.
+              // This is only relevant for very large update operations.
+              const splitEvent = originalChangeDocument?.splitEvent;
 
-          const ns = 'ns' in changeDocument && 'coll' in changeDocument.ns ? changeDocument.ns : undefined;
-
-          if (ns?.coll == CHECKPOINTS_COLLECTION) {
-            /**
-             * 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')
-              );
+              if (splitDocument == null) {
+                splitDocument = originalChangeDocument;
+              } else {
+                splitDocument = Object.assign(splitDocument, originalChangeDocument);
+              }
+
+              if (splitEvent.fragment == splitEvent.of) {
+                // Got all fragments
+                changeDocument = splitDocument;
+                splitDocument = null;
+              } else {
+                // Wait for more fragments
+                continue;
+              }
+            } else if (splitDocument != null) {
+              // We were waiting for fragments, but got a different event
+              throw new ReplicationAssertionError(`Incomplete splitEvent: ${JSON.stringify(splitDocument.splitEvent)}`);
             }
 
             if (
-              !(
-                changeDocument.operationType == 'insert' ||
-                changeDocument.operationType == 'update' ||
-                changeDocument.operationType == 'replace'
-              )
+              !filters.multipleDatabases &&
+              'ns' in changeDocument &&
+              changeDocument.ns.db != this.defaultDb.databaseName &&
+              changeDocument.ns.db.endsWith(`_${this.defaultDb.databaseName}`)
             ) {
-              continue;
+              // 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;
+                this.logger.warn(
+                  `Incorrect DB name in change stream: ${changeDocument.ns.db}. Changed to ${this.defaultDb.databaseName}.`
+                );
+              }
             }
 
-            // We handle two types of checkpoint events:
-            // 1. "Standalone" checkpoints, typically write checkpoints. We want to process these
-            //    immediately, regardless of where they were created.
-            // 2. "Batch" checkpoints for the current stream. This is used as a form of dynamic rate
-            //    limiting of commits, so we specifically want to exclude checkpoints from other streams.
-            //
-            // It may be useful to also throttle commits due to standalone checkpoints in the future.
-            // However, these typically have a much lower rate than batch checkpoints, so we don't do that for now.
-
-            const checkpointId = changeDocument.documentKey._id as string | mongo.ObjectId;
-
-            if (checkpointId == STANDALONE_CHECKPOINT_ID) {
-              // Standalone / write checkpoint received.
-              // When we are caught up, commit immediately to keep write checkpoint latency low.
-              // Once there is already a batch checkpoint pending, or the driver has buffered more
-              // change stream events, collapse standalone checkpoints into the normal batch
-              // checkpoint flow to avoid commit churn under sustained load.
-              if (waitForCheckpointLsn != null || this.getBufferedChangeCount(stream) > 0) {
-                if (waitForCheckpointLsn == null) {
-                  waitForCheckpointLsn = await createCheckpoint(this.client, this.defaultDb, this.checkpointStreamId);
-                }
+            const ns = 'ns' in changeDocument && 'coll' in changeDocument.ns ? changeDocument.ns : undefined;
+
+            if (ns?.coll == CHECKPOINTS_COLLECTION) {
+              /**
+               * 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.
+               * This 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 `_powersync_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('_powersync_checkpoints collection was dropped')
+                );
+              }
+
+              if (
+                !(
+                  changeDocument.operationType == 'insert' ||
+                  changeDocument.operationType == 'update' ||
+                  changeDocument.operationType == 'replace'
+                )
+              ) {
                 continue;
               }
-            } else if (!this.checkpointStreamId.equals(checkpointId)) {
-              continue;
-            }
-            const { comparable: lsn } = new MongoLSN({
-              timestamp: changeDocument.clusterTime!,
-              resume_token: changeDocument._id
-            });
-            if (batch.lastCheckpointLsn != null && lsn < batch.lastCheckpointLsn) {
-              // Checkpoint out of order - should never happen with MongoDB.
-              // If it does happen, we throw an error to stop the replication - restarting should recover.
-              // Since we use batch.lastCheckpointLsn for the next resumeAfter, this should not result in an infinite loop.
-              // Originally a workaround for https://jira.mongodb.org/browse/NODE-7042.
-              // This has been fixed in the driver in the meantime, but we still keep this as a safety-check.
-              throw new ReplicationAssertionError(
-                `Change resumeToken ${(changeDocument._id as any)._data} (${timestampToDate(changeDocument.clusterTime!).toISOString()}) is less than last checkpoint LSN ${batch.lastCheckpointLsn}. Restarting replication.`
-              );
-            }
 
-            if (waitForCheckpointLsn != null && lsn >= waitForCheckpointLsn) {
-              waitForCheckpointLsn = null;
-            }
-            const { checkpointBlocked } = await batch.commit(lsn, {
-              oldestUncommittedChange: this.replicationLag.oldestUncommittedChange
-            });
+              // We handle two types of checkpoint events:
+              // 1. "Standalone" checkpoints, typically write checkpoints. We want to process these
+              //    immediately, regardless of where they were created.
+              // 2. "Batch" checkpoints for the current stream. This is used as a form of dynamic rate
+              //    limiting of commits, so we specifically want to exclude checkpoints from other streams.
+              //
+              // It may be useful to also throttle commits due to standalone checkpoints in the future.
+              // However, these typically have a much lower rate than batch checkpoints, so we don't do that for now.
+
+              const checkpointId = changeDocument.documentKey._id as string | mongo.ObjectId;
+
+              if (checkpointId == STANDALONE_CHECKPOINT_ID) {
+                // Standalone / write checkpoint received.
+                // When we are caught up, commit immediately to keep write checkpoint latency low.
+                // Once there is already a batch checkpoint pending, or the driver has buffered more
+                // change stream events, collapse standalone checkpoints into the normal batch
+                // checkpoint flow to avoid commit churn under sustained load.
+                const hasBufferedChanges = eventIndex < events.length - 1;
+                if (waitForCheckpointLsn != null || hasBufferedChanges) {
+                  if (waitForCheckpointLsn == null) {
+                    waitForCheckpointLsn = await createCheckpoint(this.client, this.defaultDb, this.checkpointStreamId);
+                  }
+                  continue;
+                }
+              } else if (!this.checkpointStreamId.equals(checkpointId)) {
+                continue;
+              }
+              const { comparable: lsn } = new MongoLSN({
+                timestamp: changeDocument.clusterTime!,
+                resume_token: changeDocument._id
+              });
+              if (batch.lastCheckpointLsn != null && lsn < batch.lastCheckpointLsn) {
+                // Checkpoint out of order - should never happen with MongoDB.
+                // If it does happen, we throw an error to stop the replication - restarting should recover.
+                // Since we use batch.lastCheckpointLsn for the next resumeAfter, this should not result in an infinite loop.
+                // Originally a workaround for https://jira.mongodb.org/browse/NODE-7042.
+                // This has been fixed in the driver in the meantime, but we still keep this as a safety-check.
+                throw new ReplicationAssertionError(
+                  `Change resumeToken ${(changeDocument._id as any)._data} (${timestampToDate(changeDocument.clusterTime!).toISOString()}) is less than last checkpoint LSN ${batch.lastCheckpointLsn}. Restarting replication.`
+                );
+              }
 
-            if (!checkpointBlocked) {
-              this.replicationLag.markCommitted();
-              changesSinceLastCheckpoint = 0;
-            }
-          } else if (
-            changeDocument.operationType == 'insert' ||
-            changeDocument.operationType == 'update' ||
-            changeDocument.operationType == 'replace' ||
-            changeDocument.operationType == 'delete'
-          ) {
-            if (waitForCheckpointLsn == null) {
-              waitForCheckpointLsn = await createCheckpoint(this.client, this.defaultDb, this.checkpointStreamId);
-            }
+              if (waitForCheckpointLsn != null && lsn >= waitForCheckpointLsn) {
+                waitForCheckpointLsn = null;
+              }
+              const { checkpointBlocked } = await batch.commit(lsn, {
+                oldestUncommittedChange: this.replicationLag.oldestUncommittedChange
+              });
 
-            const rel = getMongoRelation(changeDocument.ns);
-            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) {
-              this.replicationLag.trackUncommittedChange(
-                changeDocument.clusterTime == null ? null : timestampToDate(changeDocument.clusterTime)
-              );
+              if (!checkpointBlocked) {
+                this.replicationLag.markCommitted();
+              }
+            } else if (
+              changeDocument.operationType == 'insert' ||
+              changeDocument.operationType == 'update' ||
+              changeDocument.operationType == 'replace' ||
+              changeDocument.operationType == 'delete'
+            ) {
+              if (waitForCheckpointLsn == null) {
+                waitForCheckpointLsn = await createCheckpoint(this.client, this.defaultDb, this.checkpointStreamId);
+              }
 
-              const transactionKeyValue = transactionKey(changeDocument);
+              const rel = getMongoRelation(changeDocument.ns);
+              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 config was deployed.
+                snapshot: true
+              });
+              if (table.syncAny) {
+                this.replicationLag.trackUncommittedChange(
+                  changeDocument.clusterTime == null ? null : timestampToDate(changeDocument.clusterTime)
+                );
+
+                const transactionKeyValue = transactionKey(changeDocument);
+
+                if (transactionKeyValue == null || lastTxnKey != transactionKeyValue) {
+                  // Very crude metric for counting transactions replicated.
+                  // We ignore operations other than basic CRUD, and ignore changes to _powersync_checkpoints.
+                  // Individual writes may not have a txnNumber, in which case we count them as separate transactions.
+                  lastTxnKey = transactionKeyValue;
+                  transactionsReplicatedMetric.add(1);
+                }
 
-              if (transactionKeyValue == null || lastTxnKey != transactionKeyValue) {
-                // Very crude metric for counting transactions replicated.
-                // We ignore operations other than basic CRUD, and ignore changes to _powersync_checkpoints.
-                // Individual writes may not have a txnNumber, in which case we count them as separate transactions.
-                lastTxnKey = transactionKeyValue;
-                transactionsReplicatedMetric.add(1);
+                await this.writeChange(batch, table, changeDocument);
               }
-
-              const flushResult = await this.writeChange(batch, table, changeDocument);
-              changesSinceLastCheckpoint += 1;
-              if (flushResult != null && changesSinceLastCheckpoint >= 20_000) {
-                // When we are catching up replication after an initial snapshot, there may be a very long delay
-                // before we do a commit(). In that case, we need to periodically persist the resume LSN, so
-                // we don't restart from scratch if we restart replication.
-                // The same could apply if we need to catch up on replication after some downtime.
-                const { comparable: lsn } = new MongoLSN({
-                  timestamp: changeDocument.clusterTime!,
-                  resume_token: changeDocument._id
-                });
-                this.logger.info(`Updating resume LSN to ${lsn} after ${changesSinceLastCheckpoint} changes`);
-                await batch.setResumeLsn(lsn);
-                changesSinceLastCheckpoint = 0;
+            } else if (changeDocument.operationType == 'drop') {
+              const rel = getMongoRelation(changeDocument.ns);
+              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.relationCache.delete(table);
               }
+            } else if (changeDocument.operationType == 'rename') {
+              const relFrom = getMongoRelation(changeDocument.ns);
+              const relTo = getMongoRelation(changeDocument.to);
+              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.relationCache.delete(relFrom);
+              }
+              // Here we do need to snapshot the new table
+              const collection = await this.getCollectionInfo(relTo.schema, relTo.name);
+              await this.handleRelation(batch, relTo, {
+                // This is a new (renamed) collection, so always snapshot it.
+                snapshot: true,
+                collectionInfo: collection
+              });
             }
-          } else if (changeDocument.operationType == 'drop') {
-            const rel = getMongoRelation(changeDocument.ns);
-            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.relationCache.delete(table);
-            }
-          } else if (changeDocument.operationType == 'rename') {
-            const relFrom = getMongoRelation(changeDocument.ns);
-            const relTo = getMongoRelation(changeDocument.to);
-            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.relationCache.delete(relFrom);
-            }
-            // Here we do need to snapshot the new table
-            const collection = await this.getCollectionInfo(relTo.schema, relTo.name);
-            await this.handleRelation(batch, relTo, {
-              // This is a new (renamed) collection, so always snapshot it.
-              snapshot: true,
-              collectionInfo: collection
-            });
           }
+
+          if (splitDocument == null) {
+            // We flush and mark progress on every batch of data we receive.
+            // Batches are generally large (64MB or 6000 events, whichever comes first),
+            // so this is a good natural point to flush and mark progress.
+            // We avoid this when splitDocument is set, since we cannot resume in the middle of a split event.
+            const { comparable: lsn } = MongoLSN.fromResumeToken(resumeToken);
+            await batch.flush({ oldestUncommittedChange: this.replicationLag.oldestUncommittedChange });
+            // TODO: We should consider making this standard behavior of flush().
+            await batch.setResumeLsn(lsn);
+          }
+
+          batchSpan.end();
+          const durations = outerSpan.end();
+          const duration = batchSpan.endAt - batchSpan.startAt;
+
+          this.logger.info(
+            `Processed batch of ${events.length} changes / ${eventBatch.byteSize} bytes in ${duration}ms`,
+            {
+              count: events.length,
+              bytes: eventBatch.byteSize,
+              duration,
+              t: durations
+            }
+          );
+          outerSpan = tracer.span('batch');
         }
       }
     );
@@ -1175,32 +1174,10 @@ export class ChangeStream {
   }
 }
 
-function mapChangeStreamError(e: any) {
-  if (isMongoNetworkTimeoutError(e)) {
-    // This typically has an unhelpful message like "connection 2 to 159.41.94.47:27017 timed out".
-    // We wrap the error to make it more useful.
-    throw new DatabaseConnectionError(ErrorCode.PSYNC_S1345, `Timeout while reading MongoDB ChangeStream`, e);
-  } else if (isMongoServerError(e) && e.codeName == 'MaxTimeMSExpired') {
-    // maxTimeMS was reached. Example message:
-    // MongoServerError: Executor error during aggregate command on namespace: powersync_test_data.$cmd.aggregate :: caused by :: operation exceeded time limit
-    throw new DatabaseConnectionError(ErrorCode.PSYNC_S1345, `Timeout while reading MongoDB ChangeStream`, e);
-  } else if (
-    isMongoServerError(e) &&
-    e.codeName == 'NoMatchingDocument' &&
-    e.errmsg?.includes('post-image was not found')
-  ) {
-    throw new ChangeStreamInvalidatedError(e.errmsg, e);
-  } else if (isMongoServerError(e) && e.hasErrorLabel('NonResumableChangeStreamError')) {
-    throw new ChangeStreamInvalidatedError(e.message, e);
-  } else {
-    throw new DatabaseConnectionError(ErrorCode.PSYNC_S1346, `Error reading MongoDB ChangeStream`, e);
-  }
-}
-
 /**
  * Transaction key for a change stream event, used to detect transaction boundaries. Returns null if the event is not part of a transaction.
  */
-function transactionKey(doc: mongo.ChangeStreamDocument): string | null {
+function transactionKey(doc: Pick<mongo.ChangeStreamDocument, 'lsid' | 'txnNumber'>): string | null {
   if (doc.txnNumber == null || doc.lsid == null) {
     return null;
   }