From f640a294a4e4fa92346b6ba993443c2829159f39 Mon Sep 17 00:00:00 2001
From: FUJIWARA Shunichiro <fujiwara.shunichiro@gmail.com>
Date: Wed, 14 Mar 2018 17:41:24 +0900
Subject: [PATCH] add --deregister-task-definition to rollback command

Deregister a rolled back task definition when rollback is over successfully.
---
 cmd/ecspresso/main.go |  1 +
 ecspresso.go          | 30 +++++++++++++++++++++++-------
 options.go            |  3 ++-
 3 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/cmd/ecspresso/main.go b/cmd/ecspresso/main.go
index fea8f0df..336cd35e 100644
--- a/cmd/ecspresso/main.go
+++ b/cmd/ecspresso/main.go
@@ -38,6 +38,7 @@ func _main() int {
 	rollback := kingpin.Command("rollback", "rollback service")
 	rollbackOption := ecspresso.RollbackOption{
 		DryRun: rollback.Flag("dry-run", "dry-run").Bool(),
+		DeregisterTaskDefinition: rollback.Flag("dereginster-task-definition", "deregister rolled back task definition").Bool(),
 	}
 
 	sub := kingpin.Parse()
diff --git a/ecspresso.go b/ecspresso.go
index 5953ef20..d30f29d4 100644
--- a/ecspresso.go
+++ b/ecspresso.go
@@ -210,7 +210,7 @@ func (d *App) Deploy(opt DeployOption) error {
 		return nil
 	}
 
-	if err := d.UpdateService(ctx, tdArn, count, opt.ForceNewDeployment); err != nil {
+	if err := d.UpdateService(ctx, tdArn, *count, *opt.ForceNewDeployment); err != nil {
 		return errors.Wrap(err, "deploy failed")
 	}
 	if err := d.WaitServiceStable(ctx, time.Now()); err != nil {
@@ -230,7 +230,8 @@ func (d *App) Rollback(opt RollbackOption) error {
 	if err != nil {
 		return errors.Wrap(err, "rollback failed")
 	}
-	targetArn, err := d.FindRollbackTarget(ctx, *service.TaskDefinition)
+	currentArn := *service.TaskDefinition
+	targetArn, err := d.FindRollbackTarget(ctx, currentArn)
 	if err != nil {
 		return errors.Wrap(err, "rollback failed")
 	}
@@ -240,7 +241,7 @@ func (d *App) Rollback(opt RollbackOption) error {
 		return nil
 	}
 
-	if err := d.UpdateService(ctx, targetArn, service.DesiredCount, nil); err != nil {
+	if err := d.UpdateService(ctx, targetArn, *service.DesiredCount, false); err != nil {
 		return errors.Wrap(err, "rollback failed")
 	}
 	if err := d.WaitServiceStable(ctx, time.Now()); err != nil {
@@ -248,6 +249,21 @@ func (d *App) Rollback(opt RollbackOption) error {
 	}
 
 	d.Log("Service is stable now. Completed!")
+
+	if *opt.DeregisterTaskDefinition {
+		d.Log("Deregistering rolled back task definition", arnToName(currentArn))
+		_, err := d.ecs.DeregisterTaskDefinitionWithContext(
+			ctx,
+			&ecs.DeregisterTaskDefinitionInput{
+				TaskDefinition: &currentArn,
+			},
+		)
+		if err != nil {
+			return errors.Wrap(err, "deregister task definition failed")
+		}
+		d.Log(arnToName(currentArn), "was deregistered successfully")
+	}
+
 	return nil
 }
 
@@ -318,9 +334,9 @@ func (d *App) WaitServiceStable(ctx context.Context, startedAt time.Time) error
 	return d.ecs.WaitUntilServicesStableWithContext(ctx, d.DescribeServicesInput())
 }
 
-func (d *App) UpdateService(ctx context.Context, taskDefinitionArn string, count *int64, force *bool) error {
+func (d *App) UpdateService(ctx context.Context, taskDefinitionArn string, count int64, force bool) error {
 	msg := "Updating service"
-	if *force {
+	if force {
 		msg = msg + " with force new deployment"
 	}
 	msg = msg + "..."
@@ -332,8 +348,8 @@ func (d *App) UpdateService(ctx context.Context, taskDefinitionArn string, count
 			Service:            aws.String(d.Service),
 			Cluster:            aws.String(d.Cluster),
 			TaskDefinition:     aws.String(taskDefinitionArn),
-			DesiredCount:       count,
-			ForceNewDeployment: force,
+			DesiredCount:       &count,
+			ForceNewDeployment: &force,
 		},
 	)
 	return err
diff --git a/options.go b/options.go
index e9b8d98b..6b388a60 100644
--- a/options.go
+++ b/options.go
@@ -17,5 +17,6 @@ type StatusOption struct {
 }
 
 type RollbackOption struct {
-	DryRun *bool
+	DryRun                   *bool
+	DeregisterTaskDefinition *bool
 }