@powersync/service-module-postgres-storage 0.12.0 → 0.13.1

package/dist/storage/batch/PostgresBucketBatch.js

@@ -1,20 +1,29 @@
-import * as lib_postgres from '@powersync/lib-service-postgres';
 import { BaseObserver, container, ErrorCode, errors, ReplicationAssertionError, ServiceAssertionError, ServiceError } from '@powersync/lib-services-framework';
 import { deserializeReplicaId, storage, utils } from '@powersync/service-core';
 import * as timers from 'timers/promises';
 import * as t from 'ts-codec';
-import { CurrentData } from '../../types/models/CurrentData.js';
 import { models } from '../../types/types.js';
-import { NOTIFICATION_CHANNEL, sql } from '../../utils/db.js';
+import { NOTIFICATION_CHANNEL } from '../../utils/db.js';
 import { pick } from '../../utils/ts-codec.js';
 import { batchCreateCustomWriteCheckpoints } from '../checkpoints/PostgresWriteCheckpointAPI.js';
 import { cacheKey, encodedCacheKey, OperationBatch, RecordOperation } from './OperationBatch.js';
 import { PostgresPersistedBatch } from './PostgresPersistedBatch.js';
+import { bigint } from '../../types/codecs.js';
+import { PostgresCurrentDataStore } from '../current-data-store.js';
+import { postgresTableId } from '../table-id.js';
 /**
  * Intermediate type which helps for only watching the active sync rules
  * via the Postgres NOTIFY protocol.
  */
 const StatefulCheckpoint = models.ActiveCheckpoint.and(t.object({ state: t.Enum(storage.SyncRuleState) }));
