From 0ce45be5360250567547b8f6fbe7138da3d81e9e Mon Sep 17 00:00:00 2001
From: Tristan Partin <tristan@neon.tech>
Date: Mon, 23 Dec 2024 11:41:24 -0600
Subject: [PATCH] Refactor MigrationRunner::run_migrations() to call a helper

This will make it easier to add cluster migrations, such as that for
CVE-2024-4317.

Link: https://www.postgresql.org/support/security/CVE-2024-4317/
Signed-off-by: Tristan Partin <tristan@neon.tech>
---
 compute_tools/src/migration.rs | 71 ++++++++++++++++------------------
 1 file changed, 33 insertions(+), 38 deletions(-)

diff --git a/compute_tools/src/migration.rs b/compute_tools/src/migration.rs
index 1f3de65806a8..ce8a9466b072 100644
--- a/compute_tools/src/migration.rs
+++ b/compute_tools/src/migration.rs
@@ -81,53 +81,48 @@ impl<'m> MigrationRunner<'m> {
         Ok(())
     }
 
-    /// Run the configrured set of migrations
-    pub fn run_migrations(mut self) -> Result<()> {
-        self.prepare_database()?;
+    /// Run an individual migration
+    fn run_migration(&mut self, migration_id: i64, migration: &str) -> Result<()> {
+        if migration.starts_with("-- SKIP") {
+            info!("Skipping migration id={}", migration_id);
 
-        let mut current_migration = self.get_migration_id()? as usize;
-        while current_migration < self.migrations.len() {
-            macro_rules! migration_id {
-                ($cm:expr) => {
-                    ($cm + 1) as i64
-                };
-            }
+            // Even though we are skipping the migration, updating the
+            // migration ID should help keep logic easy to understand when
+            // trying to understand the state of a cluster.
+            self.update_migration_id(migration_id)?;
+        } else {
+            info!("Running migration id={}:\n{}\n", migration_id, migration);
 
-            let migration = self.migrations[current_migration];
+            self.client
+                .simple_query("BEGIN")
+                .context("begin migration")?;
 
-            if migration.starts_with("-- SKIP") {
-                info!("Skipping migration id={}", migration_id!(current_migration));
+            self.client
+                .simple_query(migration)
+                .with_context(|| format!("run_migrations migration id={migration_id}"))?;
 
-                // Even though we are skipping the migration, updating the
-                // migration ID should help keep logic easy to understand when
-                // trying to understand the state of a cluster.
-                self.update_migration_id(migration_id!(current_migration))?;
-            } else {
-                info!(
-                    "Running migration id={}:\n{}\n",
-                    migration_id!(current_migration),
-                    migration
-                );
+            self.update_migration_id(migration_id)?;
 
-                self.client
-                    .simple_query("BEGIN")
-                    .context("begin migration")?;
+            self.client
+                .simple_query("COMMIT")
+                .context("commit migration")?;
 
-                self.client.simple_query(migration).with_context(|| {
-                    format!(
-                        "run_migrations migration id={}",
-                        migration_id!(current_migration)
-                    )
-                })?;
+            info!("Finished migration id={}", migration_id);
+        }
 
-                self.update_migration_id(migration_id!(current_migration))?;
+        Ok(())
+    }
 
-                self.client
-                    .simple_query("COMMIT")
-                    .context("commit migration")?;
+    /// Run the configrured set of migrations
+    pub fn run_migrations(mut self) -> Result<()> {
+        self.prepare_database()?;
 
-                info!("Finished migration id={}", migration_id!(current_migration));
-            }
+        let mut current_migration = self.get_migration_id()? as usize;
+        while current_migration < self.migrations.len() {
+            self.run_migration(
+                (current_migration + 1) as i64,
+                self.migrations[current_migration],
+            )?;
 
             current_migration += 1;
         }