From 2a3c802cff2fec322b1b251f9cb3fc6726792333 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Thu, 22 Oct 2026 14:14:56 +0530 Subject: [PATCH 1/4] feat(logical): support moving a single table with COPY_DATA --- pgdog/src/admin/copy_data.rs | 126 +++- pgdog/src/api/replication.rs | 18 +- pgdog/src/api/resharding.rs | 67 +- pgdog/src/api/schema_sync.rs | 11 + .../src/backend/replication/logical/error.rs | 18 + pgdog/src/backend/replication/logical/mod.rs | 1 + .../replication/logical/schema_sync.rs | 13 +- .../backend/replication/logical/table_move.rs | 572 ++++++++++++++++++ pgdog/src/backend/replication/tests.rs | 214 ++++++- pgdog/src/backend/schema/sync/pg_dump.rs | 78 ++- pgdog/src/cli.rs | 72 ++- 11 files changed, 1135 insertions(+), 55 deletions(-) create mode 100644 pgdog/src/backend/replication/logical/table_move.rs diff --git a/pgdog/src/admin/copy_data.rs b/pgdog/src/admin/copy_data.rs index 0d9b8b3e0..76560e29b 100644 --- a/pgdog/src/admin/copy_data.rs +++ b/pgdog/src/admin/copy_data.rs @@ -4,17 +4,30 @@ use tracing::info; use crate::api::resharding::ReshardTask; use crate::api::run_task; +use crate::backend::replication::logical::Error as ReplicationError; +use crate::backend::replication::logical::table_move::{QualifiedTable, 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(QualifiedTable), +} + 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, } +fn table(name: &str) -> Result { + let table = name.parse().map_err(ReplicationError::from)?; + Ok(CopySource::Table(table)) +} + #[async_trait] impl Command for CopyData { fn name(&self) -> String { @@ -24,45 +37,67 @@ 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", name] => (from, to, table(name)?, None), + ["copy_data", from, to, "table", name, slot] => (from, to, table(name)?, Some(slot)), + ["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 +109,38 @@ 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.schema == "sales" && t.name == "orders" + )); + assert_eq!(cmd.replication_slot, None); + + let cmd = CopyData::parse("copy_data a b table orders slot").unwrap(); + assert!(matches!( + &cmd.source, + CopySource::Table(t) if t.schema == "public" && t.name == "orders" + )); + assert_eq!(cmd.replication_slot.as_deref(), Some("slot")); + assert_eq!(cmd.from_database, "a"); + assert_eq!(cmd.to_database, "b"); + + assert!(CopyData::parse("copy_data a b table a.b.c").is_err()); + assert!(CopyData::parse("copy_data a b").is_err()); + } +} diff --git a/pgdog/src/api/replication.rs b/pgdog/src/api/replication.rs index 098207b30..aca028215 100644 --- a/pgdog/src/api/replication.rs +++ b/pgdog/src/api/replication.rs @@ -46,6 +46,10 @@ pub(crate) struct ReplicationTask { /// of waiting for an operator `CUTOVER`. #[builder(default)] pub(crate) auto_cutover: bool, + /// Never cut over: replicate until the task is stopped. A cutover swaps + /// whole databases, which is wrong when moving a single table. + #[builder(default)] + pub(crate) disable_cutover: bool, pub(crate) schema_sync: SchemaSyncBuilder, pub(crate) validate: bool, } @@ -83,9 +87,11 @@ impl Task for ReplicationTask { schema_sync, validate, auto_cutover, + disable_cutover, } = self; - let mut replication = Replication::new(&ctx, state.clone(), schema_sync, validate); + let mut replication = + Replication::new(&ctx, state.clone(), schema_sync, validate, disable_cutover); // run the replication until it's stopped - it can cutover multiple times but // it's stopped eventually only by the cancellation process. let result = replication.run(auto_cutover).await; @@ -151,6 +157,7 @@ struct Replication<'a> { validation_stage: PostDataValidationStage, direction: ReplicationDirection, maintenance: MaintenanceMode, + disable_cutover: bool, } impl<'a> Replication<'a> { @@ -159,6 +166,7 @@ impl<'a> Replication<'a> { state: ReshardingState, schema_sync: SchemaSyncBuilder, validate: bool, + disable_cutover: bool, ) -> Self { let validation_stage = if validate { config().config.resharding.post_data_validation @@ -173,6 +181,7 @@ impl<'a> Replication<'a> { validation_stage, direction: ReplicationDirection::Forward, maintenance: MaintenanceMode::new(), + disable_cutover, } } @@ -225,7 +234,9 @@ impl<'a> Replication<'a> { let ctx = self.ctx; let direction = self.direction; let task_cancel = ctx.cancellation_token(); - let mut cutover = (!auto_cutover).then(|| CutoverWaiter::register(ctx.root_id())); + let disable_cutover = self.disable_cutover; + let mut cutover = + (!auto_cutover && !disable_cutover).then(|| CutoverWaiter::register(ctx.root_id())); let mut validation = pin!(self.post_data_validation(stage).fuse()); loop { @@ -244,6 +255,9 @@ impl<'a> Replication<'a> { let result = async { let mut cutover_run = pin!(async { + if disable_cutover { + return std::future::pending().await; + } if let Some(cutover) = cutover.as_ref() { cutover.requested().await; } diff --git a/pgdog/src/api/resharding.rs b/pgdog/src/api/resharding.rs index 35e124ad1..3906a2c6b 100644 --- a/pgdog/src/api/resharding.rs +++ b/pgdog/src/api/resharding.rs @@ -16,6 +16,7 @@ use crate::api::synchronize_tables::SynchronizeTablesTask; use crate::api::task::TaskContext; use crate::api::{MigrationError, 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 { @@ -64,27 +68,49 @@ impl Task for ReshardTask { let mut state = self.state; let schema_sync = SchemaSyncTask::builder() .databases(state.databases()) - .publication(state.publication.clone()); + .publication(state.publication.clone()) + .publication_tables_only(self.table_move.is_some()); + + if let Some(table_move) = &self.table_move { + table_move.prepare(&state).await?; + } // 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()?; + let prepared: Result<(), MigrationError> = Box::pin(async { + 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()?; + } + + if let Some(table_move) = &self.table_move { + table_move.check_table(&state.destination).await?; + } + + Ok(()) + }) + .await; + + if let Err(err) = prepared { + if let Some(table_move) = &self.table_move { + table_move.finish(&state, false).await; + } + return Err(err); } + let mut copied = false; let result: Result<(), MigrationError> = Box::pin(async { // Copy the data, unless replicate-only. if !self.replicate_only { @@ -96,6 +122,7 @@ impl Task for ReshardTask { .build(), ) .await?; + copied = true; } // Post-data schema sync (secondary indexes, constraints): the @@ -140,6 +167,7 @@ impl Task for ReshardTask { ReplicationTask::builder() .state(state.clone()) .auto_cutover(self.auto_cutover) + .disable_cutover(self.table_move.is_some()) .schema_sync(schema_sync.clone()) .validate(!self.skip_schema_sync && !self.replicate_only) .build(), @@ -152,7 +180,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 @@ -167,6 +196,10 @@ impl Task for ReshardTask { state.detach_slots(); } + if let Some(table_move) = &self.table_move { + table_move.finish(&state, copied).await; + } + result } } 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 2491a16cd..2d000615d 100644 --- a/pgdog/src/backend/replication/logical/error.rs +++ b/pgdog/src/backend/replication/logical/error.rs @@ -263,6 +263,24 @@ pub(crate) enum Error { op: &'static str, }, + #[error("{0}")] + InvalidTableName(#[from] super::table_move::InvalidTableName), + + #[error("source and destination are the same database \"{0}\"")] + SameDatabase(String), + + #[error("table {table} does not exist on destination shard {shard}, check the schema sync log")] + DestinationTableMissing { table: String, shard: usize }, + + #[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..982835fdb --- /dev/null +++ b/pgdog/src/backend/replication/logical/table_move.rs @@ -0,0 +1,572 @@ +//! 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 std::fmt::{self, Display}; +use std::str::FromStr; + +use tracing::{info, warn}; + +use super::Error; +use super::resharding_state::ReshardingState; +use crate::backend::Cluster; +use crate::backend::pool::Request; +use crate::frontend::router::parser::{Column, Sequence, Table}; +use crate::net::{DataRow, Format}; +use crate::util::sql::quote_literal; +use crate::util::{escape_identifier, random_string}; + +/// Table name that isn't `table` or `schema.table`. +#[derive(Debug, thiserror::Error)] +#[error("invalid table name \"{0}\", expected \"table\" or \"schema.table\"")] +pub(crate) struct InvalidTableName(String); + +/// Schema-qualified table name. As in Postgres, unquoted parts are folded +/// to lowercase. A missing schema means `public`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct QualifiedTable { + pub(crate) schema: String, + pub(crate) name: String, +} + +impl QualifiedTable { + fn regclass(&self) -> String { + format!("{}::regclass", quote_literal(&self.to_string())) + } +} + +impl Display for QualifiedTable { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "\"{}\".\"{}\"", + escape_identifier(&self.schema), + escape_identifier(&self.name) + ) + } +} + +impl FromStr for QualifiedTable { + type Err = InvalidTableName; + + fn from_str(input: &str) -> Result { + let invalid = || InvalidTableName(input.to_owned()); + let mut parts = identifiers(input).ok_or_else(invalid)?; + let name = parts.pop().ok_or_else(invalid)?; + let schema = parts.pop().unwrap_or_else(|| "public".to_owned()); + + if !parts.is_empty() { + return Err(invalid()); + } + + Ok(Self { schema, name }) + } +} + +/// Split a dotted name into identifiers, unquoting `"..."` parts. +fn identifiers(input: &str) -> Option> { + let mut parts = vec![]; + let mut chars = input.chars().peekable(); + + loop { + let mut part = String::new(); + + if chars.next_if_eq(&'"').is_some() { + loop { + match chars.next()? { + '"' if chars.next_if_eq(&'"').is_some() => part.push('"'), + '"' => break, + c => part.push(c), + } + } + } else { + while let Some(c) = chars.next_if(|c| *c != '.') { + if c == '"' || c.is_whitespace() { + return None; + } + part.push(c.to_ascii_lowercase()); + } + } + + if part.is_empty() { + return None; + } + parts.push(part); + + match chars.next() { + None => return Some(parts), + Some('.') => continue, + Some(_) => return None, + } + } +} + +/// A single table moved between databases. +#[derive(Debug, Clone)] +pub(crate) struct TableMove { + table: QualifiedTable, + publication: String, +} + +impl TableMove { + pub(crate) fn new(table: QualifiedTable) -> 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 referenced = self.referenced_tables(&state.source).await?; + if let Some((shard, missing)) = missing_tables(&state.destination, &referenced).await? { + return Err(Error::ReferencedTablesMissing { + table: self.table.to_string(), + 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 \"{}\"", + escape_identifier(&self.table.schema) + ); + for shard in state.destination.shards() { + let mut server = shard.primary(&Request::default()).await?; + server.execute_checked(create_schema.as_str()).await?; + } + + if let Err(err) = self.create_publication(&state.source).await { + self.drop_publication(&state.source).await; + return Err(err); + } + + info!( + "moving table {} [{}], publication \"{}\"", + self.table, state, self.publication + ); + + Ok(()) + } + + /// Fix the destination sequences, list what is left on the source, + /// and drop the publication. Never touches the source table itself. + pub(crate) async fn finish(&self, state: &ReshardingState, copied: bool) { + if copied { + if let Err(err) = self.sync_sequences(&state.destination).await { + warn!( + "failed to sync sequences of {} on \"{}\": {err}", + self.table, + state.destination.identifier().database + ); + } + + match self.referencing_constraints(&state.source).await { + Ok(constraints) if !constraints.is_empty() => warn!( + "foreign keys on \"{}\" still reference {}: {}; drop them once the application stops using the source table", + state.source.identifier().database, + self.table, + constraints.join(", ") + ), + Ok(_) => (), + Err(err) => warn!( + "failed to list foreign keys referencing {}: {err}", + self.table + ), + } + + 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; + } + + /// Tables the moved table references with a foreign key, other than itself. + async fn referenced_tables(&self, source: &Cluster) -> Result, Error> { + let mut server = source.primary(0, &Request::default()).await?; + Ok(server + .fetch_all(format!( + "SELECT DISTINCT format('%I.%I', n.nspname, c.relname) \ + FROM pg_constraint con \ + JOIN pg_class c ON c.oid = con.confrelid \ + JOIN pg_namespace n ON n.oid = c.relnamespace \ + WHERE con.contype = 'f' AND con.conrelid = {} AND con.confrelid <> con.conrelid", + self.table.regclass() + )) + .await?) + } + + /// The table must exist on every destination shard before the copy, + /// created by schema sync or by the user. + pub(crate) async fn check_table(&self, dest: &Cluster) -> Result<(), Error> { + let table = self.table.to_string(); + match missing_tables(dest, std::slice::from_ref(&table)).await? { + Some((shard, _)) => Err(Error::DestinationTableMissing { table, shard }), + None => Ok(()), + } + } + + async fn create_publication(&self, source: &Cluster) -> Result<(), Error> { + let sql = format!( + "CREATE PUBLICATION \"{}\" FOR TABLE {}", + escape_identifier(&self.publication), + self.table + ); + + // 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() + ); + } + } + } + + /// Logical replication doesn't move sequences. Set every sequence owned + /// by the table (`SERIAL` or `IDENTITY`) to the column's max value. + async fn sync_sequences(&self, dest: &Cluster) -> Result<(), Error> { + let sql = format!( + "SELECT a.attname, n.nspname, s.relname \ + FROM pg_depend d \ + JOIN pg_class s ON s.oid = d.objid AND s.relkind = 'S' \ + JOIN pg_namespace n ON n.oid = s.relnamespace \ + JOIN pg_attribute a ON a.attrelid = d.refobjid AND a.attnum = d.refobjsubid \ + WHERE d.classid = 'pg_class'::regclass \ + AND d.refclassid = 'pg_class'::regclass \ + AND d.refobjid = {} \ + AND d.deptype IN ('a', 'i')", + self.table.regclass() + ); + + for shard in dest.shards() { + let mut server = shard.primary(&Request::default()).await?; + let sequences: Vec = server.fetch_all(sql.as_str()).await?; + + for row in sequences { + let column: String = row.get(0, Format::Text).ok_or(Error::MissingData)?; + let schema: String = row.get(1, Format::Text).ok_or(Error::MissingData)?; + let name: String = row.get(2, Format::Text).ok_or(Error::MissingData)?; + + let setval = Sequence::from(Table { + schema: Some(&schema), + name: &name, + alias: None, + }) + .setval_from_column(&Column { + schema: Some(&self.table.schema), + table: Some(&self.table.name), + name: &column, + })?; + + server.execute_checked(setval).await?; + info!( + "sequence \"{schema}\".\"{name}\" synced to {}.\"{column}\" on shard {}", + self.table, + shard.number() + ); + } + } + + Ok(()) + } + + /// Foreign keys on other source tables that reference the moved table. + async fn referencing_constraints(&self, source: &Cluster) -> Result, Error> { + let mut server = source.primary(0, &Request::default()).await?; + Ok(server + .fetch_all(format!( + "SELECT format('%s.%I', con.conrelid::regclass, con.conname) \ + FROM pg_constraint con \ + WHERE con.contype = 'f' AND con.confrelid = {} AND con.conrelid <> con.confrelid", + self.table.regclass() + )) + .await?) + } +} + +/// 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; + + fn table(schema: &str, name: &str) -> QualifiedTable { + QualifiedTable { + schema: schema.into(), + name: name.into(), + } + } + + #[test] + fn parses_table_names() { + for (input, expected) in [ + ("orders", table("public", "orders")), + ("Orders", table("public", "orders")), + ("sales.orders", table("sales", "orders")), + ("\"Sales\".\"Orders\"", table("Sales", "Orders")), + ("sales.\"my.table\"", table("sales", "my.table")), + ("\"a\"\"b\".c", table("a\"b", "c")), + ] { + assert_eq!( + input.parse::().unwrap(), + expected, + "{input}" + ); + } + } + + #[test] + fn rejects_invalid_table_names() { + for input in [ + "", + ".", + "a.", + ".a", + "a..b", + "db.sales.orders", + "\"unterminated", + "\"\"", + "a\"b", + "\"a\"b", + "my table", + ] { + assert!(input.parse::().is_err(), "{input:?}"); + } + } + + #[test] + fn displays_quoted() { + assert_eq!(table("a\"b", "c").to_string(), "\"a\"\"b\".\"c\""); + } + + #[tokio::test] + async fn destination_checks_find_missing_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(); + + let child = TableMove::new(table("pgdog", "table_move_child")); + let referenced = child.referenced_tables(&cluster).await.unwrap(); + assert_eq!(referenced, vec!["pgdog.table_move_parent".to_owned()]); + assert_eq!(missing_tables(&cluster, &referenced).await.unwrap(), None); + assert_eq!(missing_tables(&cluster, &[]).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()])) + ); + child.check_table(&cluster).await.unwrap(); + + let absent = TableMove::new(table("pgdog", "table_move_absent")); + let err = absent + .check_table(&cluster) + .await + .expect_err("table is missing on the destination"); + assert!( + matches!(err, Error::DestinationTableMissing { .. }), + "{err}" + ); + let err = absent + .referenced_tables(&cluster) + .await + .expect_err("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(table("pgdog", "table_move_publication")); + table_move.create_publication(&cluster).await.unwrap(); + table_move.drop_publication(&cluster).await; + + cluster.shutdown(); + server + .execute_checked("DROP TABLE pgdog.table_move_publication") + .await + .unwrap(); + } + + #[tokio::test] + async fn sequences_are_synced_to_max_value() { + crate::logger(); + let mut server = test_server().await; + for ddl in [ + "DROP TABLE IF EXISTS pgdog.table_move_seq", + "CREATE TABLE pgdog.table_move_seq (id BIGSERIAL PRIMARY KEY, \ + num BIGINT GENERATED BY DEFAULT AS IDENTITY, value TEXT)", + "INSERT INTO pgdog.table_move_seq (id, num, value) VALUES (100, 7, 'a'), (250, 9, 'b')", + ] { + server.execute_checked(ddl).await.unwrap(); + } + + let cluster = Cluster::new_test(&config()); + cluster.launch(); + + let table_move = TableMove::new(table("pgdog", "table_move_seq")); + table_move.sync_sequences(&cluster).await.unwrap(); + + let next: Vec = server + .fetch_all("INSERT INTO pgdog.table_move_seq (value) VALUES ('c') RETURNING id, num") + .await + .unwrap(); + let id: i64 = next[0].get(0, Format::Text).unwrap(); + let num: i64 = next[0].get(1, Format::Text).unwrap(); + assert_eq!(id, 251); + assert_eq!(num, 10); + + cluster.shutdown(); + server + .execute_checked("DROP TABLE pgdog.table_move_seq") + .await + .unwrap(); + } + + #[tokio::test] + async fn lists_foreign_keys_referencing_the_table() { + crate::logger(); + let mut server = test_server().await; + for ddl in [ + "DROP TABLE IF EXISTS pgdog.table_move_ref_child, pgdog.table_move_ref_parent", + "CREATE TABLE pgdog.table_move_ref_parent (id BIGINT PRIMARY KEY)", + "CREATE TABLE pgdog.table_move_ref_child (id BIGINT PRIMARY KEY, \ + parent_id BIGINT CONSTRAINT child_to_parent REFERENCES pgdog.table_move_ref_parent(id))", + ] { + server.execute_checked(ddl).await.unwrap(); + } + + let cluster = Cluster::new_test(&config()); + cluster.launch(); + + let constraints = TableMove::new(table("pgdog", "table_move_ref_parent")) + .referencing_constraints(&cluster) + .await + .unwrap(); + assert_eq!(constraints.len(), 1); + assert!(constraints[0].ends_with("table_move_ref_child.child_to_parent")); + + let none = TableMove::new(table("pgdog", "table_move_ref_child")) + .referencing_constraints(&cluster) + .await + .unwrap(); + assert!(none.is_empty()); + + cluster.shutdown(); + server + .execute_checked("DROP TABLE pgdog.table_move_ref_child, pgdog.table_move_ref_parent") + .await + .unwrap(); + } +} diff --git a/pgdog/src/backend/replication/tests.rs b/pgdog/src/backend/replication/tests.rs index c1d8a01ab..c8b475516 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::{QualifiedTable, TableMove}; use crate::{ api::{ MigrationError, copy_data::{CopyDataTask, TableDataSyncTask}, - replication::{ReplicationClusterStop, ReplicationClusterTask}, + replication::{ReplicationClusterStop, ReplicationClusterTask, ReplicationTask}, resharding::ReshardTask, run_task, schema_sync::{SchemaSyncPhase, SchemaSyncTask}, @@ -940,3 +941,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(QualifiedTable { + schema: schema.into(), + name: "main".into(), + }); + 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).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(QualifiedTable { + schema: schema.into(), + name: "child".into(), + }); + 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/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 465af28c8..57ea87184 100644 --- a/pgdog/src/cli.rs +++ b/pgdog/src/cli.rs @@ -10,6 +10,7 @@ use crate::api::run_task; use crate::api::schema_sync::{SchemaSyncPhase, SchemaSyncTask}; use crate::api::tasks_storage; use crate::backend::databases::databases; +use crate::backend::replication::logical::table_move::{QualifiedTable, TableMove}; use crate::backend::replication::resharding_state::ReshardingState; use crate::backend::schema::sync::config::ShardConfig; use crate::frontend::router::cli::RouterCli; @@ -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. @@ -254,12 +260,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) @@ -270,6 +284,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"); + }; + let table = table.expect("table is set"); + assert_eq!( + (table.schema.as_str(), table.name.as_str()), + ("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", "a.b.c"], + vec!["--table", "t", "--publication", "p"], + vec!["--table", "t", "--replicate-only"], + ] { + assert!(parse(&args).is_err(), "{args:?}"); + } + } + } + mod schema_sync { use crate::cli::*; From 9aef49412f86c0d065f2c0522faec4aa77883025 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Thu, 22 Oct 2026 15:39:48 +0530 Subject: [PATCH 2/4] handle cutover enum --- pgdog/src/api/replication.rs | 61 ++++++++++--------- pgdog/src/api/resharding.rs | 19 +++++- .../backend/replication/tests/validation.rs | 19 +++--- 3 files changed, 59 insertions(+), 40 deletions(-) diff --git a/pgdog/src/api/replication.rs b/pgdog/src/api/replication.rs index aca028215..00c143461 100644 --- a/pgdog/src/api/replication.rs +++ b/pgdog/src/api/replication.rs @@ -39,17 +39,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, - /// Never cut over: replicate until the task is stopped. A cutover swaps - /// whole databases, which is wrong when moving a single table. #[builder(default)] - pub(crate) disable_cutover: bool, + pub(crate) cutover: CutoverMode, pub(crate) schema_sync: SchemaSyncBuilder, pub(crate) validate: bool, } @@ -77,7 +84,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, } } @@ -86,15 +93,13 @@ impl Task for ReplicationTask { state, schema_sync, validate, - auto_cutover, - disable_cutover, + cutover, } = self; - let mut replication = - Replication::new(&ctx, state.clone(), schema_sync, validate, disable_cutover); + let mut replication = Replication::new(&ctx, state.clone(), schema_sync, validate); // run the replication until it's stopped - it can cutover multiple times but // it's stopped eventually only by the cancellation process. - let result = replication.run(auto_cutover).await; + let result = replication.run(cutover).await; replication.resume_traffic(); let stopped_in_rollback_window = replication.cancelled() @@ -157,7 +162,6 @@ struct Replication<'a> { validation_stage: PostDataValidationStage, direction: ReplicationDirection, maintenance: MaintenanceMode, - disable_cutover: bool, } impl<'a> Replication<'a> { @@ -166,7 +170,6 @@ impl<'a> Replication<'a> { state: ReshardingState, schema_sync: SchemaSyncBuilder, validate: bool, - disable_cutover: bool, ) -> Self { let validation_stage = if validate { config().config.resharding.post_data_validation @@ -181,17 +184,16 @@ impl<'a> Replication<'a> { validation_stage, direction: ReplicationDirection::Forward, maintenance: MaintenanceMode::new(), - disable_cutover, } } - async fn run(&mut self, auto_cutover: bool) -> Result<(), Error> { + async fn run(&mut self, 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(auto_cutover, PostDataValidationStage::DuringReplication) + self.replicate_until_cutover(cutover, PostDataValidationStage::DuringReplication) .await?; self.sync_schema( self.schema_sync @@ -207,8 +209,11 @@ impl<'a> Replication<'a> { loop { self.cutover().await?; self.flip_direction(); - self.replicate_until_cutover(false, PostDataValidationStage::AfterCutover) - .await?; + self.replicate_until_cutover( + CutoverMode::Manual, + PostDataValidationStage::AfterCutover, + ) + .await?; self.validation_stage = PostDataValidationStage::Off; self.sync_schema( SchemaSyncTask::builder() @@ -227,16 +232,15 @@ impl<'a> Replication<'a> { /// the source WAL in time, resume the traffic and start the replication again. async fn replicate_until_cutover( &mut self, - auto_cutover: bool, + mode: CutoverMode, stage: PostDataValidationStage, ) -> Result<(), Error> { let wait_for_validation = stage != PostDataValidationStage::AfterCutover; let ctx = self.ctx; let direction = self.direction; let task_cancel = ctx.cancellation_token(); - let disable_cutover = self.disable_cutover; let mut cutover = - (!auto_cutover && !disable_cutover).then(|| CutoverWaiter::register(ctx.root_id())); + (mode == CutoverMode::Manual).then(|| CutoverWaiter::register(ctx.root_id())); let mut validation = pin!(self.post_data_validation(stage).fuse()); loop { @@ -255,11 +259,10 @@ impl<'a> Replication<'a> { let result = async { let mut cutover_run = pin!(async { - if disable_cutover { - return std::future::pending().await; - } - 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 3906a2c6b..ba6ba556f 100644 --- a/pgdog/src/api/resharding.rs +++ b/pgdog/src/api/resharding.rs @@ -10,7 +10,7 @@ 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; @@ -65,6 +65,7 @@ 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()) @@ -166,8 +167,7 @@ impl Task for ReshardTask { ctx.run( ReplicationTask::builder() .state(state.clone()) - .auto_cutover(self.auto_cutover) - .disable_cutover(self.table_move.is_some()) + .cutover(cutover) .schema_sync(schema_sync.clone()) .validate(!self.skip_schema_sync && !self.replicate_only) .build(), @@ -203,3 +203,16 @@ impl Task for ReshardTask { 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/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; From a1aefdc15c80f0038b6fe285e97ce42e7ff84cf6 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Thu, 22 Oct 2026 16:00:33 +0530 Subject: [PATCH 3/4] cleanup --- pgdog/src/api/resharding.rs | 21 +- .../src/backend/replication/logical/error.rs | 3 - .../backend/replication/logical/table_move.rs | 181 +----------------- 3 files changed, 18 insertions(+), 187 deletions(-) diff --git a/pgdog/src/api/resharding.rs b/pgdog/src/api/resharding.rs index ba6ba556f..e44b2be53 100644 --- a/pgdog/src/api/resharding.rs +++ b/pgdog/src/api/resharding.rs @@ -14,7 +14,7 @@ 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; @@ -96,10 +96,6 @@ impl Task for ReshardTask { state.reload()?; } - if let Some(table_move) = &self.table_move { - table_move.check_table(&state.destination).await?; - } - Ok(()) }) .await; @@ -197,6 +193,21 @@ impl Task for ReshardTask { } 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; } diff --git a/pgdog/src/backend/replication/logical/error.rs b/pgdog/src/backend/replication/logical/error.rs index 2d000615d..b88b9ca07 100644 --- a/pgdog/src/backend/replication/logical/error.rs +++ b/pgdog/src/backend/replication/logical/error.rs @@ -269,9 +269,6 @@ pub(crate) enum Error { #[error("source and destination are the same database \"{0}\"")] SameDatabase(String), - #[error("table {table} does not exist on destination shard {shard}, check the schema sync log")] - DestinationTableMissing { table: String, shard: usize }, - #[error( "table {table} references tables that don't exist on destination shard {shard}: {missing}" )] diff --git a/pgdog/src/backend/replication/logical/table_move.rs b/pgdog/src/backend/replication/logical/table_move.rs index 982835fdb..4c5e20e9d 100644 --- a/pgdog/src/backend/replication/logical/table_move.rs +++ b/pgdog/src/backend/replication/logical/table_move.rs @@ -14,8 +14,6 @@ use super::Error; use super::resharding_state::ReshardingState; use crate::backend::Cluster; use crate::backend::pool::Request; -use crate::frontend::router::parser::{Column, Sequence, Table}; -use crate::net::{DataRow, Format}; use crate::util::sql::quote_literal; use crate::util::{escape_identifier, random_string}; @@ -163,32 +161,9 @@ impl TableMove { Ok(()) } - /// Fix the destination sequences, list what is left on the source, - /// and drop the publication. Never touches the source table itself. + /// Drop the publication. Never touches the source table itself. pub(crate) async fn finish(&self, state: &ReshardingState, copied: bool) { if copied { - if let Err(err) = self.sync_sequences(&state.destination).await { - warn!( - "failed to sync sequences of {} on \"{}\": {err}", - self.table, - state.destination.identifier().database - ); - } - - match self.referencing_constraints(&state.source).await { - Ok(constraints) if !constraints.is_empty() => warn!( - "foreign keys on \"{}\" still reference {}: {}; drop them once the application stops using the source table", - state.source.identifier().database, - self.table, - constraints.join(", ") - ), - Ok(_) => (), - Err(err) => warn!( - "failed to list foreign keys referencing {}: {err}", - self.table - ), - } - info!( "table {} is left in place on \"{}\"; drop it once the application uses \"{}\"", self.table, @@ -215,16 +190,6 @@ impl TableMove { .await?) } - /// The table must exist on every destination shard before the copy, - /// created by schema sync or by the user. - pub(crate) async fn check_table(&self, dest: &Cluster) -> Result<(), Error> { - let table = self.table.to_string(); - match missing_tables(dest, std::slice::from_ref(&table)).await? { - Some((shard, _)) => Err(Error::DestinationTableMissing { table, shard }), - None => Ok(()), - } - } - async fn create_publication(&self, source: &Cluster) -> Result<(), Error> { let sql = format!( "CREATE PUBLICATION \"{}\" FOR TABLE {}", @@ -268,67 +233,6 @@ impl TableMove { } } } - - /// Logical replication doesn't move sequences. Set every sequence owned - /// by the table (`SERIAL` or `IDENTITY`) to the column's max value. - async fn sync_sequences(&self, dest: &Cluster) -> Result<(), Error> { - let sql = format!( - "SELECT a.attname, n.nspname, s.relname \ - FROM pg_depend d \ - JOIN pg_class s ON s.oid = d.objid AND s.relkind = 'S' \ - JOIN pg_namespace n ON n.oid = s.relnamespace \ - JOIN pg_attribute a ON a.attrelid = d.refobjid AND a.attnum = d.refobjsubid \ - WHERE d.classid = 'pg_class'::regclass \ - AND d.refclassid = 'pg_class'::regclass \ - AND d.refobjid = {} \ - AND d.deptype IN ('a', 'i')", - self.table.regclass() - ); - - for shard in dest.shards() { - let mut server = shard.primary(&Request::default()).await?; - let sequences: Vec = server.fetch_all(sql.as_str()).await?; - - for row in sequences { - let column: String = row.get(0, Format::Text).ok_or(Error::MissingData)?; - let schema: String = row.get(1, Format::Text).ok_or(Error::MissingData)?; - let name: String = row.get(2, Format::Text).ok_or(Error::MissingData)?; - - let setval = Sequence::from(Table { - schema: Some(&schema), - name: &name, - alias: None, - }) - .setval_from_column(&Column { - schema: Some(&self.table.schema), - table: Some(&self.table.name), - name: &column, - })?; - - server.execute_checked(setval).await?; - info!( - "sequence \"{schema}\".\"{name}\" synced to {}.\"{column}\" on shard {}", - self.table, - shard.number() - ); - } - } - - Ok(()) - } - - /// Foreign keys on other source tables that reference the moved table. - async fn referencing_constraints(&self, source: &Cluster) -> Result, Error> { - let mut server = source.primary(0, &Request::default()).await?; - Ok(server - .fetch_all(format!( - "SELECT format('%s.%I', con.conrelid::regclass, con.conname) \ - FROM pg_constraint con \ - WHERE con.contype = 'f' AND con.confrelid = {} AND con.conrelid <> con.confrelid", - self.table.regclass() - )) - .await?) - } } /// The first destination shard missing any of `tables`, with the missing ones. @@ -449,18 +353,8 @@ mod test { .unwrap(), Some((0, vec!["pgdog.table_move_absent".to_owned()])) ); - child.check_table(&cluster).await.unwrap(); - let absent = TableMove::new(table("pgdog", "table_move_absent")); - let err = absent - .check_table(&cluster) - .await - .expect_err("table is missing on the destination"); - assert!( - matches!(err, Error::DestinationTableMissing { .. }), - "{err}" - ); - let err = absent + let err = TableMove::new(table("pgdog", "table_move_absent")) .referenced_tables(&cluster) .await .expect_err("table is missing on the source"); @@ -498,75 +392,4 @@ mod test { .await .unwrap(); } - - #[tokio::test] - async fn sequences_are_synced_to_max_value() { - crate::logger(); - let mut server = test_server().await; - for ddl in [ - "DROP TABLE IF EXISTS pgdog.table_move_seq", - "CREATE TABLE pgdog.table_move_seq (id BIGSERIAL PRIMARY KEY, \ - num BIGINT GENERATED BY DEFAULT AS IDENTITY, value TEXT)", - "INSERT INTO pgdog.table_move_seq (id, num, value) VALUES (100, 7, 'a'), (250, 9, 'b')", - ] { - server.execute_checked(ddl).await.unwrap(); - } - - let cluster = Cluster::new_test(&config()); - cluster.launch(); - - let table_move = TableMove::new(table("pgdog", "table_move_seq")); - table_move.sync_sequences(&cluster).await.unwrap(); - - let next: Vec = server - .fetch_all("INSERT INTO pgdog.table_move_seq (value) VALUES ('c') RETURNING id, num") - .await - .unwrap(); - let id: i64 = next[0].get(0, Format::Text).unwrap(); - let num: i64 = next[0].get(1, Format::Text).unwrap(); - assert_eq!(id, 251); - assert_eq!(num, 10); - - cluster.shutdown(); - server - .execute_checked("DROP TABLE pgdog.table_move_seq") - .await - .unwrap(); - } - - #[tokio::test] - async fn lists_foreign_keys_referencing_the_table() { - crate::logger(); - let mut server = test_server().await; - for ddl in [ - "DROP TABLE IF EXISTS pgdog.table_move_ref_child, pgdog.table_move_ref_parent", - "CREATE TABLE pgdog.table_move_ref_parent (id BIGINT PRIMARY KEY)", - "CREATE TABLE pgdog.table_move_ref_child (id BIGINT PRIMARY KEY, \ - parent_id BIGINT CONSTRAINT child_to_parent REFERENCES pgdog.table_move_ref_parent(id))", - ] { - server.execute_checked(ddl).await.unwrap(); - } - - let cluster = Cluster::new_test(&config()); - cluster.launch(); - - let constraints = TableMove::new(table("pgdog", "table_move_ref_parent")) - .referencing_constraints(&cluster) - .await - .unwrap(); - assert_eq!(constraints.len(), 1); - assert!(constraints[0].ends_with("table_move_ref_child.child_to_parent")); - - let none = TableMove::new(table("pgdog", "table_move_ref_child")) - .referencing_constraints(&cluster) - .await - .unwrap(); - assert!(none.is_empty()); - - cluster.shutdown(); - server - .execute_checked("DROP TABLE pgdog.table_move_ref_child, pgdog.table_move_ref_parent") - .await - .unwrap(); - } } From e24d4f8263905e87014abd0e207a42ce9ae0f0d9 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Wed, 28 Oct 2026 23:53:10 +0530 Subject: [PATCH 4/4] cleanup --- pgdog/src/admin/copy_data.rs | 30 +-- pgdog/src/api/resharding.rs | 25 +- .../src/backend/replication/logical/error.rs | 3 - .../backend/replication/logical/table_move.rs | 244 +++++------------- pgdog/src/backend/replication/tests.rs | 12 +- pgdog/src/cli.rs | 11 +- 6 files changed, 89 insertions(+), 236 deletions(-) diff --git a/pgdog/src/admin/copy_data.rs b/pgdog/src/admin/copy_data.rs index 76560e29b..82b8d0624 100644 --- a/pgdog/src/admin/copy_data.rs +++ b/pgdog/src/admin/copy_data.rs @@ -4,8 +4,7 @@ use tracing::info; use crate::api::resharding::ReshardTask; use crate::api::run_task; -use crate::backend::replication::logical::Error as ReplicationError; -use crate::backend::replication::logical::table_move::{QualifiedTable, TableMove}; +use crate::backend::replication::logical::table_move::TableMove; use crate::backend::replication::resharding_state::ReshardingState; use super::prelude::*; @@ -13,7 +12,7 @@ use super::prelude::*; /// What to copy: the tables of an existing publication, or a single table. pub(crate) enum CopySource { Publication(String), - Table(QualifiedTable), + Table(String), } pub(crate) struct CopyData { @@ -23,11 +22,6 @@ pub(crate) struct CopyData { pub(crate) replication_slot: Option, } -fn table(name: &str) -> Result { - let table = name.parse().map_err(ReplicationError::from)?; - Ok(CopySource::Table(table)) -} - #[async_trait] impl Command for CopyData { fn name(&self) -> String { @@ -38,8 +32,9 @@ impl Command for CopyData { let parts = sql.split(" ").collect::>(); let (from_database, to_database, source, replication_slot) = match parts[..] { - ["copy_data", from, to, "table", name] => (from, to, table(name)?, None), - ["copy_data", from, to, "table", name, slot] => (from, to, table(name)?, Some(slot)), + ["copy_data", from, to, "table", table] => { + (from, to, CopySource::Table(table.to_owned()), None) + } ["copy_data", from, to, publication] => ( from, to, @@ -74,7 +69,7 @@ impl Command for CopyData { } CopySource::Table(table) => { info!( - r#"copy_data "{}" to "{}", table={}"#, + r#"copy_data "{}" to "{}", table="{}""#, self.from_database, self.to_database, table ); let table_move = TableMove::new(table.clone()); @@ -125,22 +120,11 @@ mod test { 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.schema == "sales" && t.name == "orders" - )); + assert!(matches!(&cmd.source, CopySource::Table(t) if t == "sales.orders")); assert_eq!(cmd.replication_slot, None); - - let cmd = CopyData::parse("copy_data a b table orders slot").unwrap(); - assert!(matches!( - &cmd.source, - CopySource::Table(t) if t.schema == "public" && t.name == "orders" - )); - assert_eq!(cmd.replication_slot.as_deref(), Some("slot")); assert_eq!(cmd.from_database, "a"); assert_eq!(cmd.to_database, "b"); - assert!(CopyData::parse("copy_data a b table a.b.c").is_err()); assert!(CopyData::parse("copy_data a b").is_err()); } } diff --git a/pgdog/src/api/resharding.rs b/pgdog/src/api/resharding.rs index 974a7f71c..c74a0587a 100644 --- a/pgdog/src/api/resharding.rs +++ b/pgdog/src/api/resharding.rs @@ -72,13 +72,13 @@ impl Task for ReshardTask { .publication(state.publication.clone()) .publication_tables_only(self.table_move.is_some()); - if let Some(table_move) = &self.table_move { - table_move.prepare(&state).await?; - } + 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. It runs before any replication - // slots exist, so it stays outside the cleanup guard below. - let prepared: Result<(), MigrationError> = Box::pin(async { + // Pre-data schema sync, unless skipped. if !self.skip_schema_sync { ctx.set_status(ReshardStatus::SchemaSync); ctx.run( @@ -96,19 +96,6 @@ impl Task for ReshardTask { state.reload()?; } - Ok(()) - }) - .await; - - if let Err(err) = prepared { - if let Some(table_move) = &self.table_move { - table_move.finish(&state, false).await; - } - return Err(err); - } - - let mut copied = false; - let result: Result<(), MigrationError> = Box::pin(async { // Copy the data, unless replicate-only. if !self.replicate_only { ctx.set_status(ReshardStatus::SyncingData); diff --git a/pgdog/src/backend/replication/logical/error.rs b/pgdog/src/backend/replication/logical/error.rs index da23b5a37..6d870cfdc 100644 --- a/pgdog/src/backend/replication/logical/error.rs +++ b/pgdog/src/backend/replication/logical/error.rs @@ -263,9 +263,6 @@ pub(crate) enum Error { op: &'static str, }, - #[error("{0}")] - InvalidTableName(#[from] super::table_move::InvalidTableName), - #[error("source and destination are the same database \"{0}\"")] SameDatabase(String), diff --git a/pgdog/src/backend/replication/logical/table_move.rs b/pgdog/src/backend/replication/logical/table_move.rs index 4c5e20e9d..db7ae5be8 100644 --- a/pgdog/src/backend/replication/logical/table_move.rs +++ b/pgdog/src/backend/replication/logical/table_move.rs @@ -5,8 +5,6 @@ //! is stopped. There is no cutover; the application switches databases itself. use std::collections::HashSet; -use std::fmt::{self, Display}; -use std::str::FromStr; use tracing::{info, warn}; @@ -14,103 +12,30 @@ 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}; -/// Table name that isn't `table` or `schema.table`. -#[derive(Debug, thiserror::Error)] -#[error("invalid table name \"{0}\", expected \"table\" or \"schema.table\"")] -pub(crate) struct InvalidTableName(String); - -/// Schema-qualified table name. As in Postgres, unquoted parts are folded -/// to lowercase. A missing schema means `public`. -#[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) struct QualifiedTable { - pub(crate) schema: String, - pub(crate) name: String, -} - -impl QualifiedTable { - fn regclass(&self) -> String { - format!("{}::regclass", quote_literal(&self.to_string())) - } -} - -impl Display for QualifiedTable { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!( - f, - "\"{}\".\"{}\"", - escape_identifier(&self.schema), - escape_identifier(&self.name) - ) - } -} - -impl FromStr for QualifiedTable { - type Err = InvalidTableName; - - fn from_str(input: &str) -> Result { - let invalid = || InvalidTableName(input.to_owned()); - let mut parts = identifiers(input).ok_or_else(invalid)?; - let name = parts.pop().ok_or_else(invalid)?; - let schema = parts.pop().unwrap_or_else(|| "public".to_owned()); - - if !parts.is_empty() { - return Err(invalid()); - } - - Ok(Self { schema, name }) - } -} - -/// Split a dotted name into identifiers, unquoting `"..."` parts. -fn identifiers(input: &str) -> Option> { - let mut parts = vec![]; - let mut chars = input.chars().peekable(); - - loop { - let mut part = String::new(); - - if chars.next_if_eq(&'"').is_some() { - loop { - match chars.next()? { - '"' if chars.next_if_eq(&'"').is_some() => part.push('"'), - '"' => break, - c => part.push(c), - } - } - } else { - while let Some(c) = chars.next_if(|c| *c != '.') { - if c == '"' || c.is_whitespace() { - return None; - } - part.push(c.to_ascii_lowercase()); - } - } - - if part.is_empty() { - return None; - } - parts.push(part); - - match chars.next() { - None => return Some(parts), - Some('.') => continue, - Some(_) => return None, - } - } -} - /// A single table moved between databases. #[derive(Debug, Clone)] pub(crate) struct TableMove { - table: QualifiedTable, + /// 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: QualifiedTable) -> Self { + pub(crate) fn new(table: String) -> Self { Self { table, publication: format!("__pgdog_move_{}", random_string(19).to_lowercase()), @@ -129,33 +54,29 @@ impl TableMove { return Err(Error::SameDatabase(source.clone())); } - let referenced = self.referenced_tables(&state.source).await?; - if let Some((shard, missing)) = missing_tables(&state.destination, &referenced).await? { + let table = self.resolve(&state.source).await?; + if let Some((shard, missing)) = + missing_tables(&state.destination, &table.referenced).await? + { return Err(Error::ReferencedTablesMissing { - table: self.table.to_string(), + 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 \"{}\"", - escape_identifier(&self.table.schema) - ); + 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?; } - if let Err(err) = self.create_publication(&state.source).await { - self.drop_publication(&state.source).await; - return Err(err); - } + self.create_publication(&state.source, &table.name).await?; info!( "moving table {} [{}], publication \"{}\"", - self.table, state, self.publication + table.name, state, self.publication ); Ok(()) @@ -165,7 +86,7 @@ impl TableMove { 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 \"{}\"", + "table \"{}\" is left in place on \"{}\"; drop it once the application uses \"{}\"", self.table, state.source.identifier().database, state.destination.identifier().database @@ -175,26 +96,39 @@ impl TableMove { self.drop_publication(&state.source).await; } - /// Tables the moved table references with a foreign key, other than itself. - async fn referenced_tables(&self, source: &Cluster) -> Result, Error> { + async fn resolve(&self, source: &Cluster) -> Result { let mut server = source.primary(0, &Request::default()).await?; - Ok(server + let rows: Vec = server .fetch_all(format!( - "SELECT DISTINCT format('%I.%I', n.nspname, c.relname) \ - FROM pg_constraint con \ - JOIN pg_class c ON c.oid = con.confrelid \ + "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 \ - WHERE con.contype = 'f' AND con.conrelid = {} AND con.confrelid <> con.conrelid", - self.table.regclass() + 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?) + .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) -> Result<(), Error> { + async fn create_publication(&self, source: &Cluster, table: &str) -> Result<(), Error> { let sql = format!( - "CREATE PUBLICATION \"{}\" FOR TABLE {}", + "CREATE PUBLICATION \"{}\" FOR TABLE {table}", escape_identifier(&self.publication), - self.table ); // Publications are per database, and shards can share one. @@ -270,57 +204,8 @@ mod test { use crate::backend::server::test::test_server; use crate::config::config; - fn table(schema: &str, name: &str) -> QualifiedTable { - QualifiedTable { - schema: schema.into(), - name: name.into(), - } - } - - #[test] - fn parses_table_names() { - for (input, expected) in [ - ("orders", table("public", "orders")), - ("Orders", table("public", "orders")), - ("sales.orders", table("sales", "orders")), - ("\"Sales\".\"Orders\"", table("Sales", "Orders")), - ("sales.\"my.table\"", table("sales", "my.table")), - ("\"a\"\"b\".c", table("a\"b", "c")), - ] { - assert_eq!( - input.parse::().unwrap(), - expected, - "{input}" - ); - } - } - - #[test] - fn rejects_invalid_table_names() { - for input in [ - "", - ".", - "a.", - ".a", - "a..b", - "db.sales.orders", - "\"unterminated", - "\"\"", - "a\"b", - "\"a\"b", - "my table", - ] { - assert!(input.parse::().is_err(), "{input:?}"); - } - } - - #[test] - fn displays_quoted() { - assert_eq!(table("a\"b", "c").to_string(), "\"a\"\"b\".\"c\""); - } - #[tokio::test] - async fn destination_checks_find_missing_tables() { + async fn resolves_the_table_and_finds_missing_referenced_tables() { crate::logger(); let mut server = test_server().await; for ddl in [ @@ -336,11 +221,18 @@ mod test { let cluster = Cluster::new_test(&config()); cluster.launch(); - let child = TableMove::new(table("pgdog", "table_move_child")); - let referenced = child.referenced_tables(&cluster).await.unwrap(); - assert_eq!(referenced, vec!["pgdog.table_move_parent".to_owned()]); - assert_eq!(missing_tables(&cluster, &referenced).await.unwrap(), None); - assert_eq!(missing_tables(&cluster, &[]).await.unwrap(), None); + // 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, @@ -354,10 +246,11 @@ mod test { Some((0, vec!["pgdog.table_move_absent".to_owned()])) ); - let err = TableMove::new(table("pgdog", "table_move_absent")) - .referenced_tables(&cluster) + let err = TableMove::new("pgdog.table_move_absent".into()) + .resolve(&cluster) .await - .expect_err("table is missing on the source"); + .err() + .expect("table is missing on the source"); assert!(err.to_string().contains("does not exist"), "{err}"); cluster.shutdown(); @@ -382,8 +275,11 @@ mod test { let cluster = Cluster::new_test(&config()); cluster.launch(); - let table_move = TableMove::new(table("pgdog", "table_move_publication")); - table_move.create_publication(&cluster).await.unwrap(); + 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(); diff --git a/pgdog/src/backend/replication/tests.rs b/pgdog/src/backend/replication/tests.rs index 59ef9772e..5a7897077 100644 --- a/pgdog/src/backend/replication/tests.rs +++ b/pgdog/src/backend/replication/tests.rs @@ -10,7 +10,7 @@ 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::{QualifiedTable, TableMove}; +use super::logical::table_move::TableMove; use crate::{ api::{ MigrationError, @@ -993,10 +993,7 @@ async fn table_move_copies_and_replicates_one_table() -> Result<(), Box, + table: Option, /// Destination database. #[arg(long)] @@ -387,11 +387,7 @@ mod test { else { panic!("not a data sync command"); }; - let table = table.expect("table is set"); - assert_eq!( - (table.schema.as_str(), table.name.as_str()), - ("sales", "orders") - ); + assert_eq!(table.as_deref(), Some("sales.orders")); assert_eq!(publication, None); let Commands::DataSync { @@ -405,7 +401,6 @@ mod test { for args in [ vec![], - vec!["--table", "a.b.c"], vec!["--table", "t", "--publication", "p"], vec!["--table", "t", "--replicate-only"], ] {