Move getParameterSets to ReplicationCheckpoint.

Author: rkistner

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

@@ -113,15 +113,9 @@ export class MongoSyncBucketStorage
       }
     );
     if (!doc?.snapshot_done) {
-      return {
-        checkpoint: 0n,
-        lsn: null
-      };
+      return new MongoReplicationCheckpoint(this, 0n, null);
     }
-    return {
-      checkpoint: doc?.last_checkpoint ?? 0n,
-      lsn: doc?.last_checkpoint_lsn ?? null
-    };
+    return new MongoReplicationCheckpoint(this, doc.last_checkpoint ?? 0n, doc.last_checkpoint_lsn ?? null);
   }
 
   async startBatch(
@@ -661,11 +655,8 @@ export class MongoSyncBucketStorage
     return new MongoCompactor(this.db, this.group_id, options).compact();
   }
 
-  private makeActiveCheckpoint(doc: SyncRuleCheckpointState | null) {
-    return {
-      checkpoint: doc?.last_checkpoint ?? 0n,
-      lsn: doc?.last_checkpoint_lsn ?? null
-    };
+  private makeActiveCheckpoint(doc: SyncRuleCheckpointState | null): ReplicationCheckpoint {
+    return new MongoReplicationCheckpoint(this, doc?.last_checkpoint ?? 0n, doc?.last_checkpoint_lsn ?? null);
   }
 
   /**
@@ -980,3 +971,15 @@ interface InternalCheckpointChanges extends CheckpointChanges {
   updatedWriteCheckpoints: Map<string, bigint>;
   invalidateWriteCheckpoints: boolean;
 }
+
+class MongoReplicationCheckpoint implements ReplicationCheckpoint {
+  constructor(
+    private storage: MongoSyncBucketStorage,
+    public readonly checkpoint: InternalOpId,
+    public readonly lsn: string | null
+  ) {}
+
+  async getParameterSets(lookups: ParameterLookup[]): Promise<SqliteJsonRow[]> {
+    return this.storage.getParameterSets(this, lookups);
+  }
+}

modules/module-postgres-storage/src/storage/PostgresSyncRulesStorage.ts

@@ -137,10 +137,11 @@ export class PostgresSyncRulesStorage
       .decoded(pick(models.SyncRules, ['last_checkpoint', 'last_checkpoint_lsn']))
       .first();
 
-    return {
-      checkpoint: checkpointRow?.last_checkpoint ?? 0n,
-      lsn: checkpointRow?.last_checkpoint_lsn ?? null
-    };
+    return new PostgresReplicationCheckpoint(
+      this,
+      checkpointRow?.last_checkpoint ?? 0n,
+      checkpointRow?.last_checkpoint_lsn ?? null
+    );
   }
 
   async resolveTable(options: storage.ResolveTableOptions): Promise<storage.ResolveTableResult> {
@@ -834,9 +835,18 @@ export class PostgresSyncRulesStorage
   }
 
   private makeActiveCheckpoint(row: models.ActiveCheckpointDecoded | null) {
-    return {
-      checkpoint: row?.last_checkpoint ?? 0n,
-      lsn: row?.last_checkpoint_lsn ?? null
-    } satisfies storage.ReplicationCheckpoint;
+    return new PostgresReplicationCheckpoint(this, row?.last_checkpoint ?? 0n, row?.last_checkpoint_lsn ?? null);
+  }
+}
+
+class PostgresReplicationCheckpoint implements storage.ReplicationCheckpoint {
+  constructor(
+    private storage: PostgresSyncRulesStorage,
+    public readonly checkpoint: utils.InternalOpId,
+    public readonly lsn: string | null
+  ) {}
+
+  getParameterSets(lookups: sync_rules.ParameterLookup[]): Promise<sync_rules.SqliteJsonRow[]> {
+    return this.storage.getParameterSets(this, lookups);
   }
 }

packages/service-core-tests/src/tests/register-data-storage-tests.ts

@@ -77,9 +77,7 @@ bucket_definitions:
     const checkpoint = await bucketStorage.getCheckpoint();
     console.log('Checkpoint:', checkpoint);
     console.log(await bucketStorage.getStatus());
-    const parameters = await bucketStorage.getParameterSets(checkpoint, [
-      ParameterLookup.normalized('mybucket', '1', ['user1'])
-    ]);
+    const parameters = await checkpoint.getParameterSets([ParameterLookup.normalized('mybucket', '1', ['user1'])]);
     expect(parameters).toEqual([
       {
         group_id: 'group1a'
@@ -127,19 +125,15 @@ bucket_definitions:
     });
     const checkpoint2 = await bucketStorage.getCheckpoint();
 
-    const parameters = await bucketStorage.getParameterSets(checkpoint2, [
-      ParameterLookup.normalized('mybucket', '1', ['user1'])
-    ]);
+    const parameters = await checkpoint2.getParameterSets([ParameterLookup.normalized('mybucket', '1', ['user1'])]);
     expect(parameters).toEqual([
       {
         group_id: 'group2'
       }
     ]);
 
     // Use the checkpoint to get older data if relevant
-    const parameters2 = await bucketStorage.getParameterSets(checkpoint1, [
-      ParameterLookup.normalized('mybucket', '1', ['user1'])
-    ]);
+    const parameters2 = await checkpoint1.getParameterSets([ParameterLookup.normalized('mybucket', '1', ['user1'])]);
     expect(parameters2).toEqual([
       {
         group_id: 'group1'
@@ -207,7 +201,7 @@ bucket_definitions:
     // There removal operation for the association of `list2`::`todo2` should not interfere with the new
     // association of `list1`::`todo2`
     const checkpoint = await bucketStorage.getCheckpoint();
-    const parameters = await bucketStorage.getParameterSets(checkpoint, [
+    const parameters = await checkpoint.getParameterSets([
       ParameterLookup.normalized('mybucket', '1', ['list1']),
       ParameterLookup.normalized('mybucket', '1', ['list2'])
     ]);
@@ -256,15 +250,15 @@ bucket_definitions:
 
     const checkpoint = await bucketStorage.getCheckpoint();
 
-    const parameters1 = await bucketStorage.getParameterSets(checkpoint, [
+    const parameters1 = await checkpoint.getParameterSets([
       ParameterLookup.normalized('mybucket', '1', [314n, 314, 3.14])
     ]);
     expect(parameters1).toEqual([TEST_PARAMS]);
-    const parameters2 = await bucketStorage.getParameterSets(checkpoint, [
+    const parameters2 = await checkpoint.getParameterSets([
       ParameterLookup.normalized('mybucket', '1', [314, 314n, 3.14])
     ]);
     expect(parameters2).toEqual([TEST_PARAMS]);
-    const parameters3 = await bucketStorage.getParameterSets(checkpoint, [
+    const parameters3 = await checkpoint.getParameterSets([
       ParameterLookup.normalized('mybucket', '1', [314n, 314, 3])
     ]);
     expect(parameters3).toEqual([]);
@@ -319,7 +313,7 @@ bucket_definitions:
 
     const checkpoint = await bucketStorage.getCheckpoint();
 
-    const parameters1 = await bucketStorage.getParameterSets(checkpoint, [
+    const parameters1 = await checkpoint.getParameterSets([
       ParameterLookup.normalized('mybucket', '1', [1152921504606846976n])
     ]);
     expect(parameters1).toEqual([TEST_PARAMS]);
@@ -424,12 +418,12 @@ bucket_definitions:
     const lookups = q1.getLookups(parameters);
     expect(lookups).toEqual([ParameterLookup.normalized('by_workspace', '1', ['u1'])]);
 
-    const parameter_sets = await bucketStorage.getParameterSets(checkpoint, lookups);
+    const parameter_sets = await checkpoint.getParameterSets(lookups);
     expect(parameter_sets).toEqual([{ workspace_id: 'workspace1' }]);
 
     const buckets = await sync_rules.getBucketParameterQuerier(parameters).queryDynamicBucketDescriptions({
       getParameterSets(lookups) {
-        return bucketStorage.getParameterSets(checkpoint, lookups);
+        return checkpoint.getParameterSets(lookups);
       }
     });
     expect(buckets).toEqual([{ bucket: 'by_workspace["workspace1"]', priority: 3 }]);
@@ -495,13 +489,13 @@ bucket_definitions:
     const lookups = q1.getLookups(parameters);
     expect(lookups).toEqual([ParameterLookup.normalized('by_public_workspace', '1', [])]);
 
-    const parameter_sets = await bucketStorage.getParameterSets(checkpoint, lookups);
+    const parameter_sets = await checkpoint.getParameterSets(lookups);
     parameter_sets.sort((a, b) => JSON.stringify(a).localeCompare(JSON.stringify(b)));
     expect(parameter_sets).toEqual([{ workspace_id: 'workspace1' }, { workspace_id: 'workspace3' }]);
 
     const buckets = await sync_rules.getBucketParameterQuerier(parameters).queryDynamicBucketDescriptions({
       getParameterSets(lookups) {
-        return bucketStorage.getParameterSets(checkpoint, lookups);
+        return checkpoint.getParameterSets(lookups);
       }
     });
     buckets.sort((a, b) => a.bucket.localeCompare(b.bucket));
@@ -585,23 +579,23 @@ bucket_definitions:
     const lookups1 = q1.getLookups(parameters);
     expect(lookups1).toEqual([ParameterLookup.normalized('by_workspace', '1', [])]);
 
-    const parameter_sets1 = await bucketStorage.getParameterSets(checkpoint, lookups1);
+    const parameter_sets1 = await checkpoint.getParameterSets(lookups1);
     parameter_sets1.sort((a, b) => JSON.stringify(a).localeCompare(JSON.stringify(b)));
     expect(parameter_sets1).toEqual([{ workspace_id: 'workspace1' }]);
 
     const q2 = sync_rules.bucketDescriptors[0].parameterQueries[1];
     const lookups2 = q2.getLookups(parameters);
     expect(lookups2).toEqual([ParameterLookup.normalized('by_workspace', '2', ['u1'])]);
 
-    const parameter_sets2 = await bucketStorage.getParameterSets(checkpoint, lookups2);
+    const parameter_sets2 = await checkpoint.getParameterSets(lookups2);
     parameter_sets2.sort((a, b) => JSON.stringify(a).localeCompare(JSON.stringify(b)));
     expect(parameter_sets2).toEqual([{ workspace_id: 'workspace3' }]);
 
     // Test final values - the important part
     const buckets = (
       await sync_rules.getBucketParameterQuerier(parameters).queryDynamicBucketDescriptions({
         getParameterSets(lookups) {
-          return bucketStorage.getParameterSets(checkpoint, lookups);
+          return checkpoint.getParameterSets(lookups);
         }
       })
     ).map((e) => e.bucket);
@@ -904,9 +898,7 @@ bucket_definitions:
 
     const checkpoint = await bucketStorage.getCheckpoint();
 
-    const parameters = await bucketStorage.getParameterSets(checkpoint, [
-      ParameterLookup.normalized('mybucket', '1', ['user1'])
-    ]);
+    const parameters = await checkpoint.getParameterSets([ParameterLookup.normalized('mybucket', '1', ['user1'])]);
     expect(parameters).toEqual([]);
   });
 

packages/service-core/src/storage/SyncRulesBucketStorage.ts

@@ -66,11 +66,6 @@ export interface SyncRulesBucketStorage
 
   getCheckpoint(): Promise<ReplicationCheckpoint>;
 
-  /**
-   * Used to resolve "dynamic" parameter queries.
-   */
-  getParameterSets(checkpoint: ReplicationCheckpoint, lookups: ParameterLookup[]): Promise<SqliteJsonRow[]>;
-
   /**
    * Given two checkpoints, return the changes in bucket data and parameters that may have occurred
    * in that period.
@@ -243,6 +238,13 @@ export interface SyncBucketDataChunk {
 export interface ReplicationCheckpoint {
   readonly checkpoint: util.InternalOpId;
   readonly lsn: string | null;
+
+  /**
+   * Used to resolve "dynamic" parameter queries.
+   *
+   * This gets parameter sets specific to this checkpoint.
+   */
+  getParameterSets(lookups: ParameterLookup[]): Promise<SqliteJsonRow[]>;
 }
 
 export interface WatchWriteCheckpointOptions {

packages/service-core/src/sync/BucketChecksumState.ts

@@ -440,7 +440,7 @@ export class BucketParameterState {
     if (hasParameterChange || this.cachedDynamicBuckets == null || this.cachedDynamicBucketSet == null) {
       dynamicBuckets = await querier.queryDynamicBucketDescriptions({
         getParameterSets(lookups) {
-          return storage.getParameterSets(checkpoint.base, lookups);
+          return checkpoint.base.getParameterSets(lookups);
         }
       });
       this.cachedDynamicBuckets = dynamicBuckets;
@@ -501,7 +501,7 @@ export interface CheckpointLine {
 }
 
 // Use a more specific type to simplify testing
-export type BucketChecksumStateStorage = Pick<storage.SyncRulesBucketStorage, 'getChecksums' | 'getParameterSets'>;
+export type BucketChecksumStateStorage = Pick<storage.SyncRulesBucketStorage, 'getChecksums'>;
 
 function limitedBuckets(buckets: string[] | { bucket: string }[], limit: number) {
   buckets = buckets.map((b) => {