syn2mas: introduce a dry-run mode

Author: sandhose

crates/cli/src/commands/syn2mas.rs

@@ -70,8 +70,17 @@ enum Subcommand {
     ///
     /// It is OK for Synapse to be online during these checks.
     Check,
+
     /// Perform a migration. Synapse must be offline during this process.
-    Migrate,
+    Migrate {
+        /// Perform a dry-run migration, which is safe to run with Synapse
+        /// running, and will restore the MAS database to an empty state.
+        ///
+        /// This still *does* write to the MAS database, making it more
+        /// realistic compared to the final migration.
+        #[clap(long)]
+        dry_run: bool,
+    },
 }
 
 /// The number of parallel writing transactions active against the MAS database.
@@ -118,11 +127,10 @@ impl Options {
             .await
             .context("could not run migrations")?;
 
-        if matches!(&self.subcommand, Subcommand::Migrate) {
+        if matches!(&self.subcommand, Subcommand::Migrate { .. }) {
             // First perform a config sync
             // This is crucial to ensure we register upstream OAuth providers
             // in the MAS database
-            //
             let config = SyncConfig::extract(figment)?;
             let clock = SystemClock::default();
             let encrypter = config.secrets.encrypter();
@@ -201,7 +209,8 @@ impl Options {
 
                 Ok(ExitCode::SUCCESS)
             }
-            Subcommand::Migrate => {
+
+            Subcommand::Migrate { dry_run } => {
                 let provider_id_mappings: HashMap<String, Uuid> = {
                     let mas_oauth2 = UpstreamOAuth2Config::extract_or_default(figment)?;
 
@@ -217,8 +226,7 @@ impl Options {
 
                 // TODO how should we handle warnings at this stage?
 
-                // TODO this dry-run flag should be set to false in real circumstances !!!
-                let reader = SynapseReader::new(&mut syn_conn, true).await?;
+                let reader = SynapseReader::new(&mut syn_conn, dry_run).await?;
                 let writer_mas_connections =
                     futures_util::future::try_join_all((0..NUM_WRITER_CONNECTIONS).map(|_| {
                         database_connection_from_config_with_options(
@@ -230,7 +238,8 @@ impl Options {
                     }))
                     .instrument(tracing::info_span!("syn2mas.mas_writer_connections"))
                     .await?;
-                let writer = MasWriter::new(mas_connection, writer_mas_connections).await?;
+                let writer =
+                    MasWriter::new(mas_connection, writer_mas_connections, dry_run).await?;
 
                 let clock = SystemClock::default();
                 // TODO is this rng ok?

crates/syn2mas/src/mas_writer/mod.rs

@@ -242,6 +242,7 @@ impl FinishCheckerHandle {
 pub struct MasWriter {
     conn: LockedMasDatabase,
     writer_pool: WriterConnectionPool,
+    dry_run: bool,
 
     indices_to_restore: Vec<IndexDescription>,
     constraints_to_restore: Vec<ConstraintDescription>,
@@ -793,6 +794,7 @@ impl MasWriter {
     pub async fn new(
         mut conn: LockedMasDatabase,
         mut writer_connections: Vec<PgConnection>,
+        dry_run: bool,
     ) -> Result<Self, Error> {
         // Given that we don't have any concurrent transactions here,
         // the READ COMMITTED isolation level is sufficient.
@@ -902,7 +904,7 @@ impl MasWriter {
 
         Ok(Self {
             conn,
-
+            dry_run,
             writer_pool: WriterConnectionPool::new(writer_connections),
             indices_to_restore,
             constraints_to_restore,
@@ -987,7 +989,6 @@ impl MasWriter {
 
         // Now all the data has been migrated, finish off by restoring indices and
         // constraints!
-
         query("BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;")
             .execute(self.conn.as_mut())
             .await
@@ -1009,6 +1010,28 @@ impl MasWriter {
             .await
             .into_database("could not revert temporary tables")?;
 
+        // If we're in dry-run mode, truncate all the tables we've written to
+        if self.dry_run {
+            warn!("Migration ran in dry-run mode, deleting all imported data");
+            let tables = MAS_TABLES_AFFECTED_BY_MIGRATION
+                .iter()
+                .map(|table| format!("\"{table}\""))
+                .collect::<Vec<_>>()
+                .join(", ");
+
+            // Note that we do that with CASCADE, because we do that *after*
+            // restoring the FK constraints.
+            //
+            // The alternative would be to list all the tables we have FK to
+            // those tables, which would be a hassle, or to do that after
+            // restoring the constraints, which would mean we wouldn't validate
+            // that we've done valid FKs in dry-run mode.
+            query(&format!("TRUNCATE TABLE {tables} CASCADE;"))
+                .execute(self.conn.as_mut())
+                .await
+                .into_database_with(|| "failed to truncate all tables")?;
+        }
+
         query("COMMIT;")
             .execute(self.conn.as_mut())
             .await
@@ -1193,7 +1216,7 @@ mod test {
             .await
             .expect("failed to lock MAS database")
             .expect_left("MAS database is already locked");
-        MasWriter::new(locked_main_conn, writer_conns)
+        MasWriter::new(locked_main_conn, writer_conns, false)
             .await
             .expect("failed to construct MasWriter")
     }