@powersync/service-core 0.8.8 → 0.9.0

package/src/storage/mongo/MongoCompactor.ts

@@ -4,6 +4,7 @@ import { addChecksums } from '../../util/utils.js';
 import { PowerSyncMongo } from './db.js';
 import { BucketDataDocument, BucketDataKey } from './models.js';
 import { CompactOptions } from '../BucketStorage.js';
+import { cacheKey } from './OperationBatch.js';
 
 interface CurrentBucketState {
   /** Bucket name */
@@ -57,7 +58,11 @@ export class MongoCompactor {
   private maxOpId: bigint | undefined;
   private buckets: string[] | undefined;
 
-  constructor(private db: PowerSyncMongo, private group_id: number, options?: MongoCompactOptions) {
+  constructor(
+    private db: PowerSyncMongo,
+    private group_id: number,
+    options?: MongoCompactOptions
+  ) {
     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;
@@ -89,17 +94,33 @@ export class MongoCompactor {
 
     let currentState: CurrentBucketState | null = null;
 
+    let bucketLower: string | MinKey;
+    let bucketUpper: string | MaxKey;
+
+    if (bucket == null) {
+      bucketLower = new MinKey();
+      bucketUpper = new MaxKey();
+    } else if (bucket.includes('[')) {
+      // Exact bucket name
+      bucketLower = bucket;
+      bucketUpper = bucket;
+    } else {
+      // Bucket definition name
+      bucketLower = `${bucket}[`;
+      bucketUpper = `${bucket}[\uFFFF`;
+    }
+
     // Constant lower bound
     const lowerBound: BucketDataKey = {
       g: this.group_id,
-      b: bucket ?? (new MinKey() as any),
+      b: bucketLower as string,
       o: new MinKey() as any
     };
 
     // Upper bound is adjusted for each batch
     let upperBound: BucketDataKey = {
       g: this.group_id,
-      b: bucket ?? (new MaxKey() as any),
+      b: bucketUpper as string,
       o: new MaxKey() as any
     };
 
@@ -168,7 +189,7 @@ export class MongoCompactor {
         let isPersistentPut = doc.op == 'PUT';
 
         if (doc.op == 'REMOVE' || doc.op == 'PUT') {
-          const key = `${doc.table}/${doc.row_id}/${doc.source_table}/${doc.source_key?.toHexString()}`;
+          const key = `${doc.table}/${doc.row_id}/${cacheKey(doc.source_table!, doc.source_key!)}`;
           const targetOp = currentState.seen.get(key);
           if (targetOp) {
             // Will convert to MOVE, so don't count as PUT

package/src/storage/mongo/MongoPersistedSyncRulesContent.ts

@@ -1,7 +1,7 @@
 import { SqlSyncRules } from '@powersync/service-sync-rules';
 import * as mongo from 'mongodb';
 
-import { PersistedSyncRulesContent } from '../BucketStorage.js';
+import { ParseSyncRulesOptions, PersistedSyncRulesContent } from '../BucketStorage.js';
 import { MongoPersistedSyncRules } from './MongoPersistedSyncRules.js';
 import { MongoSyncRulesLock } from './MongoSyncRulesLock.js';
 import { PowerSyncMongo } from './db.js';
@@ -19,7 +19,10 @@ export class MongoPersistedSyncRulesContent implements PersistedSyncRulesContent
 
   public current_lock: MongoSyncRulesLock | null = null;
 
-  constructor(private db: PowerSyncMongo, doc: mongo.WithId<SyncRuleDocument>) {
+  constructor(
+    private db: PowerSyncMongo,
+    doc: mongo.WithId<SyncRuleDocument>
+  ) {
     this.id = doc._id;
     this.sync_rules_content = doc.content;
     this.last_checkpoint_lsn = doc.last_checkpoint_lsn;
@@ -30,10 +33,10 @@ export class MongoPersistedSyncRulesContent implements PersistedSyncRulesContent
     this.last_keepalive_ts = doc.last_keepalive_ts;
   }
 
-  parsed() {
+  parsed(options: ParseSyncRulesOptions) {
     return new MongoPersistedSyncRules(
       this.id,
-      SqlSyncRules.fromYaml(this.sync_rules_content),
+      SqlSyncRules.fromYaml(this.sync_rules_content, options),
       this.last_checkpoint_lsn,
       this.slot_name
     );

package/src/storage/mongo/MongoStorageProvider.ts

@@ -0,0 +1,31 @@
+import { logger } from '@powersync/lib-services-framework';
+import * as db from '../../db/db-index.js';
+import { MongoBucketStorage } from '../MongoBucketStorage.js';
+import { ActiveStorage, BucketStorageProvider, GetStorageOptions } from '../StorageProvider.js';
+import { PowerSyncMongo } from './db.js';
+
+export class MongoStorageProvider implements BucketStorageProvider {
+  get type() {
+    return 'mongodb';
+  }
+
+  async getStorage(options: GetStorageOptions): Promise<ActiveStorage> {
+    const { resolvedConfig } = options;
+
+    const client = db.mongo.createMongoClient(resolvedConfig.storage);
+
+    const database = new PowerSyncMongo(client, { database: resolvedConfig.storage.database });
+
+    return {
+      storage: new MongoBucketStorage(database, {
+        // TODO currently need the entire resolved config due to this
+        slot_name_prefix: resolvedConfig.slot_name_prefix
+      }),
+      shutDown: () => client.close(),
+      tearDown: () => {
+        logger.info(`Tearing down storage: ${database.db.namespace}...`);
+        return database.db.dropDatabase();
+      }
+    } satisfies ActiveStorage;
+  }
+}

package/src/storage/mongo/MongoSyncBucketStorage.ts

@@ -2,8 +2,9 @@ import { SqliteJsonRow, SqliteJsonValue, SqlSyncRules } from '@powersync/service
 import * as bson from 'bson';
 import * as mongo from 'mongodb';
 
+import { DisposableObserver, logger } from '@powersync/lib-services-framework';
+import * as timers from 'timers/promises';
 import * as db from '../../db/db-index.js';
-import * as replication from '../../replication/WalStream.js';
 import * as util from '../../util/util-index.js';
 import {
   BucketDataBatchOptions,
@@ -12,24 +13,39 @@ import {
   DEFAULT_DOCUMENT_BATCH_LIMIT,
   DEFAULT_DOCUMENT_CHUNK_LIMIT_BYTES,
   FlushedResult,
+  ParseSyncRulesOptions,
+  PersistedSyncRulesContent,
+  ReplicationCheckpoint,
   ResolveTableOptions,
   ResolveTableResult,
+  StartBatchOptions,
   SyncBucketDataBatch,
   SyncRulesBucketStorage,
-  SyncRuleStatus
+  SyncRulesBucketStorageListener,
+  SyncRuleStatus,
+  TerminateOptions
 } from '../BucketStorage.js';
 import { ChecksumCache, FetchPartialBucketChecksum, PartialChecksum, PartialChecksumMap } from '../ChecksumCache.js';
 import { MongoBucketStorage } from '../MongoBucketStorage.js';
 import { SourceTable } from '../SourceTable.js';
+import {
+  BatchedCustomWriteCheckpointOptions,
+  ManagedWriteCheckpointOptions,
+  SyncStorageLastWriteCheckpointFilters,
+  WriteCheckpointAPI,
+  WriteCheckpointMode
+} from '../WriteCheckpointAPI.js';
 import { PowerSyncMongo } from './db.js';
 import { BucketDataDocument, BucketDataKey, SourceKey, SyncRuleState } from './models.js';
 import { MongoBucketBatch } from './MongoBucketBatch.js';
 import { MongoCompactor } from './MongoCompactor.js';
+import { MongoWriteCheckpointAPI } from './MongoWriteCheckpointAPI.js';
 import { BSON_DESERIALIZE_OPTIONS, idPrefixFilter, mapOpEntry, readSingleBatch, serializeLookup } from './util.js';
-import { logger } from '@powersync/lib-services-framework';
-import * as timers from 'timers/promises';
 
-export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
+export class MongoSyncBucketStorage
+  extends DisposableObserver<SyncRulesBucketStorageListener>
+  implements SyncRulesBucketStorage
+{
   private readonly db: PowerSyncMongo;
   private checksumCache = new ChecksumCache({
     fetchChecksums: (batch) => {
@@ -37,16 +53,70 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
     }
   });
 
+  private parsedSyncRulesCache: { parsed: SqlSyncRules; options: ParseSyncRulesOptions } | undefined;
+  private writeCheckpointAPI: WriteCheckpointAPI;
+
   constructor(
     public readonly factory: MongoBucketStorage,
     public readonly group_id: number,
-    public readonly sync_rules: SqlSyncRules,
-    public readonly slot_name: string
+    private readonly sync_rules: PersistedSyncRulesContent,
+    public readonly slot_name: string,
+    writeCheckpointMode: WriteCheckpointMode = WriteCheckpointMode.MANAGED
   ) {
+    super();
     this.db = factory.db;
+    this.writeCheckpointAPI = new MongoWriteCheckpointAPI({
+      db: this.db,
+      mode: writeCheckpointMode
+    });
+  }
+
+  get writeCheckpointMode() {
+    return this.writeCheckpointAPI.writeCheckpointMode;
+  }
+
+  setWriteCheckpointMode(mode: WriteCheckpointMode): void {
+    this.writeCheckpointAPI.setWriteCheckpointMode(mode);
+  }
+
+  batchCreateCustomWriteCheckpoints(checkpoints: BatchedCustomWriteCheckpointOptions[]): Promise<void> {
+    return this.writeCheckpointAPI.batchCreateCustomWriteCheckpoints(
+      checkpoints.map((checkpoint) => ({ ...checkpoint, sync_rules_id: this.group_id }))
+    );
+  }
+
+  createCustomWriteCheckpoint(checkpoint: BatchedCustomWriteCheckpointOptions): Promise<bigint> {
+    return this.writeCheckpointAPI.createCustomWriteCheckpoint({
+      ...checkpoint,
+      sync_rules_id: this.group_id
+    });
   }
 
-  async getCheckpoint() {
+  createManagedWriteCheckpoint(checkpoint: ManagedWriteCheckpointOptions): Promise<bigint> {
+    return this.writeCheckpointAPI.createManagedWriteCheckpoint(checkpoint);
+  }
+
+  lastWriteCheckpoint(filters: SyncStorageLastWriteCheckpointFilters): Promise<bigint | null> {
+    return this.writeCheckpointAPI.lastWriteCheckpoint({
+      ...filters,
+      sync_rules_id: this.group_id
+    });
+  }
+
+  getParsedSyncRules(options: ParseSyncRulesOptions): SqlSyncRules {
+    const { parsed, options: cachedOptions } = this.parsedSyncRulesCache ?? {};
+    /**
+     * Check if the cached sync rules, if present, had the same options.
+     * Parse sync rules if the options are different or if there is no cached value.
+     */
+    if (!parsed || options.defaultSchema != cachedOptions?.defaultSchema) {
+      this.parsedSyncRulesCache = { parsed: this.sync_rules.parsed(options).sync_rules, options };
+    }
+
+    return this.parsedSyncRulesCache!.parsed;
+  }
+
+  async getCheckpoint(): Promise<ReplicationCheckpoint> {
     const doc = await this.db.sync_rules.findOne(
       { _id: this.group_id },
       {
@@ -55,11 +125,14 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
     );
     return {
       checkpoint: util.timestampToOpId(doc?.last_checkpoint ?? 0n),
-      lsn: doc?.last_checkpoint_lsn ?? replication.ZERO_LSN
+      lsn: doc?.last_checkpoint_lsn ?? null
     };
   }
 
-  async startBatch(options: {}, callback: (batch: BucketStorageBatch) => Promise<void>): Promise<FlushedResult | null> {
+  async startBatch(
+    options: StartBatchOptions,
+    callback: (batch: BucketStorageBatch) => Promise<void>
+  ): Promise<FlushedResult | null> {
     const doc = await this.db.sync_rules.findOne(
       {
         _id: this.group_id
@@ -68,35 +141,36 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
     );
     const checkpoint_lsn = doc?.last_checkpoint_lsn ?? null;
 
-    const batch = new MongoBucketBatch(
-      this.db,
-      this.sync_rules,
-      this.group_id,
-      this.slot_name,
-      checkpoint_lsn,
-      doc?.no_checkpoint_before ?? null
-    );
-    try {
-      await callback(batch);
-      await batch.flush();
-      await batch.abort();
-      if (batch.last_flushed_op) {
-        return { flushed_op: String(batch.last_flushed_op) };
-      } else {
-        return null;
-      }
-    } catch (e) {
-      await batch.abort();
-      throw e;
+    await using batch = new MongoBucketBatch({
+      db: this.db,
+      syncRules: this.sync_rules.parsed(options).sync_rules,
+      groupId: this.group_id,
+      slotName: this.slot_name,
+      lastCheckpointLsn: checkpoint_lsn,
+      noCheckpointBeforeLsn: doc?.no_checkpoint_before ?? options.zeroLSN,
+      storeCurrentData: options.storeCurrentData
+    });
+    this.iterateListeners((cb) => cb.batchStarted?.(batch));
+
+    await callback(batch);
+    await batch.flush();
+    if (batch.last_flushed_op) {
+      return { flushed_op: String(batch.last_flushed_op) };
+    } else {
+      return null;
     }
   }
 
   async resolveTable(options: ResolveTableOptions): Promise<ResolveTableResult> {
-    const { group_id, connection_id, connection_tag, relation } = options;
+    const { group_id, connection_id, connection_tag, entity_descriptor } = options;
 
-    const { schema, name: table, relationId, replicationColumns } = relation;
+    const { schema, name: table, objectId, replicationColumns } = entity_descriptor;
 
-    const columns = replicationColumns.map((column) => ({ name: column.name, type_oid: column.typeOid }));
+    const columns = replicationColumns.map((column) => ({
+      name: column.name,
+      type: column.type,
+      type_oid: column.typeId
+    }));
     let result: ResolveTableResult | null = null;
     await this.db.client.withSession(async (session) => {
       const col = this.db.source_tables;
@@ -104,7 +178,7 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
         {
           group_id: group_id,
           connection_id: connection_id,
-          relation_id: relationId,
+          relation_id: objectId,
           schema_name: schema,
           table_name: table,
           replica_id_columns2: columns
@@ -116,7 +190,7 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
           _id: new bson.ObjectId(),
           group_id: group_id,
           connection_id: connection_id,
-          relation_id: relationId,
+          relation_id: objectId,
           schema_name: schema,
           table_name: table,
           replica_id_columns: null,
@@ -129,12 +203,13 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
       const sourceTable = new SourceTable(
         doc._id,
         connection_tag,
-        relationId,
+        objectId,
         schema,
         table,
         replicationColumns,
         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);
 
@@ -144,7 +219,7 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
             group_id: group_id,
             connection_id: connection_id,
             _id: { $ne: doc._id },
-            $or: [{ relation_id: relationId }, { schema_name: schema, table_name: table }]
+            $or: [{ relation_id: objectId }, { schema_name: schema, table_name: table }]
           },
           { session }
         )
@@ -159,7 +234,7 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
               doc.relation_id ?? 0,
               doc.schema_name,
               doc.table_name,
-              doc.replica_id_columns2?.map((c) => ({ name: c.name, typeOid: c.type_oid })) ?? [],
+              doc.replica_id_columns2?.map((c) => ({ name: c.name, typeOid: c.type_oid, type: c.type })) ?? [],
               doc.snapshot_done ?? true
             )
         )
@@ -398,9 +473,11 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
     );
   }
 
-  async terminate() {
-    await this.clear();
-
+  async terminate(options?: TerminateOptions) {
+    // Default is to clear the storage except when explicitly requested not to.
+    if (!options || options?.clearStorage) {
+      await this.clear();
+    }
     await this.db.sync_rules.updateOne(
       {
         _id: this.group_id

package/src/storage/mongo/MongoSyncRulesLock.ts

@@ -9,7 +9,7 @@ import { logger } from '@powersync/lib-services-framework';
  * replicates those sync rules at a time.
  */
 export class MongoSyncRulesLock implements ReplicationLock {
-  private readonly refreshInterval: NodeJS.Timer;
+  private readonly refreshInterval: NodeJS.Timeout;
 
   static async createLock(db: PowerSyncMongo, sync_rules: PersistedSyncRulesContent): Promise<MongoSyncRulesLock> {
     const lockId = crypto.randomBytes(8).toString('hex');
@@ -30,12 +30,16 @@ export class MongoSyncRulesLock implements ReplicationLock {
     );
 
     if (doc == null) {
-      throw new Error(`Replication slot ${sync_rules.slot_name} is locked by another process`);
+      throw new Error(`Sync rules: ${sync_rules.id} have been locked by another process for replication.`);
     }
     return new MongoSyncRulesLock(db, sync_rules.id, lockId);
   }
 
-  constructor(private db: PowerSyncMongo, public sync_rules_id: number, private lock_id: string) {
+  constructor(
+    private db: PowerSyncMongo,
+    public sync_rules_id: number,
+    private lock_id: string
+  ) {
     this.refreshInterval = setInterval(async () => {
       try {
         await this.refresh();

package/src/storage/mongo/MongoWriteCheckpointAPI.ts

@@ -0,0 +1,151 @@
+import * as framework from '@powersync/lib-services-framework';
+import {
+  CustomWriteCheckpointFilters,
+  CustomWriteCheckpointOptions,
+  LastWriteCheckpointFilters,
+  ManagedWriteCheckpointFilters,
+  ManagedWriteCheckpointOptions,
+  WriteCheckpointAPI,
+  WriteCheckpointMode
+} from '../WriteCheckpointAPI.js';
+import { PowerSyncMongo } from './db.js';
+
+export type MongoCheckpointAPIOptions = {
+  db: PowerSyncMongo;
+  mode: WriteCheckpointMode;
+};
+
+export class MongoWriteCheckpointAPI implements WriteCheckpointAPI {
+  readonly db: PowerSyncMongo;
+  private _mode: WriteCheckpointMode;
+
+  constructor(options: MongoCheckpointAPIOptions) {
+    this.db = options.db;
+    this._mode = options.mode;
+  }
+
+  get writeCheckpointMode() {
+    return this._mode;
+  }
+
+  setWriteCheckpointMode(mode: WriteCheckpointMode): void {
+    this._mode = mode;
+  }
+
+  async batchCreateCustomWriteCheckpoints(checkpoints: CustomWriteCheckpointOptions[]): Promise<void> {
+    return batchCreateCustomWriteCheckpoints(this.db, checkpoints);
+  }
+
+  async createCustomWriteCheckpoint(options: CustomWriteCheckpointOptions): Promise<bigint> {
+    if (this.writeCheckpointMode !== WriteCheckpointMode.CUSTOM) {
+      throw new framework.errors.ValidationError(
+        `Creating a custom Write Checkpoint when the current Write Checkpoint mode is set to "${this.writeCheckpointMode}"`
+      );
+    }
+
+    const { checkpoint, user_id, sync_rules_id } = options;
+    const doc = await this.db.custom_write_checkpoints.findOneAndUpdate(
+      {
+        user_id: user_id,
+        sync_rules_id
+      },
+      {
+        $set: {
+          checkpoint
+        }
+      },
+      { upsert: true, returnDocument: 'after' }
+    );
+    return doc!.checkpoint;
+  }
+
+  async createManagedWriteCheckpoint(checkpoint: ManagedWriteCheckpointOptions): Promise<bigint> {
+    if (this.writeCheckpointMode !== WriteCheckpointMode.MANAGED) {
+      throw new framework.errors.ValidationError(
+        `Attempting to create a managed Write Checkpoint when the current Write Checkpoint mode is set to "${this.writeCheckpointMode}"`
+      );
+    }
+
+    const { user_id, heads: lsns } = checkpoint;
+    const doc = await this.db.write_checkpoints.findOneAndUpdate(
+      {
+        user_id: user_id
+      },
+      {
+        $set: {
+          lsns
+        },
+        $inc: {
+          client_id: 1n
+        }
+      },
+      { upsert: true, returnDocument: 'after' }
+    );
+    return doc!.client_id;
+  }
+
+  async lastWriteCheckpoint(filters: LastWriteCheckpointFilters): Promise<bigint | null> {
+    switch (this.writeCheckpointMode) {
+      case WriteCheckpointMode.CUSTOM:
+        if (false == 'sync_rules_id' in filters) {
+          throw new framework.errors.ValidationError(`Sync rules ID is required for custom Write Checkpoint filtering`);
+        }
+        return this.lastCustomWriteCheckpoint(filters);
+      case WriteCheckpointMode.MANAGED:
+        if (false == 'heads' in filters) {
+          throw new framework.errors.ValidationError(
+            `Replication HEAD is required for managed Write Checkpoint filtering`
+          );
+        }
+        return this.lastManagedWriteCheckpoint(filters);
+    }
+  }
+
+  protected async lastCustomWriteCheckpoint(filters: CustomWriteCheckpointFilters) {
+    const { user_id, sync_rules_id } = filters;
+    const lastWriteCheckpoint = await this.db.custom_write_checkpoints.findOne({
+      user_id,
+      sync_rules_id
+    });
+    return lastWriteCheckpoint?.checkpoint ?? null;
+  }
+
+  protected async lastManagedWriteCheckpoint(filters: ManagedWriteCheckpointFilters) {
+    const { user_id, heads } = filters;
+    // TODO: support multiple heads when we need to support multiple connections
+    const lsn = heads['1'];
+    if (lsn == null) {
+      // Can happen if we haven't replicated anything yet.
+      return null;
+    }
+    const lastWriteCheckpoint = await this.db.write_checkpoints.findOne({
+      user_id: user_id,
+      'lsns.1': { $lte: lsn }
+    });
+    return lastWriteCheckpoint?.client_id ?? null;
+  }
+}
+
+export async function batchCreateCustomWriteCheckpoints(
+  db: PowerSyncMongo,
+  checkpoints: CustomWriteCheckpointOptions[]
+): Promise<void> {
+  if (!checkpoints.length) {
+    return;
+  }
+
+  await db.custom_write_checkpoints.bulkWrite(
+    checkpoints.map((checkpointOptions) => ({
+      updateOne: {
+        filter: { user_id: checkpointOptions.user_id, sync_rules_id: checkpointOptions.sync_rules_id },
+        update: {
+          $set: {
+            checkpoint: checkpointOptions.checkpoint,
+            sync_rules_id: checkpointOptions.sync_rules_id
+          }
+        },
+        upsert: true
+      }
+    }))
+  );
+}

package/src/storage/mongo/OperationBatch.ts

@@ -1,8 +1,9 @@
-import * as bson from 'bson';
 import { ToastableSqliteRow } from '@powersync/service-sync-rules';
+import * as bson from 'bson';
 
-import * as util from '../../util/util-index.js';
 import { SaveOptions } from '../BucketStorage.js';
+import { isUUID } from './util.js';
+import { ReplicaId } from './models.js';
 
 /**
  * Maximum number of operations in a batch.
@@ -42,7 +43,16 @@ export class OperationBatch {
     return this.batch.length >= MAX_BATCH_COUNT || this.currentSize > MAX_RECORD_BATCH_SIZE;
   }
 
-  *batched(sizes: Map<string, number>): Generator<RecordOperation[]> {
+  /**
+   *
+   * @param sizes Map of source key to estimated size of the current_data document, or undefined if current_data is not persisted.
+   *
+   */
+  *batched(sizes: Map<string, number> | undefined): Generator<RecordOperation[]> {
+    if (sizes == null) {
+      yield this.batch;
+      return;
+    }
     let currentBatch: RecordOperation[] = [];
     let currentBatchSize = 0;
     for (let op of this.batch) {
@@ -63,18 +73,15 @@ export class OperationBatch {
 }
 
 export class RecordOperation {
-  public readonly afterId: bson.UUID | null;
-  public readonly beforeId: bson.UUID;
+  public readonly afterId: ReplicaId | null;
+  public readonly beforeId: ReplicaId;
   public readonly internalBeforeKey: string;
   public readonly internalAfterKey: string | null;
   public readonly estimatedSize: number;
 
   constructor(public readonly record: SaveOptions) {
-    const after = record.after;
-    const afterId = after ? util.getUuidReplicaIdentityBson(after, record.sourceTable.replicaIdColumns!) : null;
-    const beforeId = record.before
-      ? util.getUuidReplicaIdentityBson(record.before, record.sourceTable.replicaIdColumns!)
-      : afterId!;
+    const afterId = record.afterReplicaId ?? null;
+    const beforeId = record.beforeReplicaId ?? record.afterReplicaId;
     this.afterId = afterId;
     this.beforeId = beforeId;
     this.internalBeforeKey = cacheKey(record.sourceTable.id, beforeId);
@@ -84,8 +91,17 @@ export class RecordOperation {
   }
 }
 
-export function cacheKey(table: bson.ObjectId, id: bson.UUID) {
-  return `${table.toHexString()}.${id.toHexString()}`;
+/**
+ * In-memory cache key - must not be persisted.
+ */
+export function cacheKey(table: bson.ObjectId, id: ReplicaId) {
+  if (isUUID(id)) {
+    return `${table.toHexString()}.${id.toHexString()}`;
+  } else if (typeof id == 'string') {
+    return `${table.toHexString()}.${id}`;
+  } else {
+    return `${table.toHexString()}.${(bson.serialize({ id: id }) as Buffer).toString('base64')}`;
+  }
 }
 
 /**