fix: skip replica-identity validation for sync-only data-sync (#1363)

`data_sync` validates that every published table has a usable replica
identity (a primary key or a `REPLICA IDENTITY` index) before it starts
copying. This prevents all data syncs, including `--sync-only`.

That requirement only applies to the streaming phase, which uses the
replica identity to build the per-row `UPDATE`/`DELETE` filters. The
initial `COPY` never uses it, and a `--sync-only` run never streams.

## Change

Add a `require_replica_identity` flag to `Publisher::data_sync`
(threaded through `Orchestrator::data_sync` and `CopyDataTask`).
`ReshardTask` sets it to `!sync_only`:

- **Online** (default): unchanged — tables are still validated up front,
before the potentially multi-hour copy.
- **Sync-only**: validation is skipped, so tables without a primary key
or replica identity can be copied.

Co-authored-by: Claude <noreply@anthropic.com>

Author: bhargavtheertham-cb

pgdog/src/api/copy_data.rs

@@ -15,6 +15,9 @@ use crate::backend::replication::logical::orchestrator::Orchestrator;
 #[display("copy_data {orchestrator}")]
 pub(crate) struct CopyDataTask {
     pub orchestrator: Orchestrator,
+    /// Require a usable replica identity per table. Only streaming needs it,
+    /// so a sync-only migration passes `false`. See `Publisher::data_sync`.
+    pub require_replica_identity: bool,
 }
 
 impl Task for CopyDataTask {
@@ -33,7 +36,9 @@ impl Task for CopyDataTask {
             return Err(Error::DataSyncAborted);
         }
 
-        orchestrator.data_sync(&token).await?;
+        orchestrator
+            .data_sync(&token, self.require_replica_identity)
+            .await?;
 
         Ok(orchestrator)
     }

pgdog/src/api/resharding.rs

@@ -99,7 +99,13 @@ impl Task for ReshardTask {
             if !self.replicate_only {
                 ctx.set_status(ReshardStatus::SyncingData);
                 orchestrator = ctx
-                    .run(CopyDataTask::builder().orchestrator(orchestrator).build())
+                    .run(
+                        CopyDataTask::builder()
+                            .orchestrator(orchestrator)
+                            // Only streaming needs replica identity, not a sync-only copy.
+                            .require_replica_identity(!self.sync_only)
+                            .build(),
+                    )
                     .await?;
             }
 

pgdog/src/backend/replication/logical/orchestrator.rs

@@ -147,14 +147,23 @@ impl Orchestrator {
         Ok(())
     }
 
-    pub(crate) async fn data_sync(&self, cancel: &CancellationToken) -> Result<(), Error> {
+    pub(crate) async fn data_sync(
+        &self,
+        cancel: &CancellationToken,
+        require_replica_identity: bool,
+    ) -> Result<(), Error> {
         let mut publisher = self.publisher.lock().await;
 
         orchestrator_state(OrchestratorState::DataSync);
         // Run data sync for all tables in parallel using multiple replicas,
         // if available.
         publisher
-            .data_sync(&self.source, &self.destination, cancel)
+            .data_sync(
+                &self.source,
+                &self.destination,
+                cancel,
+                require_replica_identity,
+            )
             .await?;
 
         Ok(())

pgdog/src/backend/replication/logical/publisher/publisher_impl.rs

@@ -382,21 +382,25 @@ impl Publisher {
         source: &Cluster,
         dest: &Cluster,
         cancel: &CancellationToken,
+        require_replica_identity: bool,
     ) -> Result<(), Error> {
         // Fetch schema and column metadata first — valid() depends on it.
         self.sync_tables(true, source, dest).await?;
 
-        // Validate all tables support replication before committing to
-        // what can be a multi-hour copy.  A table with no primary key or
-        // unique replica-identity index cannot be replicated correctly.
-        let validation_errors: Vec<_> = self
-            .tables
-            .values()
-            .flat_map(|t| t.iter())
-            .filter_map(|t| t.valid().err())
-            .collect();
+        // Validate replica identity up front, before a potentially multi-hour
+        // copy. Only streaming consumes it (to build the per-row UPDATE/DELETE
+        // filters); the initial COPY does not — so a sync-only run skips the
+        // gate and can copy tables that lack an identity.
+        if require_replica_identity {
+            let validation_errors: Vec<_> = self
+                .tables
+                .values()
+                .flat_map(|t| t.iter())
+                .filter_map(|t| t.valid().err())
+                .collect();
 
-        ensure_validation!(validation_errors);
+            ensure_validation!(validation_errors);
+        }
 
         // Create replication slots only after validation passes — a slot
         // created before valid() would be orphaned on validation errors.
@@ -646,7 +650,7 @@ mod test {
 
         // Validation must fire before the copy begins.
         let result = publisher
-            .data_sync(&source, &dest, &CancellationToken::new())
+            .data_sync(&source, &dest, &CancellationToken::new(), true)
             .await;
 
         let err = result.expect_err("data_sync must fail for a publication with no-pk tables");
@@ -676,6 +680,59 @@ mod test {
         }
     }
 
+    /// A sync-only copy (`require_replica_identity = false`) must not reject a
+    /// table that lacks a replica identity.
+    ///
+    /// Asserted without running the copy: the gate runs before the first
+    /// cancellation check, so a pre-cancelled token returns `DataSyncAborted`
+    /// (aborted at slot creation, past the skipped gate) — a `TableValidation`
+    /// error here would mean the gate wrongly fired.
+    #[tokio::test]
+    async fn data_sync_skips_replica_identity_validation_when_not_required() {
+        crate::logger();
+
+        let mut server = test_replication_server().await;
+        for ddl in &[
+            "CREATE TABLE IF NOT EXISTS publication_test_sync_only_no_pk (data TEXT NOT NULL)",
+            "DROP PUBLICATION IF EXISTS publication_sync_only_no_pk",
+            "CREATE PUBLICATION publication_sync_only_no_pk FOR TABLE publication_test_sync_only_no_pk",
+        ] {
+            server.execute(*ddl).await.unwrap();
+        }
+
+        let source = Cluster::new_test(&config());
+        source.launch();
+        let dest = Cluster::new_test(&config());
+
+        let mut publisher = Publisher::new(
+            "publication_sync_only_no_pk",
+            QueryParserEngine::default(),
+            "sync_only_no_pk_slot".into(),
+        );
+
+        let cancel = CancellationToken::new();
+        cancel.cancel();
+        let result = publisher.data_sync(&source, &dest, &cancel, false).await;
+
+        assert!(
+            matches!(result, Err(Error::DataSyncAborted)),
+            "sync-only copy must skip replica-identity validation and abort at slot \
+             creation, not fail validation; got: {result:?}"
+        );
+        assert!(
+            publisher.slots.is_empty(),
+            "the cancelled token aborts before any slot is created"
+        );
+
+        source.shutdown();
+        for ddl in &[
+            "DROP PUBLICATION IF EXISTS publication_sync_only_no_pk",
+            "DROP TABLE IF EXISTS publication_test_sync_only_no_pk",
+        ] {
+            server.execute(*ddl).await.unwrap();
+        }
+    }
+
     /// `REPLICA IDENTITY NOTHING` must be rejected at `data_sync` time,
     /// before any replication slot is created. This test executes against
     /// a real Postgres instance so it validates the full metadata-fetch + valid() path.
@@ -704,7 +761,7 @@ mod test {
         );
 
         let result = publisher
-            .data_sync(&source, &dest, &CancellationToken::new())
+            .data_sync(&source, &dest, &CancellationToken::new(), true)
             .await;
 
         let err = result.expect_err("data_sync must fail for REPLICA IDENTITY NOTHING table");