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

package/src/storage/implementation/MongoBucketBatch.ts

@@ -1,11 +1,10 @@
 import { mongo } from '@powersync/lib-service-mongodb';
-import { HydratedSyncRules, SqlEventDescriptor, SqliteRow, SqliteValue } from '@powersync/service-sync-rules';
+import { HydratedSyncConfig, SqlEventDescriptor, SqliteRow, SqliteValue } from '@powersync/service-sync-rules';
 import * as bson from 'bson';
 
 import {
   BaseObserver,
   container,
-  logger as defaultLogger,
   ErrorCode,
   errors,
   Logger,
@@ -14,28 +13,24 @@ import {
 } from '@powersync/lib-services-framework';
 import {
   BucketStorageMarkRecordUnavailable,
-  CheckpointResult,
   deserializeBson,
   InternalOpId,
   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 { VersionedPowerSyncMongo } from './db.js';
-import { CurrentBucket, CurrentDataDocument, SourceKey, SyncRuleDocument } from './models.js';
+import { mongoTableId } from '../../utils/util.js';
+import { BucketDefinitionMapping } from './BucketDefinitionMapping.js';
+import { PersistedBatch } from './common/PersistedBatch.js';
+import { LoadedSourceRecord, SourceRecordStore } from './common/SourceRecordStore.js';
+import type { VersionedPowerSyncMongo } from './db.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
@@ -44,49 +39,65 @@ export const MAX_ROW_SIZE = 15 * 1024 * 1024;
 // 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 interface MongoBucketBatchOptions {
   db: VersionedPowerSyncMongo;
-  syncRules: HydratedSyncRules;
+  syncRules: HydratedSyncConfig;
   groupId: number;
   slotName: string;
+  syncConfigId?: bson.ObjectId | null;
   lastCheckpointLsn: string | null;
   keepaliveOp: InternalOpId | null;
   resumeFromLsn: string | null;
   storeCurrentData: boolean;
+  mapping: BucketDefinitionMapping;
   /**
    * Set to true for initial replication.
    */
   skipExistingRows: boolean;
 
   markRecordUnavailable: BucketStorageMarkRecordUnavailable | undefined;
+  hooks: storage.StorageHooks | undefined;
 
-  logger?: Logger;
+  logger: Logger;
+  tracer?: PerformanceTracer<'storage' | 'evaluate'>;
 }
 
-export class MongoBucketBatch
+export abstract class MongoBucketBatch
   extends BaseObserver<storage.BucketBatchStorageListener>
   implements storage.BucketStorageBatch
 {
-  private logger: Logger;
+  protected readonly options: MongoBucketBatchOptions;
+  protected logger: Logger;
 
   private readonly client: mongo.MongoClient;
   public readonly db: VersionedPowerSyncMongo;
   public readonly session: mongo.ClientSession;
-  private readonly sync_rules: HydratedSyncRules;
+  protected readonly sync_rules: HydratedSyncConfig;
 
-  private readonly group_id: number;
+  protected readonly group_id: number;
 
   private readonly slot_name: string;
+  /**
+   * 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.
+   */
   private readonly storeCurrentData: boolean;
-  private readonly skipExistingRows: boolean;
+  public readonly skipExistingRows: boolean;
+  protected readonly mapping: BucketDefinitionMapping;
 
   private batch: OperationBatch | null = null;
   private write_checkpoint_batch: storage.CustomWriteCheckpointOptions[] = [];
   private markRecordUnavailable: BucketStorageMarkRecordUnavailable | undefined;
+  private hooks: storage.StorageHooks | undefined;
   private clearedError = false;
 
+  private tracer: PerformanceTracer<'storage' | 'evaluate'>;
+
   /**
    * Last LSN received associated with a checkpoint.
    *
@@ -94,9 +105,9 @@ export class MongoBucketBatch
    * 1. A commit LSN.
    * 2. A keepalive message LSN.
    */
-  private last_checkpoint_lsn: string | null = null;
+  protected last_checkpoint_lsn: string | null = null;
 
-  private persisted_op: InternalOpId | null = null;
+  protected persisted_op: InternalOpId | null = null;
 
   /**
    * Last written op, if any. This may not reflect a consistent checkpoint.
@@ -115,11 +126,10 @@ export class MongoBucketBatch
    */
   public resumeFromLsn: string | null = null;
 
-  private needsActivation = true;
-
   constructor(options: MongoBucketBatchOptions) {
     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;
@@ -129,11 +139,14 @@ export class MongoBucketBatch
     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: storage.BatchedCustomWriteCheckpointOptions): void {
@@ -147,6 +160,33 @@ export class MongoBucketBatch
     return this.last_checkpoint_lsn;
   }
 
+  abstract resolveTables(options: storage.ResolveTablesOptions): Promise<storage.ResolveTablesResult>;
+
+  protected abstract createPersistedBatch(writtenSize: number): PersistedBatch;
+
+  protected abstract get sourceRecordStore(): SourceRecordStore;
+
+  protected abstract cleanupDroppedSourceTables(sourceTables: storage.SourceTable[]): Promise<void>;
+
+  abstract commit(lsn: string, options?: storage.BucketBatchCommitOptions): Promise<storage.CheckpointResult>;
+
+  abstract keepalive(lsn: string): Promise<storage.CheckpointResult>;
+
+  abstract setResumeLsn(lsn: string): Promise<void>;
+
+  abstract getSourceTableStatus(table: storage.SourceTable): Promise<storage.SourceTable | null>;
+
+  abstract markAllSnapshotDone(no_checkpoint_before_lsn: string): Promise<void>;
+
+  abstract markSnapshotDone(no_checkpoint_before_lsn: string, options?: { throwOnConflict?: boolean }): Promise<void>;
+
+  abstract markTableSnapshotRequired(table: storage.SourceTable): Promise<void>;
+
+  abstract markTableSnapshotDone(
+    tables: storage.SourceTable[],
+    no_checkpoint_before_lsn?: string
+  ): Promise<storage.SourceTable[]>;
+
   async flush(options?: storage.BatchBucketFlushOptions): Promise<storage.FlushedResult | null> {
     let result: storage.FlushedResult | null = null;
     // One flush may be split over multiple transactions.
@@ -165,6 +205,10 @@ export class MongoBucketBatch
     let last_op: InternalOpId | null = null;
     let resumeBatch: OperationBatch | null = 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);
@@ -188,6 +232,7 @@ export class MongoBucketBatch
 
     this.persisted_op = last_op;
     this.last_flushed_op = last_op;
+    await this.hooks?.afterBatchFlush?.(this);
     return { flushed_op: last_op };
   }
 
@@ -198,8 +243,13 @@ export class MongoBucketBatch
     options?: storage.BucketBatchCommitOptions
   ): Promise<OperationBatch | null> {
     let sizes: Map<string, number> | undefined = 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,
@@ -212,33 +262,16 @@ export class MongoBucketBatch
       // (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: SourceKey[] = batch.batch.map((r) => {
-        return { g: this.group_id, t: mongoTableId(r.record.sourceTable.id), k: r.beforeId };
-      });
-
-      sizes = new Map<string, number>();
-
-      const sizeCursor: mongo.AggregationCursor<{ _id: SourceKey; size: number }> =
-        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.
@@ -256,64 +289,65 @@ export class MongoBucketBatch
         }
         continue;
       }
