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

package/src/storage/implementation/MongoBucketBatch.ts

@@ -12,6 +12,7 @@ import {
   ServiceError
 } from '@powersync/lib-services-framework';
 import {
+  BucketDefinitionMapping,
   BucketStorageMarkRecordUnavailable,
   deserializeBson,
   InternalOpId,
@@ -23,12 +24,12 @@ 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 { MongoParsedSyncConfigSet } from './MongoParsedSyncConfigSet.js';
 import { batchCreateCustomWriteCheckpoints } from './MongoWriteCheckpointAPI.js';
 import { OperationBatch, RecordOperation } from './OperationBatch.js';
 
@@ -41,15 +42,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.
    */
@@ -70,13 +81,13 @@ 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.
    *
@@ -98,26 +109,15 @@ export abstract class MongoBucketBatch
 
   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;
-
   /**
    * Last written op, if any. This may not reflect a consistent checkpoint.
    */
   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.
@@ -132,34 +132,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;
@@ -230,7 +224,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 };
@@ -462,7 +456,7 @@ export abstract class MongoBucketBatch
           {
             level: errors.ErrorSeverity.WARNING,
             metadata: {
-              replication_slot: this.slot_name,
+              replication_slot: this.replicationStreamName,
               table: record.sourceTable.qualifiedName
             }
           }
@@ -521,7 +515,7 @@ export abstract class MongoBucketBatch
             {
               level: errors.ErrorSeverity.WARNING,
               metadata: {
-                replication_slot: this.slot_name,
+                replication_slot: this.replicationStreamName,
                 table: record.sourceTable.qualifiedName
               }
             }
@@ -556,7 +550,7 @@ export abstract class MongoBucketBatch
             {
               level: errors.ErrorSeverity.WARNING,
               metadata: {
-                replication_slot: this.slot_name,
+                replication_slot: this.replicationStreamName,
                 table: record.sourceTable.qualifiedName
               }
             }
