refactor!: Properly support contexts

The new AWS SDK uses contexts for all API functions and for loading the SDK configuration. When chamber was migrated to the SDK, all of them were set to `context.TODO()`. This commit removes those temporary values and threads contexts properly throughout the codebase, starting from the chamber commands. BREAKING CHANGE: This commit changes the signature of many exported functions to add a context argument.
segmentio · Jun 5, 2024 · 3367239 · 3367239
1 parent e70bb44
commit 3367239
Show file tree

Hide file tree

Showing 24 changed files with 289 additions and 271 deletions.
diff --git a/cmd/delete.go b/cmd/delete.go
@@ -51,7 +51,7 @@ func delete(cmd *cobra.Command, args []string) error {
 				Set("backend", backend),
 		})
 	}
-	secretStore, err := getSecretStore()
+	secretStore, err := getSecretStore(cmd.Context())
 	if err != nil {
 		return fmt.Errorf("Failed to get secret store: %w", err)
 	}
@@ -60,5 +60,5 @@ func delete(cmd *cobra.Command, args []string) error {
 		Key:     key,
 	}
 
-	return secretStore.Delete(secretId)
+	return secretStore.Delete(cmd.Context(), secretId)
 }
diff --git a/cmd/env.go b/cmd/env.go
@@ -64,12 +64,12 @@ func exportEnv(cmd *cobra.Command, args []string) ([]string, error) {
 		return nil, fmt.Errorf("Failed to validate service: %w", err)
 	}
 
-	secretStore, err := getSecretStore()
+	secretStore, err := getSecretStore(cmd.Context())
 	if err != nil {
 		return nil, fmt.Errorf("Failed to get secret store: %w", err)
 	}
 
-	rawSecrets, err := secretStore.ListRaw(service)
+	rawSecrets, err := secretStore.ListRaw(cmd.Context(), service)
 	if err != nil {
 		return nil, fmt.Errorf("Failed to list store contents: %w", err)
 	}

diff --git a/cmd/exec.go b/cmd/exec.go
@@ -92,7 +92,7 @@ func execRun(cmd *cobra.Command, args []string) error {
 		}
 	}
 
-	secretStore, err := getSecretStore()
+	secretStore, err := getSecretStore(cmd.Context())
 	if err != nil {
 		return fmt.Errorf("Failed to get secret store: %w", err)
 	}
@@ -108,7 +108,7 @@ func execRun(cmd *cobra.Command, args []string) error {
 		}
 		var err error
 		env = environ.Environ(os.Environ())
-		err = env.LoadStrict(secretStore, strictValue, pristine, services...)
+		err = env.LoadStrict(cmd.Context(), secretStore, strictValue, pristine, services...)
 		if err != nil {
 			return err
 		}
@@ -120,7 +120,7 @@ func execRun(cmd *cobra.Command, args []string) error {
 			collisions := make([]string, 0)
 			var err error
 			// TODO: these interfaces should look the same as Strict*, so move pristine in there
-			err = env.Load(secretStore, service, &collisions)
+			err = env.Load(cmd.Context(), secretStore, service, &collisions)
 			if err != nil {
 				return fmt.Errorf("Failed to list store contents: %w", err)
 			}

diff --git a/cmd/export.go b/cmd/export.go
@@ -53,7 +53,7 @@ func runExport(cmd *cobra.Command, args []string) error {
 		})
 	}
 
-	secretStore, err := getSecretStore()
+	secretStore, err := getSecretStore(cmd.Context())
 	if err != nil {
 		return err
 	}
@@ -64,7 +64,7 @@ func runExport(cmd *cobra.Command, args []string) error {
 			return fmt.Errorf("Failed to validate service %s: %w", service, err)
 		}
 
-		rawSecrets, err := secretStore.ListRaw(service)
+		rawSecrets, err := secretStore.ListRaw(cmd.Context(), service)
 		if err != nil {
 			return fmt.Errorf("Failed to list store contents for service %s: %w", service, err)
 		}

diff --git a/cmd/find.go b/cmd/find.go
@@ -39,18 +39,18 @@ func find(cmd *cobra.Command, args []string) error {
 		includeSecrets = true
 	}
 
-	secretStore, err := getSecretStore()
+	secretStore, err := getSecretStore(cmd.Context())
 	if err != nil {
 		return fmt.Errorf("Failed to get secret store: %w", err)
 	}
