Skip to content

Commit

Permalink
Refactor MigrationRunner::run_migrations() to call a helper
Browse files Browse the repository at this point in the history
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>
  • Loading branch information
tristan957 committed Dec 23, 2024
1 parent 5d2e111 commit 0ce45be
Showing 1 changed file with 33 additions and 38 deletions.
71 changes: 33 additions & 38 deletions compute_tools/src/migration.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down

0 comments on commit 0ce45be

Please sign in to comment.