@powersync/service-core 0.0.0-dev-20240718134716 → 0.0.0-dev-20240918082156

package/test/src/data_storage.test.ts

@@ -1,26 +1,25 @@
+import {
+  BucketDataBatchOptions,
+  ParseSyncRulesOptions,
+  PersistedSyncRulesContent,
+  StartBatchOptions
+} from '@/storage/BucketStorage.js';
 import { RequestParameters, SqlSyncRules } from '@powersync/service-sync-rules';
-import * as bson from 'bson';
 import { describe, expect, test } from 'vitest';
-import { SourceTable } from '../../src/storage/SourceTable.js';
-import { hashData } from '../../src/util/utils.js';
-import { MONGO_STORAGE_FACTORY, StorageFactory } from './util.js';
-import { SyncBucketData } from '../../src/util/protocol-types.js';
-import { BucketDataBatchOptions } from '../../src/storage/BucketStorage.js';
-import { fromAsync } from './wal_stream_utils.js';
-
-function makeTestTable(name: string, columns?: string[] | undefined) {
-  const relId = hashData('table', name, (columns ?? ['id']).join(','));
-  const id = new bson.ObjectId('6544e3899293153fa7b38331');
-  return new SourceTable(
-    id,
-    SourceTable.DEFAULT_TAG,
-    relId,
-    SourceTable.DEFAULT_SCHEMA,
-    name,
-    (columns ?? ['id']).map((column) => ({ name: column, typeOid: 25 })),
-    true
-  );
-}
+import { fromAsync, oneFromAsync } from './stream_utils.js';
+import {
+  BATCH_OPTIONS,
+  getBatchData,
+  getBatchMeta,
+  makeTestTable,
+  MONGO_STORAGE_FACTORY,
+  PARSE_OPTIONS,
+  rid,
+  StorageFactory,
+  testRules,
+  ZERO_LSN
+} from './util.js';
+import { getUuidReplicaIdentityBson } from '@/util/util-index.js';
 
 const TEST_TABLE = makeTestTable('test', ['id']);
 