-	services, err := secretStore.ListServices(blankService, includeSecrets)
+	services, err := secretStore.ListServices(cmd.Context(), blankService, includeSecrets)
 	if err != nil {
 		return fmt.Errorf("Failed to list store contents: %w", err)
 	}
 
 	if byValue {
 		for _, service := range services {
-			allSecrets, err := secretStore.List(service, true)
+			allSecrets, err := secretStore.List(cmd.Context(), service, true)
 			if err == nil {
 				matches = append(matches, findValueMatch(allSecrets, findSecret)...)
 			}

diff --git a/cmd/history.go b/cmd/history.go
@@ -47,7 +47,7 @@ func history(cmd *cobra.Command, args []string) error {
 		})
 	}
 
-	secretStore, err := getSecretStore()
+	secretStore, err := getSecretStore(cmd.Context())
 	if err != nil {
 		return fmt.Errorf("Failed to get secret store: %w", err)
 	}
@@ -56,7 +56,7 @@ func history(cmd *cobra.Command, args []string) error {
 		Key:     key,
 	}
 
-	events, err := secretStore.History(secretId)
+	events, err := secretStore.History(cmd.Context(), secretId)
 	if err != nil {
 		return fmt.Errorf("Failed to get history: %w", err)
 	}

diff --git a/cmd/import.go b/cmd/import.go
@@ -65,7 +65,7 @@ func importRun(cmd *cobra.Command, args []string) error {
 		})
 	}
 
-	secretStore, err := getSecretStore()
+	secretStore, err := getSecretStore(cmd.Context())
 	if err != nil {
 		return fmt.Errorf("Failed to get secret store: %w", err)
 	}
@@ -78,7 +78,7 @@ func importRun(cmd *cobra.Command, args []string) error {
 			Service: service,
 			Key:     key,
 		}
-		if err := secretStore.Write(secretId, value); err != nil {
+		if err := secretStore.Write(cmd.Context(), secretId, value); err != nil {
 			return fmt.Errorf("Failed to write secret: %w", err)
 		}
 	}

diff --git a/cmd/list-services.go b/cmd/list-services.go
@@ -34,11 +34,11 @@ func listServices(cmd *cobra.Command, args []string) error {
 		service = utils.NormalizeService(args[0])
 
 	}
-	secretStore, err := getSecretStore()
+	secretStore, err := getSecretStore(cmd.Context())
 	if err != nil {
 		return fmt.Errorf("Failed to get secret store: %w", err)
 	}
-	secrets, err := secretStore.ListServices(service, includeSecretName)
+	secrets, err := secretStore.ListServices(cmd.Context(), service, includeSecretName)
 	if err != nil {
 		return fmt.Errorf("Failed to list store contents: %w", err)
 	}

diff --git a/cmd/list.go b/cmd/list.go
@@ -54,11 +54,11 @@ func list(cmd *cobra.Command, args []string) error {
 		})
 	}
 
-	secretStore, err := getSecretStore()
+	secretStore, err := getSecretStore(cmd.Context())
 	if err != nil {
 		return fmt.Errorf("Failed to get secret store: %w", err)
 	}
-	secrets, err := secretStore.List(service, withValues)
+	secrets, err := secretStore.List(cmd.Context(), service, withValues)
 	if err != nil {
 		return fmt.Errorf("Failed to list store contents: %w", err)
 	}

diff --git a/cmd/read.go b/cmd/read.go
@@ -54,7 +54,7 @@ func read(cmd *cobra.Command, args []string) error {
 		})
 	}
 
-	secretStore, err := getSecretStore()
+	secretStore, err := getSecretStore(cmd.Context())
 	if err != nil {
 		return fmt.Errorf("Failed to get secret store: %w", err)
 	}
@@ -64,7 +64,7 @@ func read(cmd *cobra.Command, args []string) error {
 		Key:     key,
 	}
 
-	secret, err := secretStore.Read(secretId, version)
+	secret, err := secretStore.Read(cmd.Context(), secretId, version)
 	if err != nil {
 		return fmt.Errorf("Failed to read: %w", err)
 	}

diff --git a/cmd/root.go b/cmd/root.go
@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"os"
@@ -133,7 +134,7 @@ func validateKey(key string) error {
 	return nil
 }
 
