@powersync/service-module-mongodb-storage 0.0.0-dev-20250214100224 → 0.0.0-dev-20250303114151

package/src/storage/implementation/MongoSyncBucketStorage.ts

@@ -1,20 +1,44 @@
 import * as lib_mongo from '@powersync/lib-service-mongodb';
 import { mongo } from '@powersync/lib-service-mongodb';
-import { DisposableObserver, logger, ServiceAssertionError } from '@powersync/lib-services-framework';
-import { storage, utils } from '@powersync/service-core';
+import {
+  BaseObserver,
+  ErrorCode,
+  logger,
+  ServiceAssertionError,
+  ServiceError
+} from '@powersync/lib-services-framework';
+import {
+  BroadcastIterable,
+  CHECKPOINT_INVALIDATE_ALL,
+  CheckpointChanges,
+  GetCheckpointChangesOptions,
+  ReplicationCheckpoint,
+  SourceTable,
+  storage,
+  utils,
+  WatchWriteCheckpointOptions
+} from '@powersync/service-core';
 import { SqliteJsonRow, SqliteJsonValue, SqlSyncRules } from '@powersync/service-sync-rules';
 import * as bson from 'bson';
+import { wrapWithAbort } from 'ix/asynciterable/operators/withabort.js';
 import * as timers from 'timers/promises';
 import { MongoBucketStorage } from '../MongoBucketStorage.js';
 import { PowerSyncMongo } from './db.js';
-import { BucketDataDocument, BucketDataKey, SourceKey } from './models.js';
+import {
+  BucketDataDocument,
+  BucketDataKey,
+  SourceKey,
+  SourceTableDocument,
+  SyncRuleCheckpointState,
+  SyncRuleDocument
+} from './models.js';
 import { MongoBucketBatch } from './MongoBucketBatch.js';
 import { MongoCompactor } from './MongoCompactor.js';
 import { MongoWriteCheckpointAPI } from './MongoWriteCheckpointAPI.js';
 import { idPrefixFilter, mapOpEntry, readSingleBatch } from './util.js';
 
 export class MongoSyncBucketStorage
