@powersync/service-module-mongodb-storage 0.0.0-dev-20260602073133 → 0.0.0-dev-20260909133214

package/src/storage/implementation/MongoBucketBatch.ts

@@ -1,5 +1,11 @@
 import { mongo } from '@powersync/lib-service-mongodb';
-import { HydratedSyncConfig, SqlEventDescriptor, SqliteRow, SqliteValue } from '@powersync/service-sync-rules';
+import {
+  EventDefinitionId,
+  HydratedEventDescriptor,
+  HydratedSyncConfig,
+  SqliteRow,
+  SqliteValue
+} from '@powersync/service-sync-rules';
 import * as bson from 'bson';
 
 import {
@@ -12,6 +18,7 @@ import {
   ServiceError
 } from '@powersync/lib-services-framework';
 import {
+  BucketDefinitionMapping,
   BucketStorageMarkRecordUnavailable,
   deserializeBson,
   InternalOpId,
@@ -23,14 +30,15 @@ import {
 } from '@powersync/service-core';
 import * as timers from 'node:timers/promises';
 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 { MongoParsedSyncConfigSet } from './MongoParsedSyncConfigSet.js';
 import { OperationBatch, RecordOperation } from './OperationBatch.js';
+import { ObjectStorage } from './v3/object-storage/ObjectStorage.js';
+import { createObjectStorageUsageWriterId } from './v3/object-storage/ObjectStorageUsage.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
@@ -41,15 +49,25 @@ const replicationMutex = new utils.Mutex();
 
 export interface MongoBucketBatchOptions {
   db: VersionedPowerSyncMongo;
-  syncRules: HydratedSyncConfig;
-  groupId: number;
-  slotName: string;
-  syncConfigId?: bson.ObjectId | null;
-  lastCheckpointLsn: string | null;
-  keepaliveOp: InternalOpId | null;
+  /**
+   * The parsed sync config set for this batch.
+   *
+   * The batch derives both the hydrated sync rules and the bucket definition mapping from
+   * this single set, so they always come from the same parse. Do not add separate
+   * syncRules/mapping options - pairing values from different parses is exactly the bug
+   * this shape prevents.
+   */
+  parsedSyncConfig: MongoParsedSyncConfigSet;
+  replicationStreamId: number;
+  replicationStreamName: string;
+  syncConfigIds?: bson.ObjectId[];
+  /**
+   * Seeds the in-memory persisted-op tracking for v1 storage. Not used by v3 storage, which
+   * tracks the persisted-op head durably on the replication stream document instead.
+   */
+  keepaliveOp?: InternalOpId | null;
   resumeFromLsn: string | null;
   storeCurrentData: boolean;
-  mapping: BucketDefinitionMapping;
   /**
    * Set to true for initial replication.
    */
@@ -60,6 +78,13 @@ export interface MongoBucketBatchOptions {
 
   logger: Logger;
   tracer?: PerformanceTracer<'storage' | 'evaluate'>;
+  /**
+   * Aborts in-flight object storage uploads when replication stops.
+   */
+  signal?: AbortSignal;
+
+  objectStorage?: ObjectStorage;
+  inlineThresholdBytes?: number;
 }
 
 export abstract class MongoBucketBatch
@@ -70,35 +95,34 @@ export abstract class MongoBucketBatch
   protected logger: Logger;
 
   private readonly client: mongo.MongoClient;
-  public readonly db: VersionedPowerSyncMongo;
+  readonly db: VersionedPowerSyncMongo;
   public readonly session: mongo.ClientSession;
   protected readonly sync_rules: HydratedSyncConfig;
 
-  protected readonly group_id: number;
+  protected readonly replicationStreamId: number;
 
-  private readonly slot_name: string;
+  private readonly replicationStreamName: 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;
   public readonly skipExistingRows: boolean;
   protected readonly mapping: BucketDefinitionMapping;
+  protected readonly objectStorageUsageWriterId = createObjectStorageUsageWriterId();
 
   private batch: OperationBatch | null = null;
-  private write_checkpoint_batch: storage.CustomWriteCheckpointOptions[] = [];
+  protected 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.
-   *
-   * This could be either:
-   * 1. A commit LSN.
-   * 2. A keepalive message LSN.
-   */
-  protected last_checkpoint_lsn: string | null = null;
-
-  protected persisted_op: InternalOpId | null = null;
+  protected tracer: PerformanceTracer<'storage' | 'evaluate'>;
 
   /**
    * Last written op, if any. This may not reflect a consistent checkpoint.
@@ -106,9 +130,9 @@ export abstract class MongoBucketBatch
   public last_flushed_op: InternalOpId | null = null;
 
   /**
-   * lastCheckpointLsn is the last consistent commit.
+   * LSN to resume replication from.
    *
-   * While that is generally a "safe" point to resume from, there are cases where we may want to resume from a different point:
+   * This is typically the last commit LSN, but there are cases where it differs:
    * 1. After an initial snapshot, we don't have a consistent commit yet, but need to resume from the snapshot LSN.
    * 2. If "no_checkpoint_before_lsn" is set far in advance, it may take a while to reach that point. We
    *    may want to resume at incremental points before that.
@@ -123,34 +147,28 @@ export abstract class MongoBucketBatch
     this.options = options;
     this.client = options.db.client;
     this.db = options.db;
-    this.group_id = options.groupId;
-    this.last_checkpoint_lsn = options.lastCheckpointLsn;
+    this.replicationStreamId = options.replicationStreamId;
     this.resumeFromLsn = options.resumeFromLsn;
     this.session = this.client.startSession();
-    this.slot_name = options.slotName;
-    this.sync_rules = options.syncRules;
+    this.replicationStreamName = options.replicationStreamName;
+    this.sync_rules = options.parsedSyncConfig.hydratedSyncConfig;
     this.storeCurrentData = options.storeCurrentData;
-    this.mapping = options.mapping;
+    this.mapping = options.parsedSyncConfig.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 {
     this.write_checkpoint_batch.push({
       ...checkpoint,
-      sync_rules_id: this.group_id
+      sync_rules_id: this.replicationStreamId
     });
   }
 
-  get lastCheckpointLsn() {
-    return this.last_checkpoint_lsn;
-  }
-
   abstract resolveTables(options: storage.ResolveTablesOptions): Promise<storage.ResolveTablesResult>;
 
   protected abstract createPersistedBatch(writtenSize: number): PersistedBatch;
@@ -178,6 +196,11 @@ export abstract class MongoBucketBatch
     no_checkpoint_before_lsn?: string
   ): Promise<storage.SourceTable[]>;
 
+  protected abstract batchCreateCustomWriteCheckpoints(session: mongo.ClientSession, opId: InternalOpId): Promise<void>;
+
+  /** Perform any version-specific setup required before writing custom checkpoints. */
+  protected async prepareCustomWriteCheckpoints(): Promise<void> {}
+
   async flush(options?: storage.BatchBucketFlushOptions): Promise<storage.FlushedResult | null> {
     let result: storage.FlushedResult | null = null;
     // One flush may be split over multiple transactions.
@@ -195,25 +218,37 @@ export abstract class MongoBucketBatch
     const batch = this.batch;
     let last_op: InternalOpId | null = null;
     let resumeBatch: OperationBatch | null = null;
+    let clearedError = false;
 
     using _ = this.tracer.span('storage', 'flush');
 
     await this.hooks?.beforeBatchFlush?.(this);
+    if (this.write_checkpoint_batch.length > 0) {
+      // Collection/index creation cannot run inside the replication transaction.
+      await this.prepareCustomWriteCheckpoints();
+    }
 
     await this.withReplicationTransaction(`Flushing ${batch?.length ?? 0} ops`, async (session, opSeq) => {
+      clearedError = false;
       if (batch != null) {
-        resumeBatch = await this.replicateBatch(session, batch, opSeq, options);
+        const result = await this.replicateBatch(session, batch, opSeq, options);
+        resumeBatch = result.resumeBatch;
+        clearedError ||= result.clearedError;
       }
 
       if (this.write_checkpoint_batch.length > 0) {
         this.logger.info(`Writing ${this.write_checkpoint_batch.length} custom write checkpoints`);
-        await batchCreateCustomWriteCheckpoints(this.db, session, this.write_checkpoint_batch, opSeq.next());
+        await this.batchCreateCustomWriteCheckpoints(session, opSeq.next());
         this.write_checkpoint_batch = [];
       }
 
       last_op = opSeq.last();
     });
 
+    if (clearedError) {
+      this.clearedError = true;
+    }
+
     // null if done, set if we need another flush
     this.batch = resumeBatch;
 
@@ -221,7 +256,7 @@ export abstract class MongoBucketBatch
       throw new ReplicationAssertionError('Unexpected last_op == null');
     }
 
-    this.persisted_op = last_op;
+    this.recordPersistedOp(last_op);
     this.last_flushed_op = last_op;
     await this.hooks?.afterBatchFlush?.(this);
     return { flushed_op: last_op };
@@ -232,11 +267,15 @@ export abstract class MongoBucketBatch
     batch: OperationBatch,
     op_seq: MongoIdSequence,
     options?: storage.BucketBatchCommitOptions
-  ): Promise<OperationBatch | null> {
+  ): Promise<{ resumeBatch: OperationBatch | null; clearedError: boolean }> {
     let sizes: Map<string, number> | undefined = 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
+    // 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,
@@ -249,10 +288,14 @@ export abstract 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 = batch.batch.map((r) => ({
-        sourceTableId: mongoTableId(r.record.sourceTable.id),
-        replicaId: r.beforeId
-      }));
+      // 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);
     }
@@ -329,12 +372,17 @@ export abstract class MongoBucketBatch
       }
     }
 
