@powersync/service-module-postgres 0.5.1 → 0.6.1

package/src/replication/WalStream.ts

@@ -1,6 +1,8 @@
 import * as lib_postgres from '@powersync/lib-service-postgres';
 import {
   container,
+  DatabaseConnectionError,
+  ErrorCode,
   errors,
   logger,
   ReplicationAbortedError,
@@ -15,10 +17,6 @@ import { PgManager } from './PgManager.js';
 import { getPgOutputRelation, getRelId } from './PgRelation.js';
 import { checkSourceConfiguration, getReplicationIdentityColumns } from './replication-utils.js';
 
-export const ZERO_LSN = '00000000/00000000';
-export const PUBLICATION_NAME = 'powersync';
-export const POSTGRES_DEFAULT_SCHEMA = 'public';
-
 export interface WalStreamOptions {
   connections: PgManager;
   storage: storage.SyncRulesBucketStorage;
@@ -32,6 +30,35 @@ interface InitResult {
   needsNewSlot: boolean;
 }
 
+export const ZERO_LSN = '00000000/00000000';
+export const PUBLICATION_NAME = 'powersync';
+export const POSTGRES_DEFAULT_SCHEMA = 'public';
+
+export const KEEPALIVE_CONTENT = 'ping';
+export const KEEPALIVE_BUFFER = Buffer.from(KEEPALIVE_CONTENT);
+export const KEEPALIVE_STATEMENT: pgwire.Statement = {
+  statement: /* sql */ `
+    SELECT
+      *
+    FROM
+      pg_logical_emit_message(FALSE, 'powersync', $1)
+  `,
+  params: [{ type: 'varchar', value: KEEPALIVE_CONTENT }]
+} as const;
+
+export const isKeepAliveMessage = (msg: pgwire.PgoutputMessage) => {
+  return (
+    msg.tag == 'message' &&
+    msg.prefix == 'powersync' &&
+    msg.content &&
+    Buffer.from(msg.content).equals(KEEPALIVE_BUFFER)
+  );
+};
+
+export const sendKeepAlive = async (db: pgwire.PgClient) => {
+  await lib_postgres.retriedQuery(db, KEEPALIVE_STATEMENT);
+};
+
 export class MissingReplicationSlotError extends Error {
   constructor(message: string) {
     super(message);
@@ -71,10 +98,7 @@ export class WalStream {
           // Ping to speed up cancellation of streaming replication
           // We're not using pg_snapshot here, since it could be in the middle of
           // an initial replication transaction.
-          const promise = lib_postgres.retriedQuery(
-            this.connections.pool,
-            `SELECT * FROM pg_logical_emit_message(false, 'powersync', 'ping')`
-          );
+          const promise = sendKeepAlive(this.connections.pool);
           promise.catch((e) => {
             // Failures here are okay - this only speeds up stopping the process.
             logger.warn('Failed to ping connection', e);
@@ -180,6 +204,7 @@ export class WalStream {
 
   async initSlot(): Promise<InitResult> {
     await checkSourceConfiguration(this.connections.pool, PUBLICATION_NAME);
+    await this.ensureStorageCompatibility();
 
     const slotName = this.slot_name;
 
@@ -382,6 +407,15 @@ WHERE  oid = $1::regclass`,
         await batch.commit(ZERO_LSN);
       }
     );
+    /**
+     * Send a keepalive message after initial replication.
+     * In some edge cases we wait for a keepalive after the initial snapshot.
+     * If we don't explicitly check the contents of keepalive messages then a keepalive is detected
+     * rather quickly after initial replication - perhaps due to other WAL events.
+     * If we do explicitly check the contents of messages, we need an actual keepalive payload in order
+     * to advance the active sync rules LSN.
+     */
+    await sendKeepAlive(db);
   }
 
   static *getQueryData(results: Iterable<DatabaseInputRow>): Generator<SqliteRow> {
@@ -599,13 +633,33 @@ WHERE  oid = $1::regclass`,
   async streamChanges(replicationConnection: pgwire.PgConnection) {
     // When changing any logic here, check /docs/wal-lsns.md.
 
+    const { createEmptyCheckpoints } = await this.ensureStorageCompatibility();
+
+    const replicationOptions: Record<string, string> = {
+      proto_version: '1',
+      publication_names: PUBLICATION_NAME
+    };
+
+    /**
+     * Viewing the contents of logical messages emitted with `pg_logical_emit_message`
+     * is only supported on Postgres >= 14.0.
+     * https://www.postgresql.org/docs/14/protocol-logical-replication.html
+     */
+    const exposesLogicalMessages = await this.checkLogicalMessageSupport();
+    if (exposesLogicalMessages) {
+      /**
+       * Only add this option if the Postgres server supports it.
+       * Adding the option to a server that doesn't support it will throw an exception when starting logical replication.
+       * Error: `unrecognized pgoutput option: messages`
+       */
+      replicationOptions['messages'] = 'true';
+    }
+
     const replicationStream = replicationConnection.logicalReplication({
       slot: this.slot_name,
-      options: {
-        proto_version: '1',
-        publication_names: PUBLICATION_NAME
-      }
+      options: replicationOptions
     });
+
     this.startedStreaming = true;
 
     // Auto-activate as soon as initial replication is done
@@ -628,6 +682,15 @@ WHERE  oid = $1::regclass`,
           // chunkLastLsn may come from normal messages in the chunk,
           // or from a PrimaryKeepalive message.
           const { messages, lastLsn: chunkLastLsn } = chunk;
+
+          /**
+           * We can check if an explicit keepalive was sent if `exposesLogicalMessages == true`.
+           * If we can't check the logical messages, we should assume a keepalive if we
+           * receive an empty array of messages in a replication event.
+           */
+          const assumeKeepAlive = !exposesLogicalMessages;
+          let keepAliveDetected = false;
+
           for (const msg of messages) {
             if (msg.tag == 'relation') {
               await this.handleRelation(batch, getPgOutputRelation(msg), true);
@@ -636,27 +699,44 @@ WHERE  oid = $1::regclass`,
             } else if (msg.tag == 'commit') {
               Metrics.getInstance().transactions_replicated_total.add(1);
               inTx = false;
-              await batch.commit(msg.lsn!);
+              await batch.commit(msg.lsn!, { createEmptyCheckpoints });
               await this.ack(msg.lsn!, replicationStream);
             } else {
               if (count % 100 == 0) {
                 logger.info(`${this.slot_name} replicating op ${count} ${msg.lsn}`);
               }
 
+              /**
+               * If we can see the contents of logical messages, then we can check if a keepalive
+               * message is present. We only perform a keepalive (below) if we explicitly detect a keepalive message.
+               * If we can't see the contents of logical messages, then we should assume a keepalive is required
+               * due to the default value of `assumeKeepalive`.
+               */
+              if (exposesLogicalMessages && isKeepAliveMessage(msg)) {
+                keepAliveDetected = true;
+              }
+
               count += 1;
               await this.writeChange(batch, msg);
             }
           }
 
           if (!inTx) {
-            // In a transaction, we ack and commit according to the transaction progress.
-            // Outside transactions, we use the PrimaryKeepalive messages to advance progress.
-            // Big caveat: This _must not_ be used to skip individual messages, since this LSN
-            // may be in the middle of the next transaction.
-            // It must only be used to associate checkpoints with LSNs.
-            if (await batch.keepalive(chunkLastLsn)) {
-              await this.ack(chunkLastLsn, replicationStream);
+            if (assumeKeepAlive || keepAliveDetected) {
+              // Reset the detection flag.
+              keepAliveDetected = false;
+
+              // In a transaction, we ack and commit according to the transaction progress.
+              // Outside transactions, we use the PrimaryKeepalive messages to advance progress.
+              // Big caveat: This _must not_ be used to skip individual messages, since this LSN
+              // may be in the middle of the next transaction.
+              // It must only be used to associate checkpoints with LSNs.
+              await batch.keepalive(chunkLastLsn);
             }
+
+            // We receive chunks with empty messages often (about each second).
+            // Acknowledging here progresses the slot past these and frees up resources.
+            await this.ack(chunkLastLsn, replicationStream);
           }
 
           Metrics.getInstance().chunks_replicated_total.add(1);
@@ -672,6 +752,53 @@ WHERE  oid = $1::regclass`,
 
     replicationStream.ack(lsn);
   }
+
+  /**
+   * Ensures that the storage is compatible with the replication connection.
+   * @throws {DatabaseConnectionError} If the storage is not compatible with the replication connection.
+   */
+  protected async ensureStorageCompatibility(): Promise<storage.ResolvedBucketBatchCommitOptions> {
+    const supportsLogicalMessages = await this.checkLogicalMessageSupport();
+
+    const storageIdentifier = await this.storage.factory.getSystemIdentifier();
+    if (storageIdentifier.type != lib_postgres.POSTGRES_CONNECTION_TYPE) {
+      return {
+        // Keep the same behaviour as before allowing Postgres storage.
+        createEmptyCheckpoints: true
+      };
+    }
+
+    const parsedStorageIdentifier = lib_postgres.utils.decodePostgresSystemIdentifier(storageIdentifier.id);
+    /**
+     * Check if the same server is being used for both the sync bucket storage and the logical replication.
+     */
+    const replicationIdentifier = await lib_postgres.utils.queryPostgresSystemIdentifier(this.connections.pool);
+
+    if (!supportsLogicalMessages && replicationIdentifier.server_id == parsedStorageIdentifier.server_id) {
+      throw new DatabaseConnectionError(
+        ErrorCode.PSYNC_S1144,
+        `Separate Postgres servers are required for the replication source and sync bucket storage when using Postgres versions below 14.0.`,
+        new Error('Postgres version is below 14')
+      );
+    }
+
+    return {
+      /**
+       * Don't create empty checkpoints if the same Postgres database is used for the data source
+       * and sync bucket storage. Creating empty checkpoints will cause WAL feedback loops.
+       */
+      createEmptyCheckpoints: replicationIdentifier.database_name != parsedStorageIdentifier.database_name
+    };
+  }
+
+  /**
+   * Check if the replication connection Postgres server supports
+   * viewing the contents of logical replication messages.
+   */
+  protected async checkLogicalMessageSupport() {
+    const version = await this.connections.getServerVersion();
+    return version ? version.compareMain('14.0.0') >= 0 : false;
+  }
 }
 
 async function touch() {

package/src/replication/WalStreamReplicationJob.ts

@@ -1,6 +1,6 @@
 import { container } from '@powersync/lib-services-framework';
 import { PgManager } from './PgManager.js';
-import { MissingReplicationSlotError, WalStream } from './WalStream.js';
+import { MissingReplicationSlotError, sendKeepAlive, WalStream } from './WalStream.js';
 
 import { replication } from '@powersync/service-core';
 import { ConnectionManagerFactory } from './ConnectionManagerFactory.js';
@@ -37,7 +37,7 @@ export class WalStreamReplicationJob extends replication.AbstractReplicationJob
    */
   async keepAlive() {
     try {
-      await this.connectionManager.pool.query(`SELECT * FROM pg_logical_emit_message(false, 'powersync', 'ping')`);
+      await sendKeepAlive(this.connectionManager.pool);
     } catch (e) {
       this.logger.warn(`KeepAlive failed, unable to post to WAL`, e);
     }

package/test/src/storage_combination.test.ts

@@ -0,0 +1,35 @@
+import * as postgres_storage from '@powersync/service-module-postgres-storage';
+import { describe, expect, test } from 'vitest';
+import { env } from './env.js';
+import { WalStreamTestContext } from './wal_stream_utils.js';
+
+describe.skipIf(!env.TEST_POSTGRES_STORAGE)('replication storage combination - postgres', function () {
+  test('should allow the same Postgres cluster to be used for data and storage', async () => {
+    // Use the same cluster for the storage as the data source
+    await using context = await WalStreamTestContext.open(
+      postgres_storage.PostgresTestStorageFactoryGenerator({
+        url: env.PG_TEST_URL
+      }),
+      { doNotClear: false }
+    );
+
+    await context.updateSyncRules(/* yaml */
+    ` bucket_definitions:
+        global:
+          data:
+            - SELECT * FROM "test_data" `);
+
+    const { pool, connectionManager } = context;
+
+    const sourceVersion = await connectionManager.getServerVersion();
+
+    await pool.query(`CREATE TABLE test_data(id text primary key, description text, other text)`);
+
+    if (sourceVersion!.compareMain('14.0.0') < 0) {
+      await expect(context.replicateSnapshot()).rejects.toThrow();
+    } else {
+      // Should resolve
+      await context.replicateSnapshot();
+    }
+  });
+});