@powersync/service-module-mongodb-storage 0.15.4 → 0.17.0

package/dist/storage/implementation/MongoBucketBatch.js

@@ -1,25 +1,21 @@
 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, 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 {
+    options;
     logger;
     client;
     db;
@@ -27,12 +23,24 @@ export class MongoBucketBatch extends BaseObserver {
     sync_rules;
     group_id;
     slot_name;
+    /**
+     * Source-level setting for whether raw row data should be stored in current_data.
+     *
+     * Some sources always send complete rows (MongoDB, MySQL with binlog_row_image=full),
+     * in which case this is false for the whole batch. For sources where it depends on the
+     * table (Postgres REPLICA IDENTITY), this is true and the decision is refined per-table
+     * via SourceTable.storeCurrentData. The effective per-record value is the conjunction of
+     * the two.
+     */
     storeCurrentData;
     skipExistingRows;
+    mapping;
     batch = null;
     write_checkpoint_batch = [];
     markRecordUnavailable;
+    hooks;
     clearedError = false;
+    tracer;
     /**
      * Last LSN received associated with a checkpoint.
      *
@@ -57,10 +65,10 @@ export class MongoBucketBatch extends BaseObserver {
      * This is set when creating the batch, but may not be updated afterwards.
      */
     resumeFromLsn = null;
-    needsActivation = true;
     constructor(options) {
         super();
-        this.logger = options.logger ?? defaultLogger;
+        this.logger = options.logger;
+        this.options = options;
         this.client = options.db.client;
         this.db = options.db;
         this.group_id = options.groupId;
@@ -70,10 +78,13 @@ 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.hooks = options.hooks;
         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 +111,8 @@ export class MongoBucketBatch extends BaseObserver {
         const batch = this.batch;
         let last_op = null;
         let resumeBatch = null;
+        using _ = this.tracer.span('storage', 'flush');
+        await this.hooks?.beforeBatchFlush?.(this);
         await this.withReplicationTransaction(`Flushing ${batch?.length ?? 0} ops`, async (session, opSeq) => {
             if (batch != null) {
                 resumeBatch = await this.replicateBatch(session, batch, opSeq, options);
@@ -118,12 +131,17 @@ export class MongoBucketBatch extends BaseObserver {
         }
         this.persisted_op = last_op;
         this.last_flushed_op = last_op;
+        await this.hooks?.afterBatchFlush?.(this);
         return { flushed_op: last_op };
     }
     async replicateBatch(session, batch, op_seq, options) {
         let sizes = undefined;
-        if (this.storeCurrentData && !this.skipExistingRows) {
-            // We skip this step if we don't store current_data, since the sizes will
+        using _ = this.tracer.span('storage', 'replicate_batch');
+        // Only look up current_data sizes if the batch stores current_data and at least one
+        // table in it does too (per-table can disable it, e.g. Postgres REPLICA IDENTITY FULL).
+        const anyTableStoresCurrentData = this.storeCurrentData && batch.batch.some((r) => r.record.sourceTable.storeCurrentData);
+        if (anyTableStoresCurrentData && !this.skipExistingRows) {
+            // We skip this step if no tables store current_data, since the sizes will
             // always be small in that case.
             // With skipExistingRows, we don't load the full documents into memory,
             // so we can also skip the size lookup step.
@@ -134,27 +152,15 @@ 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);
-            }
+            // Within this branch this.storeCurrentData is true, so the per-table flag is the
+            // effective value - only look up sizes for tables that actually store current_data.
+            const sizeLookups = batch.batch
+                .filter((r) => r.record.sourceTable.storeCurrentData)
+                .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 +175,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,33 +214,41 @@ 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;
         let after = record.after;
         const sourceTable = record.sourceTable;
+        // Effective per-record flag: store current_data only if both the batch (source-level,
+        // e.g. Postgres) and the table (e.g. non-FULL replica identity) require it.
+        const storeCurrentData = this.storeCurrentData && sourceTable.storeCurrentData;
         let existing_buckets = [];
         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,12 +261,12 @@ 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 = [];
                 existing_lookups = [];
-                if (!isCompleteRow(this.storeCurrentData, after)) {
+                if (!isCompleteRow(storeCurrentData, after)) {
                     if (this.markRecordUnavailable != null) {
                         // This will trigger a "resnapshot" of the record.
                         // This is not relevant if storeCurrentData is false, since we'll get the full row
@@ -270,19 +282,21 @@ export class MongoBucketBatch extends BaseObserver {
             else {
                 existing_buckets = result.buckets;
                 existing_lookups = result.lookups;
-                if (this.storeCurrentData) {
+                if (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 = [];
                 existing_lookups = [];
-                // Log to help with debugging if there was a consistency issue
+                // Log to help with debugging if there was a consistency issue.
+                // Gate on the batch-level flag: FULL tables (per-record flag false) still get a
+                // current_data entry, so a missing record on DELETE is meaningful for them too.
                 if (this.storeCurrentData && this.markRecordUnavailable == null) {
                     this.logger.warn(`Cannot find previous record for delete on ${record.sourceTable.qualifiedName}: ${beforeId} / ${record.before?.id}`);
                 }
@@ -292,9 +306,9 @@ export class MongoBucketBatch extends BaseObserver {
                 existing_lookups = result.lookups;
             }
         }
-        let afterData;
-        if (afterId != null && !this.storeCurrentData) {
-            afterData = EMPTY_DATA;
+        let afterData = null;
+        if (afterId != null && !storeCurrentData) {
+            afterData = null;
         }
         else if (afterId != null) {
             try {
@@ -353,13 +367,15 @@ export class MongoBucketBatch extends BaseObserver {
         // However, it will be valid by the end of the transaction.
         //
         // In this case, we don't save the op, but we do save the current data.
-        if (afterId && after && utils.isCompleteRow(this.storeCurrentData, after)) {
+        if (afterId && after && utils.isCompleteRow(storeCurrentData, after)) {
             // Insert or update
             if (sourceTable.syncData) {
-                const { results: evaluated, errors: syncErrors } = this.sync_rules.evaluateRowWithErrors({
+                const { results, errors: syncErrors } = this.sync_rules.evaluateRowWithErrors({
                     record: after,
-                    sourceTable
+                    sourceTable: sourceTable.ref,
+                    bucketDataSources: sourceTable.bucketDataSources
                 });
+                const evaluated = results;
                 for (let error of syncErrors) {
                     container.reporter.captureMessage(`Failed to evaluate data query on ${record.sourceTable.qualifiedName}.${record.after?.id}: ${error.error}`, {
                         level: errors.ErrorSeverity.WARNING,
@@ -378,17 +394,11 @@ 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
-                const { results: paramEvaluated, errors: paramErrors } = this.sync_rules.evaluateParameterRowWithErrors(sourceTable, after);
+                const { results: paramEvaluated, errors: paramErrors } = this.sync_rules.evaluateParameterRowWithErrors(sourceTable.ref, after, { parameterLookupSources: sourceTable.parameterLookupSources });
                 for (let error of paramErrors) {
                     container.reporter.captureMessage(`Failed to evaluate parameter query on ${record.sourceTable.qualifiedName}.${record.after?.id}: ${error.error}`, {
                         level: errors.ErrorSeverity.WARNING,
@@ -406,26 +416,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 +444,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 +463,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 });
@@ -516,227 +531,17 @@ export class MongoBucketBatch extends BaseObserver {
     async dispose() {
         await this[Symbol.asyncDispose]();
     }
-    lastWaitingLogThottled = 0;
-    async commit(lsn, options) {
-        const { createEmptyCheckpoints } = { ...storage.DEFAULT_BUCKET_BATCH_COMMIT_OPTIONS, ...options };
-        await this.flush(options);
-        const now = new Date();
-        // Mark relevant write checkpoints as "processed".
-        // This makes it easier to identify write checkpoints that are "valid" in order.
-        await this.db.write_checkpoints.updateMany({
-            processed_at_lsn: null,
-            'lsns.1': { $lte: lsn }
-        }, {
-            $set: {
-                processed_at_lsn: lsn
-            }
-        }, {
-            session: this.session
-        });
-        const can_checkpoint = {
-            $and: [
-                { $eq: ['$snapshot_done', true] },
-                {
-                    $or: [{ $eq: ['$last_checkpoint_lsn', null] }, { $lte: ['$last_checkpoint_lsn', { $literal: lsn }] }]
-                },
-                {
-                    $or: [{ $eq: ['$no_checkpoint_before', null] }, { $lte: ['$no_checkpoint_before', { $literal: lsn }] }]
-                }
-            ]
-        };
-        const new_keepalive_op = {
-            $cond: [
-                can_checkpoint,
-                { $literal: null },
-                {
-                    $toString: {
-                        $max: [{ $toLong: '$keepalive_op' }, { $literal: this.persisted_op }, 0n]
-                    }
-                }
-            ]
-        };
-        const new_last_checkpoint = {
-            $cond: [
-                can_checkpoint,
-                {
-                    $max: ['$last_checkpoint', { $literal: this.persisted_op }, { $toLong: '$keepalive_op' }, 0n]
-                },
-                '$last_checkpoint'
-            ]
-        };
-        // For this query, we need to handle multiple cases, depending on the state:
-        // 1. Normal commit - advance last_checkpoint to this.persisted_op.
-        // 2. Commit delayed by no_checkpoint_before due to snapshot. In this case we only advance keepalive_op.
-        // 3. Commit with no new data - here may may set last_checkpoint = keepalive_op, if a delayed commit is relevant.
-        // We want to do as much as possible in a single atomic database operation, which makes this somewhat complex.
-        let preUpdateDocument = await this.db.sync_rules.findOneAndUpdate({ _id: this.group_id }, [
-            {
-                $set: {
-                    _can_checkpoint: can_checkpoint,
-                    _not_empty: createEmptyCheckpoints
-                        ? true
-                        : {
-                            $or: [
-                                { $literal: createEmptyCheckpoints },
-                                { $ne: ['$keepalive_op', new_keepalive_op] },
-                                { $ne: ['$last_checkpoint', new_last_checkpoint] }
-                            ]
-                        }
-                }
-            },
-            {
-                $set: {
-                    last_checkpoint_lsn: {
-                        $cond: [{ $and: ['$_can_checkpoint', '$_not_empty'] }, { $literal: lsn }, '$last_checkpoint_lsn']
-                    },
-                    last_checkpoint_ts: {
-                        $cond: [{ $and: ['$_can_checkpoint', '$_not_empty'] }, { $literal: now }, '$last_checkpoint_ts']
-                    },
-                    last_keepalive_ts: { $literal: now },
-                    last_fatal_error: { $literal: null },
-                    last_fatal_error_ts: { $literal: null },
-                    keepalive_op: new_keepalive_op,
-                    last_checkpoint: new_last_checkpoint,
-                    // Unset snapshot_lsn on checkpoint
-                    snapshot_lsn: {
-                        $cond: [{ $and: ['$_can_checkpoint', '$_not_empty'] }, { $literal: null }, '$snapshot_lsn']
-                    }
-                }
-            },
-            {
-                $unset: ['_can_checkpoint', '_not_empty']
-            }
-        ], {
-            session: this.session,
-            // We return the before document, so that we can check the previous state to determine if a checkpoint was actually created or if we were blocked by snapshot/no_checkpoint_before.
-            returnDocument: 'before',
-            projection: {
-                snapshot_done: 1,
-                last_checkpoint_lsn: 1,
-                no_checkpoint_before: 1,
-                keepalive_op: 1,
-                last_checkpoint: 1
-            }
-        });
-        if (preUpdateDocument == null) {
-            throw new ReplicationAssertionError('Failed to update checkpoint - no matching sync_rules document for _id: ' + this.group_id);
-        }
-        // This re-implements the same logic as in the pipeline, to determine what was actually updated.
-        // Unfortunately we cannot return these from the pipeline directly, so we need to re-implement the logic.
-        const canCheckpoint = preUpdateDocument.snapshot_done === true &&
-            (preUpdateDocument.last_checkpoint_lsn == null || preUpdateDocument.last_checkpoint_lsn <= lsn) &&
-            (preUpdateDocument.no_checkpoint_before == null || preUpdateDocument.no_checkpoint_before <= lsn);
-        const keepaliveOp = preUpdateDocument.keepalive_op == null ? null : BigInt(preUpdateDocument.keepalive_op);
-        const maxKeepalive = [keepaliveOp ?? 0n, this.persisted_op ?? 0n, 0n].reduce((a, b) => (a > b ? a : b));
-        const newKeepaliveOp = canCheckpoint ? null : maxKeepalive.toString();
-        const newLastCheckpoint = canCheckpoint
-            ? [preUpdateDocument.last_checkpoint ?? 0n, this.persisted_op ?? 0n, keepaliveOp ?? 0n, 0n].reduce((a, b) => a > b ? a : b)
-            : preUpdateDocument.last_checkpoint;
-        const notEmpty = createEmptyCheckpoints ||
-            preUpdateDocument.keepalive_op !== newKeepaliveOp ||
-            preUpdateDocument.last_checkpoint !== newLastCheckpoint;
-        const checkpointCreated = canCheckpoint && notEmpty;
-        const checkpointBlocked = !canCheckpoint;
-        if (checkpointBlocked) {
-            // Failed on snapshot_done or no_checkpoint_before.
-            if (Date.now() - this.lastWaitingLogThottled > 5_000) {
-                this.logger.info(`Waiting before creating checkpoint, currently at ${lsn} / ${preUpdateDocument.keepalive_op}. Current state: ${JSON.stringify({
-                    snapshot_done: preUpdateDocument.snapshot_done,
-                    last_checkpoint_lsn: preUpdateDocument.last_checkpoint_lsn,
-                    no_checkpoint_before: preUpdateDocument.no_checkpoint_before
-                })}`);
-                this.lastWaitingLogThottled = Date.now();
-            }
-        }
-        else {
-            if (checkpointCreated) {
-                this.logger.debug(`Created checkpoint at ${lsn} / ${newLastCheckpoint}`);
-            }
-            await this.autoActivate(lsn);
-            await this.db.notifyCheckpoint();
-            this.persisted_op = null;
-            this.last_checkpoint_lsn = lsn;
-            if (this.db.storageConfig.softDeleteCurrentData && newLastCheckpoint != null) {
-                await this.cleanupCurrentData(newLastCheckpoint);
-            }
-        }
-        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.
-     *
-     * Called on new commits.
-     */
-    async autoActivate(lsn) {
-        if (!this.needsActivation) {
-            return;
-        }
-        // Activate the batch, so it can start processing.
-        // This is done automatically when the first save() is called.
-        const session = this.session;
-        let activated = false;
-        await session.withTransaction(async () => {
-            const doc = await this.db.sync_rules.findOne({ _id: this.group_id }, { session });
-            if (doc && doc.state == SyncRuleState.PROCESSING && doc.snapshot_done && doc.last_checkpoint != null) {
-                await this.db.sync_rules.updateOne({
-                    _id: this.group_id
-                }, {
-                    $set: {
-                        state: storage.SyncRuleState.ACTIVE
-                    }
-                }, { session });
-                await this.db.sync_rules.updateMany({
-                    _id: { $ne: this.group_id },
-                    state: { $in: [storage.SyncRuleState.ACTIVE, storage.SyncRuleState.ERRORED] }
-                }, {
-                    $set: {
-                        state: storage.SyncRuleState.STOP
-                    }
-                }, { session });
-                activated = true;
-            }
-            else if (doc?.state != SyncRuleState.PROCESSING) {
-                this.needsActivation = false;
-            }
-        });
-        if (activated) {
-            this.logger.info(`Activated new sync rules at ${lsn}`);
-            await this.db.notifyCheckpoint();
-            this.needsActivation = false;
-        }
-    }
-    async keepalive(lsn) {
-        return await this.commit(lsn, { createEmptyCheckpoints: true });
-    }
-    async setResumeLsn(lsn) {
-        const update = {
-            snapshot_lsn: lsn
-        };
-        await this.db.sync_rules.updateOne({
-            _id: this.group_id
-        }, {
-            $set: update
-        }, { session: this.session });
-    }
     async save(record) {
         const { after, before, sourceTable, tag } = record;
+        const storeCurrentData = this.storeCurrentData && sourceTable.storeCurrentData;
         for (const event of this.getTableEvents(sourceTable)) {
             this.iterateListeners((cb) => cb.replicationEvent?.({
                 batch: this,
                 table: sourceTable,
                 data: {
                     op: tag,
-                    after: after && utils.isCompleteRow(this.storeCurrentData, after) ? after : undefined,
-                    before: before && utils.isCompleteRow(this.storeCurrentData, before) ? before : undefined
+                    after: after && utils.isCompleteRow(storeCurrentData, after) ? after : undefined,
+                    before: before && utils.isCompleteRow(storeCurrentData, before) ? before : undefined
                 },
                 event
             }));
@@ -766,9 +571,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 +601,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 +641,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),
@@ -858,59 +653,6 @@ export class MongoBucketBatch extends BaseObserver {
         });
         return copy;
     }
-    async markAllSnapshotDone(no_checkpoint_before_lsn) {
-        await this.db.sync_rules.updateOne({
-            _id: this.group_id
-        }, {
-            $set: {
-                snapshot_done: true,
-                last_keepalive_ts: new Date()
-            },
-            $max: {
-                no_checkpoint_before: no_checkpoint_before_lsn
-            }
-        }, { session: this.session });
-    }
-    async markTableSnapshotRequired(table) {
-        await this.db.sync_rules.updateOne({
-            _id: this.group_id
-        }, {
-            $set: {
-                snapshot_done: false
-            }
-        }, { session: this.session });
-    }
-    async markTableSnapshotDone(tables, no_checkpoint_before_lsn) {
-        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 } }, {
-                $set: {
-                    snapshot_done: true
-                },
-                $unset: {
-                    snapshot_status: 1
-                }
-            }, { session });
-            if (no_checkpoint_before_lsn != null) {
-                await this.db.sync_rules.updateOne({
-                    _id: this.group_id
-                }, {
-                    $set: {
-                        last_keepalive_ts: new Date()
-                    },
-                    $max: {
-                        no_checkpoint_before: no_checkpoint_before_lsn
-                    }
-                }, { session: this.session });
-            }
-        });
-        return tables.map((table) => {
-            const copy = table.clone();
-            copy.snapshotComplete = true;
-            return copy;
-        });
-    }
     async clearError() {
         // No need to clear an error more than once per batch, since an error would always result in restarting the batch.
         if (this.clearedError) {
@@ -930,10 +672,7 @@ export class MongoBucketBatch extends BaseObserver {
      * Gets relevant {@link SqlEventDescriptor}s for the given {@link SourceTable}
      */
     getTableEvents(table) {
-        return this.sync_rules.eventDescriptors.filter((evt) => [...evt.getSourceTables()].some((sourceTable) => sourceTable.matches(table)));
+        return this.sync_rules.eventDescriptors.filter((evt) => [...evt.getSourceTables()].some((sourceTable) => sourceTable.matches(table.ref)));
     }
 }
-export function currentBucketKey(b) {
-    return `${b.bucket}/${b.table}/${b.id}`;
-}
 //# sourceMappingURL=MongoBucketBatch.js.map