+const CheckpointWithStatus = StatefulCheckpoint.and(t.object({
+    snapshot_done: t.boolean,
+    no_checkpoint_before: t.string.or(t.Null),
+    can_checkpoint: t.boolean,
+    keepalive_op: bigint.or(t.Null),
+    new_last_checkpoint: bigint.or(t.Null),
+    created_checkpoint: t.boolean
+}));
 /**
  * 15MB. Currently matches MongoDB.
  * This could be increased in future.
@@ -28,7 +37,6 @@ export class PostgresBucketBatch extends BaseObserver {
     db;
     group_id;
     last_checkpoint_lsn;
-    no_checkpoint_before_lsn;
     persisted_op;
     write_checkpoint_batch;
     sync_rules;
@@ -37,6 +45,8 @@ export class PostgresBucketBatch extends BaseObserver {
     markRecordUnavailable;
     needsActivation = true;
     clearedError = false;
+    storageConfig;
+    currentDataStore;
     constructor(options) {
         super();
         this.options = options;
@@ -44,13 +54,14 @@ export class PostgresBucketBatch extends BaseObserver {
         this.db = options.db;
         this.group_id = options.group_id;
         this.last_checkpoint_lsn = options.last_checkpoint_lsn;
-        this.no_checkpoint_before_lsn = options.no_checkpoint_before_lsn;
         this.resumeFromLsn = options.resumeFromLsn;
         this.write_checkpoint_batch = [];
         this.sync_rules = options.sync_rules;
         this.markRecordUnavailable = options.markRecordUnavailable;
         this.batch = null;
         this.persisted_op = null;
+        this.storageConfig = options.storageConfig;
+        this.currentDataStore = new PostgresCurrentDataStore(this.storageConfig);
         if (options.keep_alive_op) {
             this.persisted_op = options.keep_alive_op;
         }
@@ -58,12 +69,18 @@ export class PostgresBucketBatch extends BaseObserver {
     get lastCheckpointLsn() {
         return this.last_checkpoint_lsn;
     }
-    get noCheckpointBeforeLsn() {
-        return this.no_checkpoint_before_lsn;
-    }
     async [Symbol.asyncDispose]() {
+        if (this.batch != null || this.write_checkpoint_batch.length > 0) {
+            // We don't error here, since:
+            // 1. In error states, this is expected (we can't distinguish between disposing after success or error).
+            // 2. SuppressedError is messy to deal with.
+            this.logger.warn('Disposing writer with unflushed changes');
+        }
         super.clearListeners();
     }
+    async dispose() {
+        await this[Symbol.asyncDispose]();
+    }
     async save(record) {
         // TODO maybe share with abstract class
         const { after, before, sourceTable, tag } = record;
@@ -118,30 +135,22 @@ export class PostgresBucketBatch extends BaseObserver {
         const BATCH_LIMIT = 2000;
         let lastBatchCount = BATCH_LIMIT;
         let processedCount = 0;
-        const codec = pick(models.CurrentData, ['buckets', 'lookups', 'source_key']);
         while (lastBatchCount == BATCH_LIMIT) {
             lastBatchCount = 0;
             await this.withReplicationTransaction(async (db) => {
                 const persistedBatch = new PostgresPersistedBatch({
                     group_id: this.group_id,
+                    storageConfig: this.storageConfig,
                     ...this.options.batch_limits
                 });
-                for await (const rows of db.streamRows(sql `
-          SELECT
-            buckets,
-            lookups,
-            source_key
-          FROM
-            current_data
-          WHERE
-            group_id = ${{ type: 'int4', value: this.group_id }}
-            AND source_table = ${{ type: 'varchar', value: sourceTable.id }}
-          LIMIT
-            ${{ type: 'int4', value: BATCH_LIMIT }}
-        `)) {
+                for await (const rows of this.currentDataStore.streamTruncateRows(db, {
+                    groupId: this.group_id,
+                    sourceTableId: postgresTableId(sourceTable.id),
+                    limit: BATCH_LIMIT
+                })) {
                     lastBatchCount += rows.length;
                     processedCount += rows.length;
-                    const decodedRows = rows.map((row) => codec.decode(row));
+                    const decodedRows = rows.map((row) => this.currentDataStore.decodeTruncateRow(row));
                     for (const value of decodedRows) {
                         const source_key = deserializeReplicaId(value.source_key);
                         persistedBatch.saveBucketData({
@@ -159,7 +168,9 @@ export class PostgresBucketBatch extends BaseObserver {
                         persistedBatch.deleteCurrentData({
                             // This is serialized since we got it from a DB query
                             serialized_source_key: value.source_key,
-                            source_table_id: sourceTable.id
+                            source_table_id: postgresTableId(sourceTable.id),
+                            // No need for soft delete, since this is not streaming replication
+                            soft: false
                         });
                     }
                 }
@@ -226,136 +237,223 @@ export class PostgresBucketBatch extends BaseObserver {
         return { flushed_op: lastOp };
     }
     async commit(lsn, options) {
-        const { createEmptyCheckpoints } = { ...storage.DEFAULT_BUCKET_BATCH_COMMIT_OPTIONS, ...options };
+        const createEmptyCheckpoints = options?.createEmptyCheckpoints ?? true;
         await this.flush();
-        if (this.last_checkpoint_lsn != null && lsn < this.last_checkpoint_lsn) {
-            // When re-applying transactions, don't create a new checkpoint until
-            // we are past the last transaction.
-            this.logger.info(`Re-applied transaction ${lsn} - skipping checkpoint`);
-            // Cannot create a checkpoint yet - return false
-            return false;
-        }
-        if (lsn < this.no_checkpoint_before_lsn) {
-            if (Date.now() - this.lastWaitingLogThrottled > 5_000) {
-                this.logger.info(`Waiting until ${this.no_checkpoint_before_lsn} before creating checkpoint, currently at ${lsn}. Persisted op: ${this.persisted_op}`);
-                this.lastWaitingLogThrottled = Date.now();
-            }
-            // Edge case: During initial replication, we have a no_checkpoint_before_lsn set,
-            // and don't actually commit the snapshot.
-            // The first commit can happen from an implicit keepalive message.
-            // That needs the persisted_op to get an accurate checkpoint, so
-            // we persist that in keepalive_op.
-            await this.db.sql `
-        UPDATE sync_rules
-        SET
-          keepalive_op = ${{ type: 'int8', value: this.persisted_op }}
-        WHERE
-          id = ${{ type: 'int4', value: this.group_id }}
-      `.execute();
-            // Cannot create a checkpoint yet - return false
-            return false;
-        }
-        // Don't create a checkpoint if there were no changes
-        if (!createEmptyCheckpoints && this.persisted_op == null) {
-            // Nothing to commit - return true
-            await this.autoActivate(lsn);
-            return true;
-        }
         const now = new Date().toISOString();
-        const update = {
-            last_checkpoint_lsn: lsn,
-            last_checkpoint_ts: now,
-            last_keepalive_ts: now,
-            snapshot_done: true,
-            last_fatal_error: null,
-            keepalive_op: null
-        };
-        if (this.persisted_op != null) {
-            update.last_checkpoint = this.persisted_op.toString();
-        }
-        const doc = await this.db.sql `
-      UPDATE sync_rules
-      SET
-        keepalive_op = ${{ type: 'int8', value: update.keepalive_op }},
-        last_fatal_error = ${{ type: 'varchar', value: update.last_fatal_error }},
-        snapshot_done = ${{ type: 'bool', value: update.snapshot_done }},
-        snapshot_lsn = NULL,
-        last_keepalive_ts = ${{ type: 1184, value: update.last_keepalive_ts }},
-        last_checkpoint = COALESCE(
-          ${{ type: 'int8', value: update.last_checkpoint }},
-          last_checkpoint
+        const persisted_op = this.persisted_op ?? null;
+        const result = await this.db.sql `
+      WITH
+        selected AS (
+          SELECT
+            id,
+            state,
+            last_checkpoint,
+            last_checkpoint_lsn,
+            snapshot_done,
+            no_checkpoint_before,
+            keepalive_op,
+            (
+              snapshot_done = TRUE
+              AND (
+                last_checkpoint_lsn IS NULL
+                OR last_checkpoint_lsn <= ${{ type: 'varchar', value: lsn }}
+              )
+              AND (
+                no_checkpoint_before IS NULL
+                OR no_checkpoint_before <= ${{ type: 'varchar', value: lsn }}
+              )
+            ) AS can_checkpoint
+          FROM
+            sync_rules
+          WHERE
+            id = ${{ type: 'int4', value: this.group_id }}
+          FOR UPDATE
         ),
-        last_checkpoint_ts = ${{ type: 1184, value: update.last_checkpoint_ts }},
-        last_checkpoint_lsn = ${{ type: 'varchar', value: update.last_checkpoint_lsn }}
-      WHERE
-        id = ${{ type: 'int4', value: this.group_id }}
-      RETURNING
+        computed AS (
+          SELECT
+            selected.*,
+            CASE
+              WHEN selected.can_checkpoint THEN GREATEST(
+                selected.last_checkpoint,
+                ${{ type: 'int8', value: persisted_op }},
+                selected.keepalive_op,
+                0
+              )
+              ELSE selected.last_checkpoint
+            END AS new_last_checkpoint,
+            CASE
+              WHEN selected.can_checkpoint THEN NULL
+              ELSE GREATEST(
+                selected.keepalive_op,
+                ${{ type: 'int8', value: persisted_op }},
+                0
+              )
+            END AS new_keepalive_op
+          FROM
+            selected
+        ),
+        updated AS (
+          UPDATE sync_rules AS sr
+          SET
+            last_checkpoint_lsn = CASE
+              WHEN computed.can_checkpoint THEN ${{ type: 'varchar', value: lsn }}
+              ELSE sr.last_checkpoint_lsn
+            END,
+            last_checkpoint_ts = CASE
+              WHEN computed.can_checkpoint THEN ${{ type: 1184, value: now }}
+              ELSE sr.last_checkpoint_ts
+            END,
+            last_keepalive_ts = ${{ type: 1184, value: now }},
+            last_fatal_error = CASE
+              WHEN computed.can_checkpoint THEN NULL
+              ELSE sr.last_fatal_error
+            END,
+            keepalive_op = computed.new_keepalive_op,
+            last_checkpoint = computed.new_last_checkpoint,
+            snapshot_lsn = CASE
+              WHEN computed.can_checkpoint THEN NULL
+              ELSE sr.snapshot_lsn
+            END
+          FROM
+            computed
+          WHERE
+            sr.id = computed.id
+            AND (
+              sr.keepalive_op IS DISTINCT FROM computed.new_keepalive_op
+              OR sr.last_checkpoint IS DISTINCT FROM computed.new_last_checkpoint
+              OR ${{ type: 'bool', value: createEmptyCheckpoints }}
+            )
+          RETURNING
+            sr.id,
+            sr.state,
+            sr.last_checkpoint,
+            sr.last_checkpoint_lsn,
+            sr.snapshot_done,
+            sr.no_checkpoint_before,
+            computed.can_checkpoint,
+            computed.keepalive_op,
+            computed.new_last_checkpoint
+        )
+      SELECT
         id,
         state,
         last_checkpoint,
-        last_checkpoint_lsn
+        last_checkpoint_lsn,
+        snapshot_done,
+        no_checkpoint_before,
+        can_checkpoint,
+        keepalive_op,
+        new_last_checkpoint,
+        TRUE AS created_checkpoint
+      FROM
+        updated
+      UNION ALL
+      SELECT
+        id,
+        state,
+        new_last_checkpoint AS last_checkpoint,
+        last_checkpoint_lsn,
+        snapshot_done,
+        no_checkpoint_before,
+        can_checkpoint,
+        keepalive_op,
+        new_last_checkpoint,
+        FALSE AS created_checkpoint
+      FROM
+        computed
+      WHERE
+        NOT EXISTS (
+          SELECT
+            1
+          FROM
+            updated
+        )
     `
-            .decoded(StatefulCheckpoint)
+            .decoded(CheckpointWithStatus)
             .first();
+        if (result == null) {
+            throw new ReplicationAssertionError('Failed to update sync_rules during checkpoint');
+        }
+        if (!result.can_checkpoint) {
+            if (Date.now() - this.lastWaitingLogThrottled > 5_000) {
+                this.logger.info(`Waiting before creating checkpoint, currently at ${lsn}. Last op: ${result.keepalive_op}. Current state: ${JSON.stringify({
+                    snapshot_done: result.snapshot_done,
+                    last_checkpoint_lsn: result.last_checkpoint_lsn,
+                    no_checkpoint_before: result.no_checkpoint_before
+                })}`);
+                this.lastWaitingLogThrottled = Date.now();
+            }
+            return { checkpointBlocked: true, checkpointCreated: false };
+        }
+        if (result.created_checkpoint) {
+            this.logger.debug(`Created checkpoint at ${lsn}. Last op: ${result.last_checkpoint}`);
+            if (result.last_checkpoint != null) {
+                await this.currentDataStore.cleanupPendingDeletes(this.db, {
+                    groupId: this.group_id,
+                    lastCheckpoint: result.last_checkpoint
+                });
+            }
+        }
         await this.autoActivate(lsn);
-        await notifySyncRulesUpdate(this.db, doc);
+        await notifySyncRulesUpdate(this.db, {
+            id: result.id,
+            state: result.state,
+            last_checkpoint: result.last_checkpoint,
+            last_checkpoint_lsn: result.last_checkpoint_lsn
+        });
         this.persisted_op = null;
         this.last_checkpoint_lsn = lsn;
-        return true;
+        // Even if created_checkpoint is false, if can_checkpoint is true, we need to return not blocked.
+        return { checkpointBlocked: false, checkpointCreated: result.created_checkpoint };
     }
     async keepalive(lsn) {
-        if (this.last_checkpoint_lsn != null && lsn < this.last_checkpoint_lsn) {
-            // No-op
-            return false;
-        }
-        if (lsn < this.no_checkpoint_before_lsn) {
-            return false;
-        }
-        if (this.persisted_op != null) {
-            // The commit may have been skipped due to "no_checkpoint_before_lsn".
-            // Apply it now if relevant
-            this.logger.info(`Commit due to keepalive at ${lsn} / ${this.persisted_op}`);
-            return await this.commit(lsn);
-        }
-        const updated = await this.db.sql `
+        return await this.commit(lsn, { createEmptyCheckpoints: true });
+    }
+    async setResumeLsn(lsn) {
+        await this.db.sql `
       UPDATE sync_rules
       SET
-        snapshot_done = ${{ type: 'bool', value: true }},
-        snapshot_lsn = NULL,
-        last_checkpoint_lsn = ${{ type: 'varchar', value: lsn }},
-        last_fatal_error = ${{ type: 'varchar', value: null }},
-        last_keepalive_ts = ${{ type: 1184, value: new Date().toISOString() }}
+        snapshot_lsn = ${{ type: 'varchar', value: lsn }}
       WHERE
         id = ${{ type: 'int4', value: this.group_id }}
-      RETURNING
-        id,
-        state,
-        last_checkpoint,
-        last_checkpoint_lsn
-    `
-            .decoded(StatefulCheckpoint)
-            .first();
-        await this.autoActivate(lsn);
-        await notifySyncRulesUpdate(this.db, updated);
-        this.last_checkpoint_lsn = lsn;
-        return true;
+    `.execute();
     }
-    async setResumeLsn(lsn) {
+    async markAllSnapshotDone(no_checkpoint_before_lsn) {
+        await this.db.transaction(async (db) => {
+            await db.sql `
+        UPDATE sync_rules
+        SET
+          snapshot_done = TRUE,
+          last_keepalive_ts = ${{ type: 1184, value: new Date().toISOString() }},
+          no_checkpoint_before = CASE
+            WHEN no_checkpoint_before IS NULL
+            OR no_checkpoint_before < ${{ type: 'varchar', value: no_checkpoint_before_lsn }} THEN ${{
+                type: 'varchar',
+                value: no_checkpoint_before_lsn
+            }}
+            ELSE no_checkpoint_before
+          END
+        WHERE
+          id = ${{ type: 'int4', value: this.group_id }}
+      `.execute();
+        });
+    }
+    async markTableSnapshotRequired(table) {
         await this.db.sql `
       UPDATE sync_rules
       SET
-        snapshot_lsn = ${{ type: 'varchar', value: lsn }}
+        snapshot_done = FALSE
       WHERE
         id = ${{ type: 'int4', value: this.group_id }}
     `.execute();
     }
-    async markSnapshotDone(tables, no_checkpoint_before_lsn) {
+    async markTableSnapshotDone(tables, no_checkpoint_before_lsn) {
         const ids = tables.map((table) => table.id.toString());
         await this.db.transaction(async (db) => {
             await db.sql `
         UPDATE source_tables
         SET
-          snapshot_done = ${{ type: 'bool', value: true }},
+          snapshot_done = TRUE,
           snapshot_total_estimated_count = NULL,
           snapshot_replicated_count = NULL,
           snapshot_last_key = NULL
@@ -367,30 +465,27 @@ export class PostgresBucketBatch extends BaseObserver {
               jsonb_array_elements(${{ type: 'jsonb', value: ids }}) AS value
           );
       `.execute();
-            if (no_checkpoint_before_lsn > this.no_checkpoint_before_lsn) {
-                this.no_checkpoint_before_lsn = no_checkpoint_before_lsn;
+            if (no_checkpoint_before_lsn != null) {
                 await db.sql `
           UPDATE sync_rules
           SET
-            no_checkpoint_before = ${{ type: 'varchar', value: no_checkpoint_before_lsn }},
-            last_keepalive_ts = ${{ type: 1184, value: new Date().toISOString() }}
+            last_keepalive_ts = ${{ type: 1184, value: new Date().toISOString() }},
+            no_checkpoint_before = CASE
+              WHEN no_checkpoint_before IS NULL
+              OR no_checkpoint_before < ${{ type: 'varchar', value: no_checkpoint_before_lsn }} THEN ${{
+                    type: 'varchar',
+                    value: no_checkpoint_before_lsn
+                }}
+              ELSE no_checkpoint_before
+            END
           WHERE
             id = ${{ type: 'int4', value: this.group_id }}
         `.execute();
             }
         });
         return tables.map((table) => {
-            const copy = new storage.SourceTable({
-                id: table.id,
-                connectionTag: table.connectionTag,
-                objectId: table.objectId,
-                schema: table.schema,
-                name: table.name,
-                replicaIdColumns: table.replicaIdColumns,
-                snapshotComplete: table.snapshotComplete
-            });
-            copy.syncData = table.syncData;
-            copy.syncParameters = table.syncParameters;
+            const copy = table.clone();
+            copy.snapshotComplete = true;
             return copy;
         });
     }
@@ -430,7 +525,7 @@ export class PostgresBucketBatch extends BaseObserver {
             // exceeding memory limits.
             const sizeLookups = batch.batch.map((r) => {
                 return {
-                    source_table: r.record.sourceTable.id.toString(),
+                    source_table: postgresTableId(r.record.sourceTable.id),
                     /**
                      * Encode to hex in order to pass a jsonb
                      */
