Merge pull request #36 from honeycombio/ssm-retries

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

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")

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,

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")

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,

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() {

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,

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,
 	}