-      const lookups: SourceKey[] = b.map((r) => {
-        return { g: this.group_id, t: mongoTableId(r.record.sourceTable.id), k: r.beforeId };
-      });
-      let current_data_lookup = new Map<string, CurrentDataDocument>();
-      // 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: PersistedBatch | null = 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: PersistedBatch | null = 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;
           // Computing our current progress is a little tricky here, since
           // 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();
     }
 
@@ -323,7 +357,7 @@ export class MongoBucketBatch
   private saveOperation(
     batch: PersistedBatch,
     operation: RecordOperation,
-    current_data: CurrentDataDocument | null,
+    sourceRecord: LoadedSourceRecord | null,
     opSeq: MongoIdSequence
   ) {
     const record = operation.record;
@@ -331,17 +365,20 @@ export class MongoBucketBatch
     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: CurrentBucket[] = [];
-    let new_buckets: CurrentBucket[] = [];
-    let existing_lookups: bson.Binary[] = [];
-    let new_lookups: bson.Binary[] = [];
+    let existing_buckets: LoadedSourceRecord['buckets'] = [];
+    let new_buckets: LoadedSourceRecord['buckets'] = [];
+    let existing_lookups: LoadedSourceRecord['lookups'] = [];
+    let new_lookups: LoadedSourceRecord['lookups'] = [];
 
-    const before_key: SourceKey = { 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.
@@ -354,12 +391,12 @@ export class MongoBucketBatch
     }
 
     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
@@ -375,18 +412,20 @@ export class MongoBucketBatch
       } else {
         existing_buckets = result.buckets;
         existing_lookups = result.lookups;
-        if (this.storeCurrentData) {
-          const data = deserializeBson((result.data as mongo.Binary).buffer) as SqliteRow;
+        if (storeCurrentData && result.data != null) {
+          const data = deserializeBson(result.data.buffer) as SqliteRow;
           after = storage.mergeToast<SqliteValue>(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}`
@@ -398,9 +437,9 @@ export class MongoBucketBatch
       }
     }
 
-    let afterData: bson.Binary | undefined;
-    if (afterId != null && !this.storeCurrentData) {
-      afterData = EMPTY_DATA;
+    let afterData: bson.Binary | null = null;
+    if (afterId != null && !storeCurrentData) {
+      afterData = null;
     } else if (afterId != null) {
       try {
         // This will fail immediately if the record is > 16MB.
@@ -466,13 +505,15 @@ export class MongoBucketBatch
     // 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(
@@ -498,20 +539,15 @@ export class MongoBucketBatch
           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
+          sourceTable.ref,
+          after,
+          { parameterLookupSources: sourceTable.parameterLookupSources }
         );
 
         for (let error of paramErrors) {
@@ -537,28 +573,29 @@ export class MongoBucketBatch
           evaluated: paramEvaluated,
           existing_lookups
         });
-        new_lookups = paramEvaluated.map((p) => {
-          return storage.serializeLookup(p.lookup);
-        });
+        new_lookups = this.sourceRecordStore.mapParameterLookups(paramEvaluated);
       }
     }
 
-    let result: CurrentDataDocument | null = null;
+    let result: LoadedSourceRecord | null = null;
 
     // 5. TOAST: Update current data and bucket list.
     if (afterId) {
       // Insert or update
-      const after_key: SourceKey = { 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,
-        data: afterData!,
+        sourceTableId,
+        replicaId: afterId,
+        data: afterData,
         buckets: new_buckets,
-        lookups: new_lookups
+        lookups: new_lookups,
+        cacheKey: operation.internalAfterKey!
       };
     }
 
@@ -567,13 +604,15 @@ export class MongoBucketBatch
       // 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;
   }
 
-  private async withTransaction(cb: () => Promise<void>) {
+  protected async withTransaction(cb: () => Promise<void>) {
+    using lockSpan = this.tracer.span('storage', 'internal_lock');
     await replicationMutex.exclusiveLock(async () => {
+      lockSpan.end();
       await this.session.withTransaction(
         async () => {
           try {
@@ -584,7 +623,9 @@ export class MongoBucketBatch
             } else {
               this.logger.warn('Transaction error', e as Error);
             }
-            await timers.setTimeout(Math.random() * 50);
+            const delay = Math.random() * 50;
+            using _ = this.tracer.span('storage', 'retry_delay');
+            await timers.setTimeout(delay);
             throw e;
           }
         },
@@ -678,271 +719,9 @@ export class MongoBucketBatch
     await this[Symbol.asyncDispose]();
   }
 
-  private lastWaitingLogThottled = 0;
-
-  async commit(lsn: string, options?: storage.BucketBatchCommitOptions): Promise<CheckpointResult> {
-    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 };
-  }
-
-  private async cleanupCurrentData(lastCheckpoint: bigint) {
-    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.
-   */
-  private async autoActivate(lsn: string) {
-    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: string): Promise<CheckpointResult> {
-    return await this.commit(lsn, { createEmptyCheckpoints: true });
-  }
-
-  async setResumeLsn(lsn: string): Promise<void> {
-    const update: Partial<SyncRuleDocument> = {
-      snapshot_lsn: lsn
-    };
-
-    await this.db.sync_rules.updateOne(
-      {
-        _id: this.group_id
-      },
-      {
-        $set: update
-      },
-      { session: this.session }
-    );
-  }
-
   async save(record: storage.SaveOptions): Promise<storage.FlushedResult | null> {
     const { after, before, sourceTable, tag } = record;
+    const storeCurrentData = this.storeCurrentData && sourceTable.storeCurrentData;
     for (const event of this.getTableEvents(sourceTable)) {
       this.iterateListeners((cb) =>
         cb.replicationEvent?.({
@@ -950,8 +729,8 @@ export class MongoBucketBatch
           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
         })
@@ -988,9 +767,11 @@ export class MongoBucketBatch
 
     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;
   }
 
@@ -1022,24 +803,10 @@ export class MongoBucketBatch
     let lastBatchCount = BATCH_LIMIT;
     while (lastBatchCount == BATCH_LIMIT) {
       await this.withReplicationTransaction(`Truncate ${sourceTable.qualifiedName}`, async (session, opSeq) => {
-        const current_data_filter: mongo.Filter<CurrentDataDocument> = {
-          _id: idPrefixFilter<SourceKey>({ 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({
@@ -1047,19 +814,22 @@ export class MongoBucketBatch
             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;
 
@@ -1083,7 +853,7 @@ export class MongoBucketBatch
     copy.snapshotStatus = snapshotStatus;
 
     await this.withTransaction(async () => {
-      await this.db.source_tables.updateOne(
+      await this.db.commonSourceTables(this.group_id).updateOne(
         { _id: mongoTableId(table.id) },
         {
           $set: {
@@ -1101,80 +871,6 @@ export class MongoBucketBatch
     return copy;
   }
 
-  async markAllSnapshotDone(no_checkpoint_before_lsn: string) {
-    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: storage.SourceTable): Promise<void> {
-    await this.db.sync_rules.updateOne(
-      {
-        _id: this.group_id
-      },
-      {
-        $set: {
-          snapshot_done: false
-        }
-      },
-      { session: this.session }
-    );
-  }
-
-  async markTableSnapshotDone(tables: storage.SourceTable[], no_checkpoint_before_lsn?: string) {
-    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;
-    });
-  }
-
   protected async clearError(): Promise<void> {
     // No need to clear an error more than once per batch, since an error would always result in restarting the batch.
     if (this.clearedError) {
@@ -1200,11 +896,7 @@ export class MongoBucketBatch
    */
   protected getTableEvents(table: storage.SourceTable): SqlEventDescriptor[] {
     return this.sync_rules.eventDescriptors.filter((evt) =>
-      [...evt.getSourceTables()].some((sourceTable) => sourceTable.matches(table))
+      [...evt.getSourceTables()].some((sourceTable) => sourceTable.matches(table.ref))
     );
   }
 }
-
-export function currentBucketKey(b: CurrentBucket) {
-  return `${b.bucket}/${b.table}/${b.id}`;
-}