@powersync/service-module-mongodb-storage 0.0.0-dev-20250827091123 → 0.0.0-dev-20250828134335

package/src/storage/implementation/MongoSyncBucketStorage.ts

@@ -2,19 +2,28 @@ import * as lib_mongo from '@powersync/lib-service-mongodb';
 import { mongo } from '@powersync/lib-service-mongodb';
 import {
   BaseObserver,
+  DatabaseQueryError,
+  ErrorCode,
   logger,
   ReplicationAbortedError,
   ServiceAssertionError
 } from '@powersync/lib-services-framework';
 import {
+  addBucketChecksums,
+  addPartialChecksums,
   BroadcastIterable,
+  BucketChecksum,
   CHECKPOINT_INVALIDATE_ALL,
   CheckpointChanges,
+  CompactOptions,
   deserializeParameterLookup,
   GetCheckpointChangesOptions,
   InternalOpId,
   internalToExternalOpId,
+  isPartialChecksum,
   maxLsn,
+  PartialChecksum,
+  PartialOrFullChecksum,
   ProtocolOpId,
   ReplicationCheckpoint,
   storage,
@@ -31,9 +40,16 @@ import { PowerSyncMongo } from './db.js';
 import { BucketDataDocument, BucketDataKey, BucketStateDocument, SourceKey, SourceTableDocument } from './models.js';
 import { MongoBucketBatch } from './MongoBucketBatch.js';
 import { MongoCompactor } from './MongoCompactor.js';
-import { MongoWriteCheckpointAPI } from './MongoWriteCheckpointAPI.js';
-import { idPrefixFilter, mapOpEntry, readSingleBatch, setSessionSnapshotTime } from '../../utils/util.js';
 import { MongoParameterCompactor } from './MongoParameterCompactor.js';
+import { MongoWriteCheckpointAPI } from './MongoWriteCheckpointAPI.js';
+import {
+  CHECKSUM_QUERY_GROUP_STAGE,
+  checksumFromAggregate,
+  idPrefixFilter,
+  mapOpEntry,
+  readSingleBatch,
+  setSessionSnapshotTime
+} from './util.js';
 
 export class MongoSyncBucketStorage
   extends BaseObserver<storage.SyncRulesBucketStorageListener>
@@ -490,11 +506,96 @@ export class MongoSyncBucketStorage
     return this.checksumCache.getChecksumMap(checkpoint, buckets);
   }
 