-  extends DisposableObserver<storage.SyncRulesBucketStorageListener>
+  extends BaseObserver<storage.SyncRulesBucketStorageListener>
   implements storage.SyncRulesBucketStorage
 {
   private readonly db: PowerSyncMongo;
@@ -147,17 +171,17 @@ export class MongoSyncBucketStorage
     let result: storage.ResolveTableResult | null = null;
     await this.db.client.withSession(async (session) => {
       const col = this.db.source_tables;
-      let doc = await col.findOne(
-        {
-          group_id: group_id,
-          connection_id: connection_id,
-          relation_id: objectId,
-          schema_name: schema,
-          table_name: table,
-          replica_id_columns2: columns
-        },
-        { session }
-      );
+      let filter: Partial<SourceTableDocument> = {
+        group_id: group_id,
+        connection_id: connection_id,
+        schema_name: schema,
+        table_name: table,
+        replica_id_columns2: columns
+      };
+      if (objectId != null) {
+        filter.relation_id = objectId;
+      }
+      let doc = await col.findOne(filter, { session });
       if (doc == null) {
         doc = {
           _id: new bson.ObjectId(),
@@ -186,31 +210,40 @@ export class MongoSyncBucketStorage
       sourceTable.syncData = options.sync_rules.tableSyncsData(sourceTable);
       sourceTable.syncParameters = options.sync_rules.tableSyncsParameters(sourceTable);
 
+      let dropTables: storage.SourceTable[] = [];
+      // Detect tables that are either renamed, or have different replica_id_columns
+      let truncateFilter = [{ schema_name: schema, table_name: table }] as any[];
+      if (objectId != null) {
+        // Only detect renames if the source uses relation ids.
+        truncateFilter.push({ relation_id: objectId });
+      }
       const truncate = await col
         .find(
           {
             group_id: group_id,
             connection_id: connection_id,
             _id: { $ne: doc._id },
-            $or: [{ relation_id: objectId }, { schema_name: schema, table_name: table }]
+            $or: truncateFilter
           },
           { session }
         )
         .toArray();
+      dropTables = truncate.map(
+        (doc) =>
+          new storage.SourceTable(
+            doc._id,
+            connection_tag,
+            doc.relation_id,
+            doc.schema_name,
+            doc.table_name,
+            doc.replica_id_columns2?.map((c) => ({ name: c.name, typeOid: c.type_oid, type: c.type })) ?? [],
+            doc.snapshot_done ?? true
+          )
+      );
+
       result = {
         table: sourceTable,
-        dropTables: truncate.map(
-          (doc) =>
-            new storage.SourceTable(
-              doc._id,
-              connection_tag,
-              doc.relation_id ?? 0,
-              doc.schema_name,
-              doc.table_name,
-              doc.replica_id_columns2?.map((c) => ({ name: c.name, typeOid: c.type_oid, type: c.type })) ?? [],
-              doc.snapshot_done ?? true
-            )
-        )
+        dropTables: dropTables
       };
     });
     return result!;
@@ -297,13 +330,7 @@ export class MongoSyncBucketStorage
         // 1. We can calculate the document size accurately without serializing again.
         // 2. We can delay parsing the results until it's needed.
         // We manually use bson.deserialize below
-        raw: true,
-
-        // Since we're using raw: true and parsing ourselves later, we don't need bigint
-        // support here.
-        // Disabling due to https://jira.mongodb.org/browse/NODE-6165, and the fact that this
-        // is one of our most common queries.
-        useBigInt64: false
+        raw: true
       }
     ) as unknown as mongo.FindCursor<Buffer>;
 
@@ -324,7 +351,7 @@ export class MongoSyncBucketStorage
 
     // Ordered by _id, meaning buckets are grouped together
     for (let rawData of data) {
-      const row = bson.deserialize(rawData, storage.BSON_DESERIALIZE_OPTIONS) as BucketDataDocument;
+      const row = bson.deserialize(rawData, storage.BSON_DESERIALIZE_INTERNAL_OPTIONS) as BucketDataDocument;
       const bucket = row._id.b;
 
       if (currentBatch == null || currentBatch.bucket != bucket || batchSize >= sizeLimit) {
@@ -577,7 +604,7 @@ export class MongoSyncBucketStorage
           await this.db.sync_rules.updateMany(
             {
               _id: { $ne: this.group_id },
-              state: storage.SyncRuleState.ACTIVE
+              state: { $in: [storage.SyncRuleState.ACTIVE, storage.SyncRuleState.ERRORED] }
             },
             {
               $set: {
@@ -608,4 +635,232 @@ export class MongoSyncBucketStorage
   async compact(options?: storage.CompactOptions) {
     return new MongoCompactor(this.db, this.group_id, options).compact();
   }
+
+  private makeActiveCheckpoint(doc: SyncRuleCheckpointState | null) {
+    return {
+      checkpoint: utils.timestampToOpId(doc?.last_checkpoint ?? 0n),
+      lsn: doc?.last_checkpoint_lsn ?? null
+    };
+  }
+
+  /**
+   * Instance-wide watch on the latest available checkpoint (op_id + lsn).
+   */
+  private async *watchActiveCheckpoint(signal: AbortSignal): AsyncIterable<ReplicationCheckpoint> {
+    // Use this form instead of (doc: SyncRuleCheckpointState | null = null),
+    // otherwise we get weird "doc: never" issues.
+    let doc = null as SyncRuleCheckpointState | null;
+    let clusterTime = null as mongo.Timestamp | null;
+    const syncRulesId = this.group_id;
+
+    await this.db.client.withSession(async (session) => {
+      doc = await this.db.sync_rules.findOne(
+        {
+          _id: syncRulesId,
+          state: { $in: [storage.SyncRuleState.ACTIVE, storage.SyncRuleState.ERRORED] }
+        },
+        {
+          session,
+          sort: { _id: -1 },
+          limit: 1,
+          projection: {
+            _id: 1,
+            state: 1,
+            last_checkpoint: 1,
+            last_checkpoint_lsn: 1
+          }
+        }
+      );
+      const time = session.clusterTime?.clusterTime ?? null;
+      clusterTime = time;
+    });
+    if (clusterTime == null) {
+      throw new ServiceError(ErrorCode.PSYNC_S2401, 'Could not get clusterTime');
+    }
+
+    if (signal.aborted) {
+      return;
+    }
+
+    if (doc == null) {
+      // Sync rules not present or not active.
+      // Abort the connections - clients will have to retry later.
+      // Should this error instead?
+      return;
+    }
+
+    yield this.makeActiveCheckpoint(doc);
+
+    // We only watch changes to the active sync rules.
+    // If it changes to inactive, we abort and restart with the new sync rules.
+
+    const pipeline = this.getChangeStreamPipeline();
+
+    const stream = this.db.sync_rules.watch(pipeline, {
+      // Start at the cluster time where we got the initial doc, to make sure
+      // we don't skip any updates.
+      // This may result in the first operation being a duplicate, but we filter
+      // it out anyway.
+      startAtOperationTime: clusterTime
+    });
+
+    signal.addEventListener(
+      'abort',
+      () => {
+        stream.close();
+      },
+      { once: true }
+    );
+
+    let lastOp: storage.ReplicationCheckpoint | null = null;
+    let lastDoc: SyncRuleCheckpointState | null = doc;
+
+    for await (const update of stream.stream()) {
+      if (signal.aborted) {
+        break;
+      }
+      if (update.operationType != 'insert' && update.operationType != 'update' && update.operationType != 'replace') {
+        continue;
+      }
+
+      const doc = await this.getOperationDoc(lastDoc, update as lib_mongo.mongo.ChangeStreamDocument<SyncRuleDocument>);
+      if (doc == null) {
+        // Irrelevant update
+        continue;
+      }
+      if (doc.state != storage.SyncRuleState.ACTIVE && doc.state != storage.SyncRuleState.ERRORED) {
+        // Sync rules have changed - abort and restart.
+        // Should this error instead?
+        break;
+      }
+
+      lastDoc = doc;
+
+      const op = this.makeActiveCheckpoint(doc);
+      // Check for LSN / checkpoint changes - ignore other metadata changes
+      if (lastOp == null || op.lsn != lastOp.lsn || op.checkpoint != lastOp.checkpoint) {
+        lastOp = op;
+        yield op;
+      }
+    }
+  }
+
+  // Nothing is done here until a subscriber starts to iterate
+  private readonly sharedIter = new BroadcastIterable((signal) => {
+    return this.watchActiveCheckpoint(signal);
+  });
+
+  /**
+   * User-specific watch on the latest checkpoint and/or write checkpoint.
+   */
+  async *watchWriteCheckpoint(options: WatchWriteCheckpointOptions): AsyncIterable<storage.StorageCheckpointUpdate> {
+    const { user_id, signal } = options;
+    let lastCheckpoint: utils.OpId | null = null;
+    let lastWriteCheckpoint: bigint | null = null;
+
+    const iter = wrapWithAbort(this.sharedIter, signal);
+    for await (const event of iter) {
+      const { checkpoint, lsn } = event;
+
+      // lsn changes are not important by itself.
+      // What is important is:
+      // 1. checkpoint (op_id) changes.
+      // 2. write checkpoint changes for the specific user
+
+      const lsnFilters: Record<string, string> = lsn ? { 1: lsn } : {};
+
+      const currentWriteCheckpoint = await this.lastWriteCheckpoint({
+        user_id,
+        heads: {
+          ...lsnFilters
+        }
+      });
+
+      if (currentWriteCheckpoint == lastWriteCheckpoint && checkpoint == lastCheckpoint) {
+        // No change - wait for next one
+        // In some cases, many LSNs may be produced in a short time.
+        // Add a delay to throttle the write checkpoint lookup a bit.
+        await timers.setTimeout(20 + 10 * Math.random());
+        continue;
+      }
+
+      const updates: CheckpointChanges =
+        lastCheckpoint == null
+          ? {
+              invalidateDataBuckets: true,
+              invalidateParameterBuckets: true,
+              updatedDataBuckets: [],
+              updatedParameterBucketDefinitions: []
+            }
+          : await this.getCheckpointChanges({
+              lastCheckpoint: lastCheckpoint,
+              nextCheckpoint: checkpoint
+            });
+
+      lastWriteCheckpoint = currentWriteCheckpoint;
+      lastCheckpoint = checkpoint;
+
+      yield {
+        base: event,
+        writeCheckpoint: currentWriteCheckpoint,
+        update: updates
+      };
+    }
+  }
+
+  private async getOperationDoc(
+    lastDoc: SyncRuleCheckpointState,
+    update: lib_mongo.mongo.ChangeStreamDocument<SyncRuleDocument>
+  ): Promise<SyncRuleCheckpointState | null> {
+    if (update.operationType == 'insert' || update.operationType == 'replace') {
+      return update.fullDocument;
+    } else if (update.operationType == 'update') {
+      const updatedFields = update.updateDescription.updatedFields ?? {};
+      if (lastDoc._id != update.documentKey._id) {
+        throw new ServiceAssertionError(`Sync rules id mismatch: ${lastDoc._id} != ${update.documentKey._id}`);
+      }
+
+      const mergedDoc: SyncRuleCheckpointState = {
+        _id: lastDoc._id,
+        last_checkpoint: updatedFields.last_checkpoint ?? lastDoc.last_checkpoint,
+        last_checkpoint_lsn: updatedFields.last_checkpoint_lsn ?? lastDoc.last_checkpoint_lsn,
+        state: updatedFields.state ?? lastDoc.state
+      };
+
+      return mergedDoc;
+    } else {
+      // Unknown event type
+      return null;
+    }
+  }
+
+  private getChangeStreamPipeline() {
+    const syncRulesId = this.group_id;
+    const pipeline: mongo.Document[] = [
+      {
+        $match: {
+          'documentKey._id': syncRulesId,
+          operationType: { $in: ['insert', 'update', 'replace'] }
+        }
+      },
+      {
+        $project: {
+          operationType: 1,
+          'documentKey._id': 1,
+          'updateDescription.updatedFields.state': 1,
+          'updateDescription.updatedFields.last_checkpoint': 1,
+          'updateDescription.updatedFields.last_checkpoint_lsn': 1,
+          'fullDocument._id': 1,
+          'fullDocument.state': 1,
+          'fullDocument.last_checkpoint': 1,
+          'fullDocument.last_checkpoint_lsn': 1
+        }
+      }
+    ];
+    return pipeline;
+  }
+
+  async getCheckpointChanges(options: GetCheckpointChangesOptions): Promise<CheckpointChanges> {
+    return CHECKPOINT_INVALIDATE_ALL;
+  }
 }

package/src/storage/implementation/MongoWriteCheckpointAPI.ts

@@ -1,7 +1,6 @@
 import * as framework from '@powersync/lib-services-framework';
 import { storage } from '@powersync/service-core';
 import { PowerSyncMongo } from './db.js';
-import { safeBulkWrite } from './util.js';
 
 export type MongoCheckpointAPIOptions = {
   db: PowerSyncMongo;
@@ -127,8 +126,7 @@ export async function batchCreateCustomWriteCheckpoints(
     return;
   }
 
-  await safeBulkWrite(
-    db.custom_write_checkpoints,
+  await db.custom_write_checkpoints.bulkWrite(
     checkpoints.map((checkpointOptions) => ({
       updateOne: {
         filter: { user_id: checkpointOptions.user_id, sync_rules_id: checkpointOptions.sync_rules_id },

package/src/storage/implementation/OperationBatch.ts

@@ -41,6 +41,10 @@ export class OperationBatch {
     return this.batch.length >= MAX_BATCH_COUNT || this.currentSize > MAX_RECORD_BATCH_SIZE;
   }
 
+  hasData() {
+    return this.length > 0;
+  }
+
   /**
    *
    * @param sizes Map of source key to estimated size of the current_data document, or undefined if current_data is not persisted.

package/src/storage/implementation/PersistedBatch.ts

@@ -15,7 +15,7 @@ import {
   CurrentDataDocument,
   SourceKey
 } from './models.js';
-import { replicaIdToSubkey, safeBulkWrite } from './util.js';
+import { replicaIdToSubkey } from './util.js';
 
 /**
  * Maximum size of operations we write in a single transaction.
@@ -245,33 +245,34 @@ export class PersistedBatch {
   }
 
   async flush(db: PowerSyncMongo, session: mongo.ClientSession) {
+    const startAt = performance.now();
     if (this.bucketData.length > 0) {
-      // calculate total size
-      await safeBulkWrite(db.bucket_data, this.bucketData, {
+      await db.bucket_data.bulkWrite(this.bucketData, {
         session,
         // inserts only - order doesn't matter
         ordered: false
       });
     }
     if (this.bucketParameters.length > 0) {
-      await safeBulkWrite(db.bucket_parameters, this.bucketParameters, {
+      await db.bucket_parameters.bulkWrite(this.bucketParameters, {
         session,
         // inserts only - order doesn't matter
         ordered: false
       });
     }
     if (this.currentData.length > 0) {
-      await safeBulkWrite(db.current_data, this.currentData, {
+      await db.current_data.bulkWrite(this.currentData, {
         session,
         // may update and delete data within the same batch - order matters
         ordered: true
       });
     }
 
+    const duration = performance.now() - startAt;
     logger.info(
       `powersync_${this.group_id} Flushed ${this.bucketData.length} + ${this.bucketParameters.length} + ${
         this.currentData.length
-      } updates, ${Math.round(this.currentSize / 1024)}kb. Last op_id: ${this.debugLastOpId}`
+      } updates, ${Math.round(this.currentSize / 1024)}kb in ${duration.toFixed(0)}ms. Last op_id: ${this.debugLastOpId}`
     );
 
     this.bucketData = [];

package/src/storage/implementation/db.ts

@@ -41,7 +41,7 @@ export class PowerSyncMongo {
     this.client = client;
 
     const db = client.db(options?.database, {
-      ...storage.BSON_DESERIALIZE_OPTIONS
+      ...storage.BSON_DESERIALIZE_INTERNAL_OPTIONS
     });
     this.db = db;
 

package/src/storage/implementation/models.ts

@@ -141,6 +141,11 @@ export interface SyncRuleDocument {
   content: string;
 }
 
+export type SyncRuleCheckpointState = Pick<
+  SyncRuleDocument,
+  'last_checkpoint' | 'last_checkpoint_lsn' | '_id' | 'state'
+>;
+
 export interface CustomWriteCheckpointDocument {
   _id: bson.ObjectId;
   user_id: string;

package/src/storage/implementation/util.ts

@@ -124,48 +124,3 @@ export const connectMongoForTests = (url: string, isCI: boolean) => {
   });
   return new PowerSyncMongo(client);
 };
-
-/**
- * MongoDB bulkWrite internally splits the operations into batches
- * so that no batch exceeds 16MB. However, there are cases where
- * the batch size is very close to 16MB, where additional metadata
- * on the server pushes it over the limit, resulting in this error
- * from the server:
- *
- * > MongoBulkWriteError: BSONObj size: 16814023 (0x1008FC7) is invalid. Size must be between 0 and 16793600(16MB) First element: insert: "bucket_data"
- *
- * We work around the issue by doing our own batching, limiting the
- * batch size to 15MB. This does add additional overhead with
- * BSON.calculateObjectSize.
- */
-export async function safeBulkWrite<T extends mongo.Document>(
-  collection: mongo.Collection<T>,
-  operations: mongo.AnyBulkWriteOperation<T>[],
-  options: mongo.BulkWriteOptions
-) {
-  // Must be below 16MB.
-  // We could probably go a little closer, but 15MB is a safe threshold.
-  const BULK_WRITE_LIMIT = 15 * 1024 * 1024;
-
-  let batch: mongo.AnyBulkWriteOperation<T>[] = [];
-  let currentSize = 0;
-  // Estimated overhead per operation, should be smaller in reality.
-  const keySize = 8;
-  for (let op of operations) {
-    const bsonSize =
-      mongo.BSON.calculateObjectSize(op, {
-        checkKeys: false,
-        ignoreUndefined: true
-      } as any) + keySize;
-    if (batch.length > 0 && currentSize + bsonSize > BULK_WRITE_LIMIT) {
-      await collection.bulkWrite(batch, options);
-      currentSize = 0;
-      batch = [];
-    }
-    batch.push(op);
-    currentSize += bsonSize;
-  }
-  if (batch.length > 0) {
-    await collection.bulkWrite(batch, options);
-  }
-}