WalStream.ts

import * as lib_postgres from '@powersync/lib-service-postgres';
import {
  container,
  DatabaseConnectionError,
  logger as defaultLogger,
  ErrorCode,
  Logger,
  ReplicationAbortedError,
  ReplicationAssertionError
} from '@powersync/lib-services-framework';
import {
  BucketStorageBatch,
  getUuidReplicaIdentityBson,
  MetricsEngine,
  RelationCache,
  ReplicationLagTracker,
  SaveUpdate,
  SourceEntityDescriptor,
  SourceTable,
  storage
} from '@powersync/service-core';
import * as pgwire from '@powersync/service-jpgwire';
import {
  applyValueContext,
  CompatibilityContext,
  HydratedSyncRules,
  SqliteInputRow,
  SqliteInputValue,
  SqliteRow,
  TablePattern,
  ToastableSqliteRow,
  toSyncRulesValue
} from '@powersync/service-sync-rules';

import { ReplicationMetric } from '@powersync/service-types';
import { PostgresTypeResolver } from '../types/resolver.js';
import { MissingReplicationSlotError } from './MissingReplicationSlotError.js';
import { PgManager } from './PgManager.js';
import { getPgOutputRelation, getRelId, referencedColumnTypeIds } from './PgRelation.js';
import { checkSourceConfiguration, checkTableRls, getReplicationIdentityColumns } from './replication-utils.js';
import {
  ChunkedSnapshotQuery,
  IdSnapshotQuery,
  MissingRow,
  PrimaryKeyValue,
  SimpleSnapshotQuery,
  SnapshotQuery
} from './SnapshotQuery.js';
import { computeWalBudgetReport, formatBytes, formatWalBudgetLine } from './wal-budget-utils.js';

export interface WalStreamOptions {
  logger?: Logger;
  connections: PgManager;
  storage: storage.SyncRulesBucketStorage;
  metrics: MetricsEngine;
  abort_signal: AbortSignal;

  /**
   * Override snapshot chunk length (number of rows), for testing.
   *
   * Defaults to 10_000.
   *
   * Note that queries are streamed, so we don't actually keep that much data in memory.
   */
  snapshotChunkLength?: number;

  /**
   * Called after each snapshot chunk is flushed, for testing.
   * This allows tests to perform actions (like generating WAL) synchronously
   * with the snapshot's chunk processing.
   */
  onSnapshotChunkFlushed?: () => Promise<void>;

  /**
   * Interval for slot health checks during snapshot, in milliseconds.
   * Set to 0 for every-chunk checking (useful for testing).
   * Defaults to 120_000 (2 minutes).
   */
  slotHealthCheckIntervalMs?: number;
}

