Re-apply some postgres changes.

Author: rkistner

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

@@ -144,6 +144,8 @@ export class WalStream {
    */
   private isStartingReplication = true;
 
+  private initialSnapshotPromise: Promise<void> | null = null;
+
   constructor(options: WalStreamOptions) {
     this.logger = options.logger ?? defaultLogger;
     this.storage = options.storage;
@@ -442,6 +444,24 @@ WHERE  oid = $1::regclass`,
         // This makes sure we don't skip any changes applied before starting this snapshot,
         // 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.
+        const rs = await db.query(`select pg_current_wal_lsn() as lsn`);
+        const noCommitBefore = rs.rows[0][0];
+
+        await batch.markAllSnapshotDone(noCommitBefore);
         await batch.commit(ZERO_LSN);
       }
     );
@@ -482,27 +502,11 @@ WHERE  oid = $1::regclass`,
     // replication afterwards.
     await db.query('BEGIN');
     try {
-      let tableLsnNotBefore: string;
       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`);
-      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.markSnapshotDone([table], tableLsnNotBefore);
+      const [resultTable] = await batch.markTableSnapshotDone([table]);
       this.relationCache.update(resultTable);
       return resultTable;
     } catch (e) {
@@ -818,7 +822,8 @@ WHERE  oid = $1::regclass`,
       // If anything errors here, the entire replication process is halted, and
       // all connections automatically closed, including this one.
       const initReplicationConnection = await this.connections.replicationConnection();
-      await this.initReplication(initReplicationConnection);
+      this.initialSnapshotPromise = this.initReplication(initReplicationConnection);
+      await this.initialSnapshotPromise;
       await initReplicationConnection.end();
 
       // At this point, the above connection has often timed out, so we start a new one
@@ -831,6 +836,18 @@ WHERE  oid = $1::regclass`,
     }
   }
 
+  /**
+   * After calling replicate(), call this to wait for the initial snapshot to complete.
+   *
+   * For tests only.
+   */
+  async waitForInitialSnapshot() {
+    if (this.initialSnapshotPromise == null) {
+      throw new ReplicationAssertionError(`Initial snapshot not started yet`);
+    }
+    return this.initialSnapshotPromise;
+  }
+
   async initReplication(replicationConnection: pgwire.PgConnection) {
     const result = await this.initSlot();
     if (result.needsInitialSync) {

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

@@ -1,18 +1,16 @@
+import { WalStream } from '@module/replication/WalStream.js';
+import { PostgresTypeResolver } from '@module/types/resolver.js';
 import * as pgwire from '@powersync/service-jpgwire';
 import {
   applyRowContext,
   CompatibilityContext,
-  SqliteInputRow,
+  CompatibilityEdition,
   DateTimeValue,
-  TimeValue,
-  CompatibilityEdition
+  SqliteInputRow,
+  TimeValue
 } from '@powersync/service-sync-rules';
 import { describe, expect, test } from 'vitest';
 import { clearTestDb, connectPgPool, connectPgWire, TEST_URI } from './util.js';
-import { WalStream } from '@module/replication/WalStream.js';
-import { PostgresTypeResolver } from '@module/types/resolver.js';
-import { CustomTypeRegistry } from '@module/types/registry.js';
-import { PostgresSnapshotter } from '@module/replication/PostgresSnapshotter.js';
 
 describe('pg data types', () => {
   async function setupTable(db: pgwire.PgClient) {
@@ -304,7 +302,7 @@ VALUES(10, ARRAY['null']::TEXT[]);
       await insert(db);
 
       const transformed = [
-        ...PostgresSnapshotter.getQueryData(pgwire.pgwireRows(await db.query(`SELECT * FROM test_data ORDER BY id`)))
+        ...WalStream.getQueryData(pgwire.pgwireRows(await db.query(`SELECT * FROM test_data ORDER BY id`)))
       ];
 
       checkResults(transformed);
@@ -323,7 +321,7 @@ VALUES(10, ARRAY['null']::TEXT[]);
       await insert(db);
 
       const transformed = [
-        ...PostgresSnapshotter.getQueryData(
+        ...WalStream.getQueryData(
           pgwire.pgwireRows(
             await db.query({
               statement: `SELECT * FROM test_data WHERE $1 ORDER BY id`,
@@ -347,9 +345,7 @@ VALUES(10, ARRAY['null']::TEXT[]);
       await insertArrays(db);
 
       const transformed = [
-        ...PostgresSnapshotter.getQueryData(
-          pgwire.pgwireRows(await db.query(`SELECT * FROM test_data_arrays ORDER BY id`))
-        )
+        ...WalStream.getQueryData(pgwire.pgwireRows(await db.query(`SELECT * FROM test_data_arrays ORDER BY id`)))
       ].map((e) => applyRowContext(e, CompatibilityContext.FULL_BACKWARDS_COMPATIBILITY));
 
       checkResultArrays(transformed);
@@ -452,7 +448,7 @@ INSERT INTO test_data(id, time, timestamp, timestamptz) VALUES (1, '17:42:01.12'
 `);
 
       const [row] = [
-        ...PostgresSnapshotter.getQueryData(
+        ...WalStream.getQueryData(
           pgwire.pgwireRows(await db.query(`SELECT time, timestamp, timestamptz FROM test_data`))
         )
       ];

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

@@ -2,7 +2,7 @@ import { PostgresRouteAPIAdapter } from '@module/api/PostgresRouteAPIAdapter.js'
 import * as types from '@module/types/types.js';
 import * as lib_postgres from '@powersync/lib-service-postgres';
 import { logger } from '@powersync/lib-services-framework';
-import { BucketStorageFactory, InternalOpId, TestStorageFactory } from '@powersync/service-core';
+import { BucketStorageFactory, InternalOpId, TestStorageConfig, TestStorageFactory } from '@powersync/service-core';
 import * as pgwire from '@powersync/service-jpgwire';
 import * as mongo_storage from '@powersync/service-module-mongodb-storage';
 import * as postgres_storage from '@powersync/service-module-postgres-storage';
@@ -16,11 +16,11 @@ export const INITIALIZED_MONGO_STORAGE_FACTORY = mongo_storage.test_utils.mongoT
   isCI: env.CI
 });
 
-export const INITIALIZED_POSTGRES_STORAGE_FACTORY = postgres_storage.test_utils.postgresTestStorageFactoryGenerator({
+export const INITIALIZED_POSTGRES_STORAGE_FACTORY = postgres_storage.test_utils.postgresTestSetup({
   url: env.PG_STORAGE_TEST_URL
 });
 
-export function describeWithStorage(options: TestOptions, fn: (factory: TestStorageFactory) => void) {
+export function describeWithStorage(options: TestOptions, fn: (factory: TestStorageConfig) => void) {
   describe.skipIf(!env.TEST_MONGO_STORAGE)(`mongodb storage`, options, function () {
     fn(INITIALIZED_MONGO_STORAGE_FACTORY);
   });