diff --git a/cmd/exec.go b/cmd/exec.go index 2a3735ee..8766a251 100644 --- 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 index 5c05527c..2019d64c 100644 --- 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 index b6247cb4..715fc13e 100644 --- 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 index 79b71fe0..8e6911f7 100644 --- 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 index 31c0769e..717463a5 100644 --- 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 index 999acd98..d15b1211 100644 --- 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 index 8cb0781f..1cb99b0a 100644 --- 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, }