interface InitResult {
  /** True if initial snapshot is not yet done. */
  needsInitialSync: boolean;
  /** True if snapshot must be started from scratch with a new slot. */
  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 WalStream {
  sync_rules: HydratedSyncRules;
  group_id: number;

  connection_id = 1;

  private logger: Logger;

  private readonly storage: storage.SyncRulesBucketStorage;
  private readonly metrics: MetricsEngine;
  private readonly slot_name: string;

  private connections: PgManager;

  private abort_signal: AbortSignal;

  private relationCache = new RelationCache((relation: number | SourceTable) => {
    if (typeof relation == 'number') {
      return relation;
    }
    return relation.objectId!;
  });

  private startedStreaming = false;

  private snapshotChunkLength: number;
  private onSnapshotChunkFlushed?: () => Promise<void>;

  private replicationLag = new ReplicationLagTracker();

  private initialSnapshotPromise: Promise<void> | null = null;

  private slotHealthCheckIntervalMs: number;

  /** Cached max_slot_wal_keep_size in bytes. null = not yet queried. */
  private maxSlotWalKeepSize: number | null = null;
  /** Whether max_slot_wal_keep_size has been queried (distinguishes "not
   *  queried" from "queried and found unlimited"). */
  private maxSlotWalKeepSizeQueried = false;

  /** Previous safe_wal_size sample for rate calculation. */
  private prevWalBudgetSample: { safeWalSize: number; timestamp: number } | null = null;
  /** Timestamp of last slot health check. */
  private lastSlotHealthCheckTime = 0;

  constructor(options: WalStreamOptions) {
    this.logger = options.logger ?? defaultLogger;
    this.storage = options.storage;
    this.metrics = options.metrics;
    this.sync_rules = options.storage.getParsedSyncRules({ defaultSchema: POSTGRES_DEFAULT_SCHEMA });
    this.group_id = options.storage.group_id;
    this.slot_name = options.storage.slot_name;
    this.connections = options.connections;
    this.snapshotChunkLength = options.snapshotChunkLength ?? 10_000;
    this.onSnapshotChunkFlushed = options.onSnapshotChunkFlushed;
    this.slotHealthCheckIntervalMs = options.slotHealthCheckIntervalMs ?? 120_000;

    this.abort_signal = options.abort_signal;
    this.abort_signal.addEventListener(
      'abort',
      () => {
        if (this.startedStreaming) {
          // 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 = sendKeepAlive(this.connections.pool);
          promise.catch((e) => {
            // Failures here are okay - this only speeds up stopping the process.
            this.logger.warn('Failed to ping connection', e);
          });
        } else {
          // If we haven't started streaming yet, it could be due to something like
          // and invalid password. In that case, don't attempt to ping.
        }
      },
      { once: true }
    );
  }

  get stopped() {
    return this.abort_signal.aborted;
  }

  async getQualifiedTableNames(
    batch: storage.BucketStorageBatch,
    db: pgwire.PgConnection,
    tablePattern: TablePattern
  ): Promise<storage.SourceTable[]> {
    const schema = tablePattern.schema;
    if (tablePattern.connectionTag != this.connections.connectionTag) {
      return [];
    }

    let tableRows: any[];
    const prefix = tablePattern.isWildcard ? tablePattern.tablePrefix : undefined;

    {
      let query = `
      SELECT
        c.oid AS relid,
        c.relname AS table_name,
        (SELECT 
          json_agg(DISTINCT a.atttypid)
          FROM pg_attribute a
          WHERE a.attnum > 0 AND NOT a.attisdropped AND a.attrelid = c.oid) 
        AS column_types
      FROM pg_class c
      JOIN pg_namespace n ON n.oid = c.relnamespace
      WHERE n.nspname = $1
      AND c.relkind = 'r'`;

      if (tablePattern.isWildcard) {
        query += ' AND c.relname LIKE $2';
      } else {
        query += ' AND c.relname = $2';
      }

      const result = await db.query({
        statement: query,
        params: [
          { type: 'varchar', value: schema },
          { type: 'varchar', value: tablePattern.tablePattern }
        ]
      });

      tableRows = pgwire.pgwireRows(result);
    }

    let result: storage.SourceTable[] = [];

    for (let row of tableRows) {
      const name = row.table_name as string;
      if (typeof row.relid != 'bigint') {
        throw new ReplicationAssertionError(`Missing relid for ${name}`);
      }
      const relid = Number(row.relid as bigint);

      if (prefix && !name.startsWith(prefix)) {
        continue;
      }

      const rs = await db.query({
        statement: `SELECT 1 FROM pg_publication_tables WHERE pubname = $1 AND schemaname = $2 AND tablename = $3`,
        params: [
          { type: 'varchar', value: PUBLICATION_NAME },
          { type: 'varchar', value: tablePattern.schema },
          { type: 'varchar', value: name }
        ]
      });
      if (rs.rows.length == 0) {
        this.logger.info(`Skipping ${tablePattern.schema}.${name} - not part of ${PUBLICATION_NAME} publication`);
        continue;
      }

      try {
        const result = await checkTableRls(db, relid);
        if (!result.canRead) {
          // We log the message, then continue anyway, since the check does not cover all cases.
          this.logger.warn(result.message!);
        }
      } catch (e) {
        // It's possible that we just don't have permission to access pg_roles - log the error and continue.
        this.logger.warn(`Could not check RLS access for ${tablePattern.schema}.${name}`, e);
      }

      const cresult = await getReplicationIdentityColumns(db, relid);

      const columnTypes = (JSON.parse(row.column_types) as string[]).map((e) => Number(e));
      const table = await this.handleRelation({
        batch,
        descriptor: {
          name,
          schema,
          objectId: relid,
          replicaIdColumns: cresult.replicationColumns
        } as SourceEntityDescriptor,
        snapshot: false,
        referencedTypeIds: columnTypes
      });

      result.push(table);
    }
    return result;
  }

  async initSlot(): Promise<InitResult> {
    await checkSourceConfiguration(this.connections.pool, PUBLICATION_NAME);
    await this.ensureStorageCompatibility();

    const slotName = this.slot_name;

    const status = await this.storage.getStatus();
    const snapshotDone = status.snapshot_done && status.checkpoint_lsn != null;
    if (snapshotDone) {
      // Snapshot is done, but we still need to check the replication slot status
      this.logger.info(`Initial replication already done`);
    }

    // Check if replication slot exists
    const slot = pgwire.pgwireRows(
      await this.connections.pool.query({
        // We specifically want wal_status and invalidation_reason, but it's not available on older versions,
        // so we just query *.
        statement: 'SELECT * FROM pg_replication_slots WHERE slot_name = $1',
        params: [{ type: 'varchar', value: slotName }]
      })
    )[0];

    // Previously we also used pg_catalog.pg_logical_slot_peek_binary_changes to confirm that we can query the slot.
    // However, there were some edge cases where the query times out, repeating the query, ultimately
    // causing high load on the source database and never recovering automatically.
    // We now instead jump straight to replication if the wal_status is not "lost", rather detecting those
    // errors during streaming replication, which is a little more robust.

    // We can have:
    //   1. needsInitialSync: true, lost slot -> MissingReplicationSlotError (starts new sync rules version).
    //      Theoretically we could handle this the same as (2).
    //   2. needsInitialSync: true, no slot -> create new slot
    //   3. needsInitialSync: true, valid slot -> resume initial sync
    //   4. needsInitialSync: false, lost slot -> MissingReplicationSlotError (starts new sync rules version)
    //   5. needsInitialSync: false, no slot -> MissingReplicationSlotError (starts new sync rules version)
    //   6. needsInitialSync: false, valid slot -> resume streaming replication
    // The main advantage of MissingReplicationSlotError are:
    // 1. If there was a complete snapshot already (cases 4/5), users can still sync from that snapshot while
    //    we do the reprocessing under a new slot name.
    // 2. If there was a partial snapshot (case 1), we can start with the new slot faster by not waiting for
    //    the partial data to be cleared.
    if (slot != null) {
      // This checks that the slot is still valid

      // wal_status is present in postgres 13+
      // invalidation_reason is present in postgres 17+
      const lost = slot.wal_status == 'lost';
      if (lost) {
        // Case 1 / 4
        const fixGuidance =
          slot.invalidation_reason === 'idle_timeout'
            ? `Increase idle_replication_slot_timeout on the source database.`
            : `Increase max_slot_wal_keep_size on the source database and delete the existing slot to recover.`;
        throw new MissingReplicationSlotError(
          `[PSYNC_S1146] Replication slot ${slotName} was invalidated ` +
            `(reason: ${slot.invalidation_reason ?? 'unknown'}). ` +
            `${fixGuidance}`,
          {
            walStatus: 'lost',
            phase: snapshotDone ? 'streaming' : 'snapshot',
            invalidationReason: slot.invalidation_reason ?? undefined
          }
        );
      }
      // Case 3 / 6
      return {
        needsInitialSync: !snapshotDone,
        needsNewSlot: false
      };
    } else {
      if (snapshotDone) {
        // Case 5
        // This will create a new slot, while keeping the current sync rules active
        throw new MissingReplicationSlotError(`Replication slot ${slotName} is missing`, {
          walStatus: 'missing',
          phase: 'streaming'
        });
      }
      // Case 2
      // This will clear data (if any) and re-create the same slot
      return { needsInitialSync: true, needsNewSlot: true };
    }
  }

  private async queryMaxSlotWalKeepSize(): Promise<number | null> {
    if (this.maxSlotWalKeepSizeQueried) {
      return this.maxSlotWalKeepSize;
    }
    this.maxSlotWalKeepSizeQueried = true;

    try {
      const rows = pgwire.pgwireRows(
        await this.connections.pool.query({
          statement: `SELECT setting, unit FROM pg_settings WHERE name = 'max_slot_wal_keep_size'`
        })
      );
      if (rows.length === 0) {
        // PG < 13 or setting doesn't exist
        return null;
      }
      const setting = Number(rows[0].setting);
      if (setting < 0) {
        // -1 = unlimited
        this.maxSlotWalKeepSize = null;
        return null;
      }
      // setting is in MB, convert to bytes
      const unit = rows[0].unit;
      const multiplier = unit === 'kB' ? 1024 : unit === '8kB' ? 8192 : 1024 * 1024; // default MB
      this.maxSlotWalKeepSize = setting * multiplier;
      return this.maxSlotWalKeepSize;
    } catch (e) {
      // Non-fatal — budget reporting is best-effort
      this.logger.warn(`Could not query max_slot_wal_keep_size`, e);
      return null;
    }
  }

  private async logWalBudget(slot: { safe_wal_size?: any; wal_status?: string }, now: number): Promise<void> {
    const maxSize = await this.queryMaxSlotWalKeepSize();

    // No limit configured
    if (maxSize == null) {
      this.logger.info(`WAL budget: no limit configured (max_slot_wal_keep_size is unlimited).`);
      return;
    }

    const safeWalSize = slot.safe_wal_size != null ? Number(slot.safe_wal_size) : null;
    if (safeWalSize == null) {
      // safe_wal_size is null on PG < 13, or when no limit is set
      return;
    }

    const report = computeWalBudgetReport({
      safeWalSize,
      maxSize,
      walStatus: slot.wal_status ?? 'unknown',
      prevSample: this.prevWalBudgetSample,
      now
    });

    // Update sample for next rate calculation
    this.prevWalBudgetSample = { safeWalSize, timestamp: now };

    const budgetLine = formatWalBudgetLine(report);

    if (report.isWarning) {
      this.logger.warn(budgetLine);
      this.logger.warn(
        `Replication slot may be invalidated before snapshot completes. ` +
          `Increase max_slot_wal_keep_size on the source database.`
      );
    } else {
      this.logger.info(budgetLine);
    }
  }

  /**
   * Check if the replication slot is still valid. Called after each chunk
   * flush during snapshot to detect slot invalidation early.
   *
   * The query hits pg_replication_slots (shared memory, not a table scan)
   * and costs ~1-2ms per round-trip — negligible next to the per-chunk
   * storage flush.
   */
  private async checkSlotHealth(): Promise<void> {
    // Ensure maxSlotWalKeepSize is populated for diagnostic context in error messages
    await this.queryMaxSlotWalKeepSize();

    const rows = pgwire.pgwireRows(
      await this.connections.pool.query({
        statement: 'SELECT * FROM pg_replication_slots WHERE slot_name = $1',
        params: [{ type: 'varchar', value: this.slot_name }]
      })
    );

    if (rows.length === 0) {
      // Slot row gone from pg_replication_slots — dropped externally.
      // Possible causes: pg_drop_replication_slot() call (operator, management tool, cleanup cron)
      throw new MissingReplicationSlotError(
        `[PSYNC_S1146] Replication slot ${this.slot_name} disappeared during snapshot.`,
        { walStatus: 'missing', phase: 'snapshot' }
      );
    }

    const walStatus = rows[0].wal_status;
    if (walStatus === 'lost') {
      // Postgres marked the slot invalid. Possible invalidation_reason values (PG 14+):
      // - wal_removed: WAL growth exceeded max_slot_wal_keep_size (the primary case)
      // - wal_level_insufficient: wal_level changed away from 'logical'
      // - idle_timeout (PG 18+): idle_replication_slot_timeout expired
      // - rows_removed: catalog rows needed by the slot were vacuumed (safe to retry — fresh slot won't need them)
      throw new MissingReplicationSlotError(
        `[PSYNC_S1146] Replication slot ${this.slot_name} was invalidated during snapshot` +
          `${this.formatWalBudgetContext()}. ` +
          `Increase max_slot_wal_keep_size on the source database.`,
        { walStatus: 'lost', phase: 'snapshot', invalidationReason: rows[0].invalidation_reason ?? undefined }
      );
    }

    await this.logWalBudget(rows[0], performance.now());
  }

  private formatWalBudgetContext(): string {
    if (this.maxSlotWalKeepSize == null) {
      return '';
    }
    return ` (limit: ${formatBytes(this.maxSlotWalKeepSize)})`;
  }

  async estimatedCountNumber(db: pgwire.PgConnection, table: storage.SourceTable): Promise<number> {
    const results = await db.query({
      statement: `SELECT reltuples::bigint AS estimate
FROM   pg_class
WHERE  oid = $1::regclass`,
      params: [{ value: table.qualifiedName, type: 'varchar' }]
    });
    const row = results.rows[0];
    return Number(row?.decodeWithoutCustomTypes(0) ?? -1n);
  }

  /**
   * Start initial replication.
   *
   * If (partial) replication was done before on this slot, this clears the state
   * and starts again from scratch.
   */
  async startInitialReplication(replicationConnection: pgwire.PgConnection, status: InitResult) {
    // If anything here errors, the entire replication process is aborted,
    // and all connections are closed, including this one.
    const db = await this.connections.snapshotConnection();

    const slotName = this.slot_name;

    if (status.needsNewSlot) {
      // This happens when there is no existing replication slot, or if the
      // existing one is unhealthy.
      // In those cases, we have to start replication from scratch.
      // If there is an existing healthy slot, we can skip this and continue
      // initial replication where we left off.
      await this.storage.clear({ signal: this.abort_signal });

      await db.query({
        statement: 'SELECT pg_drop_replication_slot(slot_name) FROM pg_replication_slots WHERE slot_name = $1',
        params: [{ type: 'varchar', value: slotName }]
      });

      // We use the replication connection here, not a pool.
      // The replication slot must be created before we start snapshotting tables.
      await replicationConnection.query(`CREATE_REPLICATION_SLOT ${slotName} LOGICAL pgoutput`);

      this.logger.info(`Created replication slot ${slotName}`);
    }

    await this.initialReplication(db);
  }

  async initialReplication(db: pgwire.PgConnection) {
    const sourceTables = this.sync_rules.getSourceTables();
    const flushResults = await this.storage.startBatch(
      {
        logger: this.logger,
        zeroLSN: ZERO_LSN,
        defaultSchema: POSTGRES_DEFAULT_SCHEMA,
        storeCurrentData: true,
        skipExistingRows: true
      },
      async (batch) => {
        let tablesWithStatus: SourceTable[] = [];
        for (let tablePattern of sourceTables) {
          const tables = await this.getQualifiedTableNames(batch, db, tablePattern);
          // Pre-get counts
          for (let table of tables) {
            if (table.snapshotComplete) {
              this.logger.info(`Skipping ${table.qualifiedName} - snapshot already done`);
              continue;
            }
            const count = await this.estimatedCountNumber(db, table);
            table = await batch.updateTableProgress(table, { totalEstimatedCount: count });
            this.relationCache.update(table);
            tablesWithStatus.push(table);

            this.logger.info(`To replicate: ${table.qualifiedName} ${table.formatSnapshotProgress()}`);
          }
        }

        for (let table of tablesWithStatus) {
          await this.snapshotTableInTx(batch, db, table);
          this.touch();
        }

        // Always commit the initial snapshot at zero.
        // 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 for the snapshot.
        // We could also use the LSN from the last table snapshot.
        const rs = await db.query(`select pg_current_wal_lsn() as lsn`);
        const noCommitBefore = rs.rows[0].decodeWithoutCustomTypes(0);

        await batch.markAllSnapshotDone(noCommitBefore);
        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);

    const lastOp = flushResults?.flushed_op;
    if (lastOp != null) {
      // Populate the cache _after_ initial replication, but _before_ we switch to this sync rules.
      await this.storage.populatePersistentChecksumCache({
        // No checkpoint yet, but we do have the opId.
        maxOpId: lastOp,
        signal: this.abort_signal
      });
    }
  }

  static decodeRow(row: pgwire.PgRow, types: PostgresTypeResolver): SqliteInputRow {
    let result: SqliteInputRow = {};

    row.raw.forEach((rawValue, i) => {
      const column = row.columns[i];
      let mappedValue: SqliteInputValue;

      if (typeof rawValue == 'string') {
        mappedValue = toSyncRulesValue(types.registry.decodeDatabaseValue(rawValue, column.typeOid), false, true);
      } else {
        // Binary format, expose as-is.
        mappedValue = rawValue;
      }

      result[column.name] = mappedValue;
    });
    return result;
  }

  private async snapshotTableInTx(
    batch: storage.BucketStorageBatch,
    db: pgwire.PgConnection,
    table: storage.SourceTable,
    limited?: PrimaryKeyValue[]
  ): Promise<storage.SourceTable> {
    // Note: We use the default "Read Committed" isolation level here, not snapshot isolation.
    // The data may change during the transaction, but that is compensated for in the streaming
    // replication afterwards.
    await db.query('BEGIN');
    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].decodeWithoutCustomTypes(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], tableLsnNotBefore);
      this.relationCache.update(resultTable);
      return resultTable;
    } catch (e) {
      await db.query('ROLLBACK');
      throw e;
    }
  }

  private async snapshotTable(
    batch: storage.BucketStorageBatch,
    db: pgwire.PgConnection,
    table: storage.SourceTable,
    limited?: PrimaryKeyValue[]
  ) {
    let totalEstimatedCount = table.snapshotStatus?.totalEstimatedCount;
    let at = table.snapshotStatus?.replicatedCount ?? 0;
    let lastCountTime = 0;
    let q: SnapshotQuery;
    // We do streaming on two levels:
    // 1. Coarse level: DELCARE CURSOR, FETCH 10000 at a time.
    // 2. Fine level: Stream chunks from each fetch call.
    if (limited) {
      q = new IdSnapshotQuery(db, table, limited);
    } else if (ChunkedSnapshotQuery.supports(table)) {
      // Single primary key - we can use the primary key for chunking
      const orderByKey = table.replicaIdColumns[0];
      q = new ChunkedSnapshotQuery(db, table, this.snapshotChunkLength, table.snapshotStatus?.lastKey ?? null);
      if (table.snapshotStatus?.lastKey != null) {
        this.logger.info(
          `Replicating ${table.qualifiedName} ${table.formatSnapshotProgress()} - resuming from ${orderByKey.name} > ${(q as ChunkedSnapshotQuery).lastKey}`
        );
      } else {
        this.logger.info(`Replicating ${table.qualifiedName} ${table.formatSnapshotProgress()} - resumable`);
      }
    } else {
      // Fallback case - query the entire table
      this.logger.info(`Replicating ${table.qualifiedName} ${table.formatSnapshotProgress()} - not resumable`);
      q = new SimpleSnapshotQuery(db, table, this.snapshotChunkLength);
      at = 0;
    }
    await q.initialize();

    let hasRemainingData = true;
    while (hasRemainingData) {
      // Fetch 10k at a time.
      // The balance here is between latency overhead per FETCH call,
      // and not spending too much time on each FETCH call.
      // We aim for a couple of seconds on each FETCH call.
      const cursor = q.nextChunk();
      hasRemainingData = false;
      // pgwire streams rows in chunks.
      // These chunks can be quite small (as little as 16KB), so we don't flush chunks automatically.
      // There are typically 100-200 rows per chunk.
      for await (let chunk of cursor) {
        if (chunk.tag == 'RowDescription') {
          continue;
        }

        if (chunk.rows.length > 0) {
          hasRemainingData = true;
        }

        for (const rawRow of chunk.rows) {
          const record = this.sync_rules.applyRowContext<never>(WalStream.decodeRow(rawRow, this.connections.types));

          // This auto-flushes when the batch reaches its size limit
          await batch.save({
            tag: storage.SaveOperationTag.INSERT,
            sourceTable: table,
            before: undefined,
            beforeReplicaId: undefined,
            after: record,
            afterReplicaId: getUuidReplicaIdentityBson(record, table.replicaIdColumns)
          });
        }

        at += chunk.rows.length;
        this.metrics.getCounter(ReplicationMetric.ROWS_REPLICATED).add(chunk.rows.length);

        this.touch();
      }

      // Important: flush before marking progress
      await batch.flush();
      if (this.onSnapshotChunkFlushed) {
        await this.onSnapshotChunkFlushed();
      }
      const now = performance.now();
      if (now - this.lastSlotHealthCheckTime >= this.slotHealthCheckIntervalMs) {
        this.lastSlotHealthCheckTime = now;
        await this.checkSlotHealth();
      }
      if (limited == null) {
        let lastKey: Uint8Array | undefined;
        if (q instanceof ChunkedSnapshotQuery) {
          lastKey = q.getLastKeySerialized();
        }
        if (lastCountTime < performance.now() - 10 * 60 * 1000) {
          // Even though we're doing the snapshot inside a transaction, the transaction uses
          // the default "Read Committed" isolation level. This means we can get new data
          // within the transaction, so we re-estimate the count every 10 minutes when replicating
          // large tables.
          totalEstimatedCount = await this.estimatedCountNumber(db, table);
          lastCountTime = performance.now();
        }
        table = await batch.updateTableProgress(table, {
          lastKey: lastKey,
          replicatedCount: at,
          totalEstimatedCount: totalEstimatedCount
        });
        this.relationCache.update(table);

        this.logger.info(`Replicating ${table.qualifiedName} ${table.formatSnapshotProgress()}`);
      } else {
        this.logger.info(`Replicating ${table.qualifiedName} ${at}/${limited.length} for resnapshot`);
      }

      if (this.abort_signal.aborted) {
        // We only abort after flushing
        throw new ReplicationAbortedError(`Initial replication interrupted`);
      }
    }
  }

  async handleRelation(options: {
    batch: storage.BucketStorageBatch;
    descriptor: SourceEntityDescriptor;
    snapshot: boolean;
    referencedTypeIds: number[];
  }) {
    const { batch, descriptor, snapshot, referencedTypeIds } = options;

    if (!descriptor.objectId && typeof descriptor.objectId != 'number') {
      throw new ReplicationAssertionError(`objectId expected, got ${typeof descriptor.objectId}`);
    }
    const result = await this.storage.resolveTable({
      group_id: this.group_id,
      connection_id: this.connection_id,
      connection_tag: this.connections.connectionTag,
      entity_descriptor: descriptor,
      sync_rules: this.sync_rules
    });
    this.relationCache.update(result.table);

    // Drop conflicting tables. This includes for example renamed tables.
    await batch.drop(result.dropTables);

    // Ensure we have a description for custom types referenced in the table.
    await this.connections.types.fetchTypes(referencedTypeIds);

    // Snapshot if:
    // 1. Snapshot is requested (false for initial snapshot, since that process handles it elsewhere)
    // 2. Snapshot is not already done, AND:
    // 3. The table is used in sync rules.
    const shouldSnapshot = snapshot && !result.table.snapshotComplete && result.table.syncAny;

    if (shouldSnapshot) {
      // Truncate this table, in case a previous snapshot was interrupted.
      await batch.truncate([result.table]);

      // Start the snapshot inside a transaction.
      // We use a dedicated connection for this.
      const db = await this.connections.snapshotConnection();
      try {
        const table = await this.snapshotTableInTx(batch, db, result.table);
        // After the table snapshot, we wait for replication to catch up.
        // To make sure there is actually something to replicate, we send a keepalive
        // message.
        await sendKeepAlive(db);
        return table;
      } finally {
        await db.end();
      }
    }

    return result.table;
  }

  /**
   * Process rows that have missing TOAST values.
   *
   * This can happen during edge cases in the chunked intial snapshot process.
   *
   * We handle this similar to an inline table snapshot, but limited to the specific
   * set of rows.
   */
  private async resnapshot(batch: BucketStorageBatch, rows: MissingRow[]) {
    const byTable = new Map<number, MissingRow[]>();
    for (let row of rows) {
      const relId = row.table.objectId as number; // always a number for postgres
      if (!byTable.has(relId)) {
        byTable.set(relId, []);
      }
      byTable.get(relId)!.push(row);
    }
    const db = await this.connections.snapshotConnection();
    try {
      for (let rows of byTable.values()) {
        const table = rows[0].table;
        await this.snapshotTableInTx(
          batch,
          db,
          table,
          rows.map((r) => r.key)
        );
      }
      // Even with resnapshot, we need to wait until we get a new consistent checkpoint
      // after the snapshot, so we need to send a keepalive message.
      await sendKeepAlive(db);
    } finally {
      await db.end();
    }
  }

  private getTable(relationId: number): storage.SourceTable {
    const table = this.relationCache.get(relationId);
    if (table == null) {
      // We should always receive a replication message before the relation is used.
      // If we can't find it, it's a bug.
      throw new ReplicationAssertionError(`Missing relation cache for ${relationId}`);
    }
    return table;
  }

  private syncRulesRecord(row: SqliteInputRow): SqliteRow;
  private syncRulesRecord(row: SqliteInputRow | undefined): SqliteRow | undefined;

  private syncRulesRecord(row: SqliteInputRow | undefined): SqliteRow | undefined {
    if (row == null) {
      return undefined;
    }
    return this.sync_rules.applyRowContext<never>(row);
  }

  private toastableSyncRulesRecord(row: ToastableSqliteRow<SqliteInputValue>): ToastableSqliteRow {
    return this.sync_rules.applyRowContext(row);
  }

  async writeChange(
    batch: storage.BucketStorageBatch,
    msg: pgwire.PgoutputMessage
  ): Promise<storage.FlushedResult | null> {
    if (msg.lsn == null) {
      return null;
    }
    if (msg.tag == 'insert' || msg.tag == 'update' || msg.tag == 'delete') {
      const table = this.getTable(getRelId(msg.relation));
      if (!table.syncAny) {
        this.logger.debug(`Table ${table.qualifiedName} not used in sync rules - skipping`);
        return null;
      }

      if (msg.tag == 'insert') {
        this.metrics.getCounter(ReplicationMetric.ROWS_REPLICATED).add(1);
        const baseRecord = this.syncRulesRecord(this.connections.types.constructAfterRecord(msg));
        return await batch.save({
          tag: storage.SaveOperationTag.INSERT,
          sourceTable: table,
          before: undefined,
          beforeReplicaId: undefined,
          after: baseRecord,
          afterReplicaId: getUuidReplicaIdentityBson(baseRecord, table.replicaIdColumns)
        });
      } else if (msg.tag == 'update') {
        this.metrics.getCounter(ReplicationMetric.ROWS_REPLICATED).add(1);
        // "before" may be null if the replica id columns are unchanged
        // It's fine to treat that the same as an insert.
        const before = this.syncRulesRecord(this.connections.types.constructBeforeRecord(msg));
        const after = this.toastableSyncRulesRecord(this.connections.types.constructAfterRecord(msg));
        return await batch.save({
          tag: storage.SaveOperationTag.UPDATE,
          sourceTable: table,
          before: before,
          beforeReplicaId: before ? getUuidReplicaIdentityBson(before, table.replicaIdColumns) : undefined,
          after: after,
          afterReplicaId: getUuidReplicaIdentityBson(after, table.replicaIdColumns)
        });
      } else if (msg.tag == 'delete') {
        this.metrics.getCounter(ReplicationMetric.ROWS_REPLICATED).add(1);
        const before = this.syncRulesRecord(this.connections.types.constructBeforeRecord(msg)!);

        return await batch.save({
          tag: storage.SaveOperationTag.DELETE,
          sourceTable: table,
          before: before,
          beforeReplicaId: getUuidReplicaIdentityBson(before, table.replicaIdColumns),
          after: undefined,
          afterReplicaId: undefined
        });
      }
    } else if (msg.tag == 'truncate') {
      let tables: storage.SourceTable[] = [];
      for (let relation of msg.relations) {
        const table = this.getTable(getRelId(relation));
        tables.push(table);
      }
      return await batch.truncate(tables);
    }
    return null;
  }

  async replicate() {
    try {
      // If anything errors here, the entire replication process is halted, and
      // all connections automatically closed, including this one.
      this.initialSnapshotPromise = (async () => {
        const initReplicationConnection = await this.connections.replicationConnection();
        await this.initReplication(initReplicationConnection);
        await initReplicationConnection.end();
      })();

      await this.initialSnapshotPromise;

      // At this point, the above connection has often timed out, so we start a new one
      const streamReplicationConnection = await this.connections.replicationConnection();
      await this.streamChanges(streamReplicationConnection);
      await streamReplicationConnection.end();
    } catch (e) {
      await this.storage.reportError(e);
      throw e;
    }
  }

  /**
   * 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) {
      await this.startInitialReplication(replicationConnection, result);
    }
  }

  async streamChanges(replicationConnection: pgwire.PgConnection) {
    try {
      await this.streamChangesInternal(replicationConnection);
    } catch (e) {
      if (isReplicationSlotInvalidError(e)) {
        throw new MissingReplicationSlotError(e.message, { walStatus: 'lost', phase: 'streaming', cause: e });
      }
      throw e;
    }
  }

  private async streamChangesInternal(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: replicationOptions
    });

    this.startedStreaming = true;

    let resnapshot: { table: storage.SourceTable; key: PrimaryKeyValue }[] = [];

    const markRecordUnavailable = (record: SaveUpdate) => {
      if (!IdSnapshotQuery.supports(record.sourceTable)) {
        // If it's not supported, it's also safe to ignore
        return;
      }
      let key: PrimaryKeyValue = {};
      for (let column of record.sourceTable.replicaIdColumns) {
        const name = column.name;
        const value = record.after[name];
        if (value == null) {
          // We don't expect this to actually happen.
          // The key should always be present in the "after" record.
          return;
        }
        // We just need a consistent representation of the primary key, and don't care about fixed quirks.
        key[name] = applyValueContext(value, CompatibilityContext.FULL_BACKWARDS_COMPATIBILITY);
      }
      resnapshot.push({
        table: record.sourceTable,
        key: key
      });
    };

    await this.storage.startBatch(
      {
        logger: this.logger,
        zeroLSN: ZERO_LSN,
        defaultSchema: POSTGRES_DEFAULT_SCHEMA,
        storeCurrentData: true,
        skipExistingRows: false,
        markRecordUnavailable
      },
      async (batch) => {
        // We don't handle any plain keepalive messages while we have transactions.
        // While we have transactions, we use that to advance the position.
        // Replication never starts in the middle of a transaction, so this starts as false.
        let skipKeepalive = false;
        let count = 0;

        for await (const chunk of replicationStream.pgoutputDecode()) {
          this.touch();

          if (this.abort_signal.aborted) {
            break;
          }

          // 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;
          const lastCommit = messages.findLast((msg) => msg.tag == 'commit');

          for (const msg of messages) {
            if (msg.tag == 'relation') {
              await this.handleRelation({
                batch,
                descriptor: getPgOutputRelation(msg),
                snapshot: true,
                referencedTypeIds: referencedColumnTypeIds(msg)
              });
            } else if (msg.tag == 'begin') {
              // This may span multiple transactions in the same chunk, or even across chunks.
              skipKeepalive = true;
              this.replicationLag.trackUncommittedChange(new Date(Number(msg.commitTime / 1000n)));
            } else if (msg.tag == 'commit') {
              this.metrics.getCounter(ReplicationMetric.TRANSACTIONS_REPLICATED).add(1);
              if (msg == lastCommit) {
                // Only commit if this is the last commit in the chunk.
                // This effectively lets us batch multiple transactions within the same chunk
                // into a single flush, increasing throughput for many small transactions.
                skipKeepalive = false;
                // flush() must be before the resnapshot check - that is
                // typically what reports the resnapshot records.
                await batch.flush({ oldestUncommittedChange: this.replicationLag.oldestUncommittedChange });
                // This _must_ be checked after the flush(), and before
                // commit() or ack(). We never persist the resnapshot list,
                // so we have to process it before marking our progress.
                if (resnapshot.length > 0) {
                  await this.resnapshot(batch, resnapshot);
                  resnapshot = [];
                }
                const { checkpointBlocked } = await batch.commit(msg.lsn!, {
                  createEmptyCheckpoints,
                  oldestUncommittedChange: this.replicationLag.oldestUncommittedChange
                });
                await this.ack(msg.lsn!, replicationStream);
                if (!checkpointBlocked) {
                  this.replicationLag.markCommitted();
                }
              }
            } else {
              if (count % 100 == 0) {
                this.logger.info(`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;
              const flushResult = await this.writeChange(batch, msg);
              if (flushResult != null && resnapshot.length > 0) {
                // If we have large transactions, we also need to flush the resnapshot list
                // periodically.
                // TODO: make sure this bit is actually triggered
                await this.resnapshot(batch, resnapshot);
                resnapshot = [];
              }
            }
          }

          if (!skipKeepalive) {
            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.
              const { checkpointBlocked } = await batch.keepalive(chunkLastLsn);
              if (!checkpointBlocked) {
                this.replicationLag.clearUncommittedChange();
              }

              this.replicationLag.markStarted();
            }

            // 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);
          }

          this.metrics.getCounter(ReplicationMetric.CHUNKS_REPLICATED).add(1);
        }
      }
    );
  }

  async ack(lsn: string, replicationStream: pgwire.ReplicationStream) {
    if (lsn == ZERO_LSN) {
      return;
    }

    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,
        oldestUncommittedChange: null
      };
    }

    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,
      oldestUncommittedChange: null
    };
  }

  /**
   * 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;
  }

  getReplicationLagMillis(): number | undefined {
    return this.replicationLag.getLagMillis();
  }

  private touch() {
    container.probes.touch().catch((e) => {
      this.logger.error(`Error touching probe`, e);
    });
  }
}

function isReplicationSlotInvalidError(e: any) {
  // We could access the error code from pgwire using this:
  //   e[Symbol.for('pg.ErrorCode')]
  // However, we typically get a generic code such as 42704 (undefined_object), which does not
  // help much. So we check the actual error message.
  const message = e.message ?? '';

  // Sample: record with incorrect prev-link 10000/10000 at 0/18AB778
  //   Seen during development. Some internal error, fixed by re-creating slot.
  //
  // Sample: publication "powersync" does not exist
  //   Happens when publication deleted or never created.
  //   Slot must be re-created in this case.
  return (
    /incorrect prev-link/.test(message) ||
    /replication slot.*does not exist/.test(message) ||
    /publication.*does not exist/.test(message) ||
    // Postgres 18 - exceeded max_slot_wal_keep_size
    /can no longer access replication slot/.test(message) ||
    // Postgres 17 - exceeded max_slot_wal_keep_size
    /can no longer get changes from replication slot/.test(message)
  );
}