postgres_cdc: fail the snapshot barrier on nack (CON-504)

internal/impl/postgresql/input_pg_stream.go

@@ -418,11 +418,43 @@ type pgStreamInput struct {
 	// acknowledged. The snapshot->stream handoff blocks until it drains so the
 	// replication slot is not promoted before snapshot rows are durable.
 	snapshotAckWG sync.WaitGroup
+	// snapshotNackErr records the first snapshot batch nack. auto_replay_nacks
+	// is user-toggleable, so a nack can be terminal: the barrier must fail
+	// rather than promote the replication slot over undelivered rows. Cleared
+	// per connection attempt (the input outlives reconnects).
+	snapshotNackMu  sync.Mutex
+	snapshotNackErr error
 
 	// IAM authentication fields
 	iamAuthEnabled bool
 }
 
+func (p *pgStreamInput) recordSnapshotNack(err error) {
+	p.snapshotNackMu.Lock()
+	defer p.snapshotNackMu.Unlock()
+	if p.snapshotNackErr == nil {
+		p.snapshotNackErr = err
+	}
+}
+
+func (p *pgStreamInput) snapshotNackError() error {
+	p.snapshotNackMu.Lock()
+	defer p.snapshotNackMu.Unlock()
+	return p.snapshotNackErr
+}
+
+// resetSnapshotGate clears any nack recorded by a previous connection attempt
+// so the promotion barrier judges only the current run: the input outlives
+// reconnects, and a stale error would fail every retry even after a clean
+// re-run. The WaitGroup is deliberately left untouched — batches from a
+// previous attempt that are still in flight can yet be acked or nacked, and
+// both must keep counting.
+func (p *pgStreamInput) resetSnapshotGate() {
+	p.snapshotNackMu.Lock()
+	defer p.snapshotNackMu.Unlock()
+	p.snapshotNackErr = nil
+}
+
 func (p *pgStreamInput) Connect(ctx context.Context) error {
 	// If IAM authentication is enabled, generate a new token
 	if p.iamAuthEnabled && p.streamConfig.RefreshAuthToken != nil {
@@ -431,6 +463,10 @@ func (p *pgStreamInput) Connect(ctx context.Context) error {
 		}
 	}
 
+	// The input outlives reconnects: clear any nack recorded by a previous
+	// snapshot attempt so the promotion barrier judges only this run.
+	p.resetSnapshotGate()
+
 	pgStream, err := pglogicalstream.NewPgStream(ctx, p.streamConfig)
 	if err != nil {
 		return fmt.Errorf("unable to create replication stream: %w", err)
@@ -518,6 +554,16 @@ func (p *pgStreamInput) processStream(pgStream *pglogicalstream.Stream, batcher
 				}()
 				select {
 				case <-drained:
+					if nackErr := p.snapshotNackError(); nackErr != nil {
+						// A snapshot batch was rejected downstream and
+						// auto_replay_nacks may be disabled: promoting the slot
+						// would skip the rejected rows forever. Restart instead;
+						// the temporary slot is discarded and the snapshot
+						// re-runs.
+						p.logger.Errorf("Snapshot batch was rejected downstream, restarting without promoting the replication slot (the snapshot will re-run): %v", nackErr)
+						p.stopSig.TriggerSoftStop()
+						break
+					}
 					pgStream.MarkSnapshotAcknowledged()
 				case <-p.stopSig.SoftStopChan():
 				}
@@ -605,10 +651,22 @@ func (p *pgStreamInput) flushBatch(
 	// in the read loop).
 	isSnapshot := lsn == nil
 
-	ackFn := func(ctx context.Context, _ error) error {
+	ackFn := func(ctx context.Context, err error) error {
 		if isSnapshot {
 			defer p.snapshotAckWG.Done()
 		}
+		if err != nil {
+			// auto_replay_nacks is user-toggleable, so a nack can be terminal.
+			// Never resolve: the checkpoint stays pinned before this batch so
+			// no LSN can be acknowledged past its undelivered rows. Snapshot
+			// nacks additionally fail the promotion barrier so the replication
+			// slot is not promoted and the snapshot re-runs on restart.
+			if isSnapshot {
+				p.recordSnapshotNack(err)
+			}
+			p.logger.Errorf("Batch rejected downstream (snapshot=%v): the checkpoint is now pinned before this batch and the input will stall once checkpoint_limit is reached, unless the batch is redelivered (auto_replay_nacks) or the pipeline restarts: %v", isSnapshot, err)
+			return err
+		}
 		maxOffset := resolveFn()
 		if maxOffset == nil {
 			return nil
@@ -617,7 +675,7 @@ func (p *pgStreamInput) flushBatch(
 		if maxLSN == nil {
 			return nil
 		}
-		if err = pgStream.AckLSN(ctx, *maxLSN); err != nil {
+		if err := pgStream.AckLSN(ctx, *maxLSN); err != nil {
 			return fmt.Errorf("unable to ack LSN to postgres: %w", err)
 		}
 		return nil

internal/impl/postgresql/integration_test.go

@@ -275,6 +275,97 @@ pg_stream:
 	require.NoError(t, streamOut.StopWithin(time.Second*10))
 }
 
+// TestIntegrationPostgresSnapshotNackFailsBarrier verifies that a downstream
+// rejection (nack) of a snapshot batch with auto_replay_nacks disabled fails
+// the promotion barrier instead of being treated like an ack: the replication
+// slot must not be promoted over undelivered rows. The input restarts itself,
+// re-runs the snapshot, and once a clean run is fully acked the slot is
+// promoted and every row has been delivered. See CON-504.
+func TestIntegrationPostgresSnapshotNackFailsBarrier(t *testing.T) {
+	integration.CheckSkip(t)
+	databaseURL, db, err := ResourceWithPostgreSQLVersion(t, "16")
+	require.NoError(t, err)
+
+	const rowCount = 5
+	for i := range rowCount {
+		f := GetFakeFlightRecord()
+		_, err = db.Exec(`INSERT INTO "FlightsCompositePK" ("Seq", "Name", "CreatedAt") VALUES ($1, $2, $3);`, i, f.RealAddress.City, time.Unix(f.CreatedAt, 0).Format(time.RFC3339))
+		require.NoError(t, err)
+	}
+
+	// batching.count == rowCount keeps the whole snapshot in one batch, so a
+	// single nack rejects the entire snapshot delivery.
+	template := fmt.Sprintf(`
+pg_stream:
+    dsn: %s
+    slot_name: test_slot_snapshot_nack_barrier
+    stream_snapshot: true
+    snapshot_batch_size: 1000
+    auto_replay_nacks: false
+    schema: public
+    tables:
+       - '"FlightsCompositePK"'
+    batching:
+      count: %d
+      period: 1h
+`, databaseURL, rowCount)
+
+	// The consumer nacks the FIRST snapshot batch it sees (terminal, since
+	// auto_replay_nacks is off), then acks everything afterwards. The input
+	// must restart without promoting the slot and re-run the snapshot; the
+	// second, clean run promotes and completes.
+	var (
+		mu       sync.Mutex
+		nacked   bool
+		reads    int
+		promoted = func() bool {
+			var n int
+			if err := db.QueryRow(`SELECT count(*) FROM pg_replication_slots WHERE slot_name = 'test_slot_snapshot_nack_barrier'`).Scan(&n); err != nil {
+				return false
+			}
+			return n == 1
+		}
+	)
+	builder := service.NewStreamBuilder()
+	require.NoError(t, builder.SetLoggerYAML(`level: OFF`))
+	require.NoError(t, builder.AddInputYAML(template))
+	require.NoError(t, builder.AddBatchConsumerFunc(func(_ context.Context, mb service.MessageBatch) error {
+		mu.Lock()
+		defer mu.Unlock()
+		if !nacked {
+			nacked = true
+			// The permanent slot must not exist while the snapshot is being
+			// rejected - it is only created by promotion after a fully-acked
+			// snapshot.
+			require.False(t, promoted(), "replication slot must not be promoted before the nacked snapshot is cleanly re-delivered")
+			return errors.New("simulated downstream rejection")
+		}
+		for _, msg := range mb {
+			if op, _ := msg.MetaGet("operation"); op == "read" {
+				reads++
+			}
+		}
+		return nil
+	}))
+	stream, err := builder.Build()
+	require.NoError(t, err)
+	license.InjectTestService(stream.Resources())
+	go func() { _ = stream.Run(t.Context()) }()
+
+	// The nack must not lose anything: the snapshot re-runs and a clean pass
+	// delivers every row, after which the slot is promoted.
+	assert.EventuallyWithT(t, func(c *assert.CollectT) {
+		mu.Lock()
+		r := reads
+		n := nacked
+		mu.Unlock()
+		assert.True(c, n, "the first snapshot batch was never delivered")
+		assert.Equal(c, rowCount, r, "the snapshot should have re-run and re-delivered every row after the nack")
+		assert.True(c, promoted(), "the replication slot should be promoted once a clean snapshot run is fully acked")
+	}, 5*time.Minute, 500*time.Millisecond)
+	require.NoError(t, stream.StopWithin(30*time.Second))
+}
+
 // TestIntegrationPostgresSnapshotAckBarrier verifies that a crash during the
 // snapshot->stream handoff (after snapshot rows are emitted but before they are
 // acknowledged) does not lose data: because the replication slot is only