@@ -691,7 +685,7 @@ export abstract class MongoBucketBatch
 
       await this.db.sync_rules.updateOne(
         {
-          _id: this.group_id
+          _id: this.replicationStreamId
         },
         {
           $set: {
@@ -700,12 +694,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.
@@ -722,19 +745,24 @@ export abstract class MongoBucketBatch
   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?.({
-          batch: this,
-          table: sourceTable,
-          data: {
-            op: tag,
-            after: after && utils.isCompleteRow(storeCurrentData, after) ? after : undefined,
-            before: before && utils.isCompleteRow(storeCurrentData, before) ? before : undefined
-          },
-          event
-        })
-      );
+    // syncEvent is the per-table designation from resolveTables. With v3 storage, multiple
+    // SourceTables can exist for the same ref, with a row change saved once per table -
+    // only the designated event carrier may fire events, so each event fires once per row.
+    if (sourceTable.syncEvent) {
+      for (const event 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
+          })
+        );
+      }
     }
 
     /**
@@ -767,7 +795,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) });
       }
     });
 
@@ -784,7 +812,7 @@ export abstract class MongoBucketBatch
     }
 
     if (last_op) {
-      this.persisted_op = last_op;
+      this.recordPersistedOp(last_op);
       return {
         flushed_op: last_op
       };
@@ -853,7 +881,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: {
@@ -879,7 +907,7 @@ export abstract class MongoBucketBatch
 
     await this.db.sync_rules.updateOne(
       {
-        _id: this.group_id
+        _id: this.replicationStreamId
       },
       {
         $set: {

package/src/storage/implementation/MongoChecksums.ts

@@ -7,19 +7,16 @@ import {
   ChecksumMap,
   FetchPartialBucketChecksum,
   InternalOpId,
-  isPartialChecksum,
   PartialChecksum,
   PartialChecksumMap,
   PartialOrFullChecksum
 } from '@powersync/service-core';
 import type { VersionedPowerSyncMongo } from './db.js';
 
-import * as lib_mongo from '@powersync/lib-service-mongodb';
 import { BucketDefinitionId } from '@powersync/service-sync-rules';
-import { BucketDefinitionMapping } from './BucketDefinitionMapping.js';
-import { BucketDataDocumentBase, StorageConfig } from './models.js';
+import { StorageConfig } from './models.js';
 
-export interface FetchPartialBucketChecksumV3 {
+export interface FetchPartialBucketChecksumByDefinition {
   bucket: string;
   definitionId: BucketDefinitionId;
   start?: InternalOpId;
@@ -47,15 +44,14 @@ export interface MongoChecksumOptions {
   operationBatchLimit?: number;
 
   storageConfig: StorageConfig;
-  mapping?: BucketDefinitionMapping;
 }
 
 const DEFAULT_BUCKET_BATCH_LIMIT = 200;
-const DEFAULT_OPERATION_BATCH_LIMIT = 50_000;
+export const DEFAULT_OPERATION_BATCH_LIMIT = 50_000;
 
 export abstract class MongoChecksums {
   private _cache: ChecksumCache | undefined;
-  private readonly storageConfig: StorageConfig;
+  protected readonly storageConfig: StorageConfig;
 
   constructor(
     protected readonly db: VersionedPowerSyncMongo,
@@ -178,138 +174,10 @@ export abstract class MongoChecksums {
   protected abstract fetchPreStates(
     batch: FetchPartialBucketChecksum[]
   ): Promise<Map<string, { opId: InternalOpId; checksum: BucketChecksum }>>;
-
-  protected async computePartialChecksumsForCollection<
-    TRequest extends FetchPartialBucketChecksumByBucket,
-    TBucketDataDocument extends BucketDataDocumentBase
-  >(
-    batch: TRequest[],
-    collection: lib_mongo.mongo.Collection<TBucketDataDocument>,
-    createFilter: (request: TRequest) => any
-  ): Promise<PartialChecksumMap> {
-    const batchLimit = this.options?.operationBatchLimit ?? DEFAULT_OPERATION_BATCH_LIMIT;
-
-    // Map requests by bucket. We adjust this as we get partial results.
-    let requests = new Map<string, TRequest>();
-    for (let request of batch) {
-      requests.set(request.bucket, request);
-    }
-
-    const partialChecksums = new Map<string, PartialOrFullChecksum>();
-
-    while (requests.size > 0) {
-      const filters = Array.from(requests.values(), createFilter);
-
-      // Historically, checksum may be stored as 'int' or 'double'.
-      // More recently, this should be a 'long'.
-      // $toLong ensures that we always sum it as a long, avoiding inaccuracies in the calculations.
-      const checksumLong = this.storageConfig.longChecksums ? '$checksum' : { $toLong: '$checksum' };
-
-      // Aggregate over a max of `batchLimit` operations at a time.
-      // Let's say we have 3 buckets (A, B, C), each with 10 operations, and our batch limit is 12.
-      // Then we'll do three batches:
-      // 1. Query: A[1-end], B[1-end], C[1-end]
-      //    Returns: A[1-10], B[1-2]
-      // 2. Query: B[3-end], C[1-end]
-      //    Returns: B[3-10], C[1-4]
-      // 3. Query: C[5-end]
-      //    Returns: C[5-10]
-      const aggregate = await collection
-        .aggregate(
-          [
-            {
-              $match: {
-                $or: filters
-              }
-            },
-            // sort and limit _before_ grouping
-            { $sort: { _id: 1 } },
-            { $limit: batchLimit },
-            {
-              $group: {
-                _id: '$_id.b',
-                checksum_total: { $sum: checksumLong },
-                count: { $sum: 1 },
-                has_clear_op: {
-                  $max: {
-                    $cond: [{ $eq: ['$op', 'CLEAR'] }, 1, 0]
-                  }
-                },
-                last_op: { $max: '$_id.o' }
-              }
-            },
-            // Sort the aggregated results (100 max, so should be fast).
-            // This is important to identify which buckets we have partial data for.
-            { $sort: { _id: 1 } }
-          ],
-          { session: undefined, readConcern: 'snapshot', maxTimeMS: lib_mongo.MONGO_CHECKSUM_TIMEOUT_MS }
-        )
-        .toArray()
-        .catch((e) => {
-          throw lib_mongo.mapQueryError(e, 'while reading checksums');
-        });
-
-      let batchCount = 0;
-      let limitReached = false;
-      for (let doc of aggregate) {
-        const bucket = doc._id;
-        const checksum = checksumFromAggregate(doc);
-
-        const existing = partialChecksums.get(bucket);
-        if (existing != null) {
-          partialChecksums.set(bucket, addPartialChecksums(bucket, existing, checksum));
-        } else {
-          partialChecksums.set(bucket, checksum);
-        }
-
-        batchCount += doc.count;
-        if (batchCount == batchLimit) {
-          // Limit reached. Request more in the next batch.
-          // Note that this only affects the _last_ bucket in a batch.
-          limitReached = true;
-          const req = requests.get(bucket);
-          requests.set(bucket, {
-            ...req!,
-            start: doc.last_op
-          });
-        } else {
-          // All done for this bucket
-          requests.delete(bucket);
-        }
-      }
-      if (!limitReached) {
-        break;
-      }
-    }
-
-    return new Map<string, PartialOrFullChecksum>(
-      batch.map((request) => {
-        const bucket = request.bucket;
-        // Could be null if we got no data
-        let partialChecksum = partialChecksums.get(bucket);
-        if (partialChecksum == null) {
-          partialChecksum = {
-            bucket,
-            partialCount: 0,
-            partialChecksum: 0
-          };
-        }
-        if (request.start == null && isPartialChecksum(partialChecksum)) {
-          partialChecksum = {
-            bucket,
-            count: partialChecksum.partialCount,
-            checksum: partialChecksum.partialChecksum
-          };
-        }
-
-        return [bucket, partialChecksum];
-      })
-    );
-  }
 }
 
 export function emptyChecksumForRequest(
-  request: Pick<FetchPartialBucketChecksum | FetchPartialBucketChecksumV3, 'bucket' | 'start'>
+  request: Pick<FetchPartialBucketChecksum | FetchPartialBucketChecksumByDefinition, 'bucket' | 'start'>
 ): PartialOrFullChecksum {
   if (request.start == null) {
     return { bucket: request.bucket, count: 0, checksum: 0 };
@@ -320,7 +188,7 @@ export function emptyChecksumForRequest(
 /**
  * Convert output of the $group stage into a checksum.
  */