+  clearChecksumCache() {
+    this.checksumCache.clear();
+  }
+
   private async getChecksumsInternal(batch: storage.FetchPartialBucketChecksum[]): Promise<storage.PartialChecksumMap> {
     if (batch.length == 0) {
       return new Map();
     }
 
+    const preFilters: any[] = [];
+    for (let request of batch) {
+      if (request.start == null) {
+        preFilters.push({
+          _id: {
+            g: this.group_id,
+            b: request.bucket
+          },
+          'compacted_state.op_id': { $exists: true, $lte: request.end }
+        });
+      }
+    }
+
+    const preStates = new Map<string, { opId: InternalOpId; checksum: BucketChecksum }>();
+
+    if (preFilters.length > 0) {
+      // For un-cached bucket checksums, attempt to use the compacted state first.
+      const states = await this.db.bucket_state
+        .find({
+          $or: preFilters
+        })
+        .toArray();
+      for (let state of states) {
+        const compactedState = state.compacted_state!;
+        preStates.set(state._id.b, {
+          opId: compactedState.op_id,
+          checksum: {
+            bucket: state._id.b,
+            checksum: Number(compactedState.checksum),
+            count: compactedState.count
+          }
+        });
+      }
+    }
+
+    const mappedRequests = batch.map((request) => {
+      let start = request.start;
+      if (start == null) {
+        const preState = preStates.get(request.bucket);
+        if (preState != null) {
+          start = preState.opId;
+        }
+      }
+      return {
+        ...request,
+        start
+      };
+    });
+
+    const queriedChecksums = await this.queryPartialChecksums(mappedRequests);
+
+    return new Map<string, storage.PartialOrFullChecksum>(
+      batch.map((request) => {
+        const bucket = request.bucket;
+        // Could be null if this is either (1) a partial request, or (2) no compacted checksum was available
+        const preState = preStates.get(bucket);
+        // Could be null if we got no data
+        const partialChecksum = queriedChecksums.get(bucket);
+        const merged = addPartialChecksums(bucket, preState?.checksum ?? null, partialChecksum ?? null);
+
+        return [bucket, merged];
+      })
+    );
+  }
+
+  async queryPartialChecksums(batch: storage.FetchPartialBucketChecksum[]): Promise<storage.PartialChecksumMap> {
+    try {
+      return await this.queryPartialChecksumsInternal(batch);
+    } catch (e) {
+      if (e.codeName == 'MaxTimeMSExpired') {
+        logger.warn(`Checksum query timed out; falling back to slower version`, e);
+        // Timeout - try the slower but more robust version
+        return await this.queryPartialChecksumsFallback(batch);
+      }
+      throw lib_mongo.mapQueryError(e, 'while reading checksums');
+    }
+  }
+
+  private async queryPartialChecksumsInternal(
+    batch: storage.FetchPartialBucketChecksum[]
+  ): Promise<storage.PartialChecksumMap> {
     const filters: any[] = [];
     for (let request of batch) {
       filters.push({
@@ -502,12 +603,12 @@ export class MongoSyncBucketStorage
           $gt: {
             g: this.group_id,
             b: request.bucket,
-            o: request.start ? BigInt(request.start) : new bson.MinKey()
+            o: request.start ?? new bson.MinKey()
           },
           $lte: {
             g: this.group_id,
             b: request.bucket,
-            o: BigInt(request.end)
+            o: request.end
           }
         }
       });
@@ -521,44 +622,126 @@ export class MongoSyncBucketStorage
               $or: filters
             }
           },
-          {
-            $group: {
-              _id: '$_id.b',
-              // Historically, checksum may be stored as 'int' or 'double'.
-              // More recently, this should be a 'long'.
-              // $toLong ensures that we always sum it as a long, avoiding inaccuracies in the calculations.
-              checksum_total: { $sum: { $toLong: '$checksum' } },
-              count: { $sum: 1 },
-              has_clear_op: {
-                $max: {
-                  $cond: [{ $eq: ['$op', 'CLEAR'] }, 1, 0]
-                }
-              }
-            }
-          }
+          CHECKSUM_QUERY_GROUP_STAGE
         ],
-        { session: undefined, readConcern: 'snapshot', maxTimeMS: lib_mongo.db.MONGO_OPERATION_TIMEOUT_MS }
+        { session: undefined, readConcern: 'snapshot', maxTimeMS: lib_mongo.MONGO_CHECKSUM_TIMEOUT_MS }
       )
-      .toArray()
-      .catch((e) => {
-        throw lib_mongo.mapQueryError(e, 'while reading checksums');
-      });
+      // Don't map the error here - we want to keep timeout errors as-is
+      .toArray();
 
-    return new Map<string, storage.PartialChecksum>(
+    const partialChecksums = new Map<string, storage.PartialOrFullChecksum>(
       aggregate.map((doc) => {
-        return [
-          doc._id,
-          {
-            bucket: doc._id,
-            partialCount: doc.count,
-            partialChecksum: Number(BigInt(doc.checksum_total) & 0xffffffffn) & 0xffffffff,
-            isFullChecksum: doc.has_clear_op == 1
-          } satisfies storage.PartialChecksum
-        ];
+        const bucket = doc._id;
+        return [bucket, checksumFromAggregate(doc)];
+      })
+    );
+
+    return new Map<string, storage.PartialOrFullChecksum>(
+      batch.map((request) => {
+        const bucket = request.bucket;
+        // Could be null if we got no data
+        let partialChecksum = partialChecksums.get(bucket);
+        if (partialChecksum == null) {
+          partialChecksum = {
+            bucket,
+            partialCount: 0,
+            partialChecksum: 0
+          };
+        }
+        if (request.start == null && isPartialChecksum(partialChecksum)) {
+          partialChecksum = {
+            bucket,
+            count: partialChecksum.partialCount,
+            checksum: partialChecksum.partialChecksum
+          };
+        }
+
+        return [bucket, partialChecksum];
       })
     );
   }
 
