From fef34be5ac1f1144a826a5049007dc9adf1f064b Mon Sep 17 00:00:00 2001
From: Dmitriy Pleshevskiy <dmitriy@ideascup.me>
Date: Thu, 25 Feb 2021 00:03:51 +0300
Subject: [PATCH] feat: add number to downgrade command config

---
 migra-cli/src/commands/downgrade.rs | 16 ++++++++++------
 migra-cli/src/main.rs               |  4 ++--
 migra-cli/src/opts.rs               |  9 ++++++++-
 migra-cli/tests/commands.rs         |  9 ++-------
 4 files changed, 22 insertions(+), 16 deletions(-)

diff --git a/migra-cli/src/commands/downgrade.rs b/migra-cli/src/commands/downgrade.rs
index 1723d66..2bc6bf5 100644
--- a/migra-cli/src/commands/downgrade.rs
+++ b/migra-cli/src/commands/downgrade.rs
@@ -1,20 +1,24 @@
 use crate::config::Config;
 use crate::database::prelude::*;
 use crate::database::MigrationManager;
+use crate::opts::DowngradeCommandOpt;
 use crate::StdResult;
+use std::cmp;
 use std::convert::TryFrom;
 
-pub(crate) fn downgrade_applied_migrations(config: Config) -> StdResult<()> {
+pub(crate) fn rollback_applied_migrations(
+    config: Config,
+    opts: DowngradeCommandOpt,
+) -> StdResult<()> {
     let mut manager = MigrationManager::try_from(&config)?;
 
     let applied_migrations = manager.applied_migration_names()?;
     let migrations = config.migrations()?;
 
-    if let Some(first_applied_migration) = applied_migrations.first() {
-        if let Some(migration) = migrations
-            .iter()
-            .find(|m| m.name() == first_applied_migration)
-        {
+    let rollback_migrations_number = cmp::min(opts.migrations_number, applied_migrations.len());
+
+    for migration_name in &applied_migrations[..rollback_migrations_number] {
+        if let Some(migration) = migrations.iter().find(|m| m.name() == migration_name) {
             println!("downgrade {}...", migration.name());
             manager.downgrade(&migration)?;
         }
diff --git a/migra-cli/src/main.rs b/migra-cli/src/main.rs
index 20b9f99..321ded4 100644
--- a/migra-cli/src/main.rs
+++ b/migra-cli/src/main.rs
@@ -35,9 +35,9 @@ fn main() -> StdResult<()> {
             let config = Config::read(opt.config)?;
             commands::upgrade_pending_migrations(config)?;
         }
-        Command::Downgrade => {
+        Command::Downgrade(opts) => {
             let config = Config::read(opt.config)?;
-            commands::downgrade_applied_migrations(config)?;
+            commands::rollback_applied_migrations(config, opts)?;
         }
         Command::Completions(opts) => {
             AppOpt::clap().gen_completions_to(
diff --git a/migra-cli/src/opts.rs b/migra-cli/src/opts.rs
index db4b648..366096b 100644
--- a/migra-cli/src/opts.rs
+++ b/migra-cli/src/opts.rs
@@ -27,7 +27,7 @@ pub(crate) enum Command {
     Upgrade,
 
     #[structopt(name = "downgrade", visible_alias = "down")]
-    Downgrade,
+    Downgrade(DowngradeCommandOpt),
 
     Completions(CompletionsShell),
 }
@@ -44,6 +44,13 @@ pub(crate) struct MakeCommandOpt {
     pub migration_name: String,
 }
 
+#[derive(Debug, StructOpt)]
+pub(crate) struct DowngradeCommandOpt {
+    /// How many applied migrations do we have to rollback
+    #[structopt(long = "number", short = "n", default_value = "1")]
+    pub migrations_number: usize,
+}
+
 #[derive(Debug, StructOpt)]
 pub(crate) enum CompletionsShell {
     Bash,
diff --git a/migra-cli/tests/commands.rs b/migra-cli/tests/commands.rs
index 796574b..f02e35a 100644
--- a/migra-cli/tests/commands.rs
+++ b/migra-cli/tests/commands.rs
@@ -229,13 +229,8 @@ Pending migrations:
             .arg("-c")
             .arg(path_to_file("Migra_env.toml"))
             .arg("down")
-            .assert()
-            .success();
-
-        Command::cargo_bin("migra")?
-            .arg("-c")
-            .arg(path_to_file("Migra_env.toml"))
-            .arg("down")
+            .arg("-n")
+            .arg("2")
             .assert()
             .success();