From 36b89949336739ae2e8ddd6a3cbd787179bf4ab7 Mon Sep 17 00:00:00 2001
From: Niklas Wigertz <niklas.wigertz@axis.com>
Date: Wed, 21 Aug 2019 16:00:54 +0200
Subject: [PATCH] add pre/post update check lifecycle hooks

---
 internal/actions/update.go | 52 ++++++++++++++++++++++++++++++++++++++
 pkg/container/metadata.go  | 12 +++++++++
 2 files changed, 64 insertions(+)

diff --git a/internal/actions/update.go b/internal/actions/update.go
index 5763017..02013d6 100644
--- a/internal/actions/update.go
+++ b/internal/actions/update.go
@@ -13,6 +13,8 @@ import (
 func Update(client container.Client, params UpdateParams) error {
 	log.Debug("Checking containers for updated images")
 
+	executePreCheck(client, params)
+
 	containers, err := client.ListContainers(params.Filter)
 	if err != nil {
 		return err
@@ -36,12 +38,14 @@ func Update(client container.Client, params UpdateParams) error {
 	checkDependencies(containers)
 
 	if params.MonitorOnly {
+		executePostCheck(client, params)
 		return nil
 	}
 
 	stopContainersInReversedOrder(containers, client, params)
 	restartContainersInSortedOrder(containers, client, params)
 
+	executePostCheck(client, params)
 	return nil
 }
 
@@ -123,11 +127,58 @@ func checkDependencies(containers []container.Container) {
 	}
 }
 
+func executePreCheck(client container.Client, params UpdateParams) {
+	containers, err := client.ListContainers(params.Filter)
+	if err != nil {
+		return
+	}
+	for _, container := range containers {
+		executePreCheckCommand(client, container)
+	}
+}
+
+func executePostCheck(client container.Client, params UpdateParams) {
+	containers, err := client.ListContainers(params.Filter)
+	if err != nil {
+		return
+	}
+	for _, container := range containers {
+		executePostCheckCommand(client, container)
+	}
+}
+
+func executePreCheckCommand(client container.Client, container container.Container) {
+	command := container.GetLifecyclePreCheckCommand()
+	if len(command) == 0 {
+		log.Debug("No pre-check command supplied. Skipping")
+		return
+	}
+
+	log.Info("Executing pre-check command.")
+	if err := client.ExecuteCommand(container.ID(), command); err != nil {
+		log.Error(err)
+	}
+}
+
+func executePostCheckCommand(client container.Client, container container.Container) {
+	command := container.GetLifecyclePostCheckCommand()
+	if len(command) == 0 {
+		log.Debug("No post-check command supplied. Skipping")
+		return
+	}
+
+	log.Info("Executing post-check command.")
+	if err := client.ExecuteCommand(container.ID(), command); err != nil {
+		log.Error(err)
+	}
+}
+
 func executePreUpdateCommand(client container.Client, container container.Container) {
 
 	command := container.GetLifecyclePreUpdateCommand()
 	if len(command) == 0 {
 		log.Debug("No pre-update command supplied. Skipping")
+		return
 	}
 
 	log.Info("Executing pre-update command.")
@@ -146,6 +197,7 @@ func executePostUpdateCommand(client container.Client, newContainerID string) {
 	command := newContainer.GetLifecyclePostUpdateCommand()
 	if len(command) == 0 {
 		log.Debug("No post-update command supplied. Skipping")
+		return
 	}
 
 	log.Info("Executing post-update command.")
diff --git a/pkg/container/metadata.go b/pkg/container/metadata.go
index 3ab9ec2..0e04350 100644
--- a/pkg/container/metadata.go
+++ b/pkg/container/metadata.go
@@ -5,10 +5,22 @@ const (
 	signalLabel     = "com.centurylinklabs.watchtower.stop-signal"
 	enableLabel     = "com.centurylinklabs.watchtower.enable"
 	zodiacLabel     = "com.centurylinklabs.zodiac.original-image"
+	preCheckLabel   = "com.centurylinklabs.watchtower.lifecycle.pre-check"
+	postCheckLabel  = "com.centurylinklabs.watchtower.lifecycle.post-check"
 	preUpdateLabel  = "com.centurylinklabs.watchtower.lifecycle.pre-update"
 	postUpdateLabel = "com.centurylinklabs.watchtower.lifecycle.post-update"
 )
 
+// GetLifecyclePreCheckCommand returns the pre-check command set in the container metadata or an empty string
+func (c Container) GetLifecyclePreCheckCommand() string {
+	return c.getLabelValueOrEmpty(preCheckLabel)
+}
+
+// GetLifecyclePostCheckCommand returns the post-check command set in the container metadata or an empty string
+func (c Container) GetLifecyclePostCheckCommand() string {
+	return c.getLabelValueOrEmpty(postCheckLabel)
+}
+
 // GetLifecyclePreUpdateCommand returns the pre-update command set in the container metadata or an empty string
 func (c Container) GetLifecyclePreUpdateCommand() string {
 	return c.getLabelValueOrEmpty(preUpdateLabel)