@powersync/service-module-mongodb-storage 0.0.0-dev-20260313100403 → 0.0.0-dev-20260515144844

package/dist/storage/implementation/MongoBucketBatch.d.ts

@@ -1,14 +1,11 @@
 import { mongo } from '@powersync/lib-service-mongodb';
 import { HydratedSyncRules, SqlEventDescriptor } from '@powersync/service-sync-rules';
 import { BaseObserver, Logger } from '@powersync/lib-services-framework';
-import { BucketStorageMarkRecordUnavailable, CheckpointResult, InternalOpId, storage } from '@powersync/service-core';
-import { VersionedPowerSyncMongo } from './db.js';
-import { CurrentBucket } from './models.js';
-/**
- * 15MB
- */
-export declare const MAX_ROW_SIZE: number;
-export declare const EMPTY_DATA: mongo.Binary;
+import { BucketStorageMarkRecordUnavailable, CheckpointResult, InternalOpId, PerformanceTracer, storage } from '@powersync/service-core';
+import { BucketDefinitionMapping } from './BucketDefinitionMapping.js';
+import { PersistedBatch } from './common/PersistedBatch.js';
+import { SourceRecordStore } from './common/SourceRecordStore.js';
+import type { VersionedPowerSyncMongo } from './db.js';
 export interface MongoBucketBatchOptions {
     db: VersionedPowerSyncMongo;
     syncRules: HydratedSyncRules;
@@ -18,27 +15,31 @@ export interface MongoBucketBatchOptions {
     keepaliveOp: InternalOpId | null;
     resumeFromLsn: string | null;
     storeCurrentData: boolean;
+    mapping: BucketDefinitionMapping;
     /**
      * Set to true for initial replication.
      */
     skipExistingRows: boolean;
     markRecordUnavailable: BucketStorageMarkRecordUnavailable | undefined;
-    logger?: Logger;
+    logger: Logger;
+    tracer?: PerformanceTracer<'storage' | 'evaluate'>;
 }
-export declare class MongoBucketBatch extends BaseObserver<storage.BucketBatchStorageListener> implements storage.BucketStorageBatch {
-    private logger;
+export declare abstract class MongoBucketBatch extends BaseObserver<storage.BucketBatchStorageListener> implements storage.BucketStorageBatch {
+    protected logger: Logger;
     private readonly client;
     readonly db: VersionedPowerSyncMongo;
     readonly session: mongo.ClientSession;
     private readonly sync_rules;
-    private readonly group_id;
+    protected readonly group_id: number;
     private readonly slot_name;
     private readonly storeCurrentData;
     private readonly skipExistingRows;
+    protected readonly mapping: BucketDefinitionMapping;
     private batch;
     private write_checkpoint_batch;
     private markRecordUnavailable;
     private clearedError;
+    private tracer;
     /**
      * Last LSN received associated with a checkpoint.
      *
@@ -67,6 +68,9 @@ export declare class MongoBucketBatch extends BaseObserver<storage.BucketBatchSt
     constructor(options: MongoBucketBatchOptions);
     addCustomWriteCheckpoint(checkpoint: storage.BatchedCustomWriteCheckpointOptions): void;
     get lastCheckpointLsn(): string | null;
+    protected abstract createPersistedBatch(writtenSize: number): PersistedBatch;
+    protected abstract get sourceRecordStore(): SourceRecordStore;
+    protected abstract cleanupDroppedSourceTables(sourceTables: storage.SourceTable[]): Promise<void>;
     flush(options?: storage.BatchBucketFlushOptions): Promise<storage.FlushedResult | null>;
     private flushInner;
     private replicateBatch;
@@ -77,7 +81,6 @@ export declare class MongoBucketBatch extends BaseObserver<storage.BucketBatchSt
     dispose(): Promise<void>;
     private lastWaitingLogThottled;
     commit(lsn: string, options?: storage.BucketBatchCommitOptions): Promise<CheckpointResult>;
-    private cleanupCurrentData;
     /**
      * Switch from processing -> active if relevant.
      *
@@ -103,4 +106,3 @@ export declare class MongoBucketBatch extends BaseObserver<storage.BucketBatchSt
      */
     protected getTableEvents(table: storage.SourceTable): SqlEventDescriptor[];
 }
-export declare function currentBucketKey(b: CurrentBucket): string;

package/dist/storage/implementation/MongoBucketBatch.js

@@ -1,24 +1,19 @@
 import { mongo } from '@powersync/lib-service-mongodb';
 import * as bson from 'bson';
-import { BaseObserver, container, logger as defaultLogger, ErrorCode, errors, ReplicationAssertionError, ServiceError } from '@powersync/lib-services-framework';
-import { deserializeBson, isCompleteRow, SaveOperationTag, storage, SyncRuleState, utils } from '@powersync/service-core';
+import { BaseObserver, container, ErrorCode, errors, ReplicationAssertionError, ServiceError } from '@powersync/lib-services-framework';
+import { deserializeBson, isCompleteRow, PerformanceTracer, SaveOperationTag, storage, SyncRuleState, utils } from '@powersync/service-core';
 import * as timers from 'node:timers/promises';
-import { idPrefixFilter, mongoTableId } from '../../utils/util.js';
+import { mongoTableId } from '../../utils/util.js';
+import { MAX_ROW_SIZE } from './MongoBucketBatchShared.js';
 import { MongoIdSequence } from './MongoIdSequence.js';
 import { batchCreateCustomWriteCheckpoints } from './MongoWriteCheckpointAPI.js';
-import { cacheKey, OperationBatch, RecordOperation } from './OperationBatch.js';
-import { PersistedBatch } from './PersistedBatch.js';
-/**
- * 15MB
- */
-export const MAX_ROW_SIZE = 15 * 1024 * 1024;
+import { OperationBatch, RecordOperation } from './OperationBatch.js';
 // Currently, we can only have a single flush() at a time, since it locks the op_id sequence.
 // While the MongoDB transaction retry mechanism handles this okay, using an in-process Mutex
 // makes it more fair and has less overhead.
 //
 // In the future, we can investigate allowing multiple replication streams operating independently.
 const replicationMutex = new utils.Mutex();
-export const EMPTY_DATA = new bson.Binary(bson.serialize({}));
 export class MongoBucketBatch extends BaseObserver {
     logger;
     client;
@@ -29,10 +24,12 @@ export class MongoBucketBatch extends BaseObserver {
     slot_name;
     storeCurrentData;
     skipExistingRows;
+    mapping;
     batch = null;
     write_checkpoint_batch = [];
     markRecordUnavailable;
     clearedError = false;
+    tracer;
     /**
      * Last LSN received associated with a checkpoint.
      *
@@ -60,7 +57,7 @@ export class MongoBucketBatch extends BaseObserver {
     needsActivation = true;
     constructor(options) {
         super();
-        this.logger = options.logger ?? defaultLogger;
+        this.logger = options.logger;
         this.client = options.db.client;
         this.db = options.db;
         this.group_id = options.groupId;
@@ -70,10 +67,12 @@ export class MongoBucketBatch extends BaseObserver {
         this.slot_name = options.slotName;
         this.sync_rules = options.syncRules;
         this.storeCurrentData = options.storeCurrentData;
+        this.mapping = options.mapping;
         this.skipExistingRows = options.skipExistingRows;
         this.markRecordUnavailable = options.markRecordUnavailable;
         this.batch = new OperationBatch();
         this.persisted_op = options.keepaliveOp ?? null;
+        this.tracer = options.tracer ?? new PerformanceTracer('MongoDB storage');
     }
     addCustomWriteCheckpoint(checkpoint) {
         this.write_checkpoint_batch.push({
@@ -100,6 +99,7 @@ export class MongoBucketBatch extends BaseObserver {
         const batch = this.batch;
         let last_op = null;
         let resumeBatch = null;
+        using _ = this.tracer.span('storage', 'flush');
         await this.withReplicationTransaction(`Flushing ${batch?.length ?? 0} ops`, async (session, opSeq) => {
             if (batch != null) {
                 resumeBatch = await this.replicateBatch(session, batch, opSeq, options);
@@ -122,6 +122,7 @@ export class MongoBucketBatch extends BaseObserver {
     }
     async replicateBatch(session, batch, op_seq, options) {
         let sizes = undefined;
+        using _ = this.tracer.span('storage', 'replicate_batch');
         if (this.storeCurrentData && !this.skipExistingRows) {
             // We skip this step if we don't store current_data, since the sizes will
             // always be small in that case.
@@ -134,27 +135,11 @@ export class MongoBucketBatch extends BaseObserver {
             // (automatically limited to 48MB(?) per batch by MongoDB). The issue is that it changes
             // the order of processing, which then becomes really tricky to manage.
             // This now takes 2+ queries, but doesn't have any issues with order of operations.
-            const sizeLookups = batch.batch.map((r) => {
-                return { g: this.group_id, t: mongoTableId(r.record.sourceTable.id), k: r.beforeId };
-            });
-            sizes = new Map();
-            const sizeCursor = this.db.common_current_data.aggregate([
-                {
-                    $match: {
-                        _id: { $in: sizeLookups }
-                    }
-                },
-                {
-                    $project: {
-                        _id: 1,
-                        size: { $bsonSize: '$$ROOT' }
-                    }
-                }
-            ], { session });
-            for await (let doc of sizeCursor.stream()) {
-                const key = cacheKey(doc._id.t, doc._id.k);
-                sizes.set(key, doc.size);
-            }
+            const sizeLookups = batch.batch.map((r) => ({
+                sourceTableId: mongoTableId(r.record.sourceTable.id),
+                replicaId: r.beforeId
+            }));
+            sizes = await this.sourceRecordStore.loadSizes(session, sizeLookups);
         }
         // If set, we need to start a new transaction with this batch.
         let resumeBatch = null;
@@ -169,40 +154,38 @@ export class MongoBucketBatch extends BaseObserver {
                 }
                 continue;
             }
-            const lookups = b.map((r) => {
-                return { g: this.group_id, t: mongoTableId(r.record.sourceTable.id), k: r.beforeId };
-            });
-            let current_data_lookup = new Map();
-            // With skipExistingRows, we only need to know whether or not the row exists.
-            const projection = this.skipExistingRows ? { _id: 1 } : undefined;
-            const cursor = this.db.common_current_data.find({
-                _id: { $in: lookups }
-            }, { session, projection });
-            for await (let doc of cursor.stream()) {
-                current_data_lookup.set(cacheKey(doc._id.t, doc._id.k), doc);
-            }
-            let persistedBatch = new PersistedBatch(this.db, this.group_id, transactionSize, {
-                logger: this.logger
-            });
+            using lookupSpan = this.tracer.span('storage', 'lookup');
+            const lookups = b.map((r) => ({
+                sourceTableId: mongoTableId(r.record.sourceTable.id),
+                replicaId: r.beforeId
+            }));
+            let sourceRecordLookup = await this.sourceRecordStore.loadDocuments(session, lookups, this.skipExistingRows);
+            lookupSpan.end();
+            let persistedBatch = this.createPersistedBatch(transactionSize);
+            // The current code structure makes it tricky to cleanly split this span from the one
+            // where fluhsing. So we manually end and re-create this span whenever we flush.
+            let evalSpan = this.tracer.span('evaluate');
             for (let op of b) {
                 if (resumeBatch) {
                     resumeBatch.push(op);
                     continue;
                 }
-                const currentData = current_data_lookup.get(op.internalBeforeKey) ?? null;
-                if (currentData != null) {
+                const sourceRecord = sourceRecordLookup.get(op.internalBeforeKey) ?? null;
+                if (sourceRecord != null) {
                     // If it will be used again later, it will be set again using nextData below
-                    current_data_lookup.delete(op.internalBeforeKey);
+                    sourceRecordLookup.delete(op.internalBeforeKey);
                 }
-                const nextData = this.saveOperation(persistedBatch, op, currentData, op_seq);
+                const nextData = this.saveOperation(persistedBatch, op, sourceRecord, op_seq);
                 if (nextData != null) {
                     // Update our current_data and size cache
-                    current_data_lookup.set(op.internalAfterKey, nextData);
-                    sizes?.set(op.internalAfterKey, nextData.data.length());
+                    sourceRecordLookup.set(op.internalAfterKey, nextData);
+                    sizes?.set(op.internalAfterKey, nextData.data?.length() ?? 0);
                 }
                 if (persistedBatch.shouldFlushTransaction()) {
+                    evalSpan.end();
                     // Transaction is getting big.
                     // Flush, and resume in a new transaction.
+                    using persistSpan = this.tracer.span('storage', 'persist_flush');
                     const { flushedAny } = await persistedBatch.flush(this.session, options);
                     didFlush ||= flushedAny;
                     persistedBatch = null;
@@ -210,20 +193,25 @@ export class MongoBucketBatch extends BaseObserver {
                     // we're stopping in the middle of a batch.
                     // We create a new batch, and push any remaining operations to it.
                     resumeBatch = new OperationBatch();
+                    persistSpan.end();
+                    evalSpan = this.tracer.span('evaluate');
                 }
             }
+            evalSpan.end();
             if (persistedBatch) {
                 transactionSize = persistedBatch.currentSize;
+                using _ = this.tracer.span('storage', 'persist_flush');
                 const { flushedAny } = await persistedBatch.flush(this.session, options);
                 didFlush ||= flushedAny;
             }
         }
         if (didFlush) {
+            using _ = this.tracer.span('storage', 'clear_error');
             await this.clearError();
         }
         return resumeBatch?.hasData() ? resumeBatch : null;
     }
-    saveOperation(batch, operation, current_data, opSeq) {
+    saveOperation(batch, operation, sourceRecord, opSeq) {
         const record = operation.record;
         const beforeId = operation.beforeId;
         const afterId = operation.afterId;
@@ -233,10 +221,10 @@ export class MongoBucketBatch extends BaseObserver {
         let new_buckets = [];
         let existing_lookups = [];
         let new_lookups = [];
-        const before_key = { g: this.group_id, t: mongoTableId(record.sourceTable.id), k: beforeId };
+        const sourceTableId = mongoTableId(record.sourceTable.id);
         if (this.skipExistingRows) {
             if (record.tag == SaveOperationTag.INSERT) {
-                if (current_data != null) {
+                if (sourceRecord != null) {
                     // Initial replication, and we already have the record.
                     // This may be a different version of the record, but streaming replication
                     // will take care of that.
@@ -249,7 +237,7 @@ export class MongoBucketBatch extends BaseObserver {
             }
         }
         if (record.tag == SaveOperationTag.UPDATE) {
-            const result = current_data;
+            const result = sourceRecord;
             if (result == null) {
                 // Not an error if we re-apply a transaction
                 existing_buckets = [];
@@ -270,14 +258,14 @@ export class MongoBucketBatch extends BaseObserver {
             else {
                 existing_buckets = result.buckets;
                 existing_lookups = result.lookups;
-                if (this.storeCurrentData) {
+                if (this.storeCurrentData && result.data != null) {
                     const data = deserializeBson(result.data.buffer);
                     after = storage.mergeToast(after, data);
                 }
             }
         }
         else if (record.tag == SaveOperationTag.DELETE) {
-            const result = current_data;
+            const result = sourceRecord;
             if (result == null) {
                 // Not an error if we re-apply a transaction
                 existing_buckets = [];
@@ -292,9 +280,9 @@ export class MongoBucketBatch extends BaseObserver {
                 existing_lookups = result.lookups;
             }
         }
-        let afterData;
+        let afterData = null;
         if (afterId != null && !this.storeCurrentData) {
-            afterData = EMPTY_DATA;
+            afterData = null;
         }
         else if (afterId != null) {
             try {
@@ -378,13 +366,7 @@ export class MongoBucketBatch extends BaseObserver {
                     table: sourceTable,
                     before_buckets: existing_buckets
                 });
-                new_buckets = evaluated.map((e) => {
-                    return {
-                        bucket: e.bucket,
-                        table: e.table,
-                        id: e.id
-                    };
-                });
+                new_buckets = this.sourceRecordStore.mapEvaluatedBuckets(evaluated);
             }
             if (sourceTable.syncParameters) {
                 // Parameters
@@ -406,26 +388,27 @@ export class MongoBucketBatch extends BaseObserver {
                     evaluated: paramEvaluated,
                     existing_lookups
                 });
-                new_lookups = paramEvaluated.map((p) => {
-                    return storage.serializeLookup(p.lookup);
-                });
+                new_lookups = this.sourceRecordStore.mapParameterLookups(paramEvaluated);
             }
         }
         let result = null;
         // 5. TOAST: Update current data and bucket list.
         if (afterId) {
             // Insert or update
-            const after_key = { g: this.group_id, t: mongoTableId(sourceTable.id), k: afterId };
-            batch.upsertCurrentData(after_key, {
+            batch.upsertCurrentData({
+                sourceTableId,
+                replicaId: afterId,
                 data: afterData,
                 buckets: new_buckets,
                 lookups: new_lookups
             });
             result = {
-                _id: after_key,
+                sourceTableId,
+                replicaId: afterId,
                 data: afterData,
                 buckets: new_buckets,
-                lookups: new_lookups
+                lookups: new_lookups,
+                cacheKey: operation.internalAfterKey
             };
         }
         if (afterId == null || !storage.replicaIdEquals(beforeId, afterId)) {
@@ -433,12 +416,14 @@ export class MongoBucketBatch extends BaseObserver {
             // Note that this is a soft delete.
             // We don't specifically need a new or unique op_id here, but it must be greater than the
             // last checkpoint, so we use next().
-            batch.softDeleteCurrentData(before_key, opSeq.next());
+            batch.softDeleteCurrentData(sourceTableId, beforeId, opSeq.next());
         }
         return result;
     }
     async withTransaction(cb) {
+        using lockSpan = this.tracer.span('storage', 'internal_lock');
         await replicationMutex.exclusiveLock(async () => {
+            lockSpan.end();
             await this.session.withTransaction(async () => {
                 try {
                     await cb();
@@ -450,7 +435,9 @@ export class MongoBucketBatch extends BaseObserver {
                     else {
                         this.logger.warn('Transaction error', e);
                     }
-                    await timers.setTimeout(Math.random() * 50);
+                    const delay = Math.random() * 50;
+                    using _ = this.tracer.span('storage', 'retry_delay');
+                    await timers.setTimeout(delay);
                     throw e;
                 }
             }, { maxCommitTimeMS: 10000 });
@@ -656,21 +643,12 @@ export class MongoBucketBatch extends BaseObserver {
             await this.db.notifyCheckpoint();
             this.persisted_op = null;
             this.last_checkpoint_lsn = lsn;
-            if (this.db.storageConfig.softDeleteCurrentData && newLastCheckpoint != null) {
-                await this.cleanupCurrentData(newLastCheckpoint);
+            if (newLastCheckpoint != null) {
+                await this.sourceRecordStore.postCommitCleanup(newLastCheckpoint, this.logger);
             }
         }
         return { checkpointBlocked, checkpointCreated };
     }
-    async cleanupCurrentData(lastCheckpoint) {
-        const result = await this.db.v3_current_data.deleteMany({
-            '_id.g': this.group_id,
-            pending_delete: { $exists: true, $lte: lastCheckpoint }
-        });
-        if (result.deletedCount > 0) {
-            this.logger.info(`Cleaned up ${result.deletedCount} pending delete current_data records for checkpoint ${lastCheckpoint}`);
-        }
-    }
     /**
      * Switch from processing -> active if relevant.
      *
@@ -709,7 +687,7 @@ export class MongoBucketBatch extends BaseObserver {
             }
         });
         if (activated) {
-            this.logger.info(`Activated new sync rules at ${lsn}`);
+            this.logger.info(`Activated new replication stream at ${lsn}`);
             await this.db.notifyCheckpoint();
             this.needsActivation = false;
         }
@@ -766,9 +744,10 @@ export class MongoBucketBatch extends BaseObserver {
         const result = await this.flush();
         await this.withTransaction(async () => {
             for (let table of sourceTables) {
-                await this.db.source_tables.deleteOne({ _id: mongoTableId(table.id) });
+                await this.db.commonSourceTables(this.group_id).deleteOne({ _id: mongoTableId(table.id) });
             }
         });
+        await this.cleanupDroppedSourceTables(sourceTables);
         return result;
     }
     async truncate(sourceTables) {
@@ -795,41 +774,30 @@ export class MongoBucketBatch extends BaseObserver {
         let lastBatchCount = BATCH_LIMIT;
         while (lastBatchCount == BATCH_LIMIT) {
             await this.withReplicationTransaction(`Truncate ${sourceTable.qualifiedName}`, async (session, opSeq) => {
-                const current_data_filter = {
-                    _id: idPrefixFilter({ g: this.group_id, t: mongoTableId(sourceTable.id) }, ['k']),
-                    // Skip soft-deleted data
-                    // Works for both v1 and v3 current_data schemas
-                    pending_delete: { $exists: false }
-                };
-                const cursor = this.db.common_current_data.find(current_data_filter, {
-                    projection: {
-                        _id: 1,
-                        buckets: 1,
-                        lookups: 1
-                    },
-                    limit: BATCH_LIMIT,
-                    session: session
-                });
-                const batch = await cursor.toArray();
-                const persistedBatch = new PersistedBatch(this.db, this.group_id, 0, { logger: this.logger });
+                using evalSpan = this.tracer.span('evaluate');
+                const sourceTableId = mongoTableId(sourceTable.id);
+                const batch = await this.sourceRecordStore.loadTruncateBatch(session, sourceTableId, BATCH_LIMIT);
+                const persistedBatch = this.createPersistedBatch(0);
                 for (let value of batch) {
                     persistedBatch.saveBucketData({
                         op_seq: opSeq,
                         before_buckets: value.buckets,
                         evaluated: [],
                         table: sourceTable,
-                        sourceKey: value._id.k
+                        sourceKey: value.replicaId
                     });
                     persistedBatch.saveParameterData({
                         op_seq: opSeq,
                         existing_lookups: value.lookups,
                         evaluated: [],
                         sourceTable: sourceTable,
-                        sourceKey: value._id.k
+                        sourceKey: value.replicaId
                     });
                     // Since this is not from streaming replication, we can do a hard delete
-                    persistedBatch.hardDeleteCurrentData(value._id);
+                    persistedBatch.hardDeleteCurrentData(sourceTableId, value.replicaId);
                 }
+                evalSpan.end();
+                using _ = this.tracer.span('storage', 'persist_flush');
                 await persistedBatch.flush(session);
                 lastBatchCount = batch.length;
                 last_op = opSeq.last();
@@ -846,7 +814,7 @@ export class MongoBucketBatch extends BaseObserver {
         };
         copy.snapshotStatus = snapshotStatus;
         await this.withTransaction(async () => {
-            await this.db.source_tables.updateOne({ _id: mongoTableId(table.id) }, {
+            await this.db.commonSourceTables(this.group_id).updateOne({ _id: mongoTableId(table.id) }, {
                 $set: {
                     snapshot_status: {
                         last_key: snapshotStatus.lastKey == null ? null : new bson.Binary(snapshotStatus.lastKey),
@@ -884,7 +852,7 @@ export class MongoBucketBatch extends BaseObserver {
         const session = this.session;
         const ids = tables.map((table) => mongoTableId(table.id));
         await this.withTransaction(async () => {
-            await this.db.source_tables.updateMany({ _id: { $in: ids } }, {
+            await this.db.commonSourceTables(this.group_id).updateMany({ _id: { $in: ids } }, {
                 $set: {
                     snapshot_done: true
                 },
@@ -933,7 +901,4 @@ export class MongoBucketBatch extends BaseObserver {
         return this.sync_rules.eventDescriptors.filter((evt) => [...evt.getSourceTables()].some((sourceTable) => sourceTable.matches(table)));
     }
 }
-export function currentBucketKey(b) {
-    return `${b.bucket}/${b.table}/${b.id}`;
-}
 //# sourceMappingURL=MongoBucketBatch.js.map