From cbd10dbf36a2f0a0a76d18b8c653d12e3d96978f Mon Sep 17 00:00:00 2001
From: wulonghui <wulonghui6666@gmail.com>
Date: Fri, 25 Sep 2015 10:55:03 +0800
Subject: [PATCH] Update rollback process

---
 pkg/kubectl/cmd/rollingupdate.go      | 70 +++++++++++++++++++--------
 pkg/kubectl/cmd/rollingupdate_test.go |  2 +-
 2 files changed, 50 insertions(+), 22 deletions(-)

diff --git a/pkg/kubectl/cmd/rollingupdate.go b/pkg/kubectl/cmd/rollingupdate.go
index 7a283c4f179..3192d27e20e 100644
--- a/pkg/kubectl/cmd/rollingupdate.go
+++ b/pkg/kubectl/cmd/rollingupdate.go
@@ -100,47 +100,62 @@ func NewCmdRollingUpdate(f *cmdutil.Factory, out io.Writer) *cobra.Command {
 	return cmd
 }
 
-func validateArguments(cmd *cobra.Command, filenames, args []string) (deploymentKey, filename, image, oldName string, err error) {
-	deploymentKey = cmdutil.GetFlagString(cmd, "deployment-label-key")
-	image = cmdutil.GetFlagString(cmd, "image")
-	filename = ""
+func validateArguments(cmd *cobra.Command, filenames, args []string) error {
+	deploymentKey := cmdutil.GetFlagString(cmd, "deployment-label-key")
+	image := cmdutil.GetFlagString(cmd, "image")
+	rollback := cmdutil.GetFlagBool(cmd, "rollback")
 
 	if len(deploymentKey) == 0 {
-		return "", "", "", "", cmdutil.UsageError(cmd, "--deployment-label-key can not be empty")
+		return cmdutil.UsageError(cmd, "--deployment-label-key can not be empty")
 	}
 	if len(filenames) > 1 {
-		return "", "", "", "", cmdutil.UsageError(cmd, "May only specificy a single filename for new controller")
-	}
-	if len(filenames) > 0 {
-		filename = filenames[0]
-	}
-	if len(filenames) == 0 && len(image) == 0 {
-		return "", "", "", "", cmdutil.UsageError(cmd, "Must specify --filename or --image for new controller")
-	}
-	if len(filenames) != 0 && len(image) != 0 {
-		return "", "", "", "", cmdutil.UsageError(cmd, "--filename and --image can not both be specified")
-	}
-	if len(args) < 1 {
-		return "", "", "", "", cmdutil.UsageError(cmd, "Must specify the controller to update")
+		return cmdutil.UsageError(cmd, "May only specify a single filename for new controller")
 	}
 
-	return deploymentKey, filename, image, args[0], nil
+	if !rollback {
+		if len(filenames) == 0 && len(image) == 0 {
+			return cmdutil.UsageError(cmd, "Must specify --filename or --image for new controller")
+		}
+		if len(filenames) != 0 && len(image) != 0 {
+			return cmdutil.UsageError(cmd, "--filename and --image can not both be specified")
+		}
+	} else {
+		if len(filenames) != 0 || len(image) != 0 {
+			return cmdutil.UsageError(cmd, "Don't specify --filename or --image on rollback")
+		}
+	}
+
+	if len(args) < 1 {
+		return cmdutil.UsageError(cmd, "Must specify the controller to update")
+	}
+
+	return nil
 }
 
 func RunRollingUpdate(f *cmdutil.Factory, out io.Writer, cmd *cobra.Command, args []string, options *RollingUpdateOptions) error {
 	if len(os.Args) > 1 && os.Args[1] == "rollingupdate" {
 		printDeprecationWarning("rolling-update", "rollingupdate")
 	}
-	deploymentKey, filename, image, oldName, err := validateArguments(cmd, options.Filenames, args)
+	err := validateArguments(cmd, options.Filenames, args)
 	if err != nil {
 		return err
 	}
+
+	deploymentKey := cmdutil.GetFlagString(cmd, "deployment-label-key")
+	filename := ""
+	image := cmdutil.GetFlagString(cmd, "image")
+	oldName := args[0]
+	rollback := cmdutil.GetFlagBool(cmd, "rollback")
 	period := cmdutil.GetFlagDuration(cmd, "update-period")
 	interval := cmdutil.GetFlagDuration(cmd, "poll-interval")
 	timeout := cmdutil.GetFlagDuration(cmd, "timeout")
 	dryrun := cmdutil.GetFlagBool(cmd, "dry-run")
 	outputFormat := cmdutil.GetFlagString(cmd, "output")
 
+	if len(options.Filenames) > 0 {
+		filename = options.Filenames[0]
+	}
+
 	cmdNamespace, enforceNamespace, err := f.DefaultNamespace()
 	if err != nil {
 		return err
@@ -238,6 +253,19 @@ func RunRollingUpdate(f *cmdutil.Factory, out io.Writer, cmd *cobra.Command, arg
 			return err
 		}
 	}
+
+	if rollback {
+		keepOldName = len(args) == 1
+		newName := findNewName(args, oldRc)
+		if newRc, err = kubectl.LoadExistingNextReplicationController(client, cmdNamespace, newName); err != nil {
+			return err
+		}
+
+		if newRc == nil {
+			return cmdutil.UsageError(cmd, "Could not find %s to rollback.\n", newName)
+		}
+	}
+
 	if oldName == newRc.Name {
 		return cmdutil.UsageError(cmd, "%s cannot have the same name as the existing ReplicationController %s",
 			filename, oldName)
@@ -295,7 +323,7 @@ func RunRollingUpdate(f *cmdutil.Factory, out io.Writer, cmd *cobra.Command, arg
 		MaxUnavailable: util.NewIntOrStringFromInt(0),
 		MaxSurge:       util.NewIntOrStringFromInt(1),
 	}
-	if cmdutil.GetFlagBool(cmd, "rollback") {
+	if rollback {
 		err = kubectl.AbortRollingUpdate(config)
 		if err != nil {
 			return err
diff --git a/pkg/kubectl/cmd/rollingupdate_test.go b/pkg/kubectl/cmd/rollingupdate_test.go
index a132a7dcee6..bf2204b2731 100644
--- a/pkg/kubectl/cmd/rollingupdate_test.go
+++ b/pkg/kubectl/cmd/rollingupdate_test.go
@@ -79,7 +79,7 @@ func TestValidateArgs(t *testing.T) {
 				cmd.Flags().Set(key, val)
 			}
 		}
-		_, _, _, _, err := validateArguments(cmd, test.filenames, test.args)
+		err := validateArguments(cmd, test.filenames, test.args)
 		if err != nil && !test.expectErr {
 			t.Errorf("unexpected error: %v (%s)", err, test.testName)
 		}