Tests and fixes for custom write checkpoints.

Author: rkistner

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

@@ -87,13 +87,6 @@ export class MongoSyncBucketStorage
     );
   }
 
-  createCustomWriteCheckpoint(checkpoint: storage.BatchedCustomWriteCheckpointOptions): Promise<bigint> {
-    return this.writeCheckpointAPI.createCustomWriteCheckpoint({
-      ...checkpoint,
-      sync_rules_id: this.group_id
-    });
-  }
-
   createManagedWriteCheckpoint(checkpoint: storage.ManagedWriteCheckpointOptions): Promise<bigint> {
     return this.writeCheckpointAPI.createManagedWriteCheckpoint(checkpoint);
   }

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

@@ -278,6 +278,11 @@ export class MongoWriteCheckpointAPI implements storage.WriteCheckpointAPI {
         lsn: null
       };
       lastId = doc.checkpoint;
+    } else {
+      yield {
+        id: null,
+        lsn: null
+      };
     }
 
     for await (let event of stream) {
@@ -324,7 +329,7 @@ export async function batchCreateCustomWriteCheckpoints(
   db: PowerSyncMongo,
   checkpoints: storage.CustomWriteCheckpointOptions[]
 ): Promise<void> {
-  if (!checkpoints.length) {
+  if (checkpoints.length == 0) {
     return;
   }
 

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

@@ -74,6 +74,7 @@ export class PowerSyncMongo {
     await this.instance.deleteOne({});
     await this.locks.deleteMany({});
     await this.bucket_state.deleteMany({});
+    await this.custom_write_checkpoints.deleteMany({});
   }
 
   /**

modules/module-postgres-storage/src/storage/checkpoints/PostgresWriteCheckpointAPI.ts

@@ -1,7 +1,7 @@
 import * as lib_postgres from '@powersync/lib-service-postgres';
 import * as framework from '@powersync/lib-services-framework';
-import { storage } from '@powersync/service-core';
-import { JSONBig } from '@powersync/service-jsonbig';
+import { storage, sync } from '@powersync/service-core';
+import { JSONBig, JsonContainer } from '@powersync/service-jsonbig';
 import { models } from '../../types/types.js';
 
 export type PostgresCheckpointAPIOptions = {
@@ -129,11 +129,22 @@ export async function batchCreateCustomWriteCheckpoints(
     return;
   }
 
+  // Needs to be encoded using plain JSON.stringify
+  const mappedCheckpoints = checkpoints.map((cp) => {
+    return {
+      user_id: cp.user_id,
+      // Cannot encode bigint directly using JSON.stringify.
+      // The ::int8 in the query below will take care of casting back to a number
+      checkpoint: String(cp.checkpoint),
+      sync_rules_id: cp.sync_rules_id
+    };
+  });
+
   await db.sql`
     WITH
       json_data AS (
         SELECT
-          jsonb_array_elements(${{ type: 'jsonb', value: JSONBig.stringify(checkpoints) }}) AS
+          jsonb_array_elements(${{ type: 'jsonb', value: mappedCheckpoints }}) AS
         CHECKPOINT
       )
     INSERT INTO

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

@@ -1560,7 +1560,7 @@ bucket_definitions:
     });
     // We have to trigger a new keepalive after the checkpoint, at least to cover postgres storage.
     // This is what is effetively triggered with RouteAPI.createReplicationHead().
-    // MongoDB storage doesn't need this.
+    // MongoDB storage doesn't explicitly need this anymore.
     await bucketStorage.startBatch(test_utils.BATCH_OPTIONS, async (batch) => {
       await batch.keepalive('6/0');
     });
@@ -1577,4 +1577,135 @@ bucket_definitions:
       }
     });
   });
+
+  test('custom write checkpoints - checkpoint after write', async (context) => {
+    await using factory = await generateStorageFactory();
+    const r = await factory.configureSyncRules({
+      content: `
+bucket_definitions:
+  mybucket:
+    data: []
+    `,
+      validate: false
+    });
+    const bucketStorage = factory.getInstance(r.persisted_sync_rules!);
+    await bucketStorage.autoActivate();
+    bucketStorage.setWriteCheckpointMode(storage.WriteCheckpointMode.CUSTOM);
+
+    const abortController = new AbortController();
+    context.onTestFinished(() => abortController.abort());
+    const iter = bucketStorage
+      .watchCheckpointChanges({ user_id: 'user1', signal: abortController.signal })
+      [Symbol.asyncIterator]();
+
+    await bucketStorage.batchCreateCustomWriteCheckpoints([
+      {
+        checkpoint: 5n,
+        user_id: 'user1'
+      }
+    ]);
+
+    await bucketStorage.startBatch(test_utils.BATCH_OPTIONS, async (batch) => {
+      await batch.keepalive('5/0');
+    });
+
+    const result = await iter.next();
+    expect(result).toMatchObject({
+      done: false,
+      value: {
+        base: {
+          checkpoint: 0n,
+          lsn: '5/0'
+        },
+        writeCheckpoint: 5n
+      }
+    });
+  });
+
+  test('custom write checkpoints - write after checkpoint', async (context) => {
+    await using factory = await generateStorageFactory();
+    const r = await factory.configureSyncRules({
+      content: `
+bucket_definitions:
+  mybucket:
+    data: []
+    `,
+      validate: false
+    });
+    const bucketStorage = factory.getInstance(r.persisted_sync_rules!);
+    await bucketStorage.autoActivate();
+    bucketStorage.setWriteCheckpointMode(storage.WriteCheckpointMode.CUSTOM);
+
+    const abortController = new AbortController();
+    context.onTestFinished(() => abortController.abort());
+    const iter = bucketStorage
+      .watchCheckpointChanges({ user_id: 'user1', signal: abortController.signal })
+      [Symbol.asyncIterator]();
+
+    await bucketStorage.startBatch(test_utils.BATCH_OPTIONS, async (batch) => {
+      await batch.keepalive('5/0');
+    });
+
+    const result = await iter.next();
+    expect(result).toMatchObject({
+      done: false,
+      value: {
+        base: {
+          checkpoint: 0n,
+          lsn: '5/0'
+        },
+        writeCheckpoint: null
+      }
+    });
+
+    await bucketStorage.batchCreateCustomWriteCheckpoints([
+      {
+        checkpoint: 6n,
+        user_id: 'user1'
+      }
+    ]);
+    // We have to trigger a new keepalive after the checkpoint, at least to cover postgres storage.
+    // This is what is effetively triggered with RouteAPI.createReplicationHead().
+    // MongoDB storage doesn't explicitly need this anymore.
+    await bucketStorage.startBatch(test_utils.BATCH_OPTIONS, async (batch) => {
+      await batch.keepalive('6/0');
+    });
+
+    const result2 = await iter.next();
+    expect(result2).toMatchObject({
+      done: false,
+      value: {
+        base: {
+          checkpoint: 0n,
+          lsn: '6/0'
+        },
+        writeCheckpoint: 6n
+      }
+    });
+
+    await bucketStorage.batchCreateCustomWriteCheckpoints([
+      {
+        checkpoint: 7n,
+        user_id: 'user1'
+      }
+    ]);
+    // We have to trigger a new keepalive after the checkpoint, at least to cover postgres storage.
+    // This is what is effetively triggered with RouteAPI.createReplicationHead().
+    // MongoDB storage doesn't explicitly need this anymore.
+    await bucketStorage.startBatch(test_utils.BATCH_OPTIONS, async (batch) => {
+      await batch.keepalive('7/0');
+    });
+
+    const result3 = await iter.next();
+    expect(result3).toMatchObject({
+      done: false,
+      value: {
+        base: {
+          checkpoint: 0n,
+          lsn: '7/0'
+        },
+        writeCheckpoint: 7n
+      }
+    });
+  });
 }