@powersync/service-module-mongodb-storage 0.18.3 → 0.20.0

package/src/storage/implementation/v1/MongoBucketBatchV1.ts

@@ -1,5 +1,6 @@
 import { ReplicationAssertionError } from '@powersync/lib-services-framework';
 import { ColumnDescriptor, InternalOpId, SourceTable, storage } from '@powersync/service-core';
+import { HydratedSyncConfig } from '@powersync/service-sync-rules';
 import * as bson from 'bson';
 import { mongoTableId } from '../../../utils/util.js';
 import { calculateCheckpointState } from '../CheckpointState.js';
@@ -9,7 +10,7 @@ import { SourceRecordStore } from '../common/SourceRecordStore.js';
 import { PersistedBatchV1 } from './PersistedBatchV1.js';
 import { SourceRecordStoreV1 } from './SourceRecordStoreV1.js';
 import { VersionedPowerSyncMongoV1 } from './VersionedPowerSyncMongoV1.js';
-import { SyncRuleDocumentV1 } from './models.js';
+import { SourceTableDocumentV1, SyncRuleDocumentV1 } from './models.js';
 
 export class MongoBucketBatchV1 extends MongoBucketBatch {
   declare readonly db: VersionedPowerSyncMongoV1;
@@ -51,6 +52,7 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
 
   async resolveTables(options: storage.ResolveTablesOptions): Promise<storage.ResolveTablesResult> {
     const syncRules = options.parsedSyncConfig?.hydratedSyncConfig ?? this.sync_rules;
+    const reconcile = options.reconcileSourceTables ?? storage.defaultSourceTableReconciler;
     const { connection_id, source } = options;
     const { schema, name, objectId, replicaIdColumns, connectionTag, sendsCompleteRows } = source;
 
@@ -63,20 +65,38 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
     let result: storage.ResolveTablesResult | null = null;
     await this.db.client.withSession(async (session) => {
       const col = this.db.sourceTablesV1(this.replicationStreamId);
-      const filter: any = {
-        group_id: this.replicationStreamId,
-        connection_id,
-        schema_name: schema,
-        table_name: name,
-        replica_id_columns2: normalizedReplicaIdColumns
-      };
+
+      // Find records that overlap by name or relation id.
+      const orClauses: Record<string, unknown>[] = [{ schema_name: schema, table_name: name }];
       if (objectId != null) {
-        filter.relation_id = objectId;
+        orClauses.push({ relation_id: objectId });
+      }
+      const candidateDocs = await col
+        .find({ group_id: this.replicationStreamId, connection_id, $or: orClauses }, { session })
+        .toArray();
+
+      const candidateTables = candidateDocs.map((doc) =>
+        this.hydrateSourceTable(doc, connectionTag, syncRules, sendsCompleteRows)
+      );
+      const candidates = candidateTables.map((table) => table.clone());
+      const resolution = await reconcile({ source, candidates });
+      storage.validateSourceTableCandidateResolution(candidates, resolution);
+
+      for (const { id, sourceMetadata } of storage.diffSourceTableUpdates(candidateTables, resolution)) {
+        await col.updateOne({ _id: mongoTableId(id) }, { $set: { source_metadata: sourceMetadata } }, { session });
       }
 
-      let doc = await col.findOne(filter, { session });
-      if (doc == null) {
-        doc = {
+      const materializedResolution = storage.materializeSourceTableResolution(candidateTables, resolution);
+      const compatibleTables = materializedResolution.compatibleTables;
+
+      // V1 keeps one record, preferring one that has already been snapshotted.
+      const reuse = compatibleTables.find((table) => table.snapshotComplete) ?? compatibleTables[0] ?? null;
+
+      let sourceTable: storage.SourceTable;
+      if (reuse != null) {
+        sourceTable = reuse;
+      } else {
+        const doc: SourceTableDocumentV1 = {
           _id: options.idGenerator ? (options.idGenerator() as bson.ObjectId) : new bson.ObjectId(),
           group_id: this.replicationStreamId,
           connection_id,
@@ -86,69 +106,18 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
           replica_id_columns: null,
           replica_id_columns2: normalizedReplicaIdColumns,
           snapshot_done: false,
-          snapshot_status: undefined
+          snapshot_status: undefined,
+          source_metadata: resolution.newTableValues.sourceMetadata
         };
         await col.insertOne(doc, { session });
+        sourceTable = this.hydrateSourceTable(doc, connectionTag, syncRules, sendsCompleteRows);
       }
 
-      const sourceTable = new storage.SourceTable({
-        id: doc._id,
-        ref: source,
-        objectId,
-        replicaIdColumns,
-        snapshotComplete: doc.snapshot_done ?? true,
-        ...syncRules.getMatchingSources(source)
-      });
-      sourceTable.syncEvent = syncRules.tableTriggersEvent(source);
-      sourceTable.syncData = sourceTable.bucketDataSources.length > 0;
-      sourceTable.syncParameters = sourceTable.parameterLookupSources.length > 0;
-      sourceTable.storeCurrentData = sendsCompleteRows !== true;
-      sourceTable.snapshotStatus =
-        doc.snapshot_status == null
-          ? undefined
-          : {
-              lastKey: doc.snapshot_status.last_key?.buffer ?? null,
-              totalEstimatedCount: doc.snapshot_status.total_estimated_count,
-              replicatedCount: doc.snapshot_status.replicated_count
-            };
-
-      const truncateFilter = [{ schema_name: schema, table_name: name }] as any[];
-      if (objectId != null) {
-        truncateFilter.push({ relation_id: objectId });
-      }
-      const truncate = await col
-        .find(
-          {
-            group_id: this.replicationStreamId,
-            connection_id,
-            _id: { $ne: doc._id },
-            $or: truncateFilter
-          },
-          { session }
-        )
-        .toArray();
-      const dropTables = truncate.map((dropDoc) => {
-        const ref = {
-          connectionTag,
-          schema: dropDoc.schema_name,
-          name: dropDoc.table_name
-        };
-        const table = new storage.SourceTable({
-          id: dropDoc._id,
-          ref,
-          objectId: dropDoc.relation_id,
-          replicaIdColumns:
-            dropDoc.replica_id_columns2?.map(
-              (c) => ({ name: c.name, typeId: c.type_oid, type: c.type }) satisfies ColumnDescriptor
-            ) ?? [],
-          snapshotComplete: dropDoc.snapshot_done ?? true,
-          ...syncRules.getMatchingSources(ref)
-        });
-        table.syncEvent = syncRules.tableTriggersEvent(ref);
-        table.syncData = table.bucketDataSources.length > 0;
-        table.syncParameters = table.parameterLookupSources.length > 0;
-        return table;
-      });
+      const dropTables = [
+        ...materializedResolution.incompatibleTables,
+        // V1 only keeps one SourceTable per physical table.
+        ...compatibleTables.filter((table) => !storage.sourceTableIdEquals(table.id, sourceTable.id))
+      ];
 
       result = { tables: [sourceTable], dropTables };
     });
@@ -156,38 +125,33 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
     return result!;
   }
 
-  async getSourceTableStatus(table: storage.SourceTable): Promise<storage.SourceTable | null> {
-    const doc = await this.db.sourceTablesV1(this.replicationStreamId).findOne(
-      {
-        group_id: this.replicationStreamId,
-        _id: mongoTableId(table.id)
-      },
-      { session: this.session }
-    );
-    if (doc == null) {
-      return null;
-    }
-
-    const ref = {
-      connectionTag: table.ref.connectionTag,
-      schema: doc.schema_name,
-      name: doc.table_name
-    };
-    const sourceTable = new storage.SourceTable({
+  /**
+   * Build a SourceTable from its persisted v1 record.
+   */
+  private hydrateSourceTable(
+    doc: SourceTableDocumentV1,
+    connectionTag: string,
+    syncRules: HydratedSyncConfig,
+    sendsCompleteRows: boolean | undefined
+  ): storage.SourceTable {
+    const ref = { connectionTag, schema: doc.schema_name, name: doc.table_name };
+    const table = new storage.SourceTable({
       id: doc._id,
       ref,
       objectId: doc.relation_id,
       replicaIdColumns:
         doc.replica_id_columns2?.map(
-          (c) => ({ name: c.name, typeId: c.type_oid, type: c.type }) satisfies ColumnDescriptor
+          (column) => ({ name: column.name, typeId: column.type_oid, type: column.type }) satisfies ColumnDescriptor
         ) ?? [],
       snapshotComplete: doc.snapshot_done ?? true,
-      ...this.sync_rules.getMatchingSources(ref)
+      sourceMetadata: doc.source_metadata ?? null,
+      ...syncRules.getMatchingSources(ref)
     });
-    sourceTable.syncEvent = this.sync_rules.tableTriggersEvent(ref);
-    sourceTable.syncData = sourceTable.bucketDataSources.length > 0;
-    sourceTable.syncParameters = sourceTable.parameterLookupSources.length > 0;
-    sourceTable.snapshotStatus =
+    table.syncEvent = syncRules.tableTriggersEvent(ref);
+    table.syncData = table.bucketDataSources.length > 0;
+    table.syncParameters = table.parameterLookupSources.length > 0;
+    table.storeCurrentData = sendsCompleteRows !== true;
+    table.snapshotStatus =
       doc.snapshot_status == null
         ? undefined
         : {
@@ -195,7 +159,23 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
             totalEstimatedCount: doc.snapshot_status.total_estimated_count,
             replicatedCount: doc.snapshot_status.replicated_count
           };
-    return sourceTable;
+    return table;
+  }
+
+  async getSourceTableStatus(table: storage.SourceTable): Promise<storage.SourceTable | null> {
+    const doc = await this.db.sourceTablesV1(this.replicationStreamId).findOne(
+      {
+        group_id: this.replicationStreamId,
+        _id: mongoTableId(table.id)
+      },
+      { session: this.session }
+    );
+    if (doc == null) {
+      return null;
+    }
+
+    // Use the default storeCurrentData value when reading outside resolution.
+    return this.hydrateSourceTable(doc, table.ref.connectionTag, this.sync_rules, undefined);
   }
 
   async commit(lsn: string, options?: storage.BucketBatchCommitOptions): Promise<storage.CheckpointResult> {

package/src/storage/implementation/v1/MongoCompactorV1.ts

@@ -1,20 +1,25 @@
 import { mongo } from '@powersync/lib-service-mongodb';
 import { ReplicationAssertionError } from '@powersync/lib-services-framework';
-import { storage } from '@powersync/service-core';
+import { addChecksums, storage, utils } from '@powersync/service-core';
 import { BucketDefinitionId } from '@powersync/service-sync-rules';
-import { SingleBucketStore } from '../common/SingleBucketStore.js';
+import { BucketDataDoc } from '../common/BucketDataDoc.js';
 import { BucketStateDocumentBase, LEGACY_BUCKET_DATA_DEFINITION_ID } from '../models.js';
-import { DirtyBucket, MongoCompactor } from '../MongoCompactor.js';
-import { BucketStateDocumentV1 } from './models.js';
+import { CurrentBucketState, DirtyBucket, MongoCompactor } from '../MongoCompactor.js';
+import { cacheKey } from '../OperationBatch.js';
+import { BucketDataDocumentV1, BucketStateDocumentV1 } from './models.js';
 import type { MongoSyncBucketStorageV1 } from './MongoSyncBucketStorageV1.js';
 import { SingleBucketStoreV1 } from './SingleBucketStoreV1.js';
 import { VersionedPowerSyncMongoV1 } from './VersionedPowerSyncMongoV1.js';
 
+type CompactClearProperties = 'op' | 'checksum' | 'target_op';
+
 export class MongoCompactorV1 extends MongoCompactor {
   // Override types to the more specific ones
   declare protected readonly db: VersionedPowerSyncMongoV1;
   declare protected readonly storage: MongoSyncBucketStorageV1;
 
+  private updates: mongo.AnyBulkWriteOperation<BucketDataDocumentV1>[] = [];
+
   public async *dirtyBucketBatches(options: {
     minBucketChanges: number;
     minChangeRatio: number;
@@ -81,14 +86,293 @@ export class MongoCompactorV1 extends MongoCompactor {
     };
   }
 
-  protected async getBucketDataContext(
-    bucket: string,
-    _definitionId: BucketDefinitionId | null
-  ): Promise<SingleBucketStore | null> {
+  private getBucketDataContext(bucket: string): SingleBucketStoreV1 {
     return new SingleBucketStoreV1(this.db, {
       replicationStreamId: this.group_id,
       definitionId: LEGACY_BUCKET_DATA_DEFINITION_ID,
       bucket
     });
   }
+
+  protected async compactSingleBucket(bucket: string) {
+    // Do not carry queued writes from a failed attempt into the rescan
+    this.updates = [];
+
+    const idLimitBytes = this.idLimitBytes;
+    const bucketContext = this.getBucketDataContext(bucket);
+    const currentState: CurrentBucketState = {
+      bucket,
+      definitionId: bucketContext.key.definitionId,
+      seen: new Map(),
+      trackingSize: 0,
+      lastNotPut: null,
+      opsSincePut: 0,
+      checksum: 0,
+      opCount: 0,
+      opBytes: 0
+    };
+
+    // Constant lower bound.
+    const lowerBound = bucketContext.minId;
+    // Upper bound is adjusted for each batch.
+    let upperBound = bucketContext.maxId;
+
+    // Paginate through bucket data in batches to avoid cursor timeouts.
+    while (true) {
+      this.signal?.throwIfAborted();
+
+      // Query one batch at a time, to avoid cursor timeouts.
+      const pipeline = [
+        {
+          $match: {
+            _id: {
+              $gte: lowerBound,
+              $lt: upperBound
+            }
+          }
+        },
+        { $sort: { _id: -1 } },
+        { $limit: this.moveBatchQueryLimit },
+        {
+          $project: {
+            _id: 1,
+            op: 1,
+            table: 1,
+            row_id: 1,
+            source_table: 1,
+            source_key: 1,
+            checksum: 1,
+            size: { $bsonSize: '$$ROOT' }
+          }
+        }
+      ];
+
+      const cursor = bucketContext.collection.aggregate<BucketDataDocumentV1 & { size: number | bigint }>(pipeline, {
+        // batchSize is 1 more than limit to auto-close the cursor.
+        // See https://github.com/mongodb/node-mongodb-native/pull/4580
+        batchSize: this.moveBatchQueryLimit + 1
+      });
+      // We don't limit to a single batch here, since that often causes MongoDB to scan through more than it returns.
+      // Instead, we load up to the limit.
+      const rawBatch = await cursor.toArray();
+      const batch = rawBatch.map((document) => {
+        const { size, ...rest } = document;
+        return {
+          doc: bucketContext.fromPersistedDocument(rest),
+          size
+        };
+      });
+
+      if (batch.length == 0) {
+        // We've reached the end.
+        break;
+      }
+
+      // Reuse the exact collection _id value from Mongo for the next bound.
+      upperBound = rawBatch[rawBatch.length - 1]._id;
+
+      for (const { doc, size } of batch) {
+        if (doc.o > this.maxOpId) {
+          continue;
+        }
+
+        currentState.checksum = addChecksums(currentState.checksum, Number(doc.checksum));
+        currentState.opCount += 1;
+
+        let isPersistentPut = doc.op == 'PUT';
+
+        currentState.opBytes += Number(size);
+        if (doc.op == 'REMOVE' || doc.op == 'PUT') {
+          const key = `${doc.table}/${doc.row_id}/${cacheKey(doc.source_table!, doc.source_key!)}`;
+          const targetOp = currentState.seen.get(key);
+          if (targetOp) {
+            // Will convert to MOVE, so don't count as PUT.
+            isPersistentPut = false;
+
+            this.updates.push({
+              updateOne: {
+                filter: { _id: bucketContext.docId(doc.o) },
+                update: {
+                  $set: {
+                    op: 'MOVE',
+                    target_op: targetOp
+                  },
+                  $unset: {
+                    source_table: 1,
+                    source_key: 1,
+                    table: 1,
+                    row_id: 1,
+                    data: 1
+                  }
+                } satisfies mongo.UpdateFilter<BucketDataDocumentV1>
+              }
+            });
+
+            // TODO: better estimate for this.
+            currentState.opBytes += 200 - Number(size);
+          } else if (currentState.trackingSize < idLimitBytes) {
+            // flatstr reduces the memory usage by flattening the string.
+            currentState.seen.set(utils.flatstr(key), doc.o);
+            // length + 16 for the string
+            // 24 for the bigint
+            // 50 for map overhead
+            // 50 for additional overhead
+            currentState.trackingSize += key.length + 140;
+          }
+        }
+
+        if (isPersistentPut) {
+          currentState.lastNotPut = null;
+          currentState.opsSincePut = 0;
+        } else if (doc.op != 'CLEAR') {
+          if (currentState.lastNotPut == null) {
+            currentState.lastNotPut = doc.o;
+          }
+          currentState.opsSincePut += 1;
+        }
+
+        if (this.updates.length + this.bucketStateUpdates.length >= this.moveBatchLimit) {
+          await this.flush(bucketContext);
+        }
+      }
+
+      this.logger.info(`Processed batch of length ${batch.length} current bucket: ${bucket}`);
+    }
+
+    // Free memory before clearing the bucket.
+    currentState.seen.clear();
+    if (currentState.lastNotPut != null && currentState.opsSincePut >= 1) {
+      this.logger.info(
+        `Inserting CLEAR at ${this.group_id}:${bucket}:${currentState.lastNotPut} to remove ${currentState.opsSincePut} operations`
+      );
+      // Need flush() before clear().
+      await this.flush(bucketContext);
+      await this.clearBucket(currentState, bucketContext);
+    }
+
+    // Do this after clearBucket so we have accurate counts.
+    this.updateBucketChecksums(currentState, this.maxOpId);
+    // Need another flush after updateBucketChecksums().
+    await this.flush(bucketContext);
+  }
+
+  private async flush(bucketContext: SingleBucketStoreV1) {
+    if (this.updates.length > 0) {
+      this.logger.info(`Compacting ${this.updates.length} ops`);
+      await bucketContext.collection.bulkWrite(this.updates, {
+        // Order is not important. Since checksums are not affected, these operations can happen in any order,
+        // and it's fine if the operations are partially applied. Each individual operation is atomic.
+        ordered: false
+      });
+      this.updates = [];
+    }
+
+    await this.flushBucketStateUpdates();
+  }
+
+  /**
+   * Perform a CLEAR compact for a bucket.
+   *
+   * @param currentState tracks the last non-PUT op, which will be converted to CLEAR.
+   */
+  private async clearBucket(currentState: CurrentBucketState, bucketContext: SingleBucketStoreV1) {
+    const clearOp = currentState.lastNotPut!;
+
+    const opFilter = {
+      _id: {
+        $gte: bucketContext.minId,
+        $lte: bucketContext.docId(clearOp)
+      }
+    };
+
+    const session = this.db.client.startSession();
+    try {
+      let done = false;
+      while (!done) {
+        this.signal?.throwIfAborted();
+        let opCountDiff = 0;
+        // Do the CLEAR operation in batches, with each batch a separate transaction.
+        // The state after each batch is fully consistent.
+        // We need a transaction per batch to make sure checksums stay consistent.
+        await session.withTransaction(
+          async () => {
+            const query = bucketContext.collection.find<Pick<BucketDataDocumentV1, '_id' | CompactClearProperties>>(
+              opFilter,
+              {
+                session,
+                sort: { _id: 1 },
+                projection: {
+                  _id: 1,
+                  op: 1,
+                  checksum: 1,
+                  target_op: 1
+                },
+                limit: this.clearBatchLimit
+              }
+            );
+            let checksum = 0;
+            let lastOp: Pick<BucketDataDoc, 'o' | CompactClearProperties> | null = null;
+            let targetOp: bigint | null = null;
+            let gotAnOp = false;
+            let numberOfOpsToClear = 0;
+            for await (const rawOp of query.stream()) {
+              const op = bucketContext.fromPartialPersistedDocument(rawOp);
+
+              if (op.op == 'MOVE' || op.op == 'REMOVE' || op.op == 'CLEAR') {
+                checksum = utils.addChecksums(checksum, Number(op.checksum));
+                lastOp = op;
+                numberOfOpsToClear += 1;
+                if (op.op != 'CLEAR') {
+                  gotAnOp = true;
+                }
+                if (op.target_op != null && (targetOp == null || op.target_op > targetOp)) {
+                  targetOp = op.target_op;
+                }
+              } else {
+                throw new ReplicationAssertionError(`Unexpected ${op.op} operation at ${this.formatBucketDataKey(op)}`);
+              }
+            }
+            if (!gotAnOp) {
+              done = true;
+              return;
+            }
+
+            this.logger.info(`Flushing CLEAR for ${numberOfOpsToClear} ops at ${lastOp?.o}`);
+            await bucketContext.collection.deleteMany(
+              {
+                _id: {
+                  $gte: bucketContext.minId,
+                  $lte: bucketContext.docId(lastOp!.o)
+                }
+              },
+              { session }
+            );
+
+            const op = bucketContext.toPersistedDocument({
+              o: lastOp!.o,
+              op: 'CLEAR',
+              checksum: BigInt(checksum),
+              data: null,
+              target_op: targetOp
+            });
+            await bucketContext.collection.insertOne(op, { session });
+
+            opCountDiff = -numberOfOpsToClear + 1;
+          },
+          {
+            writeConcern: { w: 'majority' },
+            readConcern: { level: 'snapshot' }
+          }
+        );
+        // Update outside the transaction, since the transaction can be retried multiple times.
+        currentState.opCount += opCountDiff;
+      }
+    } finally {
+      await session.endSession();
+    }
+  }
+
+  private formatBucketDataKey(doc: Pick<BucketDataDoc, 'bucketKey' | 'o'>) {
+    return `${doc.bucketKey.replicationStreamId}:${doc.bucketKey.bucket}:${doc.o}`;
+  }
 }