add --retries/-r root flag and configure the SSM client to automatica…

…lly retry in the face of throttling
segmentio · Oct 10, 2017 · aa024d6 · aa024d6
1 parent d4d20a4
commit aa024d6
Show file tree

Hide file tree

Showing 7 changed files with 16 additions and 7 deletions.
diff --git a/cmd/exec.go b/cmd/exec.go
@@ -51,7 +51,7 @@ func execRun(cmd *cobra.Command, args []string) error {
 	}
 
 	env := environ(os.Environ())
-	secretStore := store.NewSSMStore()
+	secretStore := store.NewSSMStore(numRetries)
 	for _, service := range args {
 		if err := validateService(service); err != nil {
 			return errors.Wrap(err, "Failed to validate service")

diff --git a/cmd/history.go b/cmd/history.go
@@ -40,7 +40,7 @@ func history(cmd *cobra.Command, args []string) error {
 		return errors.Wrap(err, "Failed to validate key")
 	}
 
-	secretStore := store.NewSSMStore()
+	secretStore := store.NewSSMStore(numRetries)
 	secretId := store.SecretId{
 		Service: service,
 		Key:     key,

diff --git a/cmd/list.go b/cmd/list.go
@@ -35,7 +35,7 @@ func list(cmd *cobra.Command, args []string) error {
 		return errors.Wrap(err, "Failed to validate service")
 	}
 
-	secretStore := store.NewSSMStore()
+	secretStore := store.NewSSMStore(numRetries)
 	secrets, err := secretStore.List(service, false)
 	if err != nil {
 		return errors.Wrap(err, "Failed to list store contents")

diff --git a/cmd/read.go b/cmd/read.go
@@ -47,7 +47,7 @@ func read(cmd *cobra.Command, args []string) error {
 		return errors.Wrap(err, "Failed to validate key")
 	}
 
-	secretStore := store.NewSSMStore()
+	secretStore := store.NewSSMStore(numRetries)
 	secretId := store.SecretId{
 		Service: service,
 		Key:     key,

diff --git a/cmd/root.go b/cmd/root.go
@@ -12,11 +12,16 @@ import (
 var (
 	validKeyFormat     = regexp.MustCompile(`^[A-Za-z0-9-_]+$`)
 	validServiceFormat = regexp.MustCompile(`^[A-Za-z0-9-_]+$`)
+
+	numRetries int
 )
 
 const (
 	// ShortTimeFormat is a short format for printing timestamps
 	ShortTimeFormat = "01-02 15:04:05"
+
+	// DefaultNumRetries is the default for the number of retries we'll use for our SSM client
+	DefaultNumRetries = 10
 )
 
 // RootCmd represents the base command when called without any subcommands
@@ -27,6 +32,10 @@ var RootCmd = &cobra.Command{
 	SilenceErrors: true,
 }
 
+func init() {
+	RootCmd.PersistentFlags().IntVarP(&numRetries, "retries", "r", DefaultNumRetries, "For SSM, the number of retries we'll make before giving up")
+}
+
 // Execute adds all child commands to the root command sets flags appropriately.
 // This is called by main.main(). It only needs to happen once to the rootCmd.
 func Execute() {

diff --git a/cmd/write.go b/cmd/write.go
@@ -54,7 +54,7 @@ func write(cmd *cobra.Command, args []string) error {
 		value = string(v)
 	}
 
-	secretStore := store.NewSSMStore()
+	secretStore := store.NewSSMStore(numRetries)
 	secretId := store.SecretId{
 		Service: service,
 		Key:     key,

diff --git a/store/ssmstore.go b/store/ssmstore.go
@@ -29,7 +29,7 @@ type SSMStore struct {
 }
 
 // NewSSMStore creates a new SSMStore
-func NewSSMStore() *SSMStore {
+func NewSSMStore(numRetries int) *SSMStore {
 	region, ok := os.LookupEnv("AWS_REGION")
 	if !ok {
 		// If region is not set, attempt to determine it via ec2 metadata API
@@ -41,7 +41,7 @@ func NewSSMStore() *SSMStore {
 	ssmSession := session.Must(session.NewSession(&aws.Config{
 		Region: aws.String(region),
 	}))
-	svc := ssm.New(ssmSession)
+	svc := ssm.New(ssmSession, &aws.Config{MaxRetries: aws.Int(numRetries)})
 	return &SSMStore{
 		svc: svc,
 	}