Allow checking for specific keys to be present during update

command/cmdHelp.go

@@ -45,6 +45,7 @@ func (cmd *HelpCommand) Apply(m *manager.Manager) error {
 			update
 				-t|-timeout (duration) [optional = 5*time.Second]
 				-r|-retry (duration) [optional = 0*time.Second]
+				-k|-keys (string) [optional = ""]
 
 			ls
 				-r|-regex (bool) [optional = false]

command/cmdUpdate.go

@@ -2,6 +2,7 @@ package command
 
 import (
 	"flag"
+	"strings"
 	"time"
 
 	"github.com/StackAdapt/systags/manager"
@@ -11,6 +12,7 @@ type UpdateCommand struct {
 	baseCommand
 	timeout time.Duration
 	retry   time.Duration
+	keys    string
 }
 
 func NewUpdateCommand() *UpdateCommand {
@@ -25,6 +27,8 @@ func NewUpdateCommand() *UpdateCommand {
 	cmd.flagSet.DurationVar(&cmd.timeout, "timeout", 5*time.Second, "")
 	cmd.flagSet.DurationVar(&cmd.retry, "r", 0*time.Second, "") // TODO: Does this get overwritten by `--retry` default?
 	cmd.flagSet.DurationVar(&cmd.retry, "retry", 0*time.Second, "")
+	cmd.flagSet.StringVar(&cmd.keys, "k", "", "")
+	cmd.flagSet.StringVar(&cmd.keys, "keys", "", "")
 
 	// Don't print unneeded usage
 	cmd.flagSet.Usage = func() {}
@@ -39,7 +43,12 @@ func (cmd *UpdateCommand) Apply(m *manager.Manager) error {
 		return err
 	}
 
-	err = m.UpdateRemote(cmd.timeout, cmd.retry)
+	var keys []string
+	if cmd.keys != "" {
+		keys = strings.Split(cmd.keys, ",")
+	}
+
+	err = m.UpdateRemote(cmd.timeout, cmd.retry, keys)
 	if err != nil {
 		return err
 	}

manager/manager.go

@@ -271,7 +271,11 @@ func (m *Manager) SaveFiles() error {
 // retry duration is provided, this function will auto
 // retry for the duration if empty tags are returned.
 // A bounded exponential backoff strategy is employed.
-func (m *Manager) UpdateRemote(timeout time.Duration, retry time.Duration) error {
+// If the required keys slice is provided, the function
+// will keep retrying until all the required keys are
+// present in the fetched tags or until the retry
+// duration is reached.
+func (m *Manager) UpdateRemote(timeout time.Duration, retry time.Duration, requiredKeys []string) error {
 
 	// TODO:
 	// At the moment, only AWS is supported, but if you want
@@ -292,12 +296,28 @@ func (m *Manager) UpdateRemote(timeout time.Duration, retry time.Duration) error
 	startTime := time.Now()
 	untilTime := startTime.Add(retry)
 
+	hasRequiredKeys := func(tags Tags) bool {
+
+		for _, k := range requiredKeys {
+			if _, ok := tags[k]; !ok {
+				return false
+			}
+		}
+
+		return true
+	}
+
 	for {
 		res, err = getAwsTags(m.GetLogger(), timeout)
 		if err != nil {
 			return err
 		}
 
+		// Check whether all the required keys have values
+		if len(requiredKeys) > 0 && hasRequiredKeys(res) {
+			break
+		}
+
 		// Tags are not empty or we have reached time limit
 		if len(res) > 0 || time.Since(startTime) > retry {
 			break
@@ -318,6 +338,8 @@ func (m *Manager) UpdateRemote(timeout time.Duration, retry time.Duration) error
 		}
 	}
 
+	// TODO: Should we return an error if required keys and len(res) condition not met?
+
 	m.remote = res
 	return nil
 }