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

package/src/storage/implementation/MongoSyncBucketStorage.ts

@@ -2,17 +2,12 @@ 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,
@@ -20,10 +15,7 @@ import {
   GetCheckpointChangesOptions,
   InternalOpId,
   internalToExternalOpId,
-  isPartialChecksum,
   maxLsn,
-  PartialChecksum,
-  PartialOrFullChecksum,
   ProtocolOpId,
   ReplicationCheckpoint,
   storage,
@@ -39,28 +31,19 @@ import { MongoBucketStorage } from '../MongoBucketStorage.js';
 import { PowerSyncMongo } from './db.js';
 import { BucketDataDocument, BucketDataKey, BucketStateDocument, SourceKey, SourceTableDocument } from './models.js';
 import { MongoBucketBatch } from './MongoBucketBatch.js';
+import { MongoChecksums } from './MongoChecksums.js';
 import { MongoCompactor } from './MongoCompactor.js';
 import { MongoParameterCompactor } from './MongoParameterCompactor.js';
 import { MongoWriteCheckpointAPI } from './MongoWriteCheckpointAPI.js';
-import {
-  CHECKSUM_QUERY_GROUP_STAGE,
-  checksumFromAggregate,
-  idPrefixFilter,
-  mapOpEntry,
-  readSingleBatch,
-  setSessionSnapshotTime
-} from './util.js';
+import { idPrefixFilter, mapOpEntry, readSingleBatch, setSessionSnapshotTime } from '../../utils/util.js';
+
 
 export class MongoSyncBucketStorage
   extends BaseObserver<storage.SyncRulesBucketStorageListener>
   implements storage.SyncRulesBucketStorage
 {
   private readonly db: PowerSyncMongo;
-  private checksumCache = new storage.ChecksumCache({
-    fetchChecksums: (batch) => {
-      return this.getChecksumsInternal(batch);
-    }
-  });
+  readonly checksums: MongoChecksums;
 
   private parsedSyncRulesCache: { parsed: SqlSyncRules; options: storage.ParseSyncRulesOptions } | undefined;
   private writeCheckpointAPI: MongoWriteCheckpointAPI;
@@ -74,6 +57,7 @@ export class MongoSyncBucketStorage
   ) {
     super();
     this.db = factory.db;
+    this.checksums = new MongoChecksums(this.db, this.group_id);
     this.writeCheckpointAPI = new MongoWriteCheckpointAPI({
       db: this.db,
       mode: writeCheckpointMode,
@@ -503,243 +487,11 @@ export class MongoSyncBucketStorage
   }
 
   async getChecksums(checkpoint: utils.InternalOpId, buckets: string[]): Promise<utils.ChecksumMap> {
-    return this.checksumCache.getChecksumMap(checkpoint, buckets);
+    return this.checksums.getChecksums(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({
-        _id: {
-          $gt: {
-            g: this.group_id,
-            b: request.bucket,
-            o: request.start ?? new bson.MinKey()
-          },
-          $lte: {
-            g: this.group_id,
-            b: request.bucket,
-            o: request.end
-          }
-        }
-      });
-    }
-
-    const aggregate = await this.db.bucket_data
-      .aggregate(
-        [
-          {
-            $match: {
-              $or: filters
-            }
-          },
-          CHECKSUM_QUERY_GROUP_STAGE
-        ],
-        { session: undefined, readConcern: 'snapshot', maxTimeMS: lib_mongo.MONGO_CHECKSUM_TIMEOUT_MS }
-      )
-      // Don't map the error here - we want to keep timeout errors as-is
-      .toArray();
-
-    const partialChecksums = new Map<string, storage.PartialOrFullChecksum>(
-      aggregate.map((doc) => {
-        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;
+    this.checksums.clearCache();
   }
 
   async terminate(options?: storage.TerminateOptions) {

package/src/storage/implementation/PersistedBatch.ts

@@ -16,7 +16,7 @@ import {
   CurrentDataDocument,
   SourceKey
 } from './models.js';
-import { replicaIdToSubkey } from './util.js';
+import { replicaIdToSubkey } from '../../utils/util.js';
 
 /**
  * Maximum size of operations we write in a single transaction.

package/src/storage/implementation/db.ts

@@ -8,6 +8,7 @@ import {
   BucketParameterDocument,
   BucketStateDocument,
   CheckpointEventDocument,
+  ClientConnectionDocument,
   CurrentDataDocument,
   CustomWriteCheckpointDocument,
   IdSequenceDocument,
@@ -37,6 +38,7 @@ 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;
@@ -61,6 +63,7 @@ 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');
   }
 
   /**
@@ -78,6 +81,7 @@ export class PowerSyncMongo {
     await this.locks.deleteMany({});
     await this.bucket_state.deleteMany({});
     await this.custom_write_checkpoints.deleteMany({});
+    await this.connection_report_events.deleteMany({});
   }
 
   /**
@@ -127,6 +131,20 @@ 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,6 +1,7 @@
 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.
@@ -106,7 +107,7 @@ export interface BucketStateDocument {
     op_id: InternalOpId;
     count: number;
     checksum: bigint;
-    bytes: number;
+    bytes: number | null;
   };
 
   estimate_since_compact?: {
@@ -234,3 +235,5 @@ export interface InstanceDocument {
   // The instance UUID
   _id: string;
 }
+
+export interface ClientConnectionDocument extends event_types.ClientConnection {}

package/src/storage/storage-index.ts

@@ -7,8 +7,9 @@ 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 './implementation/util.js';
+export * from '../utils/util.js';
 export * from './MongoBucketStorage.js';
+export * from './MongoReportStorage.js';
+export * as test_utils from '../utils/test-utils.js';

package/src/utils/test-utils.ts

@@ -0,0 +1,55 @@
+import { mongo } from '@powersync/lib-service-mongodb';
+import { PowerSyncMongo } from '../storage/implementation/db.js';
+import { TestStorageOptions } from '@powersync/service-core';
+import { MongoReportStorage } from '../storage/MongoReportStorage.js';
+import { MongoBucketStorage } from '../storage/MongoBucketStorage.js';
+
+export type MongoTestStorageOptions = {
+  url: string;
+  isCI: boolean;
+};
+
+export function 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_' });
+  };
+}
+
+export function mongoTestReportStorageFactoryGenerator(factoryOptions: MongoTestStorageOptions) {
+  return async (options?: TestStorageOptions) => {
+    const db = connectMongoForTests(factoryOptions.url, factoryOptions.isCI);
+
+    await db.createConnectionReportingCollection();
+
+    if (!options?.doNotClear) {
+      await db.clear();
+    }
+
+    return new MongoReportStorage(db);
+  };
+}
+
+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);
+};

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

@@ -3,11 +3,9 @@ import * as crypto from 'crypto';
 import * as uuid from 'uuid';
 
 import { mongo } from '@powersync/lib-service-mongodb';
-import { BucketChecksum, PartialChecksum, PartialOrFullChecksum, storage, utils } from '@powersync/service-core';
-
-import { PowerSyncMongo } from './db.js';
-import { BucketDataDocument } from './models.js';
+import { storage, utils } from '@powersync/service-core';
 import { ServiceAssertionError } from '@powersync/lib-services-framework';
+import { BucketDataDocument } from '../storage/implementation/models.js';
 
 export function idPrefixFilter<T>(prefix: Partial<T>, rest: (keyof T)[]): mongo.Condition<T> {
   let filter = {
@@ -105,20 +103,6 @@ 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) {
@@ -130,44 +114,3 @@ 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/utils/utils-index.ts

@@ -0,0 +1,2 @@
+export * as test_utils from './test-utils.js';
+export * from './util.js';