diff --git a/pgdog/src/admin/copy_data.rs b/pgdog/src/admin/copy_data.rs index 0d9b8b3e0..82b8d0624 100644 --- a/pgdog/src/admin/copy_data.rs +++ b/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, } @@ -24,45 +31,68 @@ impl Command for CopyData { fn parse(sql: &str) -> Result { let parts = sql.split(" ").collect::>(); - 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, 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()); + } +} diff --git a/pgdog/src/api/replication.rs b/pgdog/src/api/replication.rs index 9abac91a5..e1412f3f3 100644 --- a/pgdog/src/api/replication.rs +++ b/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 { ReplicationDefinition { databases: self.state.databases(), - auto_cutover: self.auto_cutover, + auto_cutover: self.cutover == CutoverMode::Auto, } } @@ -119,7 +130,7 @@ impl ReplicationTask { struct Replication { state: ReshardingState, schema_sync: SchemaSyncBuilder, - auto_cutover: bool, + cutover: CutoverMode, validation_stage: PostDataValidationStage, direction: ReplicationDirection, maintenance: MaintenanceMode, @@ -130,7 +141,7 @@ impl>> Replication { fn new(task: ReplicationTask) -> Self { let ReplicationTask { state, - auto_cutover, + cutover, schema_sync, validate, } = task; @@ -143,7 +154,7 @@ impl>> Replication { Self { state, schema_sync, - auto_cutover, + cutover, validation_stage, direction: ReplicationDirection::Forward, maintenance: MaintenanceMode::new(), @@ -159,7 +170,7 @@ impl>> Replication { /// forward phase returns [`Error::ReplicationAborted`]. async fn run(mut self, ctx: &TaskContext) -> 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>> Replication { async fn replicate_loop( &mut self, ctx: &TaskContext, - 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>> Replication { 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>> Replication { async fn replicate_until_cutover( &mut self, ctx: &TaskContext, - 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>> Replication { 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 diff --git a/pgdog/src/api/resharding.rs b/pgdog/src/api/resharding.rs index 03251b32f..c74a0587a 100644 --- a/pgdog/src/api/resharding.rs +++ b/pgdog/src/api/resharding.rs @@ -10,12 +10,13 @@ use std::time::Duration; use tracing::warn; use crate::api::copy_data::CopyDataTask; -use crate::api::replication::ReplicationTask; +use crate::api::replication::{CutoverMode, ReplicationTask}; use crate::api::schema_sync::{SchemaSyncPhase, SchemaSyncTask}; use crate::api::synchronize_tables::SynchronizeTablesTask; use crate::api::task::TaskContext; -use crate::api::{MigrationError, Task}; +use crate::api::{MigrationError, Task, run_task}; use crate::backend::replication::logical::resharding_state::ReshardingState; +use crate::backend::replication::logical::table_move::TableMove; use crate::config::config; use pgdog_stats::{ReshardDefinition, ReshardStatus, TaskDefinition}; @@ -38,6 +39,9 @@ pub(crate) struct ReshardTask { /// waiting for an operator `CUTOVER`. Set by the reshard flow. #[builder(default)] pub(crate) auto_cutover: bool, + /// Move this single table instead of the tables in a user publication. + /// Schema sync covers only this table, and replication never cuts over. + pub(crate) table_move: Option, } impl Task for ReshardTask { @@ -61,31 +65,37 @@ impl Task for ReshardTask { async fn run(self, ctx: TaskContext) -> Result<(), MigrationError> { let cancel = ctx.cancellation_token(); + let cutover = self.cutover_mode(); let mut state = self.state; let schema_sync = SchemaSyncTask::builder() .databases(state.databases()) - .publication(state.publication.clone()); - - // Pre-data schema sync, unless skipped. It runs before any replication - // slots exist, so it stays outside the cleanup guard below. - if !self.skip_schema_sync { - ctx.set_status(ReshardStatus::SchemaSync); - ctx.run( - schema_sync - .clone() - .phase(SchemaSyncPhase::Pre) - .ignore_errors(true) - .build(), - ) - .await?; - - // The pre-data sync changed the destination's schema, so its pools - // reloaded. `SchemaSync::reload_destination` refreshes only its own - // cluster refs, so the state still holds stale ones. - state.reload()?; - } + .publication(state.publication.clone()) + .publication_tables_only(self.table_move.is_some()); + let mut copied = false; let result: Result<(), MigrationError> = Box::pin(async { + if let Some(table_move) = &self.table_move { + table_move.prepare(&state).await?; + } + + // Pre-data schema sync, unless skipped. + if !self.skip_schema_sync { + ctx.set_status(ReshardStatus::SchemaSync); + ctx.run( + schema_sync + .clone() + .phase(SchemaSyncPhase::Pre) + .ignore_errors(true) + .build(), + ) + .await?; + + // The pre-data sync changed the destination's schema, so its pools + // reloaded. `SchemaSync::reload_destination` refreshes only its own + // cluster refs, so the state still holds stale ones. + state.reload()?; + } + // Copy the data, unless replicate-only. if !self.replicate_only { ctx.set_status(ReshardStatus::SyncingData); @@ -96,6 +106,7 @@ impl Task for ReshardTask { .build(), ) .await?; + copied = true; } // Post-data schema sync (secondary indexes, constraints): the @@ -138,7 +149,7 @@ impl Task for ReshardTask { Box::pin( ReplicationTask::builder() .state(state.clone()) - .auto_cutover(self.auto_cutover) + .cutover(cutover) .schema_sync(schema_sync.clone()) .validate(!self.skip_schema_sync && !self.replicate_only) .build() @@ -151,7 +162,8 @@ impl Task for ReshardTask { }) .await; - if result.is_err() || cancel.is_cancelled() { + // A table move drops its publication, so its slots can't be reused. + if result.is_err() || cancel.is_cancelled() || self.table_move.is_some() { // on error or cancellation make sure we drop the slots, // but only the slots we're actually created during this run. // If the slots were created outside or by calling pgdog @@ -166,6 +178,38 @@ impl Task for ReshardTask { state.detach_slots(); } + if let Some(table_move) = &self.table_move { + if copied { + // A table move never cuts over, so run the cutover phase + // here to set its sequences. It's a separate task because + // a stop cancels this one and every subtask. + let sequences = run_task( + schema_sync + .phase(SchemaSyncPhase::Cutover) + .ignore_errors(true) + .build(), + ) + .await; + if let Err(err) = sequences { + warn!("failed to sync sequences after the table move: {err}"); + } + } + table_move.finish(&state, copied).await; + } + result } } + +impl ReshardTask { + /// A table move never cuts over: a cutover swaps whole databases. + fn cutover_mode(&self) -> CutoverMode { + if self.table_move.is_some() { + CutoverMode::Disabled + } else if self.auto_cutover { + CutoverMode::Auto + } else { + CutoverMode::Manual + } + } +} diff --git a/pgdog/src/api/schema_sync.rs b/pgdog/src/api/schema_sync.rs index 18b3b740b..cdbd560ca 100644 --- a/pgdog/src/api/schema_sync.rs +++ b/pgdog/src/api/schema_sync.rs @@ -49,6 +49,8 @@ impl From for SyncState { pub(crate) struct SchemaSyncTask { #[builder(field)] dump: SchemaDump, + #[builder(field)] + publication_tables_only: bool, databases: Databases, publication: String, phase: SchemaSyncPhase, @@ -62,6 +64,14 @@ pub(crate) type SchemaSyncBuilder = SchemaSyncTaskBuilder< schema_sync_task_builder::SetPublication, >; +impl SchemaSyncTaskBuilder { + /// Sync only the publication's tables instead of the whole source schema. + pub(crate) fn publication_tables_only(mut self, publication_tables_only: bool) -> Self { + self.publication_tables_only = publication_tables_only; + self + } +} + impl Task for SchemaSyncTask { type Status = SchemaSyncStatus; type Output = (); @@ -89,6 +99,7 @@ impl Task for SchemaSyncTask { &self.databases.source, &self.databases.destination, &self.publication, + self.publication_tables_only, )?; ctx.set_status(SchemaSyncStatus::LoadingSchema); diff --git a/pgdog/src/backend/replication/logical/error.rs b/pgdog/src/backend/replication/logical/error.rs index 67eaa314c..6d870cfdc 100644 --- a/pgdog/src/backend/replication/logical/error.rs +++ b/pgdog/src/backend/replication/logical/error.rs @@ -263,6 +263,18 @@ pub(crate) enum Error { op: &'static str, }, + #[error("source and destination are the same database \"{0}\"")] + SameDatabase(String), + + #[error( + "table {table} references tables that don't exist on destination shard {shard}: {missing}" + )] + ReferencedTablesMissing { + table: String, + shard: usize, + missing: String, + }, + #[error("2pc commit failed for {transaction}: {source}")] TwoPcCleanupPending { transaction: TwoPcTransaction, diff --git a/pgdog/src/backend/replication/logical/mod.rs b/pgdog/src/backend/replication/logical/mod.rs index 1a3f41f93..45f8405ff 100644 --- a/pgdog/src/backend/replication/logical/mod.rs +++ b/pgdog/src/backend/replication/logical/mod.rs @@ -5,6 +5,7 @@ pub(crate) mod publisher; pub(crate) mod resharding_state; pub(crate) mod schema_sync; pub(crate) mod subscriber; +pub(crate) mod table_move; pub(crate) mod tables_sync; pub(crate) use copy_statement::CopyStatement; diff --git a/pgdog/src/backend/replication/logical/schema_sync.rs b/pgdog/src/backend/replication/logical/schema_sync.rs index c345916d8..b9a4325df 100644 --- a/pgdog/src/backend/replication/logical/schema_sync.rs +++ b/pgdog/src/backend/replication/logical/schema_sync.rs @@ -22,6 +22,7 @@ pub(crate) struct SchemaSync { source: Cluster, destination: Cluster, publication: String, + publication_tables_only: bool, } impl SchemaSync { @@ -29,11 +30,13 @@ impl SchemaSync { source: &str, destination: &str, publication: &str, + publication_tables_only: bool, ) -> Result { Ok(Self { source: databases().schema_owner(source)?, destination: databases().schema_owner(destination)?, publication: publication.to_owned(), + publication_tables_only, }) } @@ -70,9 +73,13 @@ impl SchemaSync { .min() .unwrap_or(0); Ok(Arc::new( - PgDump::new(&self.source, &self.publication) - .dump(destination_version) - .await?, + PgDump::new( + &self.source, + &self.publication, + self.publication_tables_only, + ) + .dump(destination_version) + .await?, )) } diff --git a/pgdog/src/backend/replication/logical/table_move.rs b/pgdog/src/backend/replication/logical/table_move.rs new file mode 100644 index 000000000..db7ae5be8 --- /dev/null +++ b/pgdog/src/backend/replication/logical/table_move.rs @@ -0,0 +1,291 @@ +//! Move a single table between databases. +//! +//! Reuses the resharding machinery: the table is wrapped in a publication +//! that PgDog creates and drops, then copied and replicated until the task +//! is stopped. There is no cutover; the application switches databases itself. + +use std::collections::HashSet; + +use tracing::{info, warn}; + +use super::Error; +use super::resharding_state::ReshardingState; +use crate::backend::Cluster; +use crate::backend::pool::Request; +use crate::net::DataRow; +use crate::util::sql::quote_literal; +use crate::util::{escape_identifier, random_string}; + +/// A single table moved between databases. +#[derive(Debug, Clone)] +pub(crate) struct TableMove { + /// As the user wrote it; Postgres resolves it like any table name. + table: String, + publication: String, +} + +/// The moved table as resolved on the source. +struct SourceTable { + /// Quoted `schema.table`. + name: String, + /// Quoted schema. + schema: String, + /// Quoted names of the tables it references with a foreign key. + referenced: Vec, +} + +impl TableMove { + pub(crate) fn new(table: String) -> Self { + Self { + table, + publication: format!("__pgdog_move_{}", random_string(19).to_lowercase()), + } + } + + pub(crate) fn publication(&self) -> &str { + &self.publication + } + + /// Validate the move, create the table's schema on the destination + /// and the publication on every source shard. + pub(crate) async fn prepare(&self, state: &ReshardingState) -> Result<(), Error> { + let source = &state.source.identifier().database; + if *source == state.destination.identifier().database { + return Err(Error::SameDatabase(source.clone())); + } + + let table = self.resolve(&state.source).await?; + if let Some((shard, missing)) = + missing_tables(&state.destination, &table.referenced).await? + { + return Err(Error::ReferencedTablesMissing { + table: table.name, + shard, + missing: missing.join(", "), + }); + } + + // `pg_dump --table` doesn't dump the schema the table is in. + let create_schema = format!("CREATE SCHEMA IF NOT EXISTS {}", table.schema); + for shard in state.destination.shards() { + let mut server = shard.primary(&Request::default()).await?; + server.execute_checked(create_schema.as_str()).await?; + } + + self.create_publication(&state.source, &table.name).await?; + + info!( + "moving table {} [{}], publication \"{}\"", + table.name, state, self.publication + ); + + Ok(()) + } + + /// Drop the publication. Never touches the source table itself. + pub(crate) async fn finish(&self, state: &ReshardingState, copied: bool) { + if copied { + info!( + "table \"{}\" is left in place on \"{}\"; drop it once the application uses \"{}\"", + self.table, + state.source.identifier().database, + state.destination.identifier().database + ); + } + + self.drop_publication(&state.source).await; + } + + async fn resolve(&self, source: &Cluster) -> Result { + let mut server = source.primary(0, &Request::default()).await?; + let rows: Vec = server + .fetch_all(format!( + "SELECT DISTINCT format('%I.%I', n.nspname, c.relname), format('%I', n.nspname), \ + CASE WHEN r.oid IS NULL THEN '' ELSE format('%I.%I', rn.nspname, r.relname) END \ + FROM pg_class c \ + JOIN pg_namespace n ON n.oid = c.relnamespace \ + LEFT JOIN pg_constraint con \ + ON con.conrelid = c.oid AND con.contype = 'f' AND con.confrelid <> c.oid \ + LEFT JOIN pg_class r ON r.oid = con.confrelid \ + LEFT JOIN pg_namespace rn ON rn.oid = r.relnamespace \ + WHERE c.oid = {}::regclass", + quote_literal(&self.table) + )) + .await?; + + let first = rows.first().ok_or(Error::MissingData)?; + Ok(SourceTable { + name: first.get_text(0).ok_or(Error::MissingData)?, + schema: first.get_text(1).ok_or(Error::MissingData)?, + referenced: rows + .iter() + .filter_map(|row| row.get_text(2)) + .filter(|name| !name.is_empty()) + .collect(), + }) + } + + async fn create_publication(&self, source: &Cluster, table: &str) -> Result<(), Error> { + let sql = format!( + "CREATE PUBLICATION \"{}\" FOR TABLE {table}", + escape_identifier(&self.publication), + ); + + // Publications are per database, and shards can share one. + let mut created = HashSet::new(); + for shard in source.shards() { + let mut server = shard.primary(&Request::default()).await?; + let addr = server.addr(); + if created.insert((addr.host.clone(), addr.port, addr.database_name.clone())) { + server.execute_checked(sql.as_str()).await?; + } + } + + Ok(()) + } + + async fn drop_publication(&self, source: &Cluster) { + let sql = format!( + "DROP PUBLICATION IF EXISTS \"{}\"", + escape_identifier(&self.publication) + ); + + for shard in source.shards() { + let result: Result<(), Error> = async { + let mut server = shard.primary(&Request::default()).await?; + server.execute_checked(sql.as_str()).await?; + Ok(()) + } + .await; + + if let Err(err) = result { + warn!( + "failed to drop publication \"{}\" on shard {}: {err}", + self.publication, + shard.number() + ); + } + } + } +} + +/// The first destination shard missing any of `tables`, with the missing ones. +async fn missing_tables( + dest: &Cluster, + tables: &[String], +) -> Result)>, Error> { + if tables.is_empty() { + return Ok(None); + } + + let names = tables + .iter() + .map(|name| quote_literal(name)) + .collect::>() + .join(", "); + let sql = format!( + "SELECT name FROM unnest(ARRAY[{names}]::text[]) AS name WHERE to_regclass(name) IS NULL" + ); + + for shard in dest.shards() { + let mut server = shard.primary(&Request::default()).await?; + let missing: Vec = server.fetch_all(sql.as_str()).await?; + if !missing.is_empty() { + return Ok(Some((shard.number(), missing))); + } + } + + Ok(None) +} + +#[cfg(test)] +mod test { + use super::*; + use crate::backend::server::test::test_server; + use crate::config::config; + + #[tokio::test] + async fn resolves_the_table_and_finds_missing_referenced_tables() { + crate::logger(); + let mut server = test_server().await; + for ddl in [ + "DROP TABLE IF EXISTS pgdog.table_move_child, pgdog.table_move_parent", + "CREATE TABLE pgdog.table_move_parent (id BIGINT PRIMARY KEY)", + "CREATE TABLE pgdog.table_move_child (id BIGINT PRIMARY KEY, \ + parent_id BIGINT REFERENCES pgdog.table_move_parent(id), \ + self_id BIGINT REFERENCES pgdog.table_move_child(id))", + ] { + server.execute_checked(ddl).await.unwrap(); + } + + let cluster = Cluster::new_test(&config()); + cluster.launch(); + + // Unquoted names fold to lowercase, as in SQL. + let child = TableMove::new("PGDOG.Table_Move_Child".into()) + .resolve(&cluster) + .await + .unwrap(); + assert_eq!(child.name, "pgdog.table_move_child"); + assert_eq!(child.schema, "pgdog"); + assert_eq!(child.referenced, ["pgdog.table_move_parent"]); + assert_eq!( + missing_tables(&cluster, &child.referenced).await.unwrap(), + None + ); + assert_eq!( + missing_tables( + &cluster, + &[ + "pgdog.table_move_parent".to_owned(), + "pgdog.table_move_absent".to_owned() + ] + ) + .await + .unwrap(), + Some((0, vec!["pgdog.table_move_absent".to_owned()])) + ); + + let err = TableMove::new("pgdog.table_move_absent".into()) + .resolve(&cluster) + .await + .err() + .expect("table is missing on the source"); + assert!(err.to_string().contains("does not exist"), "{err}"); + + cluster.shutdown(); + server + .execute_checked("DROP TABLE pgdog.table_move_child, pgdog.table_move_parent") + .await + .unwrap(); + } + + #[tokio::test] + async fn publication_is_created_once_when_shards_share_a_database() { + crate::logger(); + let mut server = test_server().await; + server + .execute_checked( + "CREATE TABLE IF NOT EXISTS pgdog.table_move_publication (id BIGINT PRIMARY KEY)", + ) + .await + .unwrap(); + + // Both test shards point at the same database. + let cluster = Cluster::new_test(&config()); + cluster.launch(); + + let table_move = TableMove::new("pgdog.table_move_publication".into()); + table_move + .create_publication(&cluster, "pgdog.table_move_publication") + .await + .unwrap(); + table_move.drop_publication(&cluster).await; + + cluster.shutdown(); + server + .execute_checked("DROP TABLE pgdog.table_move_publication") + .await + .unwrap(); + } +} diff --git a/pgdog/src/backend/replication/tests.rs b/pgdog/src/backend/replication/tests.rs index 8081b4bb2..5a7897077 100644 --- a/pgdog/src/backend/replication/tests.rs +++ b/pgdog/src/backend/replication/tests.rs @@ -10,11 +10,12 @@ use super::logical::Error; use super::logical::publisher::Table; use super::logical::publisher::replication_progress::ReplicationProgress; use super::logical::resharding_state::ReshardingState; +use super::logical::table_move::TableMove; use crate::{ api::{ MigrationError, copy_data::{CopyDataTask, TableDataSyncTask}, - replication::{ReplicationClusterStop, ReplicationClusterTask}, + replication::{ReplicationClusterStop, ReplicationClusterTask, ReplicationTask}, resharding::ReshardTask, run_task, schema_sync::{SchemaSyncPhase, SchemaSyncTask}, @@ -968,3 +969,214 @@ async fn post_data_fails_when_index_rejects_copied_rows() -> Result<(), Box Result<(), Box> { + crate::logger(); + let schema = "table_move_e2e"; + let destination = "table_move_e2e_dest"; + let original_config = config(); + let mut admin = test_server().await; + let result = async { + setup_replication_test(&mut admin, schema, destination).await?; + let source = databases::databases().schema_owner(schema)?; + let dest = databases::databases().schema_owner(destination)?; + let mut server = source.primary(0, &Request::default()).await?; + server + .execute_checked(format!( + "CREATE SCHEMA {schema}; \ + CREATE TABLE {schema}.main (id BIGSERIAL PRIMARY KEY, payload TEXT); \ + CREATE INDEX main_payload_idx ON {schema}.main (payload); \ + CREATE TABLE {schema}.other (id BIGINT PRIMARY KEY); \ + INSERT INTO {schema}.main (payload) SELECT 'row' FROM generate_series(1, 10)" + )) + .await?; + let mut dest_server = dest.primary(0, &Request::default()).await?; + + let table_move = TableMove::new(format!("{schema}.main")); + let publication = table_move.publication().to_owned(); + let state = ReshardingState::builder() + .source(schema) + .destination(destination) + .publication(&publication) + .maybe_replication_slot(Some(schema.into())) + .build()?; + let task = run_task( + ReshardTask::builder() + .state(state) + .table_move(table_move) + .build(), + ); + let id = task.id(); + + wait_for_reshard_status( + id, + ReshardStatus::Replication { + stage: ReplicationStatus::Replicating, + }, + ) + .await?; + wait_for_active_source_slot(&mut server, schema).await?; + let cutover_accepted = ReplicationTask::trigger_cutover(Some(id)); + + server + .execute_checked(format!( + "INSERT INTO {schema}.main (payload) SELECT 'live' FROM generate_series(1, 5)" + )) + .await?; + tokio::time::timeout(Duration::from_secs(30), async { + loop { + let rows: Vec = dest_server + .fetch_all(format!("SELECT count(*) FROM {schema}.main")) + .await?; + if rows == [15] { + return Ok::<_, Box>(()); + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + }) + .await??; + + tasks_storage().cancel_task(id); + let outcome = tokio::time::timeout(Duration::from_secs(90), task).await?; + + let leftover_slots = source_slots(&mut server, schema).await?; + let leftover_publications: Vec = server + .fetch_all(format!( + "SELECT count(*) FROM pg_publication WHERE pubname = '{publication}'" + )) + .await?; + let source_rows: Vec = server + .fetch_all(format!("SELECT count(*) FROM {schema}.main")) + .await?; + let next_id: Vec = dest_server + .fetch_all(format!( + "INSERT INTO {schema}.main (payload) VALUES ('after') RETURNING id" + )) + .await?; + let other_on_dest: Vec = dest_server + .fetch_all(format!( + "SELECT count(*) FROM pg_tables WHERE schemaname = '{schema}' AND tablename = 'other'" + )) + .await?; + let index_on_dest: Vec = dest_server + .fetch_all(format!( + "SELECT count(*) FROM pg_indexes WHERE schemaname = '{schema}' AND indexname = 'main_payload_idx'" + )) + .await?; + + Ok::<_, Box>(( + outcome.is_err(), + cutover_accepted, + leftover_slots, + leftover_publications, + source_rows, + next_id, + other_on_dest, + index_on_dest, + )) + } + .await; + + cleanup_replication_test(&mut admin, &original_config, [schema, destination]).await?; + + let ( + stopped, + cutover_accepted, + leftover_slots, + leftover_publications, + source_rows, + next_id, + other_on_dest, + index_on_dest, + ) = result?; + assert!(stopped, "a stopped table move reports as cancelled"); + assert!(!cutover_accepted, "a table move must never cut over"); + assert_eq!(leftover_slots, Vec::::new()); + assert_eq!(leftover_publications, [0]); + assert_eq!(source_rows, [15], "the source table is left in place"); + assert_eq!(next_id, [16], "the destination sequence follows the data"); + assert_eq!(other_on_dest, [0], "only the moved table is copied"); + assert_eq!( + index_on_dest, + [1], + "post-data schema sync adds the table's indexes" + ); + + Ok(()) +} + +#[tokio::test] +async fn table_move_fails_before_copy_when_referenced_table_is_missing() +-> Result<(), Box> { + let schema = "table_move_fk"; + let destination = "table_move_fk_dest"; + let original_config = config(); + let mut admin = test_server().await; + let result = async { + setup_replication_test(&mut admin, schema, destination).await?; + let source = databases::databases().schema_owner(schema)?; + let dest = databases::databases().schema_owner(destination)?; + let mut server = source.primary(0, &Request::default()).await?; + server + .execute_checked(format!( + "CREATE SCHEMA {schema}; \ + CREATE TABLE {schema}.parent (id BIGINT PRIMARY KEY); \ + CREATE TABLE {schema}.child (id BIGINT PRIMARY KEY, \ + parent_id BIGINT REFERENCES {schema}.parent(id))" + )) + .await?; + + let table_move = TableMove::new(format!("{schema}.child")); + let state = ReshardingState::builder() + .source(schema) + .destination(destination) + .publication(table_move.publication()) + .build()?; + let outcome = tokio::time::timeout( + Duration::from_secs(30), + run_task( + ReshardTask::builder() + .state(state) + .table_move(table_move) + .build(), + ), + ) + .await?; + + let publications: Vec = server + .fetch_all("SELECT count(*) FROM pg_publication WHERE pubname LIKE '__pgdog_move_%'") + .await?; + let slots: Vec = server + .fetch_all("SELECT count(*) FROM pg_replication_slots") + .await?; + let mut dest_server = dest.primary(0, &Request::default()).await?; + let child_on_dest: Vec = dest_server + .fetch_all(format!( + "SELECT count(*) FROM pg_tables WHERE schemaname = '{schema}' AND tablename = 'child'" + )) + .await?; + + Ok::<_, Box>((outcome, publications, slots, child_on_dest)) + } + .await; + + cleanup_replication_test(&mut admin, &original_config, [schema, destination]).await?; + + let (outcome, publications, slots, child_on_dest) = result?; + let err = outcome.expect_err("a missing referenced table fails the move"); + assert!( + err.to_string() + .contains("references tables that don't exist"), + "{err}" + ); + assert_eq!(publications, [0]); + assert_eq!(slots, [0]); + assert_eq!( + child_on_dest, + [0], + "nothing is created before the checks pass" + ); + + Ok(()) +} diff --git a/pgdog/src/backend/replication/tests/validation.rs b/pgdog/src/backend/replication/tests/validation.rs index a99ae24ba..030d05fb4 100644 --- a/pgdog/src/backend/replication/tests/validation.rs +++ b/pgdog/src/backend/replication/tests/validation.rs @@ -5,7 +5,10 @@ use pgdog_config::resharding::PostDataValidationStage; use pgdog_stats::{ReplicationStatus, SyncState, TaskDefinitionKind, TaskProgress}; use super::*; -use crate::api::{replication::ReplicationTask, schema_sync::SchemaSyncBuilder}; +use crate::api::{ + replication::{CutoverMode, ReplicationTask}, + schema_sync::SchemaSyncBuilder, +}; use crate::backend::maintenance_mode; type TestResult = Result<(), Box>; @@ -19,13 +22,13 @@ struct FkValidation { } impl FkValidation { - fn start(&self, auto_cutover: bool) -> TaskWaiter<(), Error> { + fn start(&self, cutover: CutoverMode) -> TaskWaiter<(), Error> { run_task( ReplicationTask::builder() .state(self.state.clone()) .schema_sync(self.schema_sync.clone()) .validate(true) - .auto_cutover(auto_cutover) + .cutover(cutover) .build(), ) } @@ -245,7 +248,7 @@ async fn test_replication_fk_inconsistent_validation() -> TestResult { PostDataValidationStage::DuringReplication, true, async |fixture| { - let outcome = fixture.start(false).await; + let outcome = fixture.start(CutoverMode::Manual).await; assert!( is_fk_violation(&outcome), "validation did not reject the orphan with a foreign key error: {outcome:?}" @@ -276,7 +279,7 @@ async fn post_data_validation_off_does_not_validate() -> TestResult { PostDataValidationStage::Off, true, async |fixture| { - let task = fixture.start(false); + let task = fixture.start(CutoverMode::Manual); let id = task.id(); wait_for_replication_status(id, ReplicationStatus::Replicating).await?; tokio::time::sleep(Duration::from_secs(2)).await; @@ -304,7 +307,7 @@ async fn post_data_validation_before_cutover_stops_migration() -> TestResult { PostDataValidationStage::BeforeCutover, true, async |fixture| { - let task = fixture.start(true); + let task = fixture.start(CutoverMode::Auto); let id = task.id(); let outcome = tokio::time::timeout(Duration::from_secs(120), task).await?; assert!( @@ -329,7 +332,7 @@ async fn post_data_validation_after_cutover_keeps_replicating() -> TestResult { PostDataValidationStage::AfterCutover, true, async |fixture| { - let task = fixture.start(true); + let task = fixture.start(CutoverMode::Auto); let id = task.id(); wait_for_replication_status(id, ReplicationStatus::ReverseReplicating).await?; let failed = wait_for_validation(id, |validations| { @@ -380,7 +383,7 @@ async fn post_data_validation_during_replication_gates_cutover() -> TestResult { )) .await?; - let task = fixture.start(true); + let task = fixture.start(CutoverMode::Auto); let id = task.id(); wait_for_validation(id, |validations| !validations.is_empty()).await?; tokio::time::sleep(Duration::from_secs(3)).await; diff --git a/pgdog/src/backend/schema/sync/pg_dump.rs b/pgdog/src/backend/schema/sync/pg_dump.rs index e646b0cd8..fff0d7a15 100644 --- a/pgdog/src/backend/schema/sync/pg_dump.rs +++ b/pgdog/src/backend/schema/sync/pg_dump.rs @@ -110,6 +110,7 @@ use tokio::process::Command; pub(crate) struct PgDump { source: Cluster, publication: String, + publication_tables_only: bool, } fn build_pg_dump_command( @@ -138,12 +139,39 @@ fn build_pg_dump_command( command } +/// `pg_dump --table` pattern matching exactly one table. +fn table_pattern(schema: &str, name: &str) -> String { + format!( + "\"{}\".\"{}\"", + crate::util::escape_identifier(schema), + crate::util::escape_identifier(name) + ) +} + impl PgDump { - pub(crate) fn new(source: &Cluster, publication: &str) -> Self { + /// With `publication_tables_only`, dump just the publication's tables + /// (and their partition parents) instead of the whole database. + pub(crate) fn new(source: &Cluster, publication: &str, publication_tables_only: bool) -> Self { Self { source: source.clone(), publication: publication.to_string(), + publication_tables_only, + } + } + + fn table_patterns(tables: &[PublicationTable]) -> Vec { + let mut patterns = vec![]; + for table in tables { + for pattern in [ + table_pattern(table.destination_schema(), table.destination_name()), + table_pattern(&table.schema, &table.name), + ] { + if !patterns.contains(&pattern) { + patterns.push(pattern); + } + } } + patterns } fn clean(source: &str) -> String { @@ -222,6 +250,11 @@ impl PgDump { "server has no configured passwords".into(), ))?, ); + if self.publication_tables_only { + for pattern in Self::table_patterns(&comparison) { + command.arg("--table").arg(pattern); + } + } let output = command.output().await?; if !output.status.success() { @@ -989,7 +1022,48 @@ mod test { #[tokio::test] async fn test_pg_dump_execute() { let cluster = Cluster::new_test_single_shard(&config()); - let _pg_dump = PgDump::new(&cluster, "test_pg_dump_execute"); + let _pg_dump = PgDump::new(&cluster, "test_pg_dump_execute", false); + } + + #[test] + fn test_table_patterns_include_partition_parents() { + let tables = [ + PublicationTable { + schema: "sales".into(), + name: "orders".into(), + ..Default::default() + }, + PublicationTable { + schema: "sales".into(), + name: "events_2026".into(), + parent_schema: "sales".into(), + parent_name: "events".into(), + ..Default::default() + }, + PublicationTable { + schema: "sales".into(), + name: "events_2027".into(), + parent_schema: "sales".into(), + parent_name: "events".into(), + ..Default::default() + }, + PublicationTable { + schema: "My \"Schema\"".into(), + name: "T".into(), + ..Default::default() + }, + ]; + + assert_eq!( + PgDump::table_patterns(&tables), + [ + r#""sales"."orders""#, + r#""sales"."events""#, + r#""sales"."events_2026""#, + r#""sales"."events_2027""#, + r#""My ""Schema"""."T""#, + ] + ); } #[test] diff --git a/pgdog/src/cli.rs b/pgdog/src/cli.rs index 0c08f3c93..a45307f4f 100644 --- a/pgdog/src/cli.rs +++ b/pgdog/src/cli.rs @@ -9,6 +9,7 @@ use crate::api::task::TaskError; use crate::api::{MigrationError, run_task}; use crate::backend::databases::databases; use crate::backend::replication::logical::Error::ReplicationAborted; +use crate::backend::replication::logical::table_move::TableMove; use crate::backend::replication::resharding_state::ReshardingState; use crate::backend::schema::sync::config::ShardConfig; use crate::frontend::comms::comms; @@ -83,15 +84,20 @@ pub(crate) enum Commands { from_database: String, /// Publication name. + #[arg(long, required_unless_present = "table", conflicts_with = "table")] + publication: Option, + + /// Move a single table ("table" or "schema.table") instead of a publication. + /// Schema sync creates only this table. Runs until interrupted, without cutover. #[arg(long)] - publication: String, + table: Option, /// Destination database. #[arg(long)] to_database: String, /// Only replicate; skip the initial data copy. - #[arg(long, default_value = "false")] + #[arg(long, default_value = "false", conflicts_with = "table")] replicate_only: bool, /// Only copy data and synchronize tables; skip replication and cutover. @@ -230,12 +236,20 @@ pub(crate) async fn data_sync(commands: Commands) -> Result<(), Box table_move.publication().to_owned(), + (None, Some(publication)) => publication, + (None, None) => return Err("--publication or --table is required".into()), + }; + let state = ReshardingState::builder() .source(&from_database) .destination(&to_database) @@ -246,6 +260,7 @@ pub(crate) async fn data_sync(commands: Commands) -> Result<(), Box Result<(), Box Result { + let mut argv = vec![ + "pgdog", + "data-sync", + "--from-database", + "a", + "--to-database", + "b", + ]; + argv.extend_from_slice(args); + Ok(Cli::try_parse_from(argv)? + .command + .expect("data-sync is a command")) + } + + #[test] + fn test_table_or_publication() { + let Commands::DataSync { + table, publication, .. + } = parse(&["--table", "sales.orders"]).unwrap() + else { + panic!("not a data sync command"); + }; + assert_eq!(table.as_deref(), Some("sales.orders")); + assert_eq!(publication, None); + + let Commands::DataSync { + table, publication, .. + } = parse(&["--publication", "p"]).unwrap() + else { + panic!("not a data sync command"); + }; + assert_eq!(table, None); + assert_eq!(publication.as_deref(), Some("p")); + + for args in [ + vec![], + vec!["--table", "t", "--publication", "p"], + vec!["--table", "t", "--replicate-only"], + ] { + assert!(parse(&args).is_err(), "{args:?}"); + } + } + } + mod schema_sync { use crate::cli::*;