+  /**
+   * Checksums for large buckets can run over the query timeout.
+   * To avoid this, we query in batches.
+   * This version can handle larger amounts of data, but is slower, especially for many buckets.
+   */
+  async queryPartialChecksumsFallback(
+    batch: storage.FetchPartialBucketChecksum[]
+  ): Promise<storage.PartialChecksumMap> {
+    const partialChecksums = new Map<string, storage.PartialOrFullChecksum>();
+    for (let request of batch) {
+      const checksum = await this.slowChecksum(request);
+      partialChecksums.set(request.bucket, checksum);
+    }
+
+    return partialChecksums;
+  }
+
+  private async slowChecksum(request: storage.FetchPartialBucketChecksum): Promise<PartialOrFullChecksum> {
+    const batchLimit = 50_000;
+
+    let lowerBound = 0n;
+    const bucket = request.bucket;
+
+    let runningChecksum: PartialOrFullChecksum = {
+      bucket,
+      partialCount: 0,
+      partialChecksum: 0
+    };
+    if (request.start == null) {
+      runningChecksum = {
+        bucket,
+        count: 0,
+        checksum: 0
+      };
+    }
+
+    while (true) {
+      const filter = {
+        _id: {
+          $gt: {
+            g: this.group_id,
+            b: bucket,
+            o: lowerBound
+          },
+          $lte: {
+            g: this.group_id,
+            b: bucket,
+            o: request.end
+          }
+        }
+      };
+      const docs = await this.db.bucket_data
+        .aggregate(
+          [
+            {
+              $match: filter
+            },
+            // sort and limit _before_ grouping
+            { $sort: { _id: 1 } },
+            { $limit: batchLimit },
+            CHECKSUM_QUERY_GROUP_STAGE
+          ],
+          { session: undefined, readConcern: 'snapshot', maxTimeMS: lib_mongo.MONGO_CHECKSUM_TIMEOUT_MS }
+        )
+        .toArray();
+      const doc = docs[0];
+      if (doc == null) {
+        return runningChecksum;
+      }
+      const partial = checksumFromAggregate(doc);
+      runningChecksum = addPartialChecksums(bucket, runningChecksum, partial);
+      const isFinal = doc.count != batchLimit;
+      if (isFinal) {
+        break;
+      } else {
+        lowerBound = doc.last_op;
+      }
+    }
+    return runningChecksum;
+  }
+
   async terminate(options?: storage.TerminateOptions) {
     // Default is to clear the storage except when explicitly requested not to.
     if (!options || options?.clearStorage) {
@@ -701,11 +884,32 @@ export class MongoSyncBucketStorage
   }
 
   async compact(options?: storage.CompactOptions) {
-    const checkpoint = await this.getCheckpointInternal();
-    await new MongoCompactor(this.db, this.group_id, options).compact();
-    if (checkpoint != null && options?.compactParameterData) {
-      await new MongoParameterCompactor(this.db, this.group_id, checkpoint.checkpoint, options).compact();
+    let maxOpId = options?.maxOpId;
+    if (maxOpId == null) {
+      const checkpoint = await this.getCheckpointInternal();
+      maxOpId = checkpoint?.checkpoint ?? undefined;
     }
+    await new MongoCompactor(this, this.db, { ...options, maxOpId }).compact();
+
+    if (maxOpId != null && options?.compactParameterData) {
+      await new MongoParameterCompactor(this.db, this.group_id, maxOpId, options).compact();
+    }
+  }
+
+  async populatePersistentChecksumCache(options: Required<Pick<CompactOptions, 'signal' | 'maxOpId'>>): Promise<void> {
+    logger.info(`Populating persistent checksum cache...`);
+    const start = Date.now();
+    // We do a minimal compact here.
+    // We can optimize this in the future.
+    const compactor = new MongoCompactor(this, this.db, {
+      ...options,
+      // Don't track updates for MOVE compacting
+      memoryLimitMB: 0
+    });
+
+    await compactor.populateChecksums();
+    const duration = Date.now() - start;
+    logger.info(`Populated persistent checksum cache in ${(duration / 1000).toFixed(1)}s`);
   }
 
   /**

package/src/storage/implementation/MongoTestStorageFactoryGenerator.ts

@@ -0,0 +1,28 @@
+import { TestStorageOptions } from '@powersync/service-core';
+import { MongoBucketStorage } from '../MongoBucketStorage.js';
+import { connectMongoForTests } from './util.js';
+
+export type MongoTestStorageOptions = {
+  url: string;
+  isCI: boolean;
+};
+
+export const MongoTestStorageFactoryGenerator = (factoryOptions: MongoTestStorageOptions) => {
+  return async (options?: TestStorageOptions) => {
+    const db = connectMongoForTests(factoryOptions.url, factoryOptions.isCI);
+
+    // None of the tests insert data into this collection, so it was never created
+    if (!(await db.db.listCollections({ name: db.bucket_parameters.collectionName }).hasNext())) {
+      await db.db.createCollection('bucket_parameters');
+    }
+
+    // Full migrations are not currently run for tests, so we manually create this
+    await db.createCheckpointEventsCollection();
+
+    if (!options?.doNotClear) {
+      await db.clear();
+    }
+
+    return new MongoBucketStorage(db, { slot_name_prefix: 'test_' });
+  };
+};

package/src/storage/implementation/PersistedBatch.ts

@@ -16,7 +16,7 @@ import {
   CurrentDataDocument,
   SourceKey
 } from './models.js';
-import { replicaIdToSubkey } from '../../utils/util.js';
+import { replicaIdToSubkey } from './util.js';
 
 /**
  * Maximum size of operations we write in a single transaction.
@@ -71,15 +71,17 @@ export class PersistedBatch {
     this.logger = options?.logger ?? defaultLogger;
   }
 
-  private incrementBucket(bucket: string, op_id: InternalOpId) {
+  private incrementBucket(bucket: string, op_id: InternalOpId, bytes: number) {
     let existingState = this.bucketStates.get(bucket);
     if (existingState) {
       existingState.lastOp = op_id;
       existingState.incrementCount += 1;
+      existingState.incrementBytes += bytes;
     } else {
       this.bucketStates.set(bucket, {
         lastOp: op_id,
-        incrementCount: 1
+        incrementCount: 1,
+        incrementBytes: bytes
       });
     }
   }
@@ -115,7 +117,8 @@ export class PersistedBatch {
       }
 
       remaining_buckets.delete(key);
-      this.currentSize += recordData.length + 200;
+      const byteEstimate = recordData.length + 200;
+      this.currentSize += byteEstimate;
 
       const op_id = options.op_seq.next();
       this.debugLastOpId = op_id;
@@ -138,7 +141,7 @@ export class PersistedBatch {
           }
         }
       });
-      this.incrementBucket(k.bucket, op_id);
+      this.incrementBucket(k.bucket, op_id, byteEstimate);
     }
 
     for (let bd of remaining_buckets.values()) {
@@ -166,7 +169,7 @@ export class PersistedBatch {
         }
       });
       this.currentSize += 200;
-      this.incrementBucket(bd.bucket, op_id);
+      this.incrementBucket(bd.bucket, op_id, 200);
     }
   }
 
@@ -369,6 +372,10 @@ export class PersistedBatch {
           update: {
             $set: {
               last_op: state.lastOp
+            },
+            $inc: {
+              'estimate_since_compact.count': state.incrementCount,
+              'estimate_since_compact.bytes': state.incrementBytes
             }
           },
           upsert: true
@@ -381,4 +388,5 @@ export class PersistedBatch {
 interface BucketStateUpdate {
   lastOp: InternalOpId;
   incrementCount: number;
+  incrementBytes: number;
 }

package/src/storage/implementation/db.ts

@@ -8,7 +8,6 @@ import {
   BucketParameterDocument,
   BucketStateDocument,
   CheckpointEventDocument,
-  ClientConnectionDocument,
   CurrentDataDocument,
   CustomWriteCheckpointDocument,
   IdSequenceDocument,
@@ -38,7 +37,6 @@ export class PowerSyncMongo {
   readonly locks: mongo.Collection<lib_mongo.locks.Lock>;
   readonly bucket_state: mongo.Collection<BucketStateDocument>;
   readonly checkpoint_events: mongo.Collection<CheckpointEventDocument>;
-  readonly connection_report_events: mongo.Collection<ClientConnectionDocument>;
 
   readonly client: mongo.MongoClient;
   readonly db: mongo.Db;
@@ -63,7 +61,6 @@ export class PowerSyncMongo {
     this.locks = this.db.collection('locks');
     this.bucket_state = this.db.collection('bucket_state');
     this.checkpoint_events = this.db.collection('checkpoint_events');
-    this.connection_report_events = this.db.collection('connection_report_events');
   }
 
   /**
@@ -81,7 +78,6 @@ export class PowerSyncMongo {
     await this.locks.deleteMany({});
     await this.bucket_state.deleteMany({});
     await this.custom_write_checkpoints.deleteMany({});
-    await this.connection_report_events.deleteMany({});
   }
 
   /**
@@ -131,20 +127,6 @@ export class PowerSyncMongo {
       max: 50 // max number of documents
     });
   }
-
-  /**
-   * Only use in migrations and tests.
-   */
-  async createConnectionReportingCollection() {
-    const existingCollections = await this.db
-      .listCollections({ name: 'connection_report_events' }, { nameOnly: false })
-      .toArray();
-    const collection = existingCollections[0];
-    if (collection != null) {
-      return;
-    }
-    await this.db.createCollection('connection_report_events');
-  }
 }
 
 export function createPowerSyncMongo(config: MongoStorageConfig, options?: lib_mongo.MongoConnectionOptions) {

package/src/storage/implementation/models.ts

@@ -1,7 +1,6 @@
 import { InternalOpId, storage } from '@powersync/service-core';
 import { SqliteJsonValue } from '@powersync/service-sync-rules';
 import * as bson from 'bson';
-import { event_types } from '@powersync/service-types';
 
 /**
  * Replica id uniquely identifying a row on the source database.
@@ -99,6 +98,21 @@ export interface BucketStateDocument {
     b: string;
   };
   last_op: bigint;
+  /**
+   * If set, this can be treated as "cache" of a checksum at a specific point.
+   * Can be updated periodically, for example by the compact job.
+   */
+  compacted_state?: {
+    op_id: InternalOpId;
+    count: number;
+    checksum: bigint;
+    bytes: number;
+  };
+
+  estimate_since_compact?: {
+    count: number;
+    bytes: number;
+  };
 }
 
 export interface IdSequenceDocument {
@@ -220,5 +234,3 @@ export interface InstanceDocument {
   // The instance UUID
   _id: string;
 }
-
-export interface ClientConnectionDocument extends event_types.ClientConnection {}

package/src/{utils → storage/implementation}/util.ts

@@ -1,9 +1,12 @@
 import * as bson from 'bson';
 import * as crypto from 'crypto';
 import * as uuid from 'uuid';
+
 import { mongo } from '@powersync/lib-service-mongodb';
-import { storage, utils } from '@powersync/service-core';
-import { BucketDataDocument } from '../storage/implementation/models.js';
+import { BucketChecksum, PartialChecksum, PartialOrFullChecksum, storage, utils } from '@powersync/service-core';
+
+import { PowerSyncMongo } from './db.js';
+import { BucketDataDocument } from './models.js';
 import { ServiceAssertionError } from '@powersync/lib-services-framework';
 
 export function idPrefixFilter<T>(prefix: Partial<T>, rest: (keyof T)[]): mongo.Condition<T> {
@@ -40,7 +43,7 @@ export function generateSlotName(prefix: string, sync_rules_id: number) {
  *
  * For this to be effective, set batchSize = limit in the find command.
  */
-export async function readSingleBatch<T>(cursor: mongo.FindCursor<T>): Promise<{ data: T[]; hasMore: boolean }> {
+export async function readSingleBatch<T>(cursor: mongo.AbstractCursor<T>): Promise<{ data: T[]; hasMore: boolean }> {
   try {
     let data: T[];
     let hasMore = true;
@@ -102,6 +105,20 @@ export function replicaIdToSubkey(table: bson.ObjectId, id: storage.ReplicaId):
   }
 }
 
+/**
+ * Helper for unit tests
+ */
+export const connectMongoForTests = (url: string, isCI: boolean) => {
+  // Short timeout for tests, to fail fast when the server is not available.
+  // Slightly longer timeouts for CI, to avoid arbitrary test failures
+  const client = new mongo.MongoClient(url, {
+    connectTimeoutMS: isCI ? 15_000 : 5_000,
+    socketTimeoutMS: isCI ? 15_000 : 5_000,
+    serverSelectionTimeoutMS: isCI ? 15_000 : 2_500
+  });
+  return new PowerSyncMongo(client);
+};
+
 export function setSessionSnapshotTime(session: mongo.ClientSession, time: bson.Timestamp) {
   // This is a workaround for the lack of direct support for snapshot reads in the MongoDB driver.
   if (!session.snapshotEnabled) {
@@ -113,3 +130,44 @@ export function setSessionSnapshotTime(session: mongo.ClientSession, time: bson.
     throw new ServiceAssertionError(`Session snapshotTime is already set`);
   }
 }
+
+export const CHECKSUM_QUERY_GROUP_STAGE = {
+  $group: {
+    _id: '$_id.b',
+    // Historically, checksum may be stored as 'int' or 'double'.
+    // More recently, this should be a 'long'.
+    // $toLong ensures that we always sum it as a long, avoiding inaccuracies in the calculations.
+    checksum_total: { $sum: { $toLong: '$checksum' } },
+    count: { $sum: 1 },
+    has_clear_op: {
+      $max: {
+        $cond: [{ $eq: ['$op', 'CLEAR'] }, 1, 0]
+      }
+    },
+    last_op: { $max: '$_id.o' }
+  }
+};
+
+/**
+ * Convert output of CHECKSUM_QUERY_GROUP_STAGE into a checksum.
+ */
+export function checksumFromAggregate(doc: bson.Document): PartialOrFullChecksum {
+  const partialChecksum = Number(BigInt(doc.checksum_total) & 0xffffffffn) & 0xffffffff;
+  const bucket = doc._id;
+
+  if (doc.has_clear_op == 1) {
+    return {
+      // full checksum - replaces any previous one
+      bucket,
+      checksum: partialChecksum,
+      count: doc.count
+    } satisfies BucketChecksum;
+  } else {
+    return {
+      // partial checksum - is added to a previous one
+      bucket,
+      partialCount: doc.count,
+      partialChecksum
+    } satisfies PartialChecksum;
+  }
+}

package/src/storage/storage-index.ts

@@ -7,9 +7,8 @@ export * from './implementation/MongoPersistedSyncRulesContent.js';
 export * from './implementation/MongoStorageProvider.js';
 export * from './implementation/MongoSyncBucketStorage.js';
 export * from './implementation/MongoSyncRulesLock.js';
+export * from './implementation/MongoTestStorageFactoryGenerator.js';
 export * from './implementation/OperationBatch.js';
 export * from './implementation/PersistedBatch.js';
-export * from '../utils/util.js';
+export * from './implementation/util.js';
 export * from './MongoBucketStorage.js';
-export * from './MongoReportStorage.js';
-export * as test_utils from '../utils/test-utils.js';