-function checksumFromAggregate(doc: bson.Document): PartialOrFullChecksum {
+export function checksumFromAggregate(doc: bson.Document): PartialOrFullChecksum {
   const partialChecksum = Number(BigInt(doc.checksum_total) & 0xffffffffn) & 0xffffffff;
   const bucket = doc._id;
 

package/src/storage/implementation/MongoCompactor.ts

@@ -63,6 +63,7 @@ export interface MongoCompactOptions extends storage.CompactOptions {}
 const DEFAULT_CLEAR_BATCH_LIMIT = 5000;
 const DEFAULT_MOVE_BATCH_LIMIT = 2000;
 const DEFAULT_MOVE_BATCH_QUERY_LIMIT = 10_000;
+const DEFAULT_MOVE_BATCH_BYTE_LIMIT = 16 * 1024 * 1024;
 const DEFAULT_MIN_BUCKET_CHANGES = 10;
 const DEFAULT_MIN_CHANGE_RATIO = 0.1;
 const DIRTY_BUCKET_SCAN_BATCH_SIZE = 2_000;
@@ -83,6 +84,7 @@ export abstract class MongoCompactor {
   protected readonly idLimitBytes: number;
   protected readonly moveBatchLimit: number;
   protected readonly moveBatchQueryLimit: number;
+  protected readonly moveBatchByteLimit: number;
   protected readonly clearBatchLimit: number;
   protected readonly minBucketChanges: number;
   protected readonly minChangeRatio: number;
@@ -98,11 +100,15 @@ export abstract class MongoCompactor {
     protected readonly db: VersionedPowerSyncMongo,
     options: MongoCompactOptions
   ) {
-    this.group_id = storage.group_id;
+    this.group_id = storage.replicationStreamId;
     this.idLimitBytes = (options.memoryLimitMB ?? DEFAULT_MEMORY_LIMIT_MB) * 1024 * 1024;
     this.moveBatchLimit = options.moveBatchLimit ?? DEFAULT_MOVE_BATCH_LIMIT;
     this.moveBatchQueryLimit = options.moveBatchQueryLimit ?? DEFAULT_MOVE_BATCH_QUERY_LIMIT;
+    this.moveBatchByteLimit = options.moveBatchByteLimit ?? DEFAULT_MOVE_BATCH_BYTE_LIMIT;
     this.clearBatchLimit = options.clearBatchLimit ?? DEFAULT_CLEAR_BATCH_LIMIT;
+    if (this.clearBatchLimit < 2) {
+      throw new ReplicationAssertionError('clearBatchLimit must be >= 2');
+    }
     this.minBucketChanges = options.minBucketChanges ?? DEFAULT_MIN_BUCKET_CHANGES;
     this.minChangeRatio = options.minChangeRatio ?? DEFAULT_MIN_CHANGE_RATIO;
     this.maxOpId = options.maxOpId ?? 0n;
@@ -133,6 +139,7 @@ export abstract class MongoCompactor {
    */
   async populateChecksums(options: { minBucketChanges: number }): Promise<PopulateChecksumCacheResults> {
     let count = 0;
+    // Paginate through dirty buckets in batches until no more buckets meet the criteria.
     while (true) {
       this.signal?.throwIfAborted();
       const buckets = await this.dirtyBucketBatchForChecksums(options);
@@ -172,6 +179,7 @@ export abstract class MongoCompactor {
     },
     getDefinitionId: (state: TCollectionBucketState) => BucketDefinitionId | null
   ): AsyncGenerator<DirtyBucket[]> {
+    // Paginate through the bucket state collection using cursor-based scanning.
     while (true) {
       // To avoid timeouts from too many buckets not meeting the minBucketChanges criteria, use an aggregation pipeline
       // to scan a fixed batch of buckets at a time, but only return buckets that meet the criteria.
@@ -304,6 +312,7 @@ export abstract class MongoCompactor {
    */
   protected async compactSingleBucketRetried(bucket: string, definitionId: BucketDefinitionId | null = null) {
     let retryCount = 0;
+    // Retry with exponential backoff up to 3 times on MongoDB errors.
     while (true) {
       try {
         await this.compactSingleBucket(bucket, definitionId);
@@ -343,6 +352,7 @@ export abstract class MongoCompactor {
     // Upper bound is adjusted for each batch.
     let upperBound = bucketContext.maxId;
 
+    // Paginate through bucket data in batches to avoid cursor timeouts.
     while (true) {
       this.signal?.throwIfAborted();
 
@@ -353,11 +363,7 @@ export abstract class MongoCompactor {
             _id: {
               $gte: lowerBound,
               $lt: upperBound
-            },
-            // Workaround for a clustered collection bug where the $lt operator may include upperBound.
-            // Technically only needed for storage V3.
-            // https://jira.mongodb.org/browse/SERVER-121822
-            '_id.o': { $lt: upperBound.o }
+            }
           }
         },
         { $sort: { _id: -1 } },
@@ -488,13 +494,13 @@ export abstract class MongoCompactor {
     await this.flush(bucketContext);
   }
 
-  protected updateBucketChecksums(state: CurrentBucketState) {
+  protected collectBucketStateUpdates(state: CurrentBucketState): mongo.AnyBulkWriteOperation<BucketStateDocumentBase> {
     if (state.opCount < 0) {
       throw new ServiceAssertionError(
         `Invalid opCount: ${state.opCount} checksum ${state.checksum} opsSincePut: ${state.opsSincePut} maxOpId: ${this.maxOpId}`
       );
     }
-    this.bucketStateUpdates.push({
+    return {
       updateOne: {
         filter: this.bucketStateFilter(state.bucket, state.definitionId),
         update: {
@@ -517,7 +523,11 @@ export abstract class MongoCompactor {
         // We don't create new ones here, to avoid issues with the unique index on bucket_updates.
         upsert: false
       }
-    });
+    };
+  }
+
+  protected updateBucketChecksums(state: CurrentBucketState) {
+    this.bucketStateUpdates.push(this.collectBucketStateUpdates(state));
   }
 
   protected async flush(col: SingleBucketStore) {

package/src/storage/implementation/MongoHydrationState.ts

@@ -0,0 +1,61 @@
+import { ServiceAssertionError } from '@powersync/lib-services-framework';
+import { SyncConfigWithRequiredMapping } from '@powersync/service-core';
+import {
+  BucketDataScope,
+  BucketDataSource,
+  HydrationState,
+  ParameterIndexLookupCreator,
+  ParameterLookupScope
+} from '@powersync/service-sync-rules';
+
+export class MongoHydrationState implements HydrationState {
+  private bucketDataSourceSyncConfig = new WeakMap<BucketDataSource, SyncConfigWithRequiredMapping>();
+  private parameterIndexLookupSyncConfig = new WeakMap<ParameterIndexLookupCreator, SyncConfigWithRequiredMapping>();
+
+  constructor(
+    readonly syncConfigs: SyncConfigWithRequiredMapping[],
+    private readonly replicationStreamId: number
+  ) {
+    for (let syncConfig of syncConfigs) {
+      for (let source of syncConfig.syncConfig.config.bucketDataSources) {
+        this.bucketDataSourceSyncConfig.set(source, syncConfig);
+      }
+      for (let source of syncConfig.syncConfig.config.bucketParameterLookupSources) {
+        this.parameterIndexLookupSyncConfig.set(source, syncConfig);
+      }
+    }
+  }
+
+  getBucketSourceScope(source: BucketDataSource): BucketDataScope {
+    const syncConfig = this.bucketDataSourceSyncConfig.get(source);
+    if (syncConfig == null) {
+      throw new ServiceAssertionError(`No sync config found for bucket data source ${source.uniqueName}`);
+    }
+    const mapping = syncConfig.mapping;
+    const defId = mapping.bucketSourceId(source);
+    // Currently uniqueName is constant for a BucketDataSource within a replication stream. This may change in the
+    // future, in which case we need to reconsider the bucketPrefix here.
+    // However, the definition may change without changing the name, so we include the defId in the bucketPrefix.
+    // defId is not unique across replication streams, so we include the replicationStreamId as well.
+    return {
+      bucketPrefix: `${source.uniqueName}.${this.replicationStreamId.toString(16)}.${defId}`,
+      source
+    };
+  }
+
+  getParameterIndexLookupScope(source: ParameterIndexLookupCreator): ParameterLookupScope {
+    const syncConfig = this.parameterIndexLookupSyncConfig.get(source);
+    if (syncConfig == null) {
+      throw new ServiceAssertionError(
+        `No sync config found for parameter index lookup source ${source.sourceId.lookupName}#${source.sourceId.queryId}`
+      );
+    }
+    const mapping = syncConfig.mapping;
+    const defId = mapping.parameterLookupId(source);
+    return {
+      lookupName: defId,
+      queryId: '',
+      source
+    };
+  }
+}