From 0a1be86e1d61e44dc0278e50c55928bc72e75ec3 Mon Sep 17 00:00:00 2001 From: meskill <8974488+meskill@users.noreply.github.com> Date: Fri, 25 Sep 2026 09:41:26 +0000 Subject: [PATCH] feat(resharding): add validation step for schema_sync --- README.md | 2 +- docs/RESHARDING.md | 28 +-- .../admin/resharding/schema_sync.rs | 167 +++++++++++++++++- pgdog-stats/src/resharding.rs | 3 + pgdog-stats/src/task.rs | 12 ++ pgdog/src/admin/schema_sync.rs | 2 +- pgdog/src/api/schema_sync.rs | 19 +- pgdog/src/backend/schema/sync/pg_dump.rs | 92 ++++++++-- pgdog/src/cli.rs | 105 +++++++++-- 9 files changed, 370 insertions(+), 60 deletions(-) diff --git a/README.md b/README.md index cd0ad5c7f..9ec6ed361 100644 --- a/README.md +++ b/README.md @@ -566,7 +566,7 @@ The re-sharding process is done in 5 steps: 1. Create new empty cluster with the desired number of shards 2. Configure it in `pgdog.toml` and run `schema-sync` command to copy table schemas to the new databases 3. Run `data-sync` command to copy and re-shard table data with logical replication (tables are copied in parallel) -4. While keeping previous command running (it streams row updates in real-time), run `schema-sync --data-sync-complete` to create secondary indexes on the new databases (much faster to do this after data is copied) +4. While keeping previous command running (it streams row updates in real-time), run `schema-sync --phase post` to create secondary indexes on the new databases (much faster to do this after data is copied) 5. Cutover traffic to new cluster with `MAINTENANCE ON`, `RELOAD`, `MAINTENANCE OFF` command sequence Cutover can be done atomically with multiple PgDog containers because `RELOAD` doesn't resume traffic, `MAINTENANCE OFF` does, so the config is the same in all containers before queries are resumed. No complex synchronization tooling like etcd or Zookeeper is required. diff --git a/docs/RESHARDING.md b/docs/RESHARDING.md index a14b9e039..0a4c3294d 100644 --- a/docs/RESHARDING.md +++ b/docs/RESHARDING.md @@ -53,7 +53,7 @@ flowchart LR `Orchestrator::load_schema()` creates a `PgDump` ([`pgdog/src/backend/schema/sync/pg_dump.rs`](../pgdog/src/backend/schema/sync/pg_dump.rs)) with the source cluster and publication name, calls `pg_dump.dump().await`, and stores the `PgDumpOutput` on the orchestrator. This output carries pre-data (tables, types, extensions, -primary key constraints), secondary index DDL, post-cutover operations, and sequences — split +primary key constraints), secondary index DDL, validation operations, and sequences — split into `SyncState` phases so they can be applied in the right order later. --- @@ -122,9 +122,9 @@ The recorded LSN becomes the replay watermark for that table's WAL stream in Ste ## Step 4 — Post-data schema sync -`schema_sync_post()` restores `SyncState::PostData` — secondary indexes, non-PK constraints, -and any other DDL that was deferred. Deferring index creation until after the bulk copy avoids -index maintenance overhead during the high-throughput copy phase. +`schema_sync_post()` restores `SyncState::PostData` for secondary indexes and non-PK constraints. +It creates eligible foreign keys as `NOT VALID`, which avoids a blocking data scan. +Deferring index creation also avoids index maintenance during the bulk copy. --- @@ -199,8 +199,11 @@ returns `Error::DrainTimeout`. The reverse phase runs in the same task, not in a separate one. A `STOP_TASK` during the reverse phase ends the rollback window, and the task reports the migration as finished. -The cutover schema sync (`SyncState::Cutover`, then `SyncState::PostCutover`) runs as a -`SchemaSyncTask` subtask after each stream phase ends. +`SyncState::Cutover` runs as an awaited `SchemaSyncTask` before the traffic swap. +`SyncState::PostDataValidation` is not part of cutover. +Run it manually against the destination during replication or at another time. +It validates each deferred foreign key with `ALTER TABLE ... VALIDATE CONSTRAINT`. +Validation failures remain visible through `SHOW TASKS`. --- @@ -214,12 +217,10 @@ steps. A failure here leaves traffic unaffected and the source untouched, making ### Schema DDL — intentional error tolerance -`schema_sync_pre`, `schema_sync_post`, `schema_sync_cutover`, and `schema_sync_post_cutover` are -all called with `ignore_errors = true`. The `PgDumpOutput::restore()` method logs errors and -continues when this flag is set. The intent is to tolerate pre-existing objects on the destination -— a common condition when a previous reshard attempt failed mid-schema-sync and left partial DDL -behind. Re-running `RESHARD` after such a failure will not abort on `table already exists` or -similar conflicts. +The pre-data, post-data, and cutover stages use `ignore_errors = true`. +This setting tolerates objects left by an earlier failed schema sync. +The validation stage uses `ignore_errors = false`. +A foreign key violation therefore marks the validation task as failed. ### Data sync — abort propagation and cooperative cancellation @@ -280,7 +281,8 @@ Several mechanisms make it safe to replay data across a restart: | Mechanism | Where | Effect | |---|---|---| | Temporary replication slots | `Table::data_sync()` | Auto-dropped on connection close; no orphaned per-table slots | -| `ignore_errors = true` | All schema sync steps | Pre-existing DDL on destination does not abort the run | +| `ignore_errors = true` | Pre-data, post-data, and cutover schema sync | Pre-existing DDL does not abort the run | +| FK validation | Manual validation schema sync | Validates deferred foreign keys in a reported task | | LSN watermark guard | `StreamSubscriber::lsn_applied()` | Rows bulk-copied in Step 3 are skipped during WAL replay in Step 5 | | Upsert on INSERT messages | `Table::insert(upsert=true)` | `ON CONFLICT (pk) DO UPDATE SET` prevents duplicates on WAL re-delivery | | PK validation | `Table::valid()` | Fails before any data moves; restart is clean | \ No newline at end of file diff --git a/integration/rust/tests/integration/admin/resharding/schema_sync.rs b/integration/rust/tests/integration/admin/resharding/schema_sync.rs index 14ed216a8..99dc35b27 100644 --- a/integration/rust/tests/integration/admin/resharding/schema_sync.rs +++ b/integration/rust/tests/integration/admin/resharding/schema_sync.rs @@ -1,12 +1,21 @@ -use crate::setup::{admin_sqlx, connection_sqlx_direct}; +use crate::setup::{admin_sqlx, connection_sqlx_direct, connection_sqlx_direct_db}; use pgdog_stats::TaskProgress; use sqlx::{Executor, Pool, Postgres, Row}; +use std::time::Duration; +use tokio::time::{sleep, timeout}; use super::{ - TEST_PUB, TEST_SCHEMA, TEST_TABLE, assert_layout, cleanup, create_publication, + POLL, TEST_PUB, TEST_SCHEMA, TEST_TABLE, assert_layout, cleanup, create_publication, create_test_table, run_task_command, wait_for_relation_on_shards, wait_for_task_status, + with_cleanup, }; +const VALID_FK: &str = "validation_children_parent_fk"; +const SOURCE_NOT_VALID_FK: &str = "validation_unchecked_parent_fk"; +const VALIDATION_PARENT: &str = "validation_parents"; +const VALIDATION_CHILD: &str = "validation_children"; +const VALIDATION_UNCHECKED_CHILD: &str = "validation_unchecked_children"; + const SHOW_SCHEMA_SYNC_LAYOUT: &[(&str, &str)] = &[ ("parent_id", "INT8"), ("id", "INT8"), @@ -116,6 +125,94 @@ async fn assert_schema_sync_rows(admin: &Pool, task_id: i64, sync_stat } } +async fn create_validation_schema(direct: &Pool) { + let script = format!( + "CREATE SCHEMA {TEST_SCHEMA}; + + CREATE TABLE {TEST_SCHEMA}.{VALIDATION_PARENT} (id BIGINT PRIMARY KEY); + + CREATE TABLE {TEST_SCHEMA}.{VALIDATION_CHILD} ( + id BIGINT PRIMARY KEY, + parent_id BIGINT, + CONSTRAINT {VALID_FK} FOREIGN KEY (parent_id) + REFERENCES {TEST_SCHEMA}.{VALIDATION_PARENT}(id) + ); + + CREATE TABLE {TEST_SCHEMA}.{VALIDATION_UNCHECKED_CHILD} ( + id BIGINT PRIMARY KEY, + parent_id BIGINT + ); + + ALTER TABLE {TEST_SCHEMA}.{VALIDATION_UNCHECKED_CHILD} + ADD CONSTRAINT {SOURCE_NOT_VALID_FK} + FOREIGN KEY (parent_id) + REFERENCES {TEST_SCHEMA}.{VALIDATION_PARENT}(id) + NOT VALID; + + CREATE PUBLICATION {TEST_PUB} FOR TABLE + {TEST_SCHEMA}.{VALIDATION_PARENT}, + {TEST_SCHEMA}.{VALIDATION_CHILD}, + {TEST_SCHEMA}.{VALIDATION_UNCHECKED_CHILD};" + ); + + direct.execute(script.as_str()).await.unwrap(); +} + +async fn run_schema_sync_phase(admin: &Pool, phase: &str) -> i64 { + let task_id = run_task_command( + admin, + &format!("SCHEMA_SYNC {phase} pgdog pgdog_sharded {TEST_PUB}"), + ) + .await; + wait_for_task_status(admin, task_id, TaskProgress::Finished).await; + task_id +} + +async fn constraint_is_validated(database: &str, constraint: &str) -> bool { + let database = connection_sqlx_direct_db(database).await; + sqlx::query_scalar( + "SELECT constraint_row.convalidated + FROM pg_constraint AS constraint_row + JOIN pg_namespace AS namespace + ON namespace.oid = constraint_row.connamespace + WHERE namespace.nspname = $1 + AND constraint_row.conname = $2", + ) + .bind(TEST_SCHEMA) + .bind(constraint) + .fetch_one(&database) + .await + .unwrap() +} + +async fn assert_constraint_validation(constraint: &str, expected: bool) { + for database in ["shard_0", "shard_1"] { + assert_eq!( + constraint_is_validated(database, constraint).await, + expected, + "{constraint} validation state on {database}" + ); + } +} + +async fn wait_for_validation_failure(admin: &Pool, task_id: i64) -> SchemaSyncRow { + timeout(Duration::from_secs(30), async { + loop { + if let Some(row) = schema_sync_rows(admin) + .await + .into_iter() + .find(|row| row.id == task_id && row.shard.is_none()) + && row.status.is_error() + { + return row; + } + sleep(POLL).await; + } + }) + .await + .expect("validation task did not fail in time") +} + #[tokio::test] async fn test_schema_sync_pre() { let direct = connection_sqlx_direct().await; @@ -175,3 +272,69 @@ async fn test_schema_sync_post() { cleanup(&admin, &direct).await; } + +#[tokio::test] +async fn test_schema_sync_validation() { + let direct = connection_sqlx_direct().await; + let admin = admin_sqlx().await; + cleanup(&admin, &direct).await; + + with_cleanup(&admin, &direct, async { + create_validation_schema(&direct).await; + run_schema_sync_phase(&admin, "pre").await; + run_schema_sync_phase(&admin, "post").await; + + assert_constraint_validation(VALID_FK, false).await; + assert_constraint_validation(SOURCE_NOT_VALID_FK, false).await; + + let task_id = run_schema_sync_phase(&admin, "post-data-validation").await; + + assert_constraint_validation(VALID_FK, true).await; + assert_constraint_validation(SOURCE_NOT_VALID_FK, false).await; + assert_schema_sync_rows(&admin, task_id, "post_data_validation").await; + }) + .await; +} + +#[tokio::test] +async fn test_schema_sync_validation_rejects_existing_orphan() { + let direct = connection_sqlx_direct().await; + let admin = admin_sqlx().await; + cleanup(&admin, &direct).await; + + with_cleanup(&admin, &direct, async { + create_validation_schema(&direct).await; + run_schema_sync_phase(&admin, "pre").await; + + let shard = connection_sqlx_direct_db("shard_0").await; + shard + .execute( + format!( + "INSERT INTO {TEST_SCHEMA}.{VALIDATION_CHILD} (id, parent_id) VALUES (1, 42)" + ) + .as_str(), + ) + .await + .unwrap(); + + run_schema_sync_phase(&admin, "post").await; + + let task_id = run_task_command( + &admin, + &format!("SCHEMA_SYNC post-data-validation pgdog pgdog_sharded {TEST_PUB}"), + ) + .await; + let failure = wait_for_validation_failure(&admin, task_id).await; + + assert_eq!(failure.sync_state, "post_data_validation"); + assert_eq!(failure.source, "pgdog"); + assert_eq!(failure.destination, "pgdog_sharded"); + let TaskProgress::Error { message } = failure.status else { + panic!("validation task did not report an error"); + }; + assert!(message.contains("23503"), "{message}"); + assert!(message.contains(VALID_FK), "{message}"); + assert!(!constraint_is_validated("shard_0", VALID_FK).await); + }) + .await; +} diff --git a/pgdog-stats/src/resharding.rs b/pgdog-stats/src/resharding.rs index 1c0ba12ef..cda5a988e 100644 --- a/pgdog-stats/src/resharding.rs +++ b/pgdog-stats/src/resharding.rs @@ -80,6 +80,9 @@ pub enum SyncState { PreData, PostData, Cutover, + PostDataValidation, + #[serde(other)] + Unknown, } #[derive(Debug, Default, Clone, Copy, PartialEq, Serialize, Deserialize, schemars::JsonSchema)] diff --git a/pgdog-stats/src/task.rs b/pgdog-stats/src/task.rs index b9f87b442..7c3d5d9e4 100644 --- a/pgdog-stats/src/task.rs +++ b/pgdog-stats/src/task.rs @@ -543,6 +543,10 @@ mod test { SyncState::Cutover, "schema_sync(cutover) prod -> prod_sharded", ), + ( + SyncState::PostDataValidation, + "schema_sync(post_data_validation) prod -> prod_sharded", + ), ] { assert_eq!( TaskDefinition::from(SchemaSyncDefinition { @@ -791,6 +795,14 @@ mod test { .unwrap(), TaskStatus::Replication(ReplicationStatus::Other) ); + assert_eq!( + serde_json::from_str::( + r#"{"name":"schema_sync","kind":"schema_sync","databases":{"source":"prod","destination":"prod_sharded"},"sync_state":"new_phase","ignore_errors":false,"dry_run":false}"# + ) + .unwrap() + .to_string(), + "schema_sync(unknown) prod -> prod_sharded" + ); assert_eq!( serde_json::from_str::>( diff --git a/pgdog/src/admin/schema_sync.rs b/pgdog/src/admin/schema_sync.rs index 2ce5625ad..0f65b5f4c 100644 --- a/pgdog/src/admin/schema_sync.rs +++ b/pgdog/src/admin/schema_sync.rs @@ -70,7 +70,7 @@ impl Command for SchemaSync { }) .publication(self.publication.clone()) .phase(self.phase) - .ignore_errors(true) + .ignore_errors(self.phase != SchemaSyncPhase::PostDataValidation) .build(), ) .id(); diff --git a/pgdog/src/api/schema_sync.rs b/pgdog/src/api/schema_sync.rs index af44f2c10..f07af5b4d 100644 --- a/pgdog/src/api/schema_sync.rs +++ b/pgdog/src/api/schema_sync.rs @@ -1,5 +1,3 @@ -//! Schema-sync background task (pre-data, post-data, or cutover). - use std::sync::Arc; use std::time::Duration; @@ -23,14 +21,14 @@ use crate::backend::schema::sync::pg_dump::{Collapsed, PgDumpOutput, Statement, /// The source dump, taken once and shared by every phase of one migration. pub(crate) type SchemaDump = Arc>>; -#[derive(Clone, Copy, PartialEq, Eq, Debug, Display, FromStr)] +#[derive(Clone, Copy, PartialEq, Eq, Debug, Display, FromStr, clap::ValueEnum)] +#[display(rename_all = "kebab-case")] +#[from_str(rename_all = "kebab-case")] pub(crate) enum SchemaSyncPhase { - #[display("pre")] Pre, - #[display("post")] Post, - #[display("cutover")] Cutover, + PostDataValidation, } impl From for SyncState { @@ -39,6 +37,7 @@ impl From for SyncState { SchemaSyncPhase::Pre => SyncState::PreData, SchemaSyncPhase::Post => SyncState::PostData, SchemaSyncPhase::Cutover => SyncState::Cutover, + SchemaSyncPhase::PostDataValidation => SyncState::PostDataValidation, } } } @@ -268,15 +267,12 @@ mod tests { ("pre", SchemaSyncPhase::Pre), ("post", SchemaSyncPhase::Post), ("cutover", SchemaSyncPhase::Cutover), + ("post-data-validation", SchemaSyncPhase::PostDataValidation), ] { assert_eq!(text.parse::().unwrap(), phase); assert_eq!(phase.to_string(), text); } - // Parsing is case-insensitive; unknown phases are rejected. - assert_eq!( - "CUTOVER".parse::().unwrap(), - SchemaSyncPhase::Cutover - ); + assert!("CUTOVER".parse::().is_err()); assert!("bogus".parse::().is_err()); } @@ -289,6 +285,7 @@ mod tests { (SchemaSyncPhase::Pre, "pre_data"), (SchemaSyncPhase::Post, "post_data"), (SchemaSyncPhase::Cutover, "cutover"), + (SchemaSyncPhase::PostDataValidation, "post_data_validation"), ] { let definition = TaskDefinition::from(SchemaSyncDefinition { databases: Databases { diff --git a/pgdog/src/backend/schema/sync/pg_dump.rs b/pgdog/src/backend/schema/sync/pg_dump.rs index 5d18df6c4..56101bd4b 100644 --- a/pgdog/src/backend/schema/sync/pg_dump.rs +++ b/pgdog/src/backend/schema/sync/pg_dump.rs @@ -37,6 +37,23 @@ fn schema_name(relation: &nodes::RangeVar) -> &str { relation.schemaname().unwrap_or("public") } +fn validate_constraint( + relation: &nodes::RangeVar, + constraint: &str, +) -> Result { + // TODO: multiple times we create schema.table identifier, + // make it reusable + let schema = crate::util::escape_identifier(schema_name(relation)); + let table = + crate::util::escape_identifier(relation.relname().ok_or(SchemaSyncError::MissingEntity)?); + let constraint = crate::util::escape_identifier(constraint); + let only = if relation.inh { "" } else { "ONLY " }; + + Ok(Statement::new(format!( + "ALTER TABLE {only}\"{schema}\".\"{table}\" VALIDATE CONSTRAINT \"{constraint}\"" + ))) +} + fn is_integer_type(type_name: &str) -> bool { matches!( type_name, @@ -700,24 +717,37 @@ impl PgDumpOutput { } } else if cons.contype == nodes::ConstrType::CONSTR_FOREIGN { - // FK columns referencing integer PKs are - // computed from fk_columns at the end - if state == SyncState::PostData { - let sql = if cons.skip_validation - || stmt.relation().is_some_and(|relation| { - partitioned_tables - .contains(&Table::from(relation)) - }) { - original.to_owned() - } else { - format!( + let partitioned = + stmt.relation().is_some_and(|relation| { + partitioned_tables + .contains(&Table::from(relation)) + }); + let defer_validation = + !cons.skip_validation && !partitioned; + + let statement = match (state, defer_validation) { + (SyncState::PostData, true) => { + Statement::new(format!( "{} NOT VALID", original.trim_end().trim_end_matches(';') - ) - }; - result - .push(Statement::new(sql).set_skip_if_exists()); - } + )) + .set_skip_if_exists() + } + (SyncState::PostData, false) => { + Statement::new(original).set_skip_if_exists() + } + (SyncState::PostDataValidation, true) => { + let relation = stmt + .relation() + .ok_or(SchemaSyncError::MissingEntity)?; + let constraint = cons + .conname() + .ok_or(SchemaSyncError::MissingEntity)?; + validate_constraint(relation, constraint)? + } + _ => continue, + }; + result.push(statement); } else if state == SyncState::PostData { result.push( Statement::new(original).set_skip_if_exists(), @@ -1263,7 +1293,7 @@ ALTER TABLE test ADD CONSTRAINT id_pkey PRIMARY KEY (id);"#, CREATE TABLE parent (id INTEGER, name TEXT); CREATE TABLE child (id INTEGER, parent_id INTEGER); ALTER TABLE parent ADD CONSTRAINT parent_pkey PRIMARY KEY (id); -ALTER TABLE child ADD CONSTRAINT child_parent_fk FOREIGN KEY (parent_id) REFERENCES parent(id);"#, +ALTER TABLE ONLY child ADD CONSTRAINT child_parent_fk FOREIGN KEY (parent_id) REFERENCES parent(id);"#, ); let statements = output.statements(SyncState::PreData).unwrap(); @@ -1285,9 +1315,13 @@ ALTER TABLE child ADD CONSTRAINT child_parent_fk FOREIGN KEY (parent_id) REFEREN ); assert_eq!( output.statements(SyncState::PostData).unwrap()[0].sql, - "ALTER TABLE child ADD CONSTRAINT child_parent_fk FOREIGN KEY (parent_id) REFERENCES parent(id) NOT VALID" + "ALTER TABLE ONLY child ADD CONSTRAINT child_parent_fk FOREIGN KEY (parent_id) REFERENCES parent(id) NOT VALID" ); assert!(output.statements(SyncState::Cutover).unwrap().is_empty()); + assert_eq!( + output.statements(SyncState::PostDataValidation).unwrap()[0].sql, + "ALTER TABLE ONLY \"public\".\"child\" VALIDATE CONSTRAINT \"child_parent_fk\"" + ); } #[test] @@ -1302,6 +1336,12 @@ ALTER TABLE child ADD CONSTRAINT child_parent_fk FOREIGN KEY (parent_id) REFEREN assert_eq!(statements[0].sql, sql.trim_end_matches(';')); assert!(output.statements(SyncState::PreData).unwrap().is_empty()); assert!(output.statements(SyncState::Cutover).unwrap().is_empty()); + assert!( + output + .statements(SyncState::PostDataValidation) + .unwrap() + .is_empty() + ); } #[test] @@ -1335,14 +1375,16 @@ ALTER TABLE child ADD CONSTRAINT child_parent_fk FOREIGN KEY (parent_id) REFEREN CREATE TABLE child_0 PARTITION OF child FOR VALUES FROM (0) TO (100); ALTER TABLE child ADD CONSTRAINT child_parent_fk FOREIGN KEY (parent_id) REFERENCES parent(id);", ); - for (version, expected) in [ + for (version, expected, validation) in [ ( 170_009, "ALTER TABLE child ADD CONSTRAINT child_parent_fk FOREIGN KEY (parent_id) REFERENCES parent(id)", + None, ), ( 180_000, "ALTER TABLE child ADD CONSTRAINT child_parent_fk FOREIGN KEY (parent_id) REFERENCES parent(id) NOT VALID", + Some("ALTER TABLE \"public\".\"child\" VALIDATE CONSTRAINT \"child_parent_fk\""), ), ] { output.destination_version = version; @@ -1350,6 +1392,14 @@ ALTER TABLE child ADD CONSTRAINT child_parent_fk FOREIGN KEY (parent_id) REFEREN assert_eq!(statements.len(), 1); assert_eq!(statements[0].sql.trim(), expected); assert!(output.statements(SyncState::Cutover).unwrap().is_empty()); + assert_eq!( + output + .statements(SyncState::PostDataValidation) + .unwrap() + .first() + .map(|statement| statement.sql.as_str()), + validation + ); } } @@ -1373,6 +1423,10 @@ ALTER TABLE child ADD CONSTRAINT child_parent_fk FOREIGN KEY (parent_id) REFEREN statements[1].sql.trim(), "ALTER TABLE other.child ADD CONSTRAINT partitioned_fk FOREIGN KEY (parent_id) REFERENCES public.parent(id) ON DELETE CASCADE DEFERRABLE INITIALLY DEFERRED" ); + assert_eq!( + output.statements(SyncState::PostDataValidation).unwrap()[0].sql, + "ALTER TABLE \"public\".\"child\" VALIDATE CONSTRAINT \"ordinary_fk\"" + ); } #[test] diff --git a/pgdog/src/cli.rs b/pgdog/src/cli.rs index c23c6f85f..d0615b18a 100644 --- a/pgdog/src/cli.rs +++ b/pgdog/src/cli.rs @@ -128,13 +128,23 @@ pub(crate) enum Commands { #[arg(long)] ignore_errors: bool, - /// Data sync has been complete. - #[arg(long)] + #[arg( + long, + value_enum, + help = "Schema sync phase to run.", + default_value_t = SchemaSyncPhase::Pre, + conflicts_with_all = ["data_sync_complete", "cutover", "validation"] + )] + phase: SchemaSyncPhase, + + #[arg(long, hide = true, conflicts_with_all = ["cutover", "validation"])] data_sync_complete: bool, - /// Execute cutover statements. - #[arg(long)] + #[arg(long, hide = true, conflicts_with = "validation")] cutover: bool, + + #[arg(long, hide = true)] + validation: bool, }, /// For testing purposes only. @@ -271,6 +281,22 @@ pub(crate) async fn data_sync(commands: Commands) -> Result<(), Box Option { + if data_sync_complete { + Some(SchemaSyncPhase::Post) + } else if cutover { + Some(SchemaSyncPhase::Cutover) + } else if validation { + Some(SchemaSyncPhase::PostDataValidation) + } else { + None + } +} + pub(crate) async fn schema_sync(commands: Commands) -> Result<(), Box> { if let Commands::SchemaSync { from_database, @@ -278,18 +304,12 @@ pub(crate) async fn schema_sync(commands: Commands) -> Result<(), Box Result<(), Box Result<(), Box Result { + let mut argv = vec![ + "pgdog", + "schema-sync", + "--from-database", + "a", + "--to-database", + "b", + "--publication", + "p", + ]; + argv.extend_from_slice(args); + + let Some(Commands::SchemaSync { + phase, + data_sync_complete, + cutover, + validation, + .. + }) = Cli::try_parse_from(argv)?.command + else { + panic!("not a schema sync command"); + }; + + Ok(legacy_phase(data_sync_complete, cutover, validation).unwrap_or(phase)) + } + + #[test] + fn test_phase_flags() { + for (args, expected) in [ + (vec![], Some(SchemaSyncPhase::Pre)), + (vec!["--phase", "post"], Some(SchemaSyncPhase::Post)), + (vec!["--phase", "cutover"], Some(SchemaSyncPhase::Cutover)), + ( + vec!["--phase", "post-data-validation"], + Some(SchemaSyncPhase::PostDataValidation), + ), + (vec!["--data-sync-complete"], Some(SchemaSyncPhase::Post)), + (vec!["--cutover"], Some(SchemaSyncPhase::Cutover)), + ( + vec!["--validation"], + Some(SchemaSyncPhase::PostDataValidation), + ), + (vec!["--phase", "bogus"], None), + (vec!["--phase", "post", "--cutover"], None), + (vec!["--cutover", "--validation"], None), + (vec!["--data-sync-complete", "--validation"], None), + ] { + assert_eq!(parse(&args).ok(), expected, "{args:?}"); + } + } + } +}