[Incremental Reprocessing] [WIP] Proof-of-concept incremental reprocessing for MongoDB

modules/module-mongodb-storage/src/storage/MongoBucketStorage.ts

@@ -1,6 +1,6 @@
-import { SqlSyncRules } from '@powersync/service-sync-rules';
+import { BucketDataSource, ParameterIndexLookupCreator, SqlSyncRules } from '@powersync/service-sync-rules';
 
-import { GetIntanceOptions, storage } from '@powersync/service-core';
+import { GetIntanceOptions, maxLsn, CreateWriterOptions, storage } from '@powersync/service-core';
 
 import { BaseObserver, ErrorCode, logger, ServiceError } from '@powersync/lib-services-framework';
 import { v4 as uuid } from 'uuid';
@@ -13,6 +13,9 @@ import { SyncRuleDocument } from './implementation/models.js';
 import { MongoPersistedSyncRulesContent } from './implementation/MongoPersistedSyncRulesContent.js';
 import { MongoSyncBucketStorage, MongoSyncBucketStorageOptions } from './implementation/MongoSyncBucketStorage.js';
 import { generateSlotName } from '../utils/util.js';
+import { BucketDefinitionMapping } from './implementation/BucketDefinitionMapping.js';
+import { MongoBucketDataWriter } from './storage-index.js';
+import { MergedSyncRules } from './implementation/MergedSyncRules.js';
 
 export class MongoBucketStorage
   extends BaseObserver<storage.BucketStorageFactoryListener>
@@ -50,7 +53,14 @@ export class MongoBucketStorage
     if ((typeof id as any) == 'bigint') {
       id = Number(id);
     }
-    const storage = new MongoSyncBucketStorage(this, id, syncRules, slot_name, undefined, this.internalOptions);
+    const storage = new MongoSyncBucketStorage(
+      this,
+      id,
+      syncRules as MongoPersistedSyncRulesContent,
+      slot_name,
+      undefined,
+      this.internalOptions
+    );
     if (!options?.skipLifecycleHooks) {
       this.iterateListeners((cb) => cb.syncStorageCreated?.(storage));
     }
@@ -64,6 +74,48 @@ export class MongoBucketStorage
     return storage;
   }
 
