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

package/src/storage/implementation/MongoPersistedSyncRules.ts

@@ -1,3 +1,5 @@
+import * as sqlite from 'node:sqlite';
+
 import { ServiceAssertionError } from '@powersync/lib-services-framework';
 import { storage } from '@powersync/service-core';
 import {
@@ -5,9 +7,11 @@ import {
   BucketDataSource,
   CompatibilityOption,
   DEFAULT_HYDRATION_STATE,
-  HydratedSyncRules,
+  HydratedSyncConfig,
   HydrationState,
+  nodeSqlite,
   ParameterIndexLookupCreator,
+  ParameterLookupScope,
   SyncConfigWithErrors,
   versionedHydrationState
 } from '@powersync/service-sync-rules';
@@ -19,7 +23,7 @@ export class MongoPersistedSyncRules implements storage.PersistedSyncRules {
 
   constructor(
     public readonly id: number,
-    public readonly sync_rules: SyncConfigWithErrors,
+    public readonly syncConfigWithErrors: SyncConfigWithErrors,
     public readonly slot_name: string,
     private readonly mapping: BucketDefinitionMapping | null,
     private readonly storageConfig: StorageConfig
@@ -30,7 +34,7 @@ export class MongoPersistedSyncRules implements storage.PersistedSyncRules {
       }
       this.hydrationState = new MongoHydrationState(this.mapping, this.id);
     } else if (
-      !this.sync_rules.config.compatibility.isEnabled(CompatibilityOption.versionedBucketIds) &&
+      !this.syncConfigWithErrors.config.compatibility.isEnabled(CompatibilityOption.versionedBucketIds) &&
       !this.storageConfig.versionedBuckets
     ) {
       this.hydrationState = DEFAULT_HYDRATION_STATE;
@@ -39,8 +43,11 @@ export class MongoPersistedSyncRules implements storage.PersistedSyncRules {
     }
   }
 
-  hydratedSyncRules(): HydratedSyncRules {
-    return this.sync_rules.config.hydrate({ hydrationState: this.hydrationState });
+  hydratedSyncConfig(): HydratedSyncConfig {
+    return this.syncConfigWithErrors.config.hydrate({
+      hydrationState: this.hydrationState,
+      sqlite: nodeSqlite(sqlite)
+    });
   }
 }
 
@@ -54,7 +61,6 @@ class MongoHydrationState implements HydrationState {
     // Keep this aligned with versionedHydrationState() for now.
     //
     // Previous Mongo-specific behavior:
-    // const defId = this.mapping.bucketSourceId(source);
     // return {
     //   bucketPrefix: defId,
     //   source
@@ -65,7 +71,7 @@ class MongoHydrationState implements HydrationState {
     };
   }
 
-  getParameterIndexLookupScope(source: ParameterIndexLookupCreator) {
+  getParameterIndexLookupScope(source: ParameterIndexLookupCreator): ParameterLookupScope {
     const defId = this.mapping.parameterLookupId(source);
     return {
       lookupName: defId,

package/src/storage/implementation/MongoPersistedSyncRulesContent.ts

@@ -1,34 +1,31 @@
 import { mongo } from '@powersync/lib-service-mongodb';
-import { storage } from '@powersync/service-core';
+import { ServiceAssertionError } from '@powersync/lib-services-framework';
+import { storage, SyncRuleState } from '@powersync/service-core';
+import * as bson from 'bson';
+import { ReplicationStreamDocumentV3, SyncConfigDefinition } from '../storage-index.js';
 import { BucketDefinitionMapping } from './BucketDefinitionMapping.js';
 import { MongoPersistedSyncRules } from './MongoPersistedSyncRules.js';
 import { MongoSyncRulesLock } from './MongoSyncRulesLock.js';
 import { PowerSyncMongo } from './db.js';
-import { getMongoStorageConfig, SyncRuleDocument } from './models.js';
+import { getMongoStorageConfig } from './models.js';
+import { SyncRuleDocumentV1 } from './v1/models.js';
 
-export class MongoPersistedSyncRulesContent extends storage.PersistedSyncRulesContent {
+abstract class MongoPersistedSyncRulesContentBase extends storage.PersistedSyncRulesContent {
   public current_lock: MongoSyncRulesLock | null = null;
   public readonly mapping: BucketDefinitionMapping;
+  public readonly syncConfigId: bson.ObjectId | null;
 
-  constructor(
-    private db: PowerSyncMongo,
-    doc: mongo.WithId<SyncRuleDocument>
+  protected constructor(
+    protected readonly db: PowerSyncMongo,
+    options: ConstructorParameters<typeof storage.PersistedSyncRulesContent>[0] & {
+      mapping: BucketDefinitionMapping;
+      syncConfigId: bson.ObjectId | null;
+    }
   ) {
-    super({
-      id: doc._id,
-      sync_rules_content: doc.content,
-      compiled_plan: doc.serialized_plan ?? null,
-      last_checkpoint_lsn: doc.last_checkpoint_lsn,
-      // Handle legacy values
-      slot_name: doc.slot_name ?? `powersync_${doc._id}`,
-      last_fatal_error: doc.last_fatal_error,
-      last_fatal_error_ts: doc.last_fatal_error_ts,
-      last_checkpoint_ts: doc.last_checkpoint_ts,
-      last_keepalive_ts: doc.last_keepalive_ts,
-      active: doc.state == 'ACTIVE',
-      storageVersion: doc.storage_version ?? storage.LEGACY_STORAGE_VERSION
-    });
-    this.mapping = BucketDefinitionMapping.fromSyncRules(doc);
+    const { mapping, syncConfigId, ...base } = options;
+    super(base);
+    this.mapping = mapping;
+    this.syncConfigId = syncConfigId;
   }
 
   getStorageConfig() {
@@ -41,16 +38,64 @@ export class MongoPersistedSyncRulesContent extends storage.PersistedSyncRulesCo
 
     return new MongoPersistedSyncRules(
       parsed.id,
-      parsed.sync_rules,
+      parsed.syncConfigWithErrors,
       parsed.slot_name,
       storageConfig.incrementalReprocessing ? this.mapping : null,
       storageConfig
     );
   }
 
-  async lock() {
-    const lock = await MongoSyncRulesLock.createLock(this.db.versioned(this.getStorageConfig()), this);
+  async lock(session?: mongo.ClientSession) {
+    const lock = await MongoSyncRulesLock.createLock(this.db.versioned(this.getStorageConfig()), this, session);
     this.current_lock = lock;
     return lock;
   }
 }
+
+export class MongoPersistedSyncRulesContentV1 extends MongoPersistedSyncRulesContentBase {
+  constructor(db: PowerSyncMongo, doc: SyncRuleDocumentV1) {
+    super(db, {
+      id: doc._id,
+      sync_rules_content: doc.content,
+      compiled_plan: doc.serialized_plan ?? null,
+      last_checkpoint_lsn: doc.last_checkpoint_lsn,
+      // Handle legacy values
+      slot_name: doc.slot_name ?? `powersync_${doc._id}`,
+      last_fatal_error: doc.last_fatal_error,
+      last_fatal_error_ts: doc.last_fatal_error_ts,
+      last_checkpoint_ts: doc.last_checkpoint_ts,
+      last_keepalive_ts: doc.last_keepalive_ts,
+      active: doc.state == SyncRuleState.ACTIVE,
+      storageVersion: doc.storage_version ?? storage.LEGACY_STORAGE_VERSION,
+      mapping: new BucketDefinitionMapping(),
+      syncConfigId: null
+    });
+  }
+}
+
+export class MongoPersistedSyncRulesContentV3 extends MongoPersistedSyncRulesContentBase {
+  declare public readonly syncConfigId: bson.ObjectId;
+
+  constructor(db: PowerSyncMongo, doc: ReplicationStreamDocumentV3, config: SyncConfigDefinition) {
+    const state = doc.sync_configs.find((c) => c._id.equals(config._id));
+    if (state == null) {
+      throw new ServiceAssertionError(`Cannot find sync config ${config._id} in replication stream ${doc._id}`);
+    }
+    super(db, {
+      id: doc._id,
+      sync_rules_content: config.content,
+      compiled_plan: config.serialized_plan ?? null,
+
+      last_checkpoint_lsn: state?.last_checkpoint_lsn ?? null,
+      slot_name: doc.slot_name ?? `powersync_${doc._id}`,
+      last_fatal_error: doc.last_fatal_error,
+      last_fatal_error_ts: doc.last_fatal_error_ts,
+      last_checkpoint_ts: doc.last_checkpoint_ts,
+      last_keepalive_ts: doc.last_keepalive_ts,
+      active: doc.state == SyncRuleState.ACTIVE && state.state == SyncRuleState.ACTIVE,
+      storageVersion: doc.storage_version,
+      mapping: BucketDefinitionMapping.fromSyncConfig(config),
+      syncConfigId: config._id
+    });
+  }
+}

package/src/storage/implementation/MongoSyncBucketStorage.ts

@@ -13,7 +13,6 @@ import {
   CheckpointChanges,
   GetCheckpointChangesOptions,
   InternalOpId,
-  maxLsn,
   mergeAsyncIterables,
   PopulateChecksumCacheOptions,
   PopulateChecksumCacheResults,
@@ -22,7 +21,7 @@ import {
   utils,
   WatchWriteCheckpointOptions
 } from '@powersync/service-core';
-import { HydratedSyncRules, ParameterLookupRows, ScopedParameterLookup } from '@powersync/service-sync-rules';
+import { HydratedSyncConfig, ParameterLookupRows, ScopedParameterLookup } from '@powersync/service-sync-rules';
 import * as bson from 'bson';
 import { LRUCache } from 'lru-cache';
 import * as timers from 'timers/promises';
@@ -30,12 +29,12 @@ import { retryOnMongoMaxTimeMSExpired } from '../../utils/util.js';
 import { MongoBucketStorage } from '../MongoBucketStorage.js';
 import { MongoSyncBucketStorageContext } from './common/MongoSyncBucketStorageContext.js';
 import type { VersionedPowerSyncMongo } from './db.js';
-import { CommonSourceTableDocument, StorageConfig } from './models.js';
+import { StorageConfig } from './models.js';
 import { MongoBucketBatchOptions } from './MongoBucketBatch.js';
 import { MongoChecksumOptions, MongoChecksums } from './MongoChecksums.js';
 import { MongoCompactOptions, MongoCompactor } from './MongoCompactor.js';
 import { MongoParameterCompactor } from './MongoParameterCompactor.js';
-import { MongoPersistedSyncRulesContent } from './MongoPersistedSyncRulesContent.js';
+import { MongoPersistedSyncRulesContentV1 } from './MongoPersistedSyncRulesContent.js';
 import { MongoWriteCheckpointAPI } from './MongoWriteCheckpointAPI.js';
 
 export interface MongoSyncBucketStorageOptions {
@@ -48,6 +47,13 @@ interface InternalCheckpointChanges extends CheckpointChanges {
   invalidateWriteCheckpoints: boolean;
 }
 
+interface WriterSyncState {
+  lastCheckpointLsn: string | null;
+  resumeFromLsn: string | null;
+  keepaliveOp: InternalOpId | null;
+  syncConfigId?: bson.ObjectId | null;
+}
+
 /**
  * Only keep checkpoints around for a minute, before fetching a fresh one.
  *
@@ -68,21 +74,23 @@ export abstract class MongoSyncBucketStorage
 
   readonly checksums: MongoChecksums;
 
-  private parsedSyncRulesCache: { parsed: HydratedSyncRules; options: storage.ParseSyncRulesOptions } | undefined;
+  private parsedSyncRulesCache: { parsed: HydratedSyncConfig; options: storage.ParseSyncRulesOptions } | undefined;
   private writeCheckpointAPI: MongoWriteCheckpointAPI;
   public readonly logger: Logger;
+  public readonly storageConfig: StorageConfig;
   #storageInitialized = false;
 
   constructor(
     public readonly factory: MongoBucketStorage,
     public readonly group_id: number,
-    protected readonly sync_rules: MongoPersistedSyncRulesContent,
+    protected readonly sync_rules: MongoPersistedSyncRulesContentV1,
     public readonly slot_name: string,
     writeCheckpointMode: storage.WriteCheckpointMode | undefined,
     options: MongoSyncBucketStorageOptions
   ) {
     super();
-    this.db = factory.db.versioned(sync_rules.getStorageConfig());
+    this.storageConfig = options.storageConfig;
+    this.db = factory.db.versioned(this.storageConfig);
     this.checksums = this.createMongoChecksums(options);
     this.writeCheckpointAPI = new MongoWriteCheckpointAPI({
       db: this.db,
@@ -136,10 +144,10 @@ export abstract class MongoSyncBucketStorage
     });
   }
 
-  getParsedSyncRules(options: storage.ParseSyncRulesOptions): HydratedSyncRules {
+  getParsedSyncRules(options: storage.ParseSyncRulesOptions): HydratedSyncConfig {
     const { parsed, options: cachedOptions } = this.parsedSyncRulesCache ?? {};
     if (!parsed || options.defaultSchema != cachedOptions?.defaultSchema) {
-      this.parsedSyncRulesCache = { parsed: this.sync_rules.parsed(options).hydratedSyncRules(), options };
+      this.parsedSyncRulesCache = { parsed: this.sync_rules.parsed(options).hydratedSyncConfig(), options };
     }
 
     return this.parsedSyncRulesCache!.parsed;
@@ -149,16 +157,14 @@ export abstract class MongoSyncBucketStorage
     return (await this.getCheckpointInternal()) ?? new EmptyReplicationCheckpoint();
   }
 
+  protected abstract fetchCheckpointState(
+    session: mongo.ClientSession
+  ): Promise<{ checkpoint: bigint; lsn: string | null } | null>;
+
   async getCheckpointInternal(): Promise<storage.ReplicationCheckpoint | null> {
     return await this.db.client.withSession({ snapshot: true }, async (session) => {
-      const doc = await this.db.sync_rules.findOne(
-        { _id: this.group_id },
-        {
-          session,
-          projection: { _id: 1, state: 1, last_checkpoint: 1, last_checkpoint_lsn: 1, snapshot_done: 1 }
-        }
-      );
-      if (!doc?.snapshot_done || !['ACTIVE', 'ERRORED'].includes(doc.state)) {
+      const state = await this.fetchCheckpointState(session);
+      if (state == null) {
         return null;
       }
 
@@ -166,12 +172,7 @@ export abstract class MongoSyncBucketStorage
       if (snapshotTime == null) {
         throw new ServiceAssertionError('Missing snapshotTime in getCheckpoint()');
       }
-      return new MongoReplicationCheckpoint(
-        this,
-        doc.last_checkpoint ?? 0n,
-        doc.last_checkpoint_lsn ?? null,
-        snapshotTime
-      );
+      return new MongoReplicationCheckpoint(this, state.checkpoint, state.lsn, snapshotTime);
     });
   }
 
@@ -188,31 +189,28 @@ export abstract class MongoSyncBucketStorage
   }
 
   protected abstract createWriterImpl(batchOptions: MongoBucketBatchOptions): storage.BucketStorageBatch;
+  protected abstract getWriterSyncState(): Promise<WriterSyncState>;
 
   async createWriter(options: storage.CreateWriterOptions): Promise<storage.BucketStorageBatch> {
     await this.initializeStorage();
 
-    const doc = await this.db.sync_rules.findOne(
-      {
-        _id: this.group_id
-      },
-      { projection: { last_checkpoint_lsn: 1, no_checkpoint_before: 1, keepalive_op: 1, snapshot_lsn: 1 } }
-    );
-    const checkpoint_lsn = doc?.last_checkpoint_lsn ?? null;
+    const state = await this.getWriterSyncState();
 
     const batchOptions: MongoBucketBatchOptions = {
       logger: options.logger ?? this.logger,
       db: this.db,
-      syncRules: this.sync_rules.parsed(options).hydratedSyncRules(),
+      syncRules: this.sync_rules.parsed(options).hydratedSyncConfig(),
       mapping: this.sync_rules.mapping,
       groupId: this.group_id,
       slotName: this.slot_name,
-      lastCheckpointLsn: checkpoint_lsn,
-      resumeFromLsn: maxLsn(checkpoint_lsn, doc?.snapshot_lsn),
-      keepaliveOp: doc?.keepalive_op ? BigInt(doc.keepalive_op) : null,
+      lastCheckpointLsn: state.lastCheckpointLsn,
+      resumeFromLsn: state.resumeFromLsn,
+      keepaliveOp: state.keepaliveOp,
       storeCurrentData: options.storeCurrentData,
       skipExistingRows: options.skipExistingRows ?? false,
       markRecordUnavailable: options.markRecordUnavailable,
+      hooks: options.hooks,
+      syncConfigId: state.syncConfigId,
       tracer: options.tracer
     };
     const writer = this.createWriterImpl(batchOptions);
@@ -230,134 +228,6 @@ export abstract class MongoSyncBucketStorage
     return writer.last_flushed_op != null ? { flushed_op: writer.last_flushed_op } : null;
   }
 
-  protected abstract sourceTableBaseId(): Partial<CommonSourceTableDocument>;
-
-  protected abstract augmentCreatedSourceTableDocument(
-    createDoc: CommonSourceTableDocument,
-    options: storage.ResolveTableOptions,
-    candidateSourceTable: storage.SourceTable
-  ): void;
-
-  protected abstract initializeResolvedSourceRecords(sourceTableId: bson.ObjectId): Promise<void>;
-
-  async resolveTable(options: storage.ResolveTableOptions): Promise<storage.ResolveTableResult> {
-    const { group_id, connection_id, connection_tag, entity_descriptor } = options;
-
-    const { schema, name, objectId, replicaIdColumns } = entity_descriptor;
-
-    const normalizedReplicaIdColumns = replicaIdColumns.map((column) => ({
-      name: column.name,
-      type: column.type,
-      type_oid: column.typeId
-    }));
-    let result: storage.ResolveTableResult | null = null;
-    let initializeSourceRecordsFor: bson.ObjectId | null = null;
-
-    const baseId = this.sourceTableBaseId();
-    await this.db.client.withSession(async (session) => {
-      const col = this.db.commonSourceTables(group_id);
-      let filter: Partial<CommonSourceTableDocument> = {
-        ...baseId,
-        connection_id: connection_id,
-        schema_name: schema,
-        table_name: name,
-        replica_id_columns2: normalizedReplicaIdColumns
-      };
-
-      if (objectId != null) {
-        filter.relation_id = objectId;
-      }
-      let doc = await col.findOne(filter, { session });
-      if (doc == null) {
-        const candidateSourceTable = new storage.SourceTable({
-          id: new bson.ObjectId(),
-          connectionTag: connection_tag,
-          objectId: objectId,
-          schema: schema,
-          name: name,
-          replicaIdColumns: replicaIdColumns,
-          snapshotComplete: false
-        });
-        const createDoc: CommonSourceTableDocument = {
-          _id: candidateSourceTable.id as bson.ObjectId,
-          ...(baseId as any),
-          connection_id: connection_id,
-          relation_id: objectId,
-          schema_name: schema,
-          table_name: name,
-          replica_id_columns: null,
-          replica_id_columns2: normalizedReplicaIdColumns,
-          snapshot_done: false,
-          snapshot_status: undefined
-        };
-        this.augmentCreatedSourceTableDocument(createDoc, options, candidateSourceTable);
-        doc = createDoc;
-
-        await col.insertOne(doc, { session });
-        initializeSourceRecordsFor = doc._id;
-      }
-      const sourceTable = new storage.SourceTable({
-        id: doc._id,
-        connectionTag: connection_tag,
-        objectId: objectId,
-        schema: schema,
-        name: name,
-        replicaIdColumns: replicaIdColumns,
-        snapshotComplete: doc.snapshot_done ?? true
-      });
-      sourceTable.syncEvent = options.sync_rules.tableTriggersEvent(sourceTable);
-      sourceTable.syncData = options.sync_rules.tableSyncsData(sourceTable);
-      sourceTable.syncParameters = options.sync_rules.tableSyncsParameters(sourceTable);
-      sourceTable.snapshotStatus =
-        doc.snapshot_status == null
-          ? undefined
-          : {
-              lastKey: doc.snapshot_status.last_key?.buffer ?? null,
-              totalEstimatedCount: doc.snapshot_status.total_estimated_count,
-              replicatedCount: doc.snapshot_status.replicated_count
-            };
-
-      let dropTables: storage.SourceTable[] = [];
-      let truncateFilter = [{ schema_name: schema, table_name: name }] as any[];
-      if (objectId != null) {
-        truncateFilter.push({ relation_id: objectId });
-      }
-      const truncate = await col
-        .find(
-          {
-            ...baseId,
-            connection_id: connection_id,
-            _id: { $ne: doc._id },
-            $or: truncateFilter
-          },
-          { session }
-        )
-        .toArray();
-      dropTables = truncate.map(
-        (doc) =>
-          new storage.SourceTable({
-            id: doc._id,
-            connectionTag: connection_tag,
-            objectId: doc.relation_id,
-            schema: doc.schema_name,
-            name: doc.table_name,
-            replicaIdColumns:
-              doc.replica_id_columns2?.map((c) => ({ name: c.name, typeOid: c.type_oid, type: c.type })) ?? [],
-            snapshotComplete: doc.snapshot_done ?? true
-          })
-      );
-
-      result = {
-        table: sourceTable,
-        dropTables: dropTables
-      };
-    });
-    if (initializeSourceRecordsFor != null) {
-      await this.initializeResolvedSourceRecords(initializeSourceRecordsFor);
-    }
-    return result!;
-  }
-
   protected abstract getParameterSetsImpl(
     checkpoint: MongoReplicationCheckpoint,
     lookups: ScopedParameterLookup[],
@@ -397,49 +267,20 @@ export abstract class MongoSyncBucketStorage
     this.checksums.clearCache();
   }
 
+  protected abstract terminateSyncRuleState(): Promise<void>;
+
   async terminate(options?: storage.TerminateOptions) {
     if (!options || options?.clearStorage) {
       await this.clear(options);
     }
-    await this.db.sync_rules.updateOne(
-      {
-        _id: this.group_id
-      },
-      {
-        $set: {
-          state: storage.SyncRuleState.TERMINATED,
-          persisted_lsn: null,
-          snapshot_done: false
-        }
-      }
-    );
+    await this.terminateSyncRuleState();
     await this.db.notifyCheckpoint();
   }
 
-  async getStatus(): Promise<storage.SyncRuleStatus> {
-    const doc = await this.db.sync_rules.findOne(
-      {
-        _id: this.group_id
-      },
-      {
-        projection: {
-          snapshot_done: 1,
-          last_checkpoint_lsn: 1,
-          state: 1,
-          snapshot_lsn: 1
-        }
-      }
-    );
-    if (doc == null) {
-      throw new ServiceAssertionError('Cannot find replication stream status');
-    }
+  protected abstract getStatusImpl(): Promise<storage.SyncRuleStatus>;
 
-    return {
-      snapshot_done: doc.snapshot_done,
-      snapshot_lsn: doc.snapshot_lsn ?? null,
-      active: doc.state == 'ACTIVE',
-      checkpoint_lsn: doc.last_checkpoint_lsn
-    };
+  async getStatus(): Promise<storage.SyncRuleStatus> {
+    return this.getStatusImpl();
   }
 
   protected abstract clearBucketData(signal?: AbortSignal): Promise<void>;
@@ -451,6 +292,7 @@ export abstract class MongoSyncBucketStorage
   protected abstract clearBucketState(signal?: AbortSignal): Promise<void>;
 
   protected abstract clearSourceTables(signal?: AbortSignal): Promise<void>;
+  protected abstract clearSyncRuleState(): Promise<void>;
 
   async clear(options?: storage.ClearStorageOptions): Promise<void> {
     const signal = options?.signal;
@@ -459,24 +301,7 @@ export abstract class MongoSyncBucketStorage
       throw new ReplicationAbortedError('Aborted clearing data', signal.reason);
     }
 
-    await this.db.sync_rules.updateOne(
-      {
-        _id: this.group_id
-      },
-      {
-        $set: {
-          snapshot_done: false,
-          persisted_lsn: null,
-          last_checkpoint_lsn: null,
-          last_checkpoint: null,
-          no_checkpoint_before: null
-        },
-        $unset: {
-          snapshot_lsn: 1
-        }
-      },
-      { maxTimeMS: lib_mongo.db.MONGO_CLEAR_OPERATION_TIMEOUT_MS }
-    );
+    await this.clearSyncRuleState();
 
     await this.clearBucketData(signal);
     await this.clearParameterIndexes(signal);

package/src/storage/implementation/MongoSyncRulesLock.ts

@@ -1,5 +1,6 @@
 import crypto from 'crypto';
 
+import { mongo } from '@powersync/lib-service-mongodb';
 import { ErrorCode, Logger, ServiceError } from '@powersync/lib-services-framework';
 import { storage } from '@powersync/service-core';
 import { VersionedPowerSyncMongo } from './db.js';
@@ -11,9 +12,13 @@ import { VersionedPowerSyncMongo } from './db.js';
 export class MongoSyncRulesLock implements storage.ReplicationLock {
   private readonly refreshInterval: NodeJS.Timeout;
 
+  /**
+   * @param session optional session to create the lock within another transaction
+   */
   static async createLock(
     db: VersionedPowerSyncMongo,
-    sync_rules: storage.PersistedSyncRulesContent
+    sync_rules: storage.PersistedSyncRulesContent,
+    session?: mongo.ClientSession
   ): Promise<MongoSyncRulesLock> {
     const lockId = crypto.randomBytes(8).toString('hex');
     const doc = await db.sync_rules.findOneAndUpdate(
@@ -28,13 +33,14 @@ export class MongoSyncRulesLock implements storage.ReplicationLock {
       },
       {
         projection: { lock: 1 },
-        returnDocument: 'before'
+        returnDocument: 'before',
+        session
       }
     );
 
     if (doc == null) {
       // Query the existing lock to get the expiration time (best effort - it may have been released in the meantime).
-      const heldLock = await db.sync_rules.findOne({ _id: sync_rules.id }, { projection: { lock: 1 } });
+      const heldLock = await db.sync_rules.findOne({ _id: sync_rules.id }, { projection: { lock: 1 }, session });
       if (heldLock?.lock?.expires_at) {
         throw new ServiceError(
           ErrorCode.PSYNC_S1003,

package/src/storage/implementation/SyncRuleStateUpdate.ts

@@ -0,0 +1,38 @@
+import { mongo } from '@powersync/lib-service-mongodb';
+import { storage } from '@powersync/service-core';
+
+/**
+ * Update pipeline to update replication stream status, covering all storage versions.
+ *
+ * Roughly equivalent to:
+ *   $set: {
+ *     state: state,
+ *     'sync_configs.$[].state': state
+ *   }
+ *
+ * The difference is that this also handles v1 storage cases, where `sync_configs` is not present.
+ */
+export function syncRuleStateUpdatePipeline(state: storage.SyncRuleState): mongo.Document[] {
+  return [
+    {
+      $set: {
+        state,
+        sync_configs: {
+          $cond: [
+            { $isArray: '$sync_configs' },
+            {
+              $map: {
+                input: '$sync_configs',
+                as: 'config',
+                in: {
+                  $mergeObjects: ['$$config', { state }]
+                }
+              }
+            },
+            '$$REMOVE'
+          ]
+        }
+      }
+    }
+  ];
+}

package/src/storage/implementation/common/BucketDataDoc.ts

@@ -1,5 +1,5 @@
 import { InternalOpId } from '@powersync/service-core';
-import { BucketDefinitionId } from '../BucketDefinitionMapping.js';
+import { BucketDefinitionId } from '@powersync/service-sync-rules';
 import { BucketDataProperties } from '../models.js';
 
 /**

package/src/storage/implementation/common/PersistedBatch.ts

@@ -1,12 +1,12 @@
 import { mongo } from '@powersync/lib-service-mongodb';
-import { BucketDataSource, EvaluatedParameters, EvaluatedRow } from '@powersync/service-sync-rules';
+import { BucketDataSource, BucketDefinitionId, EvaluatedParameters, EvaluatedRow } from '@powersync/service-sync-rules';
 import * as bson from 'bson';
 
 import { logger as defaultLogger, Logger } from '@powersync/lib-services-framework';
 import { InternalOpId, storage, utils } from '@powersync/service-core';
 import { JSONBig } from '@powersync/service-jsonbig';
 import { mongoTableId, replicaIdToSubkey } from '../../../utils/util.js';
-import { BucketDefinitionId, BucketDefinitionMapping } from '../BucketDefinitionMapping.js';
+import { BucketDefinitionMapping } from '../BucketDefinitionMapping.js';
 import { currentBucketKey, MAX_ROW_SIZE } from '../MongoBucketBatchShared.js';
 import { MongoIdSequence } from '../MongoIdSequence.js';
 import type { VersionedPowerSyncMongo } from '../db.js';

package/src/storage/implementation/common/SourceRecordStore.ts

@@ -1,9 +1,8 @@
 import { mongo } from '@powersync/lib-service-mongodb';
 import { Logger } from '@powersync/lib-services-framework';
 import { storage } from '@powersync/service-core';
-import { EvaluatedParameters, EvaluatedRow } from '@powersync/service-sync-rules';
+import { BucketDefinitionId, EvaluatedParameters, EvaluatedRow, ParameterIndexId } from '@powersync/service-sync-rules';
 import * as bson from 'bson';
-import { BucketDefinitionId, ParameterIndexId } from '../BucketDefinitionMapping.js';
 
 export interface SourceRecordLookupEntry {
   sourceTableId: bson.ObjectId;