@powersync/service-module-postgres 0.18.0 → 0.19.0

package/test/src/wal_stream.test.ts

@@ -1,12 +1,12 @@
 import { MissingReplicationSlotError } from '@module/replication/WalStream.js';
 import { METRICS_HELPER, putOp, removeOp } from '@powersync/service-core-tests';
 import { pgwireRows } from '@powersync/service-jpgwire';
+import { JSONBig } from '@powersync/service-jsonbig';
 import { ReplicationMetric } from '@powersync/service-types';
 import * as crypto from 'crypto';
 import { describe, expect, test } from 'vitest';
 import { describeWithStorage, StorageVersionTestContext } from './util.js';
 import { WalStreamTestContext, withMaxWalSize } from './wal_stream_utils.js';
-import { JSONBig } from '@powersync/service-jsonbig';
 
 const BASIC_SYNC_RULES = `
 bucket_definitions:
@@ -105,7 +105,6 @@ bucket_definitions:
     );
 
     await context.replicateSnapshot();
-    context.startStreaming();
 
     // Must be > 8kb after compression
     const largeDescription = crypto.randomBytes(20_000).toString('hex');
@@ -212,7 +211,6 @@ bucket_definitions:
     );
 
     await context.replicateSnapshot();
-    context.startStreaming();
 
     const data = await context.getBucketData('global[]');
     expect(data).toMatchObject([putOp('test_data', { id: test_id, description: 'test1' })]);
@@ -244,8 +242,6 @@ bucket_definitions:
       params: [{ type: 'varchar', value: largeDescription }]
     });
 
-    context.startStreaming();
-
     const data = await context.getBucketData('global[]');
     expect(data.length).toEqual(1);
     const row = JSON.parse(data[0].data as string);
@@ -297,7 +293,6 @@ bucket_definitions:
         `INSERT INTO test_data(id, description) VALUES('8133cd37-903b-4937-a022-7c8294015a3a', 'test1') returning id as test_id`
       );
       await context.replicateSnapshot();
-      context.startStreaming();
 
       const data = await context.getBucketData('global[]');
 
@@ -322,15 +317,12 @@ bucket_definitions:
 
       await context.loadActiveSyncRules();
 
-      // Previously, the `replicateSnapshot` call picked up on this error.
-      // Now, we have removed that check, this only comes up when we start actually streaming.
-      // We don't get the streaming response directly here, but getCheckpoint() checks for that.
-      await context.replicateSnapshot();
-      context.startStreaming();
+      // Note: The actual error may be thrown either in replicateSnapshot(), or in getCheckpoint().
 
       if (serverVersion!.compareMain('18.0.0') >= 0) {
         // No error expected in Postres 18. Replication keeps on working depite the
         // publication being re-created.
+        await context.replicateSnapshot();
         await context.getCheckpoint();
       } else {
         // await context.getCheckpoint();
@@ -338,9 +330,9 @@ bucket_definitions:
         // In the service, this error is handled in WalStreamReplicationJob,
         // creating a new replication slot.
         await expect(async () => {
+          await context.replicateSnapshot();
           await context.getCheckpoint();
         }).rejects.toThrowError(MissingReplicationSlotError);
-        context.clearStreamError();
       }
     }
   });
@@ -362,7 +354,6 @@ bucket_definitions:
         `INSERT INTO test_data(id, description) VALUES('8133cd37-903b-4937-a022-7c8294015a3a', 'test1') returning id as test_id`
       );
       await context.replicateSnapshot();
-      context.startStreaming();
 
       const data = await context.getBucketData('global[]');
 
@@ -425,7 +416,6 @@ bucket_definitions:
         `INSERT INTO test_data(id, description) VALUES('8133cd37-903b-4937-a022-7c8294015a3a', 'test1') returning id as test_id`
       );
       await context.replicateSnapshot();
-      context.startStreaming();
 
       const data = await context.getBucketData('global[]');
 
@@ -512,7 +502,7 @@ config:
     await context.initializeReplication();
     await pool.query(`INSERT INTO test_data(id, description) VALUES ('t1', '2025-09-10 15:17:14+02')`);
 
-    const data = await context.getBucketData('1#stream|0[]');
+    const data = await context.getBucketData('stream|0[]');
     expect(data).toMatchObject([putOp('test_data', { id: 't1', description: '2025-09-10T13:17:14.000000Z' })]);
   });
 
@@ -544,7 +534,7 @@ config:
       `INSERT INTO test_data(id, description, ts) VALUES ('t2', ROW(TRUE, 2)::composite, '2025-11-17T09:12:00Z')`
     );
 
-    const data = await context.getBucketData('1#stream|0[]');
+    const data = await context.getBucketData('stream|0[]');
     expect(data).toMatchObject([
       putOp('test_data', { id: 't1', description: '{"foo":1,"bar":1}', ts: '2025-11-17T09:11:00.000000Z' }),
       putOp('test_data', { id: 't2', description: '{"foo":1,"bar":2}', ts: '2025-11-17T09:12:00.000000Z' })
@@ -571,7 +561,7 @@ config:
     await context.initializeReplication();
     await pool.query(`INSERT INTO test_data(id) VALUES ('t1')`);
 
-    const data = await context.getBucketData('1#stream|0[]');
+    const data = await context.getBucketData('stream|0[]');
     expect(data).toMatchObject([putOp('test_data', { id: 't1' })]);
   });
 
@@ -593,7 +583,6 @@ config:
     );
 
     await context.replicateSnapshot();
-    context.startStreaming();
 
     await pool.query(`UPDATE test_data SET description = 'test2' WHERE id = '${test_id}'`);
 

package/test/src/wal_stream_utils.ts

@@ -1,6 +1,6 @@
 import { PgManager } from '@module/replication/PgManager.js';
 import { PUBLICATION_NAME, WalStream, WalStreamOptions } from '@module/replication/WalStream.js';
-import { CustomTypeRegistry } from '@module/types/registry.js';
+import { ReplicationAbortedError } from '@powersync/lib-services-framework';
 import {
   BucketStorageFactory,
   createCoreReplicationMetrics,
@@ -8,23 +8,24 @@ import {
   InternalOpId,
   LEGACY_STORAGE_VERSION,
   OplogEntry,
-  STORAGE_VERSION_CONFIG,
+  settledPromise,
   storage,
+  STORAGE_VERSION_CONFIG,
   SyncRulesBucketStorage,
+  unsettledPromise,
   updateSyncRulesFromYaml
 } from '@powersync/service-core';
-import { METRICS_HELPER, test_utils } from '@powersync/service-core-tests';
+import { bucketRequest, METRICS_HELPER, test_utils } from '@powersync/service-core-tests';
 import * as pgwire from '@powersync/service-jpgwire';
 import { clearTestDb, getClientCheckpoint, TEST_CONNECTION_OPTIONS } from './util.js';
 
 export class WalStreamTestContext implements AsyncDisposable {
   private _walStream?: WalStream;
   private abortController = new AbortController();
-  private streamPromise?: Promise<void>;
   private syncRulesId?: number;
+  private syncRulesContent?: storage.PersistedSyncRulesContent;
   public storage?: SyncRulesBucketStorage;
-  private replicationConnection?: pgwire.PgConnection;
-  private snapshotPromise?: Promise<void>;
+  private settledReplicationPromise?: Promise<PromiseSettledResult<void>>;
 
   /**
    * Tests operating on the wal stream need to configure the stream and manage asynchronous
@@ -64,21 +65,10 @@ export class WalStreamTestContext implements AsyncDisposable {
     await this.dispose();
   }
 
-  /**
-   * Clear any errors from startStream, to allow for a graceful dispose when streaming errors
-   * were expected.
-   */
-  async clearStreamError() {
-    if (this.streamPromise != null) {
-      this.streamPromise = this.streamPromise.catch((e) => {});
-    }
-  }
-
   async dispose() {
     this.abortController.abort();
     try {
-      await this.snapshotPromise;
-      await this.streamPromise;
+      await this.settledReplicationPromise;
       await this.connectionManager.destroy();
       await this.factory?.[Symbol.asyncDispose]();
     } catch (e) {
@@ -108,6 +98,7 @@ export class WalStreamTestContext implements AsyncDisposable {
       updateSyncRulesFromYaml(content, { validate: true, storageVersion: this.storageVersion })
     );
     this.syncRulesId = syncRules.id;
+    this.syncRulesContent = syncRules;
     this.storage = this.factory.getInstance(syncRules);
     return this.storage!;
   }
@@ -119,6 +110,7 @@ export class WalStreamTestContext implements AsyncDisposable {
     }
 
     this.syncRulesId = syncRules.id;
+    this.syncRulesContent = syncRules;
     this.storage = this.factory.getInstance(syncRules);
     return this.storage!;
   }
@@ -130,10 +122,18 @@ export class WalStreamTestContext implements AsyncDisposable {
     }
 
     this.syncRulesId = syncRules.id;
+    this.syncRulesContent = syncRules;
     this.storage = this.factory.getInstance(syncRules);
     return this.storage!;
   }
 
+  private getSyncRulesContent(): storage.PersistedSyncRulesContent {
+    if (this.syncRulesContent == null) {
+      throw new Error('Sync rules not configured - call updateSyncRules() first');
+    }
+    return this.syncRulesContent;
+  }
+
   get walStream() {
     if (this.storage == null) {
       throw new Error('updateSyncRules() first');
@@ -157,55 +157,46 @@ export class WalStreamTestContext implements AsyncDisposable {
    */
   async initializeReplication() {
     await this.replicateSnapshot();
-    this.startStreaming();
     // Make sure we're up to date
     await this.getCheckpoint();
   }
 
+  /**
+   * Replicate the initial snapshot, and start streaming.
+   */
   async replicateSnapshot() {
-    const promise = (async () => {
-      this.replicationConnection = await this.connectionManager.replicationConnection();
-      await this.walStream.initReplication(this.replicationConnection);
-    })();
-    this.snapshotPromise = promise.catch((e) => e);
-    await promise;
-  }
-
-  startStreaming() {
-    if (this.replicationConnection == null) {
-      throw new Error('Call replicateSnapshot() before startStreaming()');
+    // Use a settledPromise to avoid unhandled rejections
+    this.settledReplicationPromise = settledPromise(this.walStream.replicate());
+    try {
+      await Promise.race([unsettledPromise(this.settledReplicationPromise), this.walStream.waitForInitialSnapshot()]);
+    } catch (e) {
+      if (e instanceof ReplicationAbortedError && e.cause != null) {
+        // Edge case for tests: replicate() can throw an error, but we'd receive the ReplicationAbortedError from
+        // waitForInitialSnapshot() first. In that case, prioritize the cause, e.g. MissingReplicationSlotError.
+        // This is not a concern for production use, since we only use waitForInitialSnapshot() in tests.
+        throw e.cause;
+      }
+      throw e;
     }
-    this.streamPromise = this.walStream.streamChanges(this.replicationConnection!);
   }
 
   async getCheckpoint(options?: { timeout?: number }) {
     let checkpoint = await Promise.race([
       getClientCheckpoint(this.pool, this.factory, { timeout: options?.timeout ?? 15_000 }),
-      this.streamPromise
+      unsettledPromise(this.settledReplicationPromise!)
     ]);
     if (checkpoint == null) {
-      // This indicates an issue with the test setup - streamingPromise completed instead
+      // This indicates an issue with the test setup - replicationPromise completed instead
       // of getClientCheckpoint()
-      throw new Error('Test failure - streamingPromise completed');
+      throw new Error('Test failure - replicationPromise completed');
     }
     return checkpoint;
   }
 
-  private resolveBucketName(bucket: string) {
-    if (!this.versionedBuckets || /^\d+#/.test(bucket)) {
-      return bucket;
-    }
-    if (this.syncRulesId == null) {
-      throw new Error('Sync rules not configured - call updateSyncRules() first');
-    }
-    return `${this.syncRulesId}#${bucket}`;
-  }
-
   async getBucketsDataBatch(buckets: Record<string, InternalOpId>, options?: { timeout?: number }) {
     let checkpoint = await this.getCheckpoint(options);
-    const map = new Map<string, InternalOpId>(
-      Object.entries(buckets).map(([bucket, opId]) => [this.resolveBucketName(bucket), opId])
-    );
+    const syncRules = this.getSyncRulesContent();
+    const map = Object.entries(buckets).map(([bucket, start]) => bucketRequest(syncRules, bucket, start));
     return test_utils.fromAsync(this.storage!.getBucketDataBatch(checkpoint, map));
   }
 
