From 5eb26748ab182fbd00b8efd26fa55cece4b8dbdb Mon Sep 17 00:00:00 2001
From: Alex V Kotelnikov <gudron@gudron.me>
Date: Thu, 13 Nov 2014 16:40:02 +0400
Subject: [PATCH] added restriction to run the rollback. small changes at
 rollback output.

---
 Mage/Command/BuiltIn/DeployCommand.php | 34 ++++++++++++++++++++++----
 1 file changed, 29 insertions(+), 5 deletions(-)

diff --git a/Mage/Command/BuiltIn/DeployCommand.php b/Mage/Command/BuiltIn/DeployCommand.php
index 7ee5235..e4e2a78 100644
--- a/Mage/Command/BuiltIn/DeployCommand.php
+++ b/Mage/Command/BuiltIn/DeployCommand.php
@@ -59,6 +59,15 @@ class DeployCommand extends AbstractCommand implements RequiresEnvironment
      */
     const IN_PROGRESS = 'in_progress';
 
+    /**
+     * Stage where possible throw Rollback Exception
+     * @var array
+     */
+    public $acceptedStagesToRollback = array(
+        AbstractTask::STAGE_POST_RELEASE,
+        AbstractTask::STAGE_POST_DEPLOY
+    );
+
     /**
      * Time the Deployment has Started
      * @var integer
@@ -430,11 +439,19 @@ class DeployCommand extends AbstractCommand implements RequiresEnvironment
         }
     }
 
-    protected function runRollbackTask(){
+    protected function runRollbackTask(AbstractTask $task){
         $this->getConfig()->reload();
         $hosts = $this->getConfig()->getHosts();
 
-        if (count($hosts) == 0) {
+        Console::output("",1,2);
+        Console::output("Starting the <bold>rollback</bold>",1,1);
+
+        if(!in_array($task->getStage(), $this->acceptedStagesToRollback ) ) {
+            $stagesString = implode(', ',$this->acceptedStagesToRollback);
+            Console::output("<light_purple>Warning!</light_purple> <bold>Rollback during deployment can be called only at the stages: $stagesString <bold>",1);
+            Console::output("<bold>Rollback:<bold> <red>ABORTING</red>",1,3);
+
+        } elseif (count($hosts) == 0) {
             Console::output('<light_purple>Warning!</light_purple> <bold>No hosts defined, unable to get releases.</bold>', 1, 3);
 
         } else {
@@ -451,7 +468,14 @@ class DeployCommand extends AbstractCommand implements RequiresEnvironment
                 $this->getConfig()->setHostConfig($hostConfig);
                 $this->getConfig()->setReleaseId(-1);
 
-                $task = Factory::get('releases/rollback', $this->getConfig());
+                $task = Factory::get(array(
+                        'name'=>'releases/rollback',
+                        'parameters' => array('inDeploy'=>true)
+                    ),
+                    $this->getConfig(),
+                    false,
+                    $task->getStage()
+                );
                 $task->init();
                 $result = $task->run() && $result;
 
@@ -495,8 +519,8 @@ class DeployCommand extends AbstractCommand implements RequiresEnvironment
                     $result = false;
                 }
             } catch (RollbackException $e) {
-                Console::output('<red>FAIL, Rollback started</red> [Message: ' . $e->getMessage() . ']', 0);
-                $this->runRollbackTask();
+                Console::output('<red>FAIL, Rollback catched</red> [Message: ' . $e->getMessage() . ']', 0);
+                $this->runRollbackTask($task);
                 $result = false;
 
             } catch (ErrorWithMessageException $e) {