+  async createCombinedWriter(
+    storages: storage.SyncRulesBucketStorage[],
+    options: CreateWriterOptions
+  ): Promise<MongoBucketDataWriter> {
+    const mongoStorages = storages as MongoSyncBucketStorage[];
+    const mappings = mongoStorages.map((s) => s.sync_rules.mapping);
+    const mergedMappings = BucketDefinitionMapping.merged(mappings);
+    const mergedProcessor = MergedSyncRules.merge(mongoStorages.map((s) => s.getParsedSyncRules(options)));
+
+    const writer = new MongoBucketDataWriter({
+      db: this.db,
+      mapping: mergedMappings,
+      markRecordUnavailable: options.markRecordUnavailable,
+      rowProcessor: mergedProcessor,
+      skipExistingRows: options.skipExistingRows ?? false,
+      slotName: '',
+      storeCurrentData: options.storeCurrentData,
+      logger: options.logger
+    });
+
+    for (let storage of mongoStorages) {
+      const doc = await this.db.sync_rules.findOne(
+        {
+          _id: storage.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 parsedSyncRules = storage.getParsedSyncRules(options);
+      const batch = writer.forSyncRules({
+        syncRules: parsedSyncRules,
+
+        lastCheckpointLsn: checkpoint_lsn,
+        resumeFromLsn: maxLsn(checkpoint_lsn, doc?.snapshot_lsn),
+        keepaliveOp: doc?.keepalive_op ? BigInt(doc.keepalive_op) : null
+      });
+      storage.iterateListeners((cb) => cb.batchStarted?.(batch));
+    }
+
+    return writer;
+  }
+
   async getSystemIdentifier(): Promise<storage.BucketStorageSystemIdentifier> {
     const { setName: id } = await this.db.db.command({
       hello: 1
@@ -184,7 +236,17 @@ export class MongoBucketStorage
         {
           state: storage.SyncRuleState.PROCESSING
         },
-        { $set: { state: storage.SyncRuleState.STOP } }
+        { $set: { state: storage.SyncRuleState.STOP } },
+        {
+          session: this.session
+        }
+      );
+
+      const activeSyncRules = await this.db.sync_rules.findOne(
+        {
+          state: storage.SyncRuleState.ACTIVE
+        },
+        { session: this.session }
       );
 
       const id_doc = await this.db.op_id_sequence.findOneAndUpdate(
@@ -198,34 +260,91 @@ export class MongoBucketStorage
         },
         {
           upsert: true,
-          returnDocument: 'after'
+          returnDocument: 'after',
+          session: this.session
         }
       );
 
       const id = Number(id_doc!.op_id);
       const slot_name = generateSlotName(this.slot_name_prefix, id);
 
+      const syncRules = SqlSyncRules.fromYaml(options.content, {
+        // No schema-based validation at this point
+        schema: undefined,
+        defaultSchema: 'not_applicable', // Not needed for validation
+        throwOnError: false
+      });
+      let bucketDefinitionMapping: Record<string, number> = {};
+      let parameterDefinitionMapping: Record<string, number> = {};
+      let bucketDefinitionId = (id << 16) + 1;
+      let parameterDefinitionId = (id << 17) + 1;
+
+      let existingMapping: BucketDefinitionMapping;
+      if (activeSyncRules != null) {
+        existingMapping = BucketDefinitionMapping.fromSyncRules(activeSyncRules);
+      } else {
+        existingMapping = new BucketDefinitionMapping({}, {});
+      }
+
+      syncRules.config.hydrate({
+        hydrationState: {
+          getBucketSourceScope(source: BucketDataSource) {
+            const existingId = existingMapping.equivalentBucketSourceId(source);
+            if (existingId != null) {
+              bucketDefinitionMapping[source.uniqueName] = existingId;
+            } else {
+              bucketDefinitionMapping[source.uniqueName] = bucketDefinitionId;
+              bucketDefinitionId += 1;
+            }
+            return {
+              // N/A
+              bucketPrefix: '',
+              source
+            };
+          },
+          getParameterIndexLookupScope(source: ParameterIndexLookupCreator) {
+            const key = `${source.defaultLookupScope.lookupName}#${source.defaultLookupScope.queryId}`;
+            const existingId = existingMapping.equivalentParameterLookupId(source);
+            if (existingId != null) {
+              parameterDefinitionMapping[key] = existingId;
+            } else {
+              parameterDefinitionMapping[key] = parameterDefinitionId;
+              parameterDefinitionId += 1;
+            }
+            // N/A
+            return source.defaultLookupScope;
+          }
+        }
+      });
+
       const doc: SyncRuleDocument = {
         _id: id,
         content: options.content,
-        last_checkpoint: null,
+        last_checkpoint: activeSyncRules?.last_checkpoint ?? null,
         last_checkpoint_lsn: null,
         no_checkpoint_before: null,
-        keepalive_op: null,
+        // HACK: copy the op from the active sync rules, if any.
+        // This specifically helps for the case of the new sync rules not replicating anything new.
+        // FIXME: Make sure this is properly sound and tested.
+        keepalive_op: activeSyncRules?.last_checkpoint ? String(activeSyncRules.last_checkpoint) : null,
         snapshot_done: false,
         snapshot_lsn: undefined,
         state: storage.SyncRuleState.PROCESSING,
         slot_name: slot_name,
         last_checkpoint_ts: null,
         last_fatal_error: null,
         last_fatal_error_ts: null,
-        last_keepalive_ts: null
+        last_keepalive_ts: null,
+        rule_mapping: {
+          definitions: bucketDefinitionMapping,
+          parameter_lookups: parameterDefinitionMapping
+        }
       };
-      await this.db.sync_rules.insertOne(doc);
+      await this.db.sync_rules.insertOne(doc, { session: this.session });
       await this.db.notifyCheckpoint();
       rules = new MongoPersistedSyncRulesContent(this.db, doc);
       if (options.lock) {
-        const lock = await rules.lock();
+        await rules.lock(this.session);
       }
     });
 
@@ -275,6 +394,8 @@ export class MongoBucketStorage
       .find({
         state: { $in: [storage.SyncRuleState.PROCESSING, storage.SyncRuleState.ACTIVE] }
       })
+      // Prioritize "ACTIVE" first
+      .sort({ state: 1, _id: 1 })
       .toArray();
 
     return docs.map((doc) => {

modules/module-mongodb-storage/src/storage/implementation/BucketDefinitionMapping.ts

@@ -0,0 +1,61 @@
+import { ServiceAssertionError } from '@powersync/lib-services-framework';
+import { BucketDataSource, ParameterIndexLookupCreator, SqlSyncRules } from '@powersync/service-sync-rules';
+import { SyncRuleDocument } from './models.js';
+
+export class BucketDefinitionMapping {
+  static fromSyncRules(doc: Pick<SyncRuleDocument, 'rule_mapping'>): BucketDefinitionMapping {
+    return new BucketDefinitionMapping(doc.rule_mapping.definitions, doc.rule_mapping.parameter_lookups);
+  }
+
+  static merged(mappings: BucketDefinitionMapping[]): BucketDefinitionMapping {
+    return mappings.reduce((acc, curr) => acc.mergeWith(curr), new BucketDefinitionMapping());
+  }
+
+  constructor(
+    private definitions: Record<string, number> = {},
+    private parameterLookupMapping: Record<string, number> = {}
+  ) {}
+
+  hasBucketSourceId(id: number) {
+    return Object.values(this.definitions).includes(id);
+  }
+
+  hasParameterLookupId(id: number) {
+    return Object.values(this.parameterLookupMapping).includes(id);
+  }
+
+  bucketSourceId(source: BucketDataSource): number {
+    const defId = this.definitions[source.uniqueName];
+    if (defId == null) {
+      throw new ServiceAssertionError(`No mapping found for bucket source ${source.uniqueName}`);
+    }
+    return defId;
+  }
+
+  parameterLookupId(source: ParameterIndexLookupCreator): number {
+    const key = `${source.defaultLookupScope.lookupName}#${source.defaultLookupScope.queryId}`;
+    const defId = this.parameterLookupMapping[key];
+    if (defId == null) {
+      throw new ServiceAssertionError(`No mapping found for parameter lookup source ${key}`);
+    }
+    return defId;
+  }
+
+  equivalentBucketSourceId(source: BucketDataSource): number | null {
+    // FIXME: Do an actual comparison, instead of just using the unique name
+    return this.definitions[source.uniqueName] ?? null;
+  }
+
+  equivalentParameterLookupId(source: ParameterIndexLookupCreator): number | null {
+    // FIXME: Do an actual comparison, instead of just using the scope
+    const key = `${source.defaultLookupScope.lookupName}#${source.defaultLookupScope.queryId}`;
+    return this.parameterLookupMapping[key] ?? null;
+  }
+
+  mergeWith(other: BucketDefinitionMapping): BucketDefinitionMapping {
+    return new BucketDefinitionMapping(
+      { ...this.definitions, ...other.definitions },
+      { ...this.parameterLookupMapping, ...other.parameterLookupMapping }
+    );
+  }
+}