@powersync/service-core 0.0.0-dev-20240725112650 → 0.0.0-dev-20240918092408

package/src/storage/MongoBucketStorage.ts

@@ -1,17 +1,19 @@
-import * as mongo from 'mongodb';
-import * as timers from 'timers/promises';
-import { LRUCache } from 'lru-cache/min';
 import { SqlSyncRules } from '@powersync/service-sync-rules';
 import { wrapWithAbort } from 'ix/asynciterable/operators/withabort.js';
+import { LRUCache } from 'lru-cache/min';
+import * as mongo from 'mongodb';
+import * as timers from 'timers/promises';
 
-import * as replication from '../replication/replication-index.js';
+import * as locks from '../locks/locks-index.js';
 import * as sync from '../sync/sync-index.js';
 import * as util from '../util/util-index.js';
-import * as locks from '../locks/locks-index.js';
 
+import { logger } from '@powersync/lib-services-framework';
+import { v4 as uuid } from 'uuid';
 import {
   ActiveCheckpoint,
   BucketStorageFactory,
+  ParseSyncRulesOptions,
   PersistedSyncRules,
   PersistedSyncRulesContent,
   StorageMetrics,
@@ -23,14 +25,13 @@ import { MongoSyncBucketStorage } from './mongo/MongoSyncBucketStorage.js';
 import { PowerSyncMongo, PowerSyncMongoOptions } from './mongo/db.js';
 import { SyncRuleDocument, SyncRuleState } from './mongo/models.js';
 import { generateSlotName } from './mongo/util.js';
-import { v4 as uuid } from 'uuid';
-import { logger } from '@powersync/lib-services-framework';
 
 export interface MongoBucketStorageOptions extends PowerSyncMongoOptions {}
 
 export class MongoBucketStorage implements BucketStorageFactory {
   private readonly client: mongo.MongoClient;
   private readonly session: mongo.ClientSession;
+  // TODO: This is still Postgres specific and needs to be reworked
   public readonly slot_name_prefix: string;
 
   private readonly storageCache = new LRUCache<number, MongoSyncBucketStorage>({
@@ -47,8 +48,7 @@ export class MongoBucketStorage implements BucketStorageFactory {
         return undefined;
       }
       const rules = new MongoPersistedSyncRulesContent(this.db, doc2);
-      const storage = this.getInstance(rules.parsed());
-      return storage;
+      return this.getInstance(rules);
     }
   });
 
@@ -61,12 +61,12 @@ export class MongoBucketStorage implements BucketStorageFactory {
     this.slot_name_prefix = options.slot_name_prefix;
   }
 
-  getInstance(options: PersistedSyncRules): MongoSyncBucketStorage {
-    let { id, sync_rules, slot_name } = options;
+  getInstance(options: PersistedSyncRulesContent): MongoSyncBucketStorage {
+    let { id, slot_name } = options;
     if ((typeof id as any) == 'bigint') {
       id = Number(id);
     }
-    return new MongoSyncBucketStorage(this, id, sync_rules, slot_name);
+    return new MongoSyncBucketStorage(this, id, options, slot_name);
   }
 
   async configureSyncRules(sync_rules: string, options?: { lock?: boolean }) {
@@ -136,7 +136,12 @@ export class MongoBucketStorage implements BucketStorageFactory {
 
   async updateSyncRules(options: UpdateSyncRulesOptions): Promise<MongoPersistedSyncRulesContent> {
     // Parse and validate before applying any changes
-    const parsed = SqlSyncRules.fromYaml(options.content);
+    const parsed = SqlSyncRules.fromYaml(options.content, {
+      // No schema-based validation at this point
+      schema: undefined,
+      defaultSchema: 'not_applicable', // Not needed for validation
+      throwOnError: true
+    });
 
     let rules: MongoPersistedSyncRulesContent | undefined = undefined;
 
@@ -204,9 +209,9 @@ export class MongoBucketStorage implements BucketStorageFactory {
     return new MongoPersistedSyncRulesContent(this.db, doc);
   }
 
-  async getActiveSyncRules(): Promise<PersistedSyncRules | null> {
+  async getActiveSyncRules(options: ParseSyncRulesOptions): Promise<PersistedSyncRules | null> {
     const content = await this.getActiveSyncRulesContent();
-    return content?.parsed() ?? null;
+    return content?.parsed(options) ?? null;
   }
 
   async getNextSyncRulesContent(): Promise<MongoPersistedSyncRulesContent | null> {
@@ -223,9 +228,9 @@ export class MongoBucketStorage implements BucketStorageFactory {
     return new MongoPersistedSyncRulesContent(this.db, doc);
   }
 
-  async getNextSyncRules(): Promise<PersistedSyncRules | null> {
+  async getNextSyncRules(options: ParseSyncRulesOptions): Promise<PersistedSyncRules | null> {
     const content = await this.getNextSyncRulesContent();
-    return content?.parsed() ?? null;
+    return content?.parsed(options) ?? null;
   }
 
   async getReplicatingSyncRules(): Promise<PersistedSyncRulesContent[]> {
@@ -294,14 +299,6 @@ export class MongoBucketStorage implements BucketStorageFactory {
   }
 
   async getStorageMetrics(): Promise<StorageMetrics> {
-    const active_sync_rules = await this.getActiveSyncRules();
-    if (active_sync_rules == null) {
-      return {
-        operations_size_bytes: 0,
-        parameters_size_bytes: 0,
-        replication_size_bytes: 0
-      };
-    }
     const operations_aggregate = await this.db.bucket_data
 
       .aggregate([
@@ -370,7 +367,7 @@ export class MongoBucketStorage implements BucketStorageFactory {
   private makeActiveCheckpoint(doc: SyncRuleDocument | null) {
     return {
       checkpoint: util.timestampToOpId(doc?.last_checkpoint ?? 0n),
-      lsn: doc?.last_checkpoint_lsn ?? replication.ZERO_LSN,
+      lsn: doc?.last_checkpoint_lsn ?? null,
       hasSyncRules() {
         return doc != null;
       },

package/src/storage/SourceEntity.ts

@@ -0,0 +1,22 @@
+export interface ColumnDescriptor {
+  name: string;
+  /**
+   *  The type of the column ie VARCHAR, INT, etc
+   */
+  type?: string;
+  /**
+   *  Some data sources have a type id that can be used to identify the type of the column
+   */
+  typeId?: number;
+}
+
+// TODO: This needs to be consolidated with SourceTable into something new.
+export interface SourceEntityDescriptor {
+  /**
+   *  The internal id of the data source structure in the database
+   */
+  objectId: number | string;
+  schema: string;
+  name: string;
+  replicationColumns: ColumnDescriptor[];
+}

package/src/storage/SourceTable.ts

@@ -1,10 +1,8 @@
-import { DEFAULT_SCHEMA, DEFAULT_TAG } from '@powersync/service-sync-rules';
-
-import * as replication from '../replication/replication-index.js';
+import { DEFAULT_TAG } from '@powersync/service-sync-rules';
 import * as util from '../util/util-index.js';
+import { ColumnDescriptor } from './SourceEntity.js';
 
 export class SourceTable {
-  static readonly DEFAULT_SCHEMA = DEFAULT_SCHEMA;
   static readonly DEFAULT_TAG = DEFAULT_TAG;
 
   /**
@@ -28,11 +26,11 @@ export class SourceTable {
   constructor(
     public readonly id: any,
     public readonly connectionTag: string,
-    public readonly relationId: number,
+    public readonly objectId: number | string,
     public readonly schema: string,
     public readonly table: string,
 
-    public readonly replicaIdColumns: replication.ReplicationColumn[],
+    public readonly replicaIdColumns: ColumnDescriptor[],
     public readonly snapshotComplete: boolean
   ) {}
 

package/src/storage/StorageEngine.ts

@@ -0,0 +1,55 @@
+import { ResolvedPowerSyncConfig } from '../util/util-index.js';
+import { BucketStorageFactory } from './BucketStorage.js';
+import { BucketStorageProvider, ActiveStorage } from './StorageProvider.js';
+import { logger } from '@powersync/lib-services-framework';
+
+export type StorageEngineOptions = {
+  configuration: ResolvedPowerSyncConfig;
+};
+
+export class StorageEngine {
+  // TODO: This will need to revisited when we actually support multiple storage providers.
+  private storageProviders: Map<string, BucketStorageProvider> = new Map();
+  private currentActiveStorage: ActiveStorage | null = null;
+
+  constructor(private options: StorageEngineOptions) {}
+
+  get activeBucketStorage(): BucketStorageFactory {
+    return this.activeStorage.storage;
+  }
+
+  get activeStorage(): ActiveStorage {
+    if (!this.currentActiveStorage) {
+      throw new Error(`No storage provider has been initialized yet.`);
+    }
+
+    return this.currentActiveStorage;
+  }
+
+  /**
+   * Register a provider which generates a {@link BucketStorageFactory}
+   * given the matching config specified in the loaded {@link ResolvedPowerSyncConfig}
+   */
+  registerProvider(provider: BucketStorageProvider) {
+    this.storageProviders.set(provider.type, provider);
+  }
+
+  public async start(): Promise<void> {
+    logger.info('Starting Storage Engine...');
+    const { configuration } = this.options;
+    this.currentActiveStorage = await this.storageProviders.get(configuration.storage.type)!.getStorage({
+      resolvedConfig: configuration
+    });
+    logger.info(`Successfully activated storage: ${configuration.storage.type}.`);
+    logger.info('Successfully started Storage Engine.');
+  }
+
+  /**
+   *  Shutdown the storage engine, safely shutting down any activated storage providers.
+   */
+  public async shutDown(): Promise<void> {
+    logger.info('Shutting down Storage Engine...');
+    await this.currentActiveStorage?.shutDown();
+    logger.info('Successfully shut down Storage Engine.');
+  }
+}

package/src/storage/StorageProvider.ts

@@ -0,0 +1,27 @@
+import { BucketStorageFactory } from './BucketStorage.js';
+import * as util from '../util/util-index.js';
+
+export interface ActiveStorage {
+  storage: BucketStorageFactory;
+  shutDown(): Promise<void>;
+
+  /**
+   *  Tear down / drop the storage permanently
+   */
+  tearDown(): Promise<boolean>;
+}
+
+export interface GetStorageOptions {
+  // TODO: This should just be the storage config. Update once the slot name prefix coupling has been removed from the storage
+  resolvedConfig: util.ResolvedPowerSyncConfig;
+}
+
+export interface BucketStorageProvider {
+  /**
+   *  The storage type that this provider provides.
+   *  The type should match the `type` field in the config.
+   */
+  type: string;
+
+  getStorage(options: GetStorageOptions): Promise<ActiveStorage>;
+}

package/src/storage/mongo/MongoBucketBatch.ts

@@ -2,9 +2,8 @@ import { SqliteRow, SqlSyncRules } from '@powersync/service-sync-rules';
 import * as bson from 'bson';
 import * as mongo from 'mongodb';
 
-import * as util from '../../util/util-index.js';
-import * as replication from '../../replication/replication-index.js';
 import { container, errors, logger } from '@powersync/lib-services-framework';
+import * as util from '../../util/util-index.js';
 import { BucketStorageBatch, FlushedResult, mergeToast, SaveOptions } from '../BucketStorage.js';
 import { SourceTable } from '../SourceTable.js';
 import { PowerSyncMongo } from './db.js';
@@ -12,7 +11,7 @@ import { CurrentBucket, CurrentDataDocument, SourceKey } from './models.js';
 import { MongoIdSequence } from './MongoIdSequence.js';
 import { cacheKey, OperationBatch, RecordOperation } from './OperationBatch.js';
 import { PersistedBatch } from './PersistedBatch.js';
-import { BSON_DESERIALIZE_OPTIONS, idPrefixFilter, serializeLookup } from './util.js';
+import { BSON_DESERIALIZE_OPTIONS, idPrefixFilter, replicaIdEquals, serializeLookup } from './util.js';
 
 /**
  * 15MB
@@ -25,6 +24,7 @@ const MAX_ROW_SIZE = 15 * 1024 * 1024;
 //
 // In the future, we can investigate allowing multiple replication streams operating independently.
 const replicationMutex = new util.Mutex();
+
 export class MongoBucketBatch implements BucketStorageBatch {
   private readonly client: mongo.MongoClient;
   public readonly db: PowerSyncMongo;
@@ -61,7 +61,7 @@ export class MongoBucketBatch implements BucketStorageBatch {
     group_id: number,
     slot_name: string,
     last_checkpoint_lsn: string | null,
-    no_checkpoint_before_lsn: string | null
+    no_checkpoint_before_lsn: string
   ) {
     this.db = db;
     this.client = db.client;
@@ -70,7 +70,7 @@ export class MongoBucketBatch implements BucketStorageBatch {
     this.slot_name = slot_name;
     this.session = this.client.startSession();
     this.last_checkpoint_lsn = last_checkpoint_lsn;
-    this.no_checkpoint_before_lsn = no_checkpoint_before_lsn ?? replication.ZERO_LSN;
+    this.no_checkpoint_before_lsn = no_checkpoint_before_lsn;
   }
 
   async flush(): Promise<FlushedResult | null> {
@@ -301,7 +301,7 @@ export class MongoBucketBatch implements BucketStorageBatch {
     }
 
     // 2. Save bucket data
-    if (beforeId != null && (afterId == null || !beforeId.equals(afterId))) {
+    if (beforeId != null && (afterId == null || !replicaIdEquals(beforeId, afterId))) {
       // Source ID updated
       if (sourceTable.syncData) {
         // Delete old record
@@ -431,7 +431,7 @@ export class MongoBucketBatch implements BucketStorageBatch {
       };
     }
 
-    if (afterId == null || !beforeId.equals(afterId)) {
+    if (afterId == null || !replicaIdEquals(beforeId, afterId)) {
       // Either a delete (afterId == null), or replaced the old replication id
       batch.deleteCurrentData(before_key);
     }
@@ -743,7 +743,7 @@ export class MongoBucketBatch implements BucketStorageBatch {
       const copy = new SourceTable(
         table.id,
         table.connectionTag,
-        table.relationId,
+        table.objectId,
         table.schema,
         table.table,
         table.replicaIdColumns,

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 */
@@ -168,7 +169,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';
@@ -30,10 +30,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 * as db from '../../db/db-index.js';
+import { MongoBucketStorage } from '../MongoBucketStorage.js';
+import { BucketStorageProvider, ActiveStorage, GetStorageOptions } from '../StorageProvider.js';
+import { PowerSyncMongo } from './db.js';
+import { logger } from '@powersync/lib-services-framework';
+
+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

@@ -3,7 +3,6 @@ import * as bson from 'bson';
 import * as mongo from 'mongodb';
 
 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,11 +11,16 @@ import {
   DEFAULT_DOCUMENT_BATCH_LIMIT,
   DEFAULT_DOCUMENT_CHUNK_LIMIT_BYTES,
   FlushedResult,
+  ParseSyncRulesOptions,
+  PersistedSyncRules,
+  PersistedSyncRulesContent,
   ResolveTableOptions,
   ResolveTableResult,
+  StartBatchOptions,
   SyncBucketDataBatch,
   SyncRulesBucketStorage,
-  SyncRuleStatus
+  SyncRuleStatus,
+  TerminateOptions
 } from '../BucketStorage.js';
 import { ChecksumCache, FetchPartialBucketChecksum } from '../ChecksumCache.js';
 import { MongoBucketStorage } from '../MongoBucketStorage.js';
@@ -35,29 +39,38 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
     }
   });
 
+  private parsedSyncRulesCache: SqlSyncRules | undefined;
+
   constructor(
     public readonly factory: MongoBucketStorage,
     public readonly group_id: number,
-    public readonly sync_rules: SqlSyncRules,
+    private readonly sync_rules: PersistedSyncRulesContent,
     public readonly slot_name: string
   ) {
     this.db = factory.db;
   }
 
+  getParsedSyncRules(options: ParseSyncRulesOptions): SqlSyncRules {
+    this.parsedSyncRulesCache ??= this.sync_rules.parsed(options).sync_rules;
+    return this.parsedSyncRulesCache;
+  }
+
   async getCheckpoint() {
     const doc = await this.db.sync_rules.findOne(
       { _id: this.group_id },
       {
-        projection: { last_checkpoint: 1, last_checkpoint_lsn: 1 }
+        projection: { last_checkpoint: 1 }
       }
     );
     return {
-      checkpoint: util.timestampToOpId(doc?.last_checkpoint ?? 0n),
-      lsn: doc?.last_checkpoint_lsn ?? replication.ZERO_LSN
+      checkpoint: util.timestampToOpId(doc?.last_checkpoint ?? 0n)
     };
   }
 
-  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,11 +81,11 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
 
     const batch = new MongoBucketBatch(
       this.db,
-      this.sync_rules,
+      this.sync_rules.parsed(options).sync_rules,
       this.group_id,
       this.slot_name,
       checkpoint_lsn,
-      doc?.no_checkpoint_before ?? null
+      doc?.no_checkpoint_before ?? options.zeroLSN
     );
     try {
       await callback(batch);
@@ -90,11 +103,15 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
   }
 
   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;
@@ -102,7 +119,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
@@ -114,7 +131,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,
@@ -127,7 +144,7 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
       const sourceTable = new SourceTable(
         doc._id,
         connection_tag,
-        relationId,
+        objectId,
         schema,
         table,
         replicationColumns,
@@ -142,7 +159,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 }
         )
@@ -157,7 +174,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
             )
         )
@@ -386,9 +403,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

@@ -30,7 +30,7 @@ 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);
   }

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.
@@ -63,18 +64,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 +82,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')}`;
+  }
 }
 
 /**

package/src/storage/mongo/PersistedBatch.ts

@@ -13,9 +13,10 @@ import {
   BucketParameterDocument,
   CurrentBucket,
   CurrentDataDocument,
-  SourceKey
+  SourceKey,
+  ReplicaId
 } from './models.js';
-import { serializeLookup } from './util.js';
+import { replicaIdToSubkey, serializeLookup } from './util.js';
 import { logger } from '@powersync/lib-services-framework';
 
 /**
@@ -59,7 +60,7 @@ export class PersistedBatch {
 
   saveBucketData(options: {
     op_seq: MongoIdSequence;
-    sourceKey: bson.UUID;
+    sourceKey: ReplicaId;
     table: SourceTable;
     evaluated: EvaluatedRow[];
     before_buckets: CurrentBucket[];
@@ -70,7 +71,7 @@ export class PersistedBatch {
       remaining_buckets.set(key, b);
     }
 
-    const dchecksum = util.hashDelete(`${options.table.id}/${options.sourceKey}`);
+    const dchecksum = util.hashDelete(replicaIdToSubkey(options.table.id, options.sourceKey));
 
     for (let k of options.evaluated) {
       const key = currentBucketKey(k);
@@ -134,7 +135,7 @@ export class PersistedBatch {
 
   saveParameterData(data: {
     op_seq: MongoIdSequence;
-    sourceKey: bson.UUID;
+    sourceKey: ReplicaId;
     sourceTable: SourceTable;
     evaluated: EvaluatedParameters[];
     existing_lookups: bson.Binary[];