Fix rollback when partially migrated

butane_cli/src/lib.rs

@@ -132,18 +132,34 @@ pub fn detach_latest_migration(base_dir: &PathBuf) -> Result<()> {
     Ok(())
 }
 
-pub fn migrate(base_dir: &PathBuf) -> Result<()> {
+pub fn migrate(base_dir: &PathBuf, name: Option<String>) -> Result<()> {
     let spec = load_connspec(base_dir)?;
     let mut conn = db::connect(&spec)?;
     let to_apply = get_migrations(base_dir)?.unapplied_migrations(&conn)?;
     println!("{} migrations to apply", to_apply.len());
     for m in to_apply {
         println!("Applying migration {}", m.name());
         m.apply(&mut conn)?;
+        if let Some(ref name) = name {
+            if name == &m.name().to_string() {
+                println!("Finishing at migration {}", m.name());
+                break;
+            }
+        }
     }
     Ok(())
 }
 
+pub fn rollback(base_dir: &PathBuf, name: Option<String>) -> Result<()> {
+    let spec = load_connspec(base_dir)?;
+    let conn = butane::db::connect(&spec)?;
+
+    match name {
+        Some(to) => rollback_to(base_dir, conn, &to),
+        None => rollback_latest(base_dir, conn),
+    }
+}
+
 pub fn rollback_to(base_dir: &Path, mut conn: Connection, to: &str) -> Result<()> {
     let ms = get_migrations(base_dir)?;
     let to_migration = match ms.get_migration(to) {
@@ -154,10 +170,35 @@ pub fn rollback_to(base_dir: &Path, mut conn: Connection, to: &str) -> Result<()
         }
     };
 
-    let to_unapply = ms.migrations_since(&to_migration)?;
+    let latest = ms
+        .last_applied_migration(&conn)
+        .unwrap_or_else(|err| {
+            eprintln!("Err: {err}");
+            std::process::exit(1);
+        })
+        .unwrap_or_else(|| {
+            eprintln!("No migrations applied!");
+            std::process::exit(1);
+        });
+
+    if to_migration == latest {
+        eprintln!("That is the latest applied migration, not rolling back to anything.");
+        std::process::exit(1);
+    }
+
+    let mut to_unapply = ms.migrations_since(&to_migration)?;
     if to_unapply.is_empty() {
         eprintln!("That is the latest migration, not rolling back to anything. If you expected something to happen, try specifying the migration to rollback to.");
+        std::process::exit(1);
+    }
+
+    if *to_unapply.last().unwrap() != latest {
+        let index = to_unapply.iter().position(|m| {
+            m.name() == latest.name()
+        }).unwrap();
+        to_unapply = to_unapply.split_at(index + 1).0.into();
     }
+
     for m in to_unapply.into_iter().rev() {
         println!("Rolling back migration {}", m.name());
         m.downgrade(&mut conn)?;
@@ -166,7 +207,7 @@ pub fn rollback_to(base_dir: &Path, mut conn: Connection, to: &str) -> Result<()
 }
 
 pub fn rollback_latest(base_dir: &Path, mut conn: Connection) -> Result<()> {
-    match get_migrations(base_dir)?.latest() {
+    match get_migrations(base_dir)?.last_applied_migration(&conn)? {
         Some(m) => {
             println!("Rolling back migration {}", m.name());
             m.downgrade(&mut conn)?;

butane_cli/src/main.rs

@@ -2,7 +2,7 @@ use std::path::PathBuf;
 
 use butane_cli::{
     base_dir, clean, clear_data, collapse_migrations, delete_table, detach_latest_migration, embed,
-    get_migrations, handle_error, init, list_migrations, make_migration, migrate, Result,
+    get_migrations, handle_error, init, list_migrations, make_migration, migrate, rollback,
 };
 use clap::{ArgAction, Parser, Subcommand};
 
@@ -42,7 +42,10 @@ However if the migration has been manually edited, it will need to be manually r
     )]
     DetachMigration,
     /// Apply migrations.
-    Migrate,
+    Migrate {
+        /// Migration to migrate to.
+        name: Option<String>,
+    },
     /// List migrations.
     List,
     /// Replace all migrations with a single migration representing the current model state.
@@ -55,7 +58,7 @@ However if the migration has been manually edited, it will need to be manually r
     /// Rollback migrations. With no arguments, undoes the latest migration. If the name of a migration is specified, rolls back until that migration is the latest applied migration.
     Rollback {
         /// Migration to roll back to.
-        name: String,
+        name: Option<String>,
     },
     /// Clear.
     Clear {
@@ -126,8 +129,8 @@ fn main() {
         )),
         Commands::MakeMigration { name } => handle_error(make_migration(&base_dir, Some(name))),
         Commands::DetachMigration => handle_error(detach_latest_migration(&base_dir)),
-        Commands::Migrate => handle_error(migrate(&base_dir)),
-        Commands::Rollback { name } => handle_error(rollback(&base_dir, Some(name))),
+        Commands::Migrate { name } => handle_error(migrate(&base_dir, name.to_owned())),
+        Commands::Rollback { name } => handle_error(rollback(&base_dir, name.to_owned())),
         Commands::Embed => handle_error(embed(&base_dir)),
         Commands::List => handle_error(list_migrations(&base_dir)),
         Commands::Collapse { name } => handle_error(collapse_migrations(&base_dir, Some(name))),
@@ -140,13 +143,3 @@ fn main() {
         Commands::Clean => handle_error(clean(&base_dir)),
     }
 }
-
-fn rollback(base_dir: &PathBuf, name: Option<&String>) -> Result<()> {
-    let spec = butane_cli::load_connspec(base_dir)?;
-    let conn = butane::db::connect(&spec)?;
-
-    match name {
-        Some(to) => butane_cli::rollback_to(base_dir, conn, to),
-        None => butane_cli::rollback_latest(base_dir, conn),
-    }
-}