-func getSecretStore() (store.Store, error) {
+func getSecretStore(ctx context.Context) (store.Store, error) {
 	rootPflags := RootCmd.PersistentFlags()
 	if backendEnvVarValue := os.Getenv(BackendEnvVar); !rootPflags.Changed("backend") && backendEnvVarValue != "" {
 		backend = backendEnvVarValue
@@ -170,7 +171,7 @@ func getSecretStore() (store.Store, error) {
 		if bucket == "" {
 			return nil, errors.New("Must set bucket for s3 backend")
 		}
-		s, err = store.NewS3StoreWithBucket(numRetries, bucket)
+		s, err = store.NewS3StoreWithBucket(ctx, numRetries, bucket)
 	case S3KMSBackend:
 		var bucket string
 		if bucketEnvVarValue := os.Getenv(BucketEnvVar); !rootPflags.Changed("backend-s3-bucket") && bucketEnvVarValue != "" {
@@ -197,9 +198,9 @@ func getSecretStore() (store.Store, error) {
 			return nil, errors.New("Must set kmsKeyAlias for S3 KMS backend")
 		}
 
-		s, err = store.NewS3KMSStore(numRetries, bucket, kmsKeyAlias)
+		s, err = store.NewS3KMSStore(ctx, numRetries, bucket, kmsKeyAlias)
 	case SecretsManagerBackend:
-		s, err = store.NewSecretsManagerStore(numRetries)
+		s, err = store.NewSecretsManagerStore(ctx, numRetries)
 	case SSMBackend:
 		if kmsKeyAliasFlag != DefaultKMSKey {
 			return nil, errors.New("Unable to use --kms-key-alias with this backend. Use CHAMBER_KMS_KEY_ALIAS instead.")
@@ -209,7 +210,7 @@ func getSecretStore() (store.Store, error) {
 		if err != nil {
 			return nil, fmt.Errorf("Invalid retry mode %s", retryMode)
 		}
-		s, err = store.NewSSMStoreWithRetryMode(numRetries, parsedRetryMode)
+		s, err = store.NewSSMStoreWithRetryMode(ctx, numRetries, parsedRetryMode)
 	default:
 		return nil, fmt.Errorf("invalid backend `%s`", backend)
 	}

diff --git a/cmd/write.go b/cmd/write.go
@@ -75,7 +75,7 @@ func write(cmd *cobra.Command, args []string) error {
 		}
 	}
 
-	secretStore, err := getSecretStore()
+	secretStore, err := getSecretStore(cmd.Context())
 	if err != nil {
 		return fmt.Errorf("Failed to get secret store: %w", err)
 	}
@@ -86,11 +86,11 @@ func write(cmd *cobra.Command, args []string) error {
 	}
 
 	if skipUnchanged {
-		currentSecret, err := secretStore.Read(secretId, -1)
+		currentSecret, err := secretStore.Read(cmd.Context(), secretId, -1)
 		if err == nil && value == *currentSecret.Value {
 			return nil
 		}
 	}
 
-	return secretStore.Write(secretId, value)
+	return secretStore.Write(cmd.Context(), secretId, value)
 }
diff --git a/environ/environ.go b/environ/environ.go
@@ -1,6 +1,7 @@
 package environ
 
 import (
+	"context"
 	"fmt"
 	"strings"
 
@@ -83,8 +84,8 @@ func normalizeEnvVarName(k string) string {
 
 // load loads environment variables into e from s given a service
 // collisions will be populated with any keys that get overwritten
-func (e *Environ) load(s store.Store, service string, collisions *[]string) error {
-	rawSecrets, err := s.ListRaw(utils.NormalizeService(service))
+func (e *Environ) load(ctx context.Context, s store.Store, service string, collisions *[]string) error {
+	rawSecrets, err := s.ListRaw(ctx, utils.NormalizeService(service))
 	if err != nil {
 		return err
 	}
@@ -104,20 +105,20 @@ func (e *Environ) load(s store.Store, service string, collisions *[]string) erro
 
 // Load loads environment variables into e from s given a service
 // collisions will be populated with any keys that get overwritten
-func (e *Environ) Load(s store.Store, service string, collisions *[]string) error {
-	return e.load(s, service, collisions)
+func (e *Environ) Load(ctx context.Context, s store.Store, service string, collisions *[]string) error {
+	return e.load(ctx, s, service, collisions)
 }
 
 // LoadStrict loads all services from s in strict mode: env vars in e with value equal to valueExpected
 // are the only ones substituted. If there are any env vars in s that are also in e, but don't have their value
 // set to valueExpected, this is an error.
-func (e *Environ) LoadStrict(s store.Store, valueExpected string, pristine bool, services ...string) error {
-	return e.loadStrict(s, valueExpected, pristine, services...)
+func (e *Environ) LoadStrict(ctx context.Context, s store.Store, valueExpected string, pristine bool, services ...string) error {
+	return e.loadStrict(ctx, s, valueExpected, pristine, services...)
 }
 
-func (e *Environ) loadStrict(s store.Store, valueExpected string, pristine bool, services ...string) error {
+func (e *Environ) loadStrict(ctx context.Context, s store.Store, valueExpected string, pristine bool, services ...string) error {
 	for _, service := range services {
-		rawSecrets, err := s.ListRaw(utils.NormalizeService(service))
+		rawSecrets, err := s.ListRaw(ctx, utils.NormalizeService(service))
 		if err != nil {
 			return err
 		}

diff --git a/store/backendbenchmarks_test.go b/store/backendbenchmarks_test.go
@@ -1,6 +1,7 @@
 package store
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"math/rand"
@@ -36,8 +37,9 @@ func RandStringRunes(n int) string {
 }
 
 func benchmarkStore(t *testing.T, store Store, services []string) {
-	setupStore(t, store, services)
-	defer cleanupStore(t, store, services)
+	ctx := context.Background()
+	setupStore(t, ctx, store, services)
+	defer cleanupStore(t, ctx, store, services)
 
 	concurrentExecs := []int{1, 10, 500, 1000}
 
@@ -47,7 +49,7 @@ func benchmarkStore(t *testing.T, store Store, services []string) {
 
 		for i := 0; i < concurrency; i++ {
 			wg.Add(1)
-			go emulateExec(t, &wg, store, services)
+			go emulateExec(t, ctx, &wg, store, services)
 
 		}
 		wg.Wait()
@@ -56,11 +58,11 @@ func benchmarkStore(t *testing.T, store Store, services []string) {
 	}
 }
 
-func emulateExec(t *testing.T, wg *sync.WaitGroup, s Store, services []string) error {
+func emulateExec(t *testing.T, ctx context.Context, wg *sync.WaitGroup, s Store, services []string) error {
 	defer wg.Done()
 	// Exec calls ListRaw once per service specified
 	for _, service := range services {
-		_, err := s.ListRaw(service)
+		_, err := s.ListRaw(ctx, service)
 		if err != nil {
 			t.Logf("Failed to execute ListRaw: %s", err)
 			return err
@@ -73,15 +75,15 @@ func TestS3StoreConcurrency(t *testing.T) {
 	if !benchmarkEnabled {
 		t.SkipNow()
 	}
-	s, _ := NewS3StoreWithBucket(10, "chamber-test")
+	s, _ := NewS3StoreWithBucket(context.Background(), 10, "chamber-test")
 	benchmarkStore(t, s, []string{"foo"})
 }
 
 func TestSecretsManagerStoreConcurrency(t *testing.T) {
 	if !benchmarkEnabled {
 		t.SkipNow()
 	}
-	s, _ := NewSecretsManagerStore(10)
+	s, _ := NewSecretsManagerStore(context.Background(), 10)
 	benchmarkStore(t, s, []string{"foo"})
 }
 
@@ -90,11 +92,11 @@ func TestSSMConcurrency(t *testing.T) {
 		t.SkipNow()
 	}
 
-	s, _ := NewSSMStore(10)
+	s, _ := NewSSMStore(context.Background(), 10)
 	benchmarkStore(t, s, []string{"foo"})
 }
 
-func setupStore(t *testing.T, store Store, services []string) {
+func setupStore(t *testing.T, ctx context.Context, store Store, services []string) {
 	// populate the store for services listed
 	for _, service := range services {
 		for i := 0; i < KeysPerService; i++ {
@@ -104,12 +106,12 @@ func setupStore(t *testing.T, store Store, services []string) {
 				Key:     key,
 			}
 
-			store.Write(id, RandStringRunes(100))
+			store.Write(ctx, id, RandStringRunes(100))
 		}
 	}
 }
 
-func cleanupStore(t *testing.T, store Store, services []string) {
+func cleanupStore(t *testing.T, ctx context.Context, store Store, services []string) {
 	for _, service := range services {
 		for i := 0; i < KeysPerService; i++ {
 			key := fmt.Sprintf("var%d", i)
@@ -118,7 +120,7 @@ func cleanupStore(t *testing.T, store Store, services []string) {
 				Key:     key,
 			}
 
-			store.Delete(id)
+			store.Delete(ctx, id)
 		}
 	}
 }