@@ -30,7 +29,7 @@ describe('store - mongodb', function () {
 
 function defineDataStorageTests(factory: StorageFactory) {
   test('save and load parameters', async () => {
-    const sync_rules = SqlSyncRules.fromYaml(`
+    const sync_rules = testRules(`
 bucket_definitions:
   mybucket:
     parameters:
@@ -38,9 +37,9 @@ bucket_definitions:
     data: [] 
     `);
 
-    const storage = (await factory()).getInstance({ id: 1, sync_rules, slot_name: 'test' });
+    const storage = (await factory()).getInstance(sync_rules);
 
-    const result = await storage.startBatch({}, async (batch) => {
+    const result = await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       await batch.save({
         sourceTable: TEST_TABLE,
         tag: 'insert',
@@ -49,7 +48,8 @@ bucket_definitions:
           id1: 'user3',
           id2: 'user4',
           group_id: 'group2a'
-        }
+        },
+        afterReplicaId: rid('t2')
       });
 
       await batch.save({
@@ -60,7 +60,8 @@ bucket_definitions:
           id1: 'user1',
           id2: 'user2',
           group_id: 'group1a'
-        }
+        },
+        afterReplicaId: rid('t1')
       });
     });
 
@@ -73,34 +74,38 @@ bucket_definitions:
   });
 
   test('it should use the latest version', async () => {
-    const sync_rules = SqlSyncRules.fromYaml(`
+    const sync_rules = testRules(
+      `
 bucket_definitions:
   mybucket:
     parameters:
       - SELECT group_id FROM test WHERE id = token_parameters.user_id
     data: [] 
-    `);
+    `
+    );
 
-    const storage = (await factory()).getInstance({ id: 1, sync_rules, slot_name: 'test' });
+    const storage = (await factory()).getInstance(sync_rules);
 
-    const result1 = await storage.startBatch({}, async (batch) => {
+    const result1 = await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       await batch.save({
         sourceTable: TEST_TABLE,
         tag: 'insert',
         after: {
           id: 'user1',
           group_id: 'group1'
-        }
+        },
+        afterReplicaId: rid('user1')
       });
     });
-    const result2 = await storage.startBatch({}, async (batch) => {
+    const result2 = await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       await batch.save({
         sourceTable: TEST_TABLE,
         tag: 'insert',
         after: {
           id: 'user1',
           group_id: 'group2'
-        }
+        },
+        afterReplicaId: rid('user1')
       });
     });
 
@@ -121,17 +126,19 @@ bucket_definitions:
   });
 
   test('save and load parameters with different number types', async () => {
-    const sync_rules = SqlSyncRules.fromYaml(`
+    const sync_rules = testRules(
+      `
 bucket_definitions:
   mybucket:
     parameters:
       - SELECT group_id FROM test WHERE n1 = token_parameters.n1 and f2 = token_parameters.f2 and f3 = token_parameters.f3
     data: []
-    `);
+    `
+    );
 
-    const storage = (await factory()).getInstance({ id: 1, sync_rules, slot_name: 'test' });
+    const storage = (await factory()).getInstance(sync_rules);
 
-    const result = await storage.startBatch({}, async (batch) => {
+    const result = await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       await batch.save({
         sourceTable: TEST_TABLE,
         tag: 'insert',
@@ -141,7 +148,8 @@ bucket_definitions:
           n1: 314n,
           f2: 314,
           f3: 3.14
-        }
+        },
+        afterReplicaId: rid('t1')
       });
     });
 
@@ -162,17 +170,19 @@ bucket_definitions:
     // This specific case tested here cannot happen with postgres in practice, but we still
     // test this to ensure correct deserialization.
 
-    const sync_rules = SqlSyncRules.fromYaml(`
+    const sync_rules = testRules(
+      `
 bucket_definitions:
   mybucket:
     parameters:
       - SELECT group_id FROM test WHERE n1 = token_parameters.n1
     data: []
-    `);
+    `
+    );
 
-    const storage = (await factory()).getInstance({ id: 1, sync_rules, slot_name: 'test' });
+    const storage = (await factory()).getInstance(sync_rules);
 
-    const result = await storage.startBatch({}, async (batch) => {
+    const result = await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       await batch.save({
         sourceTable: TEST_TABLE,
         tag: 'insert',
@@ -180,7 +190,8 @@ bucket_definitions:
           id: 't1',
           group_id: 'group1',
           n1: 1152921504606846976n // 2^60
-        }
+        },
+        afterReplicaId: rid('t1')
       });
 
       await batch.save({
@@ -192,7 +203,8 @@ bucket_definitions:
           // Simulate a TOAST value, even though it can't happen for values like this
           // in practice.
           n1: undefined
-        }
+        },
+        afterReplicaId: rid('t1')
       });
     });
 
@@ -205,15 +217,17 @@ bucket_definitions:
   });
 
   test('removing row', async () => {
-    const sync_rules = SqlSyncRules.fromYaml(`
+    const sync_rules = testRules(
+      `
 bucket_definitions:
   global:
     data:
       - SELECT id, description FROM "%"
-`);
-    const storage = (await factory()).getInstance({ id: 1, sync_rules, slot_name: 'test' });
+`
+    );
+    const storage = (await factory()).getInstance(sync_rules);
 
-    const result = await storage.startBatch({}, async (batch) => {
+    const result = await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       const sourceTable = TEST_TABLE;
 
       await batch.save({
@@ -222,21 +236,20 @@ bucket_definitions:
         after: {
           id: 'test1',
           description: 'test1'
-        }
+        },
+        afterReplicaId: rid('test1')
       });
       await batch.save({
         sourceTable,
         tag: 'delete',
-        before: {
-          id: 'test1'
-        }
+        beforeReplicaId: rid('test1')
       });
     });
 
     const checkpoint = result!.flushed_op;
 
     const batch = await fromAsync(storage.getBucketDataBatch(checkpoint, new Map([['global[]', '0']])));
-    const data = batch[0].data.map((d) => {
+    const data = batch[0].batch.data.map((d) => {
       return {
         op: d.op,
         object_id: d.object_id,
@@ -265,25 +278,29 @@ bucket_definitions:
   test('save and load parameters with workspaceId', async () => {
     const WORKSPACE_TABLE = makeTestTable('workspace', ['id']);
 
-    const sync_rules = SqlSyncRules.fromYaml(`
+    const sync_rules_content = testRules(
+      `
 bucket_definitions:
     by_workspace:
       parameters:
         - SELECT id as workspace_id FROM workspace WHERE
           workspace."userId" = token_parameters.user_id
       data: []
-    `);
+    `
+    );
+    const sync_rules = sync_rules_content.parsed(PARSE_OPTIONS).sync_rules;
 
-    const storage = (await factory()).getInstance({ id: 1, sync_rules, slot_name: 'test' });
+    const storage = (await factory()).getInstance(sync_rules_content);
 
-    const result = await storage.startBatch({}, async (batch) => {
+    const result = await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       await batch.save({
         sourceTable: WORKSPACE_TABLE,
         tag: 'insert',
         after: {
           id: 'workspace1',
           userId: 'u1'
-        }
+        },
+        afterReplicaId: rid('workspace1')
       });
     });
 
@@ -311,25 +328,29 @@ bucket_definitions:
   test('save and load parameters with dynamic global buckets', async () => {
     const WORKSPACE_TABLE = makeTestTable('workspace');
 
-    const sync_rules = SqlSyncRules.fromYaml(`
+    const sync_rules_content = testRules(
+      `
 bucket_definitions:
     by_public_workspace:
       parameters:
         - SELECT id as workspace_id FROM workspace WHERE
           workspace.visibility = 'public'
       data: []
-    `);
+    `
+    );
+    const sync_rules = sync_rules_content.parsed(PARSE_OPTIONS).sync_rules;
 
-    const storage = (await factory()).getInstance({ id: 1, sync_rules, slot_name: 'test' });
+    const storage = (await factory()).getInstance(sync_rules_content);
 
-    const result = await storage.startBatch({}, async (batch) => {
+    const result = await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       await batch.save({
         sourceTable: WORKSPACE_TABLE,
         tag: 'insert',
         after: {
           id: 'workspace1',
           visibility: 'public'
-        }
+        },
+        afterReplicaId: rid('workspace1')
       });
 
       await batch.save({
@@ -338,7 +359,8 @@ bucket_definitions:
         after: {
           id: 'workspace2',
           visibility: 'private'
-        }
+        },
+        afterReplicaId: rid('workspace2')
       });
 
       await batch.save({
@@ -347,7 +369,8 @@ bucket_definitions:
         after: {
           id: 'workspace3',
           visibility: 'public'
-        }
+        },
+        afterReplicaId: rid('workspace3')
       });
     });
 
@@ -377,7 +400,8 @@ bucket_definitions:
   test('multiple parameter queries', async () => {
     const WORKSPACE_TABLE = makeTestTable('workspace');
 
-    const sync_rules = SqlSyncRules.fromYaml(`
+    const sync_rules_content = testRules(
+      `
 bucket_definitions:
     by_workspace:
       parameters:
@@ -386,18 +410,21 @@ bucket_definitions:
         - SELECT id as workspace_id FROM workspace WHERE
             workspace.user_id = token_parameters.user_id
       data: []
-    `);
+    `
+    );
+    const sync_rules = sync_rules_content.parsed(PARSE_OPTIONS).sync_rules;
 
-    const storage = (await factory()).getInstance({ id: 1, sync_rules, slot_name: 'test' });
+    const storage = (await factory()).getInstance(sync_rules_content);
 
-    const result = await storage.startBatch({}, async (batch) => {
+    const result = await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       await batch.save({
         sourceTable: WORKSPACE_TABLE,
         tag: 'insert',
         after: {
           id: 'workspace1',
           visibility: 'public'
-        }
+        },
+        afterReplicaId: rid('workspace1')
       });
 
       await batch.save({
@@ -406,7 +433,8 @@ bucket_definitions:
         after: {
           id: 'workspace2',
           visibility: 'private'
-        }
+        },
+        afterReplicaId: rid('workspace2')
       });
 
       await batch.save({
@@ -416,7 +444,8 @@ bucket_definitions:
           id: 'workspace3',
           user_id: 'u1',
           visibility: 'private'
-        }
+        },
+        afterReplicaId: rid('workspace3')
       });
 
       await batch.save({
@@ -426,7 +455,8 @@ bucket_definitions:
           id: 'workspace4',
           user_id: 'u2',
           visibility: 'private'
-        }
+        },
+        afterReplicaId: rid('workspace4')
       });
     });
 
@@ -463,16 +493,18 @@ bucket_definitions:
   });
 
   test('changing client ids', async () => {
-    const sync_rules = SqlSyncRules.fromYaml(`
+    const sync_rules = testRules(
+      `
 bucket_definitions:
   global:
     data:
       - SELECT client_id as id, description FROM "%"
-`);
-    const storage = (await factory()).getInstance({ id: 1, sync_rules, slot_name: 'test' });
+`
+    );
+    const storage = (await factory()).getInstance(sync_rules);
 
     const sourceTable = TEST_TABLE;
-    const result = await storage.startBatch({}, async (batch) => {
+    const result = await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       await batch.save({
         sourceTable,
         tag: 'insert',
@@ -480,7 +512,8 @@ bucket_definitions:
           id: 'test1',
           client_id: 'client1a',
           description: 'test1a'
-        }
+        },
+        afterReplicaId: rid('test1')
       });
       await batch.save({
         sourceTable,
@@ -489,7 +522,8 @@ bucket_definitions:
           id: 'test1',
           client_id: 'client1b',
           description: 'test1b'
-        }
+        },
+        afterReplicaId: rid('test1')
       });
 
       await batch.save({
@@ -499,12 +533,13 @@ bucket_definitions:
           id: 'test2',
           client_id: 'client2',
           description: 'test2'
-        }
+        },
+        afterReplicaId: rid('test2')
       });
     });
     const checkpoint = result!.flushed_op;
     const batch = await fromAsync(storage.getBucketDataBatch(checkpoint, new Map([['global[]', '0']])));
-    const data = batch[0].data.map((d) => {
+    const data = batch[0].batch.data.map((d) => {
       return {
         op: d.op,
         object_id: d.object_id
@@ -520,15 +555,17 @@ bucket_definitions:
   });
 
   test('re-apply delete', async () => {
-    const sync_rules = SqlSyncRules.fromYaml(`
+    const sync_rules = testRules(
+      `
 bucket_definitions:
   global:
     data:
       - SELECT id, description FROM "%"
-`);
-    const storage = (await factory()).getInstance({ id: 1, sync_rules, slot_name: 'test' });
+`
+    );
+    const storage = (await factory()).getInstance(sync_rules);
 
-    await storage.startBatch({}, async (batch) => {
+    await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       const sourceTable = TEST_TABLE;
 
       await batch.save({
@@ -537,38 +574,35 @@ bucket_definitions:
         after: {
           id: 'test1',
           description: 'test1'
-        }
+        },
+        afterReplicaId: rid('test1')
       });
     });
 
-    await storage.startBatch({}, async (batch) => {
+    await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       const sourceTable = TEST_TABLE;
 
       await batch.save({
         sourceTable,
         tag: 'delete',
-        before: {
-          id: 'test1'
-        }
+        beforeReplicaId: rid('test1')
       });
     });
 
-    const result = await storage.startBatch({}, async (batch) => {
+    const result = await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       const sourceTable = TEST_TABLE;
 
       await batch.save({
         sourceTable,
         tag: 'delete',
-        before: {
-          id: 'test1'
-        }
+        beforeReplicaId: rid('test1')
       });
     });
 
     const checkpoint = result!.flushed_op;
 
     const batch = await fromAsync(storage.getBucketDataBatch(checkpoint, new Map([['global[]', '0']])));
-    const data = batch[0].data.map((d) => {
+    const data = batch[0].batch.data.map((d) => {
       return {
         op: d.op,
         object_id: d.object_id,
@@ -595,15 +629,17 @@ bucket_definitions:
   });
 
   test('re-apply update + delete', async () => {
-    const sync_rules = SqlSyncRules.fromYaml(`
+    const sync_rules = testRules(
+      `
 bucket_definitions:
   global:
     data:
       - SELECT id, description FROM "%"
-`);
-    const storage = (await factory()).getInstance({ id: 1, sync_rules, slot_name: 'test' });
+`
+    );
+    const storage = (await factory()).getInstance(sync_rules);
 
-    await storage.startBatch({}, async (batch) => {
+    await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       const sourceTable = TEST_TABLE;
 
       await batch.save({
@@ -612,11 +648,12 @@ bucket_definitions:
         after: {
           id: 'test1',
           description: 'test1'
-        }
+        },
+        afterReplicaId: rid('test1')
       });
     });
 
-    await storage.startBatch({}, async (batch) => {
+    await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       const sourceTable = TEST_TABLE;
 
       await batch.save({
@@ -625,7 +662,8 @@ bucket_definitions:
         after: {
           id: 'test1',
           description: undefined
-        }
+        },
+        afterReplicaId: rid('test1')
       });
 
       await batch.save({
@@ -634,19 +672,18 @@ bucket_definitions:
         after: {
           id: 'test1',
           description: undefined
-        }
+        },
+        afterReplicaId: rid('test1')
       });
 
       await batch.save({
         sourceTable,
         tag: 'delete',
-        before: {
-          id: 'test1'
-        }
+        beforeReplicaId: rid('test1')
       });
     });
 
-    const result = await storage.startBatch({}, async (batch) => {
+    const result = await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       const sourceTable = TEST_TABLE;
 
       await batch.save({
@@ -655,7 +692,8 @@ bucket_definitions:
         after: {
           id: 'test1',
           description: undefined
-        }
+        },
+        afterReplicaId: rid('test1')
       });
 
       await batch.save({
@@ -664,15 +702,14 @@ bucket_definitions:
         after: {
           id: 'test1',
           description: undefined
-        }
+        },
+        afterReplicaId: rid('test1')
       });
 
       await batch.save({
         sourceTable,
         tag: 'delete',
-        before: {
-          id: 'test1'
-        }
+        beforeReplicaId: rid('test1')
       });
     });
 
@@ -680,7 +717,7 @@ bucket_definitions:
 
     const batch = await fromAsync(storage.getBucketDataBatch(checkpoint, new Map([['global[]', '0']])));
 
-    const data = batch[0].data.map((d) => {
+    const data = batch[0].batch.data.map((d) => {
       return {
         op: d.op,
         object_id: d.object_id,
@@ -709,17 +746,19 @@ bucket_definitions:
   });
 
   test('truncate parameters', async () => {
-    const sync_rules = SqlSyncRules.fromYaml(`
+    const sync_rules = testRules(
+      `
 bucket_definitions:
   mybucket:
     parameters:
       - SELECT group_id FROM test WHERE id1 = token_parameters.user_id OR id2 = token_parameters.user_id
     data: []
-    `);
+    `
+    );
 
-    const storage = (await factory()).getInstance({ id: 1, sync_rules, slot_name: 'test' });
+    const storage = (await factory()).getInstance(sync_rules);
 
-    await storage.startBatch({}, async (batch) => {
+    await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       await batch.save({
         sourceTable: TEST_TABLE,
         tag: 'insert',
@@ -728,7 +767,8 @@ bucket_definitions:
           id1: 'user3',
           id2: 'user4',
           group_id: 'group2a'
-        }
+        },
+        afterReplicaId: rid('t2')
       });
 
       await batch.truncate([TEST_TABLE]);
@@ -749,16 +789,18 @@ bucket_definitions:
     // 1. Not getting the correct "current_data" state for each operation.
     // 2. Output order not being correct.
 
-    const sync_rules = SqlSyncRules.fromYaml(`
+    const sync_rules = testRules(
+      `
 bucket_definitions:
   global:
     data:
       - SELECT id, description FROM "test"
-`);
-    const storage = (await factory()).getInstance({ id: 1, sync_rules, slot_name: 'test' });
+`
+    );
+    const storage = (await factory()).getInstance(sync_rules);
 
     // Pre-setup
-    const result1 = await storage.startBatch({}, async (batch) => {
+    const result1 = await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       const sourceTable = TEST_TABLE;
 
       await batch.save({
@@ -767,7 +809,8 @@ bucket_definitions:
         after: {
           id: 'test1',
           description: 'test1a'
-        }
+        },
+        afterReplicaId: rid('test1')
       });
 
       await batch.save({
@@ -776,14 +819,15 @@ bucket_definitions:
         after: {
           id: 'test2',
           description: 'test2a'
-        }
+        },
+        afterReplicaId: rid('test2')
       });
     });
 
     const checkpoint1 = result1?.flushed_op ?? '0';
 
     // Test batch
-    const result2 = await storage.startBatch({}, async (batch) => {
+    const result2 = await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       const sourceTable = TEST_TABLE;
       // b
       await batch.save({
@@ -792,7 +836,8 @@ bucket_definitions:
         after: {
           id: 'test1',
           description: 'test1b'
-        }
+        },
+        afterReplicaId: rid('test1')
       });
 
       await batch.save({
@@ -801,10 +846,12 @@ bucket_definitions:
         before: {
           id: 'test1'
         },
+        beforeReplicaId: rid('test1'),
         after: {
           id: 'test2',
           description: 'test2b'
-        }
+        },
+        afterReplicaId: rid('test2')
       });
 
       await batch.save({
@@ -813,10 +860,13 @@ bucket_definitions:
         before: {
           id: 'test2'
         },
+        beforeReplicaId: rid('test2'),
         after: {
           id: 'test3',
           description: 'test3b'
-        }
+        },
+
+        afterReplicaId: rid('test3')
       });
 
       // c
@@ -826,7 +876,8 @@ bucket_definitions:
         after: {
           id: 'test2',
           description: 'test2c'
-        }
+        },
+        afterReplicaId: rid('test2')
       });
 
       // d
@@ -836,7 +887,8 @@ bucket_definitions:
         after: {
           id: 'test4',
           description: 'test4d'
-        }
+        },
+        afterReplicaId: rid('test4')
       });
 
       await batch.save({
@@ -845,17 +897,19 @@ bucket_definitions:
         before: {
           id: 'test4'
         },
+        beforeReplicaId: rid('test4'),
         after: {
           id: 'test5',
           description: 'test5d'
-        }
+        },
+        afterReplicaId: rid('test5')
       });
     });
 
     const checkpoint2 = result2!.flushed_op;
 
     const batch = await fromAsync(storage.getBucketDataBatch(checkpoint2, new Map([['global[]', checkpoint1]])));
-    const data = batch[0].data.map((d) => {
+    const data = batch[0].batch.data.map((d) => {
       return {
         op: d.op,
         object_id: d.object_id,
@@ -883,31 +937,40 @@ bucket_definitions:
   });
 
   test('changed data with replica identity full', async () => {
-    const sync_rules = SqlSyncRules.fromYaml(`
+    const sync_rules = testRules(
+      `
 bucket_definitions:
   global:
     data:
       - SELECT id, description FROM "test"
-`);
-    const storage = (await factory()).getInstance({ id: 1, sync_rules, slot_name: 'test' });
+`
+    );
+    function rid2(id: string, description: string) {
+      return getUuidReplicaIdentityBson({ id, description }, [
+        { name: 'id', type: 'VARCHAR', typeId: 25 },
+        { name: 'description', type: 'VARCHAR', typeId: 25 }
+      ]);
+    }
+    const storage = (await factory()).getInstance(sync_rules);
 
     const sourceTable = makeTestTable('test', ['id', 'description']);
 
     // Pre-setup
-    const result1 = await storage.startBatch({}, async (batch) => {
+    const result1 = await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       await batch.save({
         sourceTable,
         tag: 'insert',
         after: {
           id: 'test1',
           description: 'test1a'
-        }
+        },
+        afterReplicaId: rid2('test1', 'test1a')
       });
     });
 
     const checkpoint1 = result1?.flushed_op ?? '0';
 
-    const result2 = await storage.startBatch({}, async (batch) => {
+    const result2 = await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       // Unchanged, but has a before id
       await batch.save({
         sourceTable,
@@ -916,14 +979,16 @@ bucket_definitions:
           id: 'test1',
           description: 'test1a'
         },
+        beforeReplicaId: rid2('test1', 'test1a'),
         after: {
           id: 'test1',
           description: 'test1b'
-        }
+        },
+        afterReplicaId: rid2('test1', 'test1b')
       });
     });
 
-    const result3 = await storage.startBatch({}, async (batch) => {
+    const result3 = await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       // Delete
       await batch.save({
         sourceTable,
@@ -932,6 +997,7 @@ bucket_definitions:
           id: 'test1',
           description: 'test1b'
         },
+        beforeReplicaId: rid2('test1', 'test1b'),
         after: undefined
       });
     });
@@ -939,7 +1005,7 @@ bucket_definitions:
     const checkpoint3 = result3!.flushed_op;
 
     const batch = await fromAsync(storage.getBucketDataBatch(checkpoint3, new Map([['global[]', checkpoint1]])));
-    const data = batch[0].data.map((d) => {
+    const data = batch[0].batch.data.map((d) => {
       return {
         op: d.op,
         object_id: d.object_id,
@@ -975,31 +1041,41 @@ bucket_definitions:
   });
 
   test('unchanged data with replica identity full', async () => {
-    const sync_rules = SqlSyncRules.fromYaml(`
+    const sync_rules = testRules(
+      `
 bucket_definitions:
   global:
     data:
       - SELECT id, description FROM "test"
-`);
-    const storage = (await factory()).getInstance({ id: 1, sync_rules, slot_name: 'test' });
+`
+    );
+    function rid2(id: string, description: string) {
+      return getUuidReplicaIdentityBson({ id, description }, [
+        { name: 'id', type: 'VARCHAR', typeId: 25 },
+        { name: 'description', type: 'VARCHAR', typeId: 25 }
+      ]);
+    }
+
+    const storage = (await factory()).getInstance(sync_rules);
 
     const sourceTable = makeTestTable('test', ['id', 'description']);
 
     // Pre-setup
-    const result1 = await storage.startBatch({}, async (batch) => {
+    const result1 = await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       await batch.save({
         sourceTable,
         tag: 'insert',
         after: {
           id: 'test1',
           description: 'test1a'
-        }
+        },
+        afterReplicaId: rid2('test1', 'test1a')
       });
     });
 
     const checkpoint1 = result1?.flushed_op ?? '0';
 
-    const result2 = await storage.startBatch({}, async (batch) => {
+    const result2 = await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       // Unchanged, but has a before id
       await batch.save({
         sourceTable,
@@ -1008,14 +1084,16 @@ bucket_definitions:
           id: 'test1',
           description: 'test1a'
         },
+        beforeReplicaId: rid2('test1', 'test1a'),
         after: {
           id: 'test1',
           description: 'test1a'
-        }
+        },
+        afterReplicaId: rid2('test1', 'test1a')
       });
     });
 
-    const result3 = await storage.startBatch({}, async (batch) => {
+    const result3 = await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       // Delete
       await batch.save({
         sourceTable,
@@ -1024,6 +1102,7 @@ bucket_definitions:
           id: 'test1',
           description: 'test1a'
         },
+        beforeReplicaId: rid2('test1', 'test1a'),
         after: undefined
       });
     });
@@ -1031,7 +1110,7 @@ bucket_definitions:
     const checkpoint3 = result3!.flushed_op;
 
     const batch = await fromAsync(storage.getBucketDataBatch(checkpoint3, new Map([['global[]', checkpoint1]])));
-    const data = batch[0].data.map((d) => {
+    const data = batch[0].batch.data.map((d) => {
       return {
         op: d.op,
         object_id: d.object_id,
@@ -1064,15 +1143,17 @@ bucket_definitions:
     // but large enough in size to be split over multiple returned batches.
     // The specific batch splits is an implementation detail of the storage driver,
     // and the test will have to updated when other implementations are added.
-    const sync_rules = SqlSyncRules.fromYaml(`
+    const sync_rules = testRules(
+      `
 bucket_definitions:
   global:
     data:
       - SELECT id, description FROM "%"
-`);
-    const storage = (await factory()).getInstance({ id: 1, sync_rules, slot_name: 'test' });
+`
+    );
+    const storage = (await factory()).getInstance(sync_rules);
 
-    const result = await storage.startBatch({}, async (batch) => {
+    const result = await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       const sourceTable = TEST_TABLE;
 
       const largeDescription = '0123456789'.repeat(12_000_00);
@@ -1083,7 +1164,8 @@ bucket_definitions:
         after: {
           id: 'test1',
           description: 'test1'
-        }
+        },
+        afterReplicaId: rid('test1')
       });
 
       await batch.save({
@@ -1092,7 +1174,8 @@ bucket_definitions:
         after: {
           id: 'large1',
           description: largeDescription
-        }
+        },
+        afterReplicaId: rid('large1')
       });
 
       // Large enough to split the returned batch
@@ -1102,7 +1185,8 @@ bucket_definitions:
         after: {
           id: 'large2',
           description: largeDescription
-        }
+        },
+        afterReplicaId: rid('large2')
       });
 
       await batch.save({
@@ -1111,7 +1195,8 @@ bucket_definitions:
         after: {
           id: 'test3',
           description: 'test3'
-        }
+        },
+        afterReplicaId: rid('test3')
       });
     });
 
@@ -1133,7 +1218,7 @@ bucket_definitions:
     });
 
     const batch2 = await fromAsync(
-      storage.getBucketDataBatch(checkpoint, new Map([['global[]', batch1[0].next_after]]), options)
+      storage.getBucketDataBatch(checkpoint, new Map([['global[]', batch1[0].batch.next_after]]), options)
     );
     expect(getBatchData(batch2)).toEqual([
       { op_id: '3', op: 'PUT', object_id: 'large2', checksum: 1795508474 },
@@ -1146,7 +1231,7 @@ bucket_definitions:
     });
 
     const batch3 = await fromAsync(
-      storage.getBucketDataBatch(checkpoint, new Map([['global[]', batch2[0].next_after]]), options)
+      storage.getBucketDataBatch(checkpoint, new Map([['global[]', batch2[0].batch.next_after]]), options)
     );
     expect(getBatchData(batch3)).toEqual([]);
     expect(getBatchMeta(batch3)).toEqual(null);
@@ -1156,15 +1241,17 @@ bucket_definitions:
     // Test syncing a batch of data that is small in count,
     // but large enough in size to be split over multiple returned chunks.
     // Similar to the above test, but splits over 1MB chunks.
-    const sync_rules = SqlSyncRules.fromYaml(`
+    const sync_rules = testRules(
+      `
 bucket_definitions:
   global:
     data:
       - SELECT id, description FROM "%"
-`);
-    const storage = (await factory()).getInstance({ id: 1, sync_rules, slot_name: 'test' });
+`
+    );
+    const storage = (await factory()).getInstance(sync_rules);
 
-    const result = await storage.startBatch({}, async (batch) => {
+    const result = await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       const sourceTable = TEST_TABLE;
 
       const largeDescription = '0123456789'.repeat(2_000_00);
@@ -1175,7 +1262,8 @@ bucket_definitions:
         after: {
           id: 'test1',
           description: 'test1'
-        }
+        },
+        afterReplicaId: rid('test1')
       });
 
       await batch.save({
@@ -1184,7 +1272,8 @@ bucket_definitions:
         after: {
           id: 'large1',
           description: largeDescription
-        }
+        },
+        afterReplicaId: rid('large1')
       });
 
       // Large enough to split the returned batch
@@ -1194,7 +1283,8 @@ bucket_definitions:
         after: {
           id: 'large2',
           description: largeDescription
-        }
+        },
+        afterReplicaId: rid('large2')
       });
 
       await batch.save({
@@ -1203,7 +1293,8 @@ bucket_definitions:
         after: {
           id: 'test3',
           description: 'test3'
-        }
+        },
+        afterReplicaId: rid('test3')
       });
     });
 
@@ -1223,7 +1314,7 @@ bucket_definitions:
     });
 
     const batch2 = await fromAsync(
-      storage.getBucketDataBatch(checkpoint, new Map([['global[]', batch1[0].next_after]]), options)
+      storage.getBucketDataBatch(checkpoint, new Map([['global[]', batch1[0].batch.next_after]]), options)
     );
     expect(getBatchData(batch2)).toEqual([{ op_id: '3', op: 'PUT', object_id: 'large2', checksum: 1607205872 }]);
     expect(getBatchMeta(batch2)).toEqual({
@@ -1233,7 +1324,7 @@ bucket_definitions:
     });
 
     const batch3 = await fromAsync(
-      storage.getBucketDataBatch(checkpoint, new Map([['global[]', batch2[0].next_after]]), options)
+      storage.getBucketDataBatch(checkpoint, new Map([['global[]', batch2[0].batch.next_after]]), options)
     );
     expect(getBatchData(batch3)).toEqual([{ op_id: '4', op: 'PUT', object_id: 'test3', checksum: 1359888332 }]);
     expect(getBatchMeta(batch3)).toEqual({
@@ -1245,15 +1336,17 @@ bucket_definitions:
 
   test('long batch', async () => {
     // Test syncing a batch of data that is limited by count.
-    const sync_rules = SqlSyncRules.fromYaml(`
+    const sync_rules = testRules(
+      `
 bucket_definitions:
   global:
     data:
       - SELECT id, description FROM "%"
-`);
-    const storage = (await factory()).getInstance({ id: 1, sync_rules, slot_name: 'test' });
+`
+    );
+    const storage = (await factory()).getInstance(sync_rules);
 
-    const result = await storage.startBatch({}, async (batch) => {
+    const result = await storage.startBatch(BATCH_OPTIONS, async (batch) => {
       const sourceTable = TEST_TABLE;
 
       for (let i = 1; i <= 6; i++) {
@@ -1263,14 +1356,17 @@ bucket_definitions:
           after: {
             id: `test${i}`,
             description: `test${i}`
-          }
+          },
+          afterReplicaId: `test${i}`
         });
       }
     });
 
     const checkpoint = result!.flushed_op;
 
-    const batch1 = await fromAsync(storage.getBucketDataBatch(checkpoint, new Map([['global[]', '0']]), { limit: 4 }));
+    const batch1 = await oneFromAsync(
+      storage.getBucketDataBatch(checkpoint, new Map([['global[]', '0']]), { limit: 4 })
+    );
 
     expect(getBatchData(batch1)).toEqual([
       { op_id: '1', op: 'PUT', object_id: 'test1', checksum: 2871785649 },
@@ -1285,8 +1381,8 @@ bucket_definitions:
       next_after: '4'
     });
 
-    const batch2 = await fromAsync(
-      storage.getBucketDataBatch(checkpoint, new Map([['global[]', batch1[0].next_after]]), {
+    const batch2 = await oneFromAsync(
+      storage.getBucketDataBatch(checkpoint, new Map([['global[]', batch1.batch.next_after]]), {
         limit: 4
       })
     );
@@ -1302,7 +1398,7 @@ bucket_definitions:
     });
 
     const batch3 = await fromAsync(
-      storage.getBucketDataBatch(checkpoint, new Map([['global[]', batch2[0].next_after]]), {
+      storage.getBucketDataBatch(checkpoint, new Map([['global[]', batch2.batch.next_after]]), {
         limit: 4
       })
     );
@@ -1311,28 +1407,3 @@ bucket_definitions:
     expect(getBatchMeta(batch3)).toEqual(null);
   });
 }
-
-function getBatchData(batch: SyncBucketData[]) {
-  if (batch.length == 0) {
-    return [];
-  }
-  return batch[0].data.map((d) => {
-    return {
-      op_id: d.op_id,
-      op: d.op,
-      object_id: d.object_id,
-      checksum: d.checksum
-    };
-  });
-}
-
-function getBatchMeta(batch: SyncBucketData[]) {
-  if (batch.length == 0) {
-    return null;
-  }
-  return {
-    has_more: batch[0].has_more,
-    after: batch[0].after,
-    next_after: batch[0].next_after
-  };
-}