@@ -217,9 +208,9 @@ export class WalStreamTestContext implements AsyncDisposable {
     if (typeof start == 'string') {
       start = BigInt(start);
     }
-    const resolvedBucket = this.resolveBucketName(bucket);
+    const syncRules = this.getSyncRulesContent();
     const checkpoint = await this.getCheckpoint(options);
-    const map = new Map<string, InternalOpId>([[resolvedBucket, start]]);
+    let map = [bucketRequest(syncRules, bucket, start)];
     let data: OplogEntry[] = [];
     while (true) {
       const batch = this.storage!.getBucketDataBatch(checkpoint, map);
@@ -229,19 +220,20 @@ export class WalStreamTestContext implements AsyncDisposable {
       if (batches.length == 0 || !batches[0]!.chunkData.has_more) {
         break;
       }
-      map.set(resolvedBucket, BigInt(batches[0]!.chunkData.next_after));
+      map = [bucketRequest(syncRules, bucket, BigInt(batches[0]!.chunkData.next_after))];
     }
     return data;
   }
 
   async getChecksums(buckets: string[], options?: { timeout?: number }) {
     const checkpoint = await this.getCheckpoint(options);
-    const versionedBuckets = buckets.map((bucket) => this.resolveBucketName(bucket));
+    const syncRules = this.getSyncRulesContent();
+    const versionedBuckets = buckets.map((bucket) => bucketRequest(syncRules, bucket, 0n));
     const checksums = await this.storage!.getChecksums(checkpoint, versionedBuckets);
 
     const unversioned = new Map();
     for (let i = 0; i < buckets.length; i++) {
-      unversioned.set(buckets[i], checksums.get(versionedBuckets[i])!);
+      unversioned.set(buckets[i], checksums.get(versionedBuckets[i].bucket)!);
     }
 
     return unversioned;
@@ -260,9 +252,9 @@ export class WalStreamTestContext implements AsyncDisposable {
     if (typeof start == 'string') {
       start = BigInt(start);
     }
-    const resolvedBucket = this.resolveBucketName(bucket);
+    const syncRules = this.getSyncRulesContent();
     const { checkpoint } = await this.storage!.getCheckpoint();
-    const map = new Map<string, InternalOpId>([[resolvedBucket, start]]);
+    const map = [bucketRequest(syncRules, bucket, start)];
     const batch = this.storage!.getBucketDataBatch(checkpoint, map);
     const batches = await test_utils.fromAsync(batch);
     return batches[0]?.chunkData.data ?? [];