Re-apply more test changes.

Author: rkistner

modules/module-postgres/src/replication/WalStream.ts

@@ -445,19 +445,8 @@ WHERE  oid = $1::regclass`,
         // in the case of snapshot retries.
         // We could alternatively commit at the replication slot LSN.
 
-        // Get the current LSN.
-        // The data will only be consistent once incremental replication has passed that point.
-        // We have to get this LSN _after_ we have finished the table snapshots.
-        //
-        // There are basically two relevant LSNs here:
-        // A: The LSN before the snapshot starts. We don't explicitly record this on the PowerSync side,
-        //    but it is implicitly recorded in the replication slot.
-        // B: The LSN after the table snapshot is complete, which is what we get here.
-        // When we do the snapshot queries, the data that we get back for each chunk could match the state
-        // anywhere between A and B. To actually have a consistent state on our side, we need to:
-        // 1. Complete the snapshot.
-        // 2. Wait until logical replication has caught up with all the change between A and B.
-        // Calling `markSnapshotDone(LSN B)` covers that.
+        // Get the current LSN for hte snapshot.
+        // We could also use the LSN from the last table snapshto.
         const rs = await db.query(`select pg_current_wal_lsn() as lsn`);
         const noCommitBefore = rs.rows[0][0];
 
@@ -504,9 +493,25 @@ WHERE  oid = $1::regclass`,
     try {
       await this.snapshotTable(batch, db, table, limited);
 
+      // Get the current LSN.
+      // The data will only be consistent once incremental replication has passed that point.
+      // We have to get this LSN _after_ we have finished the table snapshot.
+      //
+      // There are basically two relevant LSNs here:
+      // A: The LSN before the snapshot starts. We don't explicitly record this on the PowerSync side,
+      //    but it is implicitly recorded in the replication slot.
+      // B: The LSN after the table snapshot is complete, which is what we get here.
+      // When we do the snapshot queries, the data that we get back for each chunk could match the state
+      // anywhere between A and B. To actually have a consistent state on our side, we need to:
+      // 1. Complete the snapshot.
+      // 2. Wait until logical replication has caught up with all the change between A and B.
+      // Calling `markSnapshotDone(LSN B)` covers that.
+      const rs = await db.query(`select pg_current_wal_lsn() as lsn`);
+      const tableLsnNotBefore = rs.rows[0][0];
+
       // Side note: A ROLLBACK would probably also be fine here, since we only read in this transaction.
       await db.query('COMMIT');
-      const [resultTable] = await batch.markTableSnapshotDone([table]);
+      const [resultTable] = await batch.markTableSnapshotDone([table], tableLsnNotBefore);
       this.relationCache.update(resultTable);
       return resultTable;
     } catch (e) {
@@ -821,10 +826,13 @@ WHERE  oid = $1::regclass`,
     try {
       // If anything errors here, the entire replication process is halted, and
       // all connections automatically closed, including this one.
-      const initReplicationConnection = await this.connections.replicationConnection();
-      this.initialSnapshotPromise = this.initReplication(initReplicationConnection);
+      this.initialSnapshotPromise = (async () => {
+        const initReplicationConnection = await this.connections.replicationConnection();
+        await this.initReplication(initReplicationConnection);
+        await initReplicationConnection.end();
+      })();
+
       await this.initialSnapshotPromise;
-      await initReplicationConnection.end();
 
       // At this point, the above connection has often timed out, so we start a new one
       const streamReplicationConnection = await this.connections.replicationConnection();

modules/module-postgres/test/src/schema_changes.test.ts

@@ -56,10 +56,17 @@ function defineTests(config: storage.TestStorageConfig) {
     // Truncate - order doesn't matter
     expect(data.slice(2, 4).sort(compareIds)).toMatchObject([REMOVE_T1, REMOVE_T2]);
 
-    expect(data.slice(4)).toMatchObject([
+    expect(data.slice(4, 5)).toMatchObject([
       // Snapshot and/or replication insert
       PUT_T3
     ]);
+
+    if (data.length > 5) {
+      expect(data.slice(5)).toMatchObject([
+        // Replicated insert (optional duplication)
+        PUT_T3
+      ]);
+    }
   });
 
   test('add table', async () => {

modules/module-postgres/test/src/wal_stream_utils.ts

@@ -6,21 +6,22 @@ import {
   initializeCoreReplicationMetrics,
   InternalOpId,
   OplogEntry,
+  settledPromise,
   storage,
-  SyncRulesBucketStorage
+  SyncRulesBucketStorage,
+  unsettledPromise
 } from '@powersync/service-core';
 import { 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';
 import { CustomTypeRegistry } from '@module/types/registry.js';
+import { ReplicationAbortedError } from '@powersync/lib-services-framework';
 
 export class WalStreamTestContext implements AsyncDisposable {
   private _walStream?: WalStream;
   private abortController = new AbortController();
-  private streamPromise?: Promise<void>;
   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
@@ -55,21 +56,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) {
@@ -143,36 +133,38 @@ 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;
   }