-    if (didFlush) {
+    const clearedError = didFlush && !this.clearedError;
+    if (clearedError) {
+      // No need to clear an error more than once per batch, since an error would always result in restarting the batch.
       using _ = this.tracer.span('storage', 'clear_error');
-      await this.clearError();
+      await this.clearError(this.session);
     }
 
-    return resumeBatch?.hasData() ? resumeBatch : null;
+    return {
+      resumeBatch: resumeBatch?.hasData() ? resumeBatch : null,
+      clearedError
+    };
   }
 
   private saveOperation(
@@ -348,6 +396,9 @@ export abstract 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: LoadedSourceRecord['buckets'] = [];
     let new_buckets: LoadedSourceRecord['buckets'] = [];
@@ -376,7 +427,7 @@ export abstract class MongoBucketBatch
         // 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
@@ -392,7 +443,7 @@ export abstract class MongoBucketBatch
       } else {
         existing_buckets = result.buckets;
         existing_lookups = result.lookups;
-        if (this.storeCurrentData && result.data != null) {
+        if (storeCurrentData && result.data != null) {
           const data = deserializeBson(result.data.buffer) as SqliteRow;
           after = storage.mergeToast<SqliteValue>(after!, data);
         }
@@ -403,7 +454,9 @@ export abstract class MongoBucketBatch
         // 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}`
@@ -416,7 +469,7 @@ export abstract class MongoBucketBatch
     }
 
     let afterData: bson.Binary | null = null;
-    if (afterId != null && !this.storeCurrentData) {
+    if (afterId != null && !storeCurrentData) {
       afterData = null;
     } else if (afterId != null) {
       try {
@@ -440,7 +493,7 @@ export abstract class MongoBucketBatch
           {
             level: errors.ErrorSeverity.WARNING,
             metadata: {
-              replication_slot: this.slot_name,
+              replication_slot: this.replicationStreamName,
               table: record.sourceTable.qualifiedName
             }
           }
@@ -483,7 +536,7 @@ export abstract 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, errors: syncErrors } = this.sync_rules.evaluateRowWithErrors({
@@ -499,7 +552,7 @@ export abstract class MongoBucketBatch
             {
               level: errors.ErrorSeverity.WARNING,
               metadata: {
-                replication_slot: this.slot_name,
+                replication_slot: this.replicationStreamName,
                 table: record.sourceTable.qualifiedName
               }
             }
@@ -534,7 +587,7 @@ export abstract class MongoBucketBatch
             {
               level: errors.ErrorSeverity.WARNING,
               metadata: {
-                replication_slot: this.slot_name,
+                replication_slot: this.replicationStreamName,
                 table: record.sourceTable.qualifiedName
               }
             }
@@ -669,7 +722,7 @@ export abstract class MongoBucketBatch
 
       await this.db.sync_rules.updateOne(
         {
-          _id: this.group_id
+          _id: this.replicationStreamId
         },
         {
           $set: {
@@ -678,12 +731,41 @@ export abstract class MongoBucketBatch
         },
         { session }
       );
+
+      // Allow subclasses to persist additional flush-time state in the same transaction
+      // (e.g. v3 advances the stream-level last_persisted_op).
+      await this.onReplicationTransactionFlush(session, opSeq.last());
       // We don't notify checkpoint here - we don't make any checkpoint updates directly
     });
   }
 
+  /**
+   * Hook called inside the replication flush transaction, after ops have been persisted.
+   *
+   * The base implementation does nothing; v3 storage overrides this to `$max` the stream-level
+   * `last_persisted_op` durably within the same transaction.
+   */
+  protected async onReplicationTransactionFlush(_session: mongo.ClientSession, _lastOp: InternalOpId): Promise<void> {
+    // No-op by default (v1 behaviour unchanged).
+  }
+
+  /**
+   * Called after a replication transaction has successfully committed, with the last persisted op id.
+   *
+   * v1 storage tracks this in memory to fold into the next checkpoint. v3 storage does not need it:
+   * the stream-level `last_persisted_op` is already advanced durably by
+   * {@link onReplicationTransactionFlush} within the same transaction, and checkpoints read it
+   * from the document.
+   *
+   * Keep calls to this adjacent to {@link withReplicationTransaction} usage - both must observe
+   * every path that persists ops.
+   */
+  protected recordPersistedOp(_lastOp: InternalOpId): void {
+    // No-op by default.
+  }
+
   async [Symbol.asyncDispose]() {
-    if (this.batch != null || this.write_checkpoint_batch.length > 0) {
+    if (this.batch?.hasData() || this.write_checkpoint_batch.length > 0) {
       // We don't error here, since:
       // 1. In error states, this is expected (we can't distinguish between disposing after success or error).
       // 2. SuppressedError is messy to deal with.
@@ -699,19 +781,27 @@ export abstract class MongoBucketBatch
 
   async save(record: storage.SaveOptions): Promise<storage.FlushedResult | null> {
     const { after, before, sourceTable, tag } = record;
-    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
-          },
-          event
-        })
-      );
+    const storeCurrentData = this.storeCurrentData && sourceTable.storeCurrentData;
+    // V3 source tables own disjoint event-definition ids for each physical table. Multiple
+    // SourceTables may evaluate different events, but each definition is evaluated through
+    // at most one record for this row change.
+    // Legacy storage leaves eventDefinitionIds undefined and selects by table ref.
+    if (sourceTable.syncEvent) {
+      for (const { event, eventId } of this.getTableEvents(sourceTable)) {
+        this.iterateListeners((cb) =>
+          cb.replicationEvent?.({
+            batch: this,
+            table: sourceTable,
+            data: {
+              op: tag,
+              after: after && utils.isCompleteRow(storeCurrentData, after) ? after : undefined,
+              before: before && utils.isCompleteRow(storeCurrentData, before) ? before : undefined
+            },
+            event,
+            event_id: eventId
+          })
+        );
+      }
     }
 
     /**
@@ -744,7 +834,7 @@ export abstract class MongoBucketBatch
 
     await this.withTransaction(async () => {
       for (let table of sourceTables) {
-        await this.db.commonSourceTables(this.group_id).deleteOne({ _id: mongoTableId(table.id) });
+        await this.db.commonSourceTables(this.replicationStreamId).deleteOne({ _id: mongoTableId(table.id) });
       }
     });
 
@@ -761,7 +851,7 @@ export abstract class MongoBucketBatch
     }
 
     if (last_op) {
-      this.persisted_op = last_op;
+      this.recordPersistedOp(last_op);
       return {
         flushed_op: last_op
       };
@@ -830,7 +920,7 @@ export abstract class MongoBucketBatch
     copy.snapshotStatus = snapshotStatus;
 
     await this.withTransaction(async () => {
-      await this.db.commonSourceTables(this.group_id).updateOne(
+      await this.db.commonSourceTables(this.replicationStreamId).updateOne(
         { _id: mongoTableId(table.id) },
         {
           $set: {
@@ -848,32 +938,47 @@ export abstract class MongoBucketBatch
     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) {
-      return;
-    }
-
+  protected async clearError(session: mongo.ClientSession): Promise<void> {
     await this.db.sync_rules.updateOne(
       {
-        _id: this.group_id
+        _id: this.replicationStreamId
       },
       {
         $set: {
           last_fatal_error: null,
           last_fatal_error_ts: null
         }
+      },
+      {
+        session
       }
     );
-    this.clearedError = true;
   }
 
   /**
-   * Gets relevant {@link SqlEventDescriptor}s for the given {@link SourceTable}
+   * Gets relevant {@link HydratedEventDescriptor}s for the given {@link SourceTable}
    */
-  protected getTableEvents(table: storage.SourceTable): SqlEventDescriptor[] {
-    return this.sync_rules.eventDescriptors.filter((evt) =>
-      [...evt.getSourceTables()].some((sourceTable) => sourceTable.matches(table.ref))
-    );
+  protected getTableEvents(
+    table: storage.SourceTable
+  ): { event: HydratedEventDescriptor; eventId?: EventDefinitionId }[] {
+    // V3 storage assigns event-definition ids to each source table, so membership is authoritative.
+    // Iterate the table's distinct ids and resolve each through the stream's deduped event map, so a
+    // definition reused across configs has one evaluator for that id. The evaluator may still return
+    // multiple payloads from matching queries. Legacy storage leaves this undefined and selects by table ref.
+    if (table.eventDefinitionIds != null) {
+      const eventById = this.options.parsedSyncConfig.eventById;
+      const events: { event: HydratedEventDescriptor; eventId: EventDefinitionId }[] = [];
+      for (const id of table.eventDefinitionIds) {
+        const event = eventById.get(id);
+        if (event != null && event.tableTriggersEvent(table.ref)) {
+          events.push({ event, eventId: id });
+        }
+      }
+      return events;
+    }
+
+    return this.sync_rules.eventDescriptors
+      .filter((event) => event.tableTriggersEvent(table.ref))
+      .map((event) => ({ event }));
   }
 }