@@ -438,26 +533,10 @@ export class PostgresBucketBatch extends BaseObserver {
                 };
             });
             sizes = new Map();
-            for await (const rows of db.streamRows(lib_postgres.sql `
-        WITH
-          filter_data AS (
-            SELECT
-              decode(FILTER ->> 'source_key', 'hex') AS source_key, -- Decoding from hex to bytea
-              (FILTER ->> 'source_table') AS source_table_id
-            FROM
-              jsonb_array_elements(${{ type: 'jsonb', value: sizeLookups }}::jsonb) AS FILTER
-          )
-        SELECT
-          octet_length(c.data) AS data_size,
-          c.source_table,
-          c.source_key
-        FROM
-          current_data c
-          JOIN filter_data f ON c.source_table = f.source_table_id
-          AND c.source_key = f.source_key
-        WHERE
-          c.group_id = ${{ type: 'int4', value: this.group_id }}
-      `)) {
+            for await (const rows of this.currentDataStore.streamSizeRows(db, {
+                groupId: this.group_id,
+                lookups: sizeLookups
+            })) {
                 for (const row of rows) {
                     const key = cacheKey(row.source_table, row.source_key);
                     sizes.set(key, row.data_size);
@@ -479,48 +558,24 @@ export class PostgresBucketBatch extends BaseObserver {
             }
             const lookups = b.map((r) => {
                 return {
-                    source_table: r.record.sourceTable.id,
+                    source_table: postgresTableId(r.record.sourceTable.id),
                     source_key: storage.serializeReplicaId(r.beforeId).toString('hex')
                 };
             });
             const current_data_lookup = new Map();
-            for await (const currentDataRows of db.streamRows({
-                statement: /* sql */ `
-          SELECT
-            ${this.options.skip_existing_rows ? `c.source_table, c.source_key` : 'c.*'}
-          FROM
-            current_data c
-            JOIN (
-              SELECT
-                decode(FILTER ->> 'source_key', 'hex') AS source_key,
-                FILTER ->> 'source_table' AS source_table_id
-              FROM
-                jsonb_array_elements($1::jsonb) AS FILTER
-            ) f ON c.source_table = f.source_table_id
-            AND c.source_key = f.source_key
-          WHERE
-            c.group_id = $2;
-        `,
-                params: [
-                    {
-                        type: 'jsonb',
-                        value: lookups
-                    },
-                    {
-                        type: 'int4',
-                        value: this.group_id
-                    }
-                ]
+            for await (const currentDataRows of this.currentDataStore.streamLookupRows(db, {
+                groupId: this.group_id,
+                lookups,
+                skipExistingRows: this.options.skip_existing_rows
             })) {
                 for (const row of currentDataRows) {
-                    const decoded = this.options.skip_existing_rows
-                        ? pick(CurrentData, ['source_key', 'source_table']).decode(row)
-                        : CurrentData.decode(row);
+                    const decoded = this.currentDataStore.decodeLookupRow(row, this.options.skip_existing_rows);
                     current_data_lookup.set(encodedCacheKey(decoded.source_table, decoded.source_key), decoded);
                 }
             }
             let persistedBatch = new PostgresPersistedBatch({
                 group_id: this.group_id,
+                storageConfig: this.storageConfig,
                 ...this.options.batch_limits
             });
             for (const op of b) {
@@ -758,17 +813,19 @@ export class PostgresBucketBatch extends BaseObserver {
                 source_key: afterId,
                 group_id: this.group_id,
                 data: afterData,
-                source_table: sourceTable.id,
+                source_table: postgresTableId(sourceTable.id),
                 buckets: newBuckets,
-                lookups: newLookups
+                lookups: newLookups,
+                pending_delete: null
             };
             persistedBatch.upsertCurrentData(result);
         }
         if (afterId == null || !storage.replicaIdEquals(beforeId, afterId)) {
             // Either a delete (afterId == null), or replaced the old replication id
             persistedBatch.deleteCurrentData({
-                source_table_id: record.sourceTable.id,
-                source_key: beforeId
+                source_table_id: postgresTableId(sourceTable.id),
+                source_key: beforeId,
+                soft: true
             });
         }
         return result;
@@ -787,40 +844,45 @@ export class PostgresBucketBatch extends BaseObserver {
         await this.db.transaction(async (db) => {
             const syncRulesRow = await db.sql `
         SELECT
-          state
+          state,
+          snapshot_done
         FROM
           sync_rules
         WHERE
           id = ${{ type: 'int4', value: this.group_id }}
+        FOR NO KEY UPDATE;
       `
-                .decoded(pick(models.SyncRules, ['state']))
+                .decoded(pick(models.SyncRules, ['state', 'snapshot_done']))
                 .first();
-            if (syncRulesRow && syncRulesRow.state == storage.SyncRuleState.PROCESSING) {
+            if (syncRulesRow && syncRulesRow.state == storage.SyncRuleState.PROCESSING && syncRulesRow.snapshot_done) {
                 await db.sql `
           UPDATE sync_rules
           SET
             state = ${{ type: 'varchar', value: storage.SyncRuleState.ACTIVE }}
           WHERE
             id = ${{ type: 'int4', value: this.group_id }}
+        `.execute();
+                await db.sql `
+          UPDATE sync_rules
+          SET
+            state = ${{ type: 'varchar', value: storage.SyncRuleState.STOP }}
+          WHERE
+            (
+              state = ${{ value: storage.SyncRuleState.ACTIVE, type: 'varchar' }}
+              OR state = ${{ value: storage.SyncRuleState.ERRORED, type: 'varchar' }}
+            )
+            AND id != ${{ type: 'int4', value: this.group_id }}
         `.execute();
                 didActivate = true;
+                this.needsActivation = false;
+            }
+            else if (syncRulesRow?.state != storage.SyncRuleState.PROCESSING) {
+                this.needsActivation = false;
             }
-            await db.sql `
-        UPDATE sync_rules
-        SET
-          state = ${{ type: 'varchar', value: storage.SyncRuleState.STOP }}
-        WHERE
-          (
-            state = ${{ value: storage.SyncRuleState.ACTIVE, type: 'varchar' }}
-            OR state = ${{ value: storage.SyncRuleState.ERRORED, type: 'varchar' }}
-          )
-          AND id != ${{ type: 'int4', value: this.group_id }}
-      `.execute();
         });
         if (didActivate) {
             this.logger.info(`Activated new sync rules at ${lsn}`);
         }
-        this.needsActivation = false;
     }
     /**
      * Gets relevant {@link SqlEventDescriptor}s for the given {@link SourceTable}
@@ -831,9 +893,29 @@ export class PostgresBucketBatch extends BaseObserver {
     }
     async withReplicationTransaction(callback) {
         try {
-            return await this.db.transaction(async (db) => {
-                return await callback(db);
-            });
+            // Try for up to a minute
+            const lastTry = Date.now() + 60_000;
+            while (true) {
+                try {
+                    return await this.db.transaction(async (db) => {
+                        // The isolation level is required to protect against concurrent updates to the same data.
+                        // In theory the "select ... for update" locks may be able to protect against this, but we
+                        // still have failing tests if we use that as the only isolation mechanism.
+                        await db.query('SET TRANSACTION ISOLATION LEVEL REPEATABLE READ;');
+                        return await callback(db);
+                    });
+                }
+                catch (err) {
+                    const code = err.cause?.code;
+                    if ((code == '40001' || code == '40P01') && Date.now() < lastTry) {
+                        // Serialization (lock) failure, retry
+                        this.logger.warn(`Serialization failure during replication transaction, retrying: ${err.message}`);
+                        await timers.setTimeout(100 + Math.random() * 200);
+                        continue;
+                    }
+                    throw err;
+                }
+            }
         }
         finally {
             await this.db.sql `