feat(logical): support moving a single table with COPY_DATA

pgdog/src/admin/copy_data.rs

@@ -4,14 +4,21 @@ use tracing::info;
 
 use crate::api::resharding::ReshardTask;
 use crate::api::run_task;
+use crate::backend::replication::logical::table_move::TableMove;
 use crate::backend::replication::resharding_state::ReshardingState;
 
 use super::prelude::*;
 
+/// What to copy: the tables of an existing publication, or a single table.
+pub(crate) enum CopySource {
+    Publication(String),
+    Table(String),
+}
+
 pub(crate) struct CopyData {
     pub(crate) from_database: String,
     pub(crate) to_database: String,
-    pub(crate) publication: String,
+    pub(crate) source: CopySource,
     pub(crate) replication_slot: Option<String>,
 }
 
@@ -24,45 +31,68 @@ impl Command for CopyData {
     fn parse(sql: &str) -> Result<Self, Error> {
         let parts = sql.split(" ").collect::<Vec<_>>();
 
-        match parts[..] {
-            ["copy_data", from_database, to_database, publication] => Ok(Self {
-                from_database: from_database.to_owned(),
-                to_database: to_database.to_owned(),
-                publication: publication.to_owned(),
-                replication_slot: None,
-            }),
-            [
-                "copy_data",
-                from_database,
-                to_database,
-                publication,
-                replication_slot,
-            ] => Ok(Self {
-                from_database: from_database.to_owned(),
-                to_database: to_database.to_owned(),
-                publication: publication.to_owned(),
-                replication_slot: Some(replication_slot.to_owned()),
-            }),
-            _ => Err(Error::Syntax),
-        }
+        let (from_database, to_database, source, replication_slot) = match parts[..] {
+            ["copy_data", from, to, "table", table] => {
+                (from, to, CopySource::Table(table.to_owned()), None)
+            }
+            ["copy_data", from, to, publication] => (
+                from,
+                to,
+                CopySource::Publication(publication.to_owned()),
+                None,
+            ),
+            ["copy_data", from, to, publication, slot] => (
+                from,
+                to,
+                CopySource::Publication(publication.to_owned()),
+                Some(slot),
+            ),
+            _ => return Err(Error::Syntax),
+        };
+
+        Ok(Self {
+            from_database: from_database.to_owned(),
+            to_database: to_database.to_owned(),
+            source,
+            replication_slot: replication_slot.map(str::to_owned),
+        })
     }
 
     async fn execute(&self) -> Result<Vec<Message>, Error> {
-        info!(
-            r#"copy_data "{}" to "{}", publication="{}""#,
-            self.from_database, self.to_database, self.publication
-        );
+        let (publication, table_move) = match &self.source {
+            CopySource::Publication(publication) => {
+                info!(
+                    r#"copy_data "{}" to "{}", publication="{}""#,
+                    self.from_database, self.to_database, publication
+                );
+                (publication.clone(), None)
+            }
+            CopySource::Table(table) => {
+                info!(
+                    r#"copy_data "{}" to "{}", table="{}""#,
+                    self.from_database, self.to_database, table
+                );
+                let table_move = TableMove::new(table.clone());
+                (table_move.publication().to_owned(), Some(table_move))
+            }
+        };
 
         let state = ReshardingState::builder()
             .source(&self.from_database)
             .destination(&self.to_database)
-            .publication(&self.publication)
+            .publication(&publication)
             .maybe_replication_slot(self.replication_slot.clone())
             .build()?;
 
         let slot_name = state.replication_slot().to_owned();
 
-        let task_id = run_task(ReshardTask::builder().state(state).build()).id();
+        let task_id = run_task(
+            ReshardTask::builder()
+                .state(state)
+                .maybe_table_move(table_move)
+                .build(),
+        )
+        .id();
 
         let mut dr = DataRow::new();
         dr.add(task_id.to_string()).add(slot_name);
@@ -74,3 +104,27 @@ impl Command for CopyData {
         ])
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn parses_publication_and_table_forms() {
+        let cmd = CopyData::parse("copy_data a b pub").unwrap();
+        assert!(matches!(&cmd.source, CopySource::Publication(p) if p == "pub"));
+        assert_eq!(cmd.replication_slot, None);
+
+        let cmd = CopyData::parse("copy_data a b pub slot").unwrap();
+        assert!(matches!(&cmd.source, CopySource::Publication(p) if p == "pub"));
+        assert_eq!(cmd.replication_slot.as_deref(), Some("slot"));
+
+        let cmd = CopyData::parse("copy_data a b table sales.orders").unwrap();
+        assert!(matches!(&cmd.source, CopySource::Table(t) if t == "sales.orders"));
+        assert_eq!(cmd.replication_slot, None);
+        assert_eq!(cmd.from_database, "a");
+        assert_eq!(cmd.to_database, "b");
+
+        assert!(CopyData::parse("copy_data a b").is_err());
+    }
+}

pgdog/src/api/replication.rs

@@ -40,13 +40,24 @@ use pgdog_stats::{
 };
 use tracing::{info, warn};
 
+/// When a replication cuts over to the destination.
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
+pub(crate) enum CutoverMode {
+    /// Wait for an operator `CUTOVER`.
+    #[default]
+    Manual,
+    /// Cut over once the destination has caught up.
+    Auto,
+    /// Never cut over; replicate until the task is stopped. A cutover swaps
+    /// whole databases, which is wrong when moving a single table.
+    Disabled,
+}
+
 #[derive(Debug, bon::Builder)]
 pub(crate) struct ReplicationTask {
     pub(crate) state: ReshardingState,
-    /// Cut over automatically once the destination has caught up, instead
-    /// of waiting for an operator `CUTOVER`.
     #[builder(default)]
-    pub(crate) auto_cutover: bool,
+    pub(crate) cutover: CutoverMode,
     pub(crate) schema_sync: SchemaSyncBuilder,
     pub(crate) validate: bool,
 }
@@ -74,7 +85,7 @@ impl Task for ReplicationTask {
     fn definition(&self) -> impl Into<TaskDefinition> {
         ReplicationDefinition {
             databases: self.state.databases(),
-            auto_cutover: self.auto_cutover,
+            auto_cutover: self.cutover == CutoverMode::Auto,
         }
     }
 
@@ -119,7 +130,7 @@ impl ReplicationTask {
 struct Replication<T> {
     state: ReshardingState,
     schema_sync: SchemaSyncBuilder,
-    auto_cutover: bool,
+    cutover: CutoverMode,
     validation_stage: PostDataValidationStage,
     direction: ReplicationDirection,
     maintenance: MaintenanceMode,
@@ -130,7 +141,7 @@ impl<T: Task<Status: From<ReplicationStatus>>> Replication<T> {
     fn new(task: ReplicationTask) -> Self {
         let ReplicationTask {
             state,
-            auto_cutover,
+            cutover,
             schema_sync,
             validate,
         } = task;
@@ -143,7 +154,7 @@ impl<T: Task<Status: From<ReplicationStatus>>> Replication<T> {
         Self {
             state,
             schema_sync,
-            auto_cutover,
+            cutover,
             validation_stage,
             direction: ReplicationDirection::Forward,
             maintenance: MaintenanceMode::new(),
@@ -159,7 +170,7 @@ impl<T: Task<Status: From<ReplicationStatus>>> Replication<T> {
     /// forward phase returns [`Error::ReplicationAborted`].
     async fn run(mut self, ctx: &TaskContext<T>) -> Result<(), Error> {
         let state = self.state.clone();
-        let result = self.replicate_loop(ctx, self.auto_cutover).await;
+        let result = self.replicate_loop(ctx, self.cutover).await;
         self.resume_traffic();
 
         let cancelled = ctx.cancellation_token().is_cancelled();
@@ -195,19 +206,15 @@ impl<T: Task<Status: From<ReplicationStatus>>> Replication<T> {
     async fn replicate_loop(
         &mut self,
         ctx: &TaskContext<T>,
-        auto_cutover: bool,
+        cutover: CutoverMode,
     ) -> Result<(), Error> {
         info!(
-            "[replication] starting {}, auto_cutover={auto_cutover}, post_data_validation={}",
+            "[replication] starting {}, cutover={cutover:?}, post_data_validation={}",
             self.state.databases(),
             self.validation_stage
         );
-        self.replicate_until_cutover(
-            ctx,
-            auto_cutover,
-            PostDataValidationStage::DuringReplication,
-        )
-        .await?;
+        self.replicate_until_cutover(ctx, cutover, PostDataValidationStage::DuringReplication)
+            .await?;
         self.sync_schema(
             ctx,
             self.schema_sync
@@ -223,8 +230,12 @@ impl<T: Task<Status: From<ReplicationStatus>>> Replication<T> {
         loop {
             self.cutover(ctx).await?;
             self.flip_direction();
-            self.replicate_until_cutover(ctx, false, PostDataValidationStage::AfterCutover)
-                .await?;
+            self.replicate_until_cutover(
+                ctx,
+                CutoverMode::Manual,
+                PostDataValidationStage::AfterCutover,
+            )
+            .await?;
             self.validation_stage = PostDataValidationStage::Off;
             self.sync_schema(
                 ctx,
@@ -245,13 +256,14 @@ impl<T: Task<Status: From<ReplicationStatus>>> Replication<T> {
     async fn replicate_until_cutover(
         &mut self,
         ctx: &TaskContext<T>,
-        auto_cutover: bool,
+        mode: CutoverMode,
         stage: PostDataValidationStage,
     ) -> Result<(), Error> {
         let wait_for_validation = stage != PostDataValidationStage::AfterCutover;
         let direction = self.direction;
         let task_cancel = ctx.cancellation_token();
-        let mut cutover = (!auto_cutover).then(|| CutoverWaiter::register(ctx.root_id()));
+        let mut cutover =
+            (mode == CutoverMode::Manual).then(|| CutoverWaiter::register(ctx.root_id()));
         let mut validation = pin!(self.post_data_validation(ctx, stage).fuse());
 
         loop {
@@ -273,8 +285,10 @@ impl<T: Task<Status: From<ReplicationStatus>>> Replication<T> {
 
             let result = async {
                 let mut cutover_run = pin!(async {
-                    if let Some(cutover) = cutover.as_ref() {
-                        cutover.requested().await;
+                    match (mode, cutover.as_ref()) {
+                        (CutoverMode::Disabled, _) => std::future::pending().await,
+                        (_, Some(cutover)) => cutover.requested().await,
+                        (_, None) => (),
                     }
                     Self::prepare_cutover(ctx, &self.state, &mut self.maintenance, progress.clone())
                         .await