semaphoreci · lyoung-confluent · Jun 26, 2024
diff --git a/cache-cli/cmd/clear.go b/cache-cli/cmd/clear.go
@@ -18,10 +18,10 @@ var clearCmd = &cobra.Command{
 }
 
 func RunClear(cmd *cobra.Command, args []string) {
-	storage, err := storage.InitStorage()
+	storage, err := storage.InitStorage(cmd.Context())
 	utils.Check(err)
 
-	err = storage.Clear()
+	err = storage.Clear(cmd.Context())
 	utils.Check(err)
 	log.Infof("Deleted all caches.")
 }

diff --git a/cache-cli/cmd/clear_test.go b/cache-cli/cmd/clear_test.go
@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"context"
 	"fmt"
 	"io/ioutil"
 	"os"
@@ -13,13 +14,14 @@ import (
 )
 
 func Test__Clear(t *testing.T) {
+	ctx := context.TODO()
 	log.SetFormatter(new(logging.CustomFormatter))
 	log.SetLevel(log.InfoLevel)
 	log.SetOutput(openLogfileForTests(t))
 
-	runTestForAllBackends(t, func(backend string, storage storage.Storage) {
+	runTestForAllBackends(ctx, t, func(backend string, storage storage.Storage) {
 		t.Run(fmt.Sprintf("%s no keys", backend), func(*testing.T) {
-			err := storage.Clear()
+			err := storage.Clear(ctx)
 			assert.Nil(t, err)
 
 			RunClear(clearCmd, []string{})
@@ -29,13 +31,13 @@ func Test__Clear(t *testing.T) {
 		})
 
 		t.Run(fmt.Sprintf("%s with keys", backend), func(*testing.T) {
-			err := storage.Clear()
+			err := storage.Clear(ctx)
 			assert.Nil(t, err)
 
 			tempFile, _ := ioutil.TempFile(os.TempDir(), "*")
-			storage.Store("abc001", tempFile.Name())
+			storage.Store(ctx, "abc001", tempFile.Name())
 
-			RunClear(hasKeyCmd, []string{})
+			RunClear(clearCmd, []string{})
 			output := readOutputFromFile(t)
 
 			assert.Contains(t, output, "Deleted all caches.")

diff --git a/cache-cli/cmd/delete.go b/cache-cli/cmd/delete.go
@@ -24,14 +24,14 @@ func RunDelete(cmd *cobra.Command, args []string) {
 		return
 	}
 
-	storage, err := storage.InitStorage()
+	storage, err := storage.InitStorage(cmd.Context())
 	utils.Check(err)
 
 	rawKey := args[0]
 	key := NormalizeKey(rawKey)
 
-	if ok, _ := storage.HasKey(key); ok {
-		err := storage.Delete(key)
+	if ok, _ := storage.HasKey(cmd.Context(), key); ok {
+		err := storage.Delete(cmd.Context(), key)
 		utils.Check(err)
 		log.Infof("Key '%s' is deleted.", key)
 	} else {

diff --git a/cache-cli/cmd/delete_test.go b/cache-cli/cmd/delete_test.go
@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"context"
 	"fmt"
 	"io/ioutil"
 	"os"
@@ -13,11 +14,12 @@ import (
 )
 
 func Test__Delete(t *testing.T) {
+	ctx := context.TODO()
 	log.SetFormatter(new(logging.CustomFormatter))
 	log.SetLevel(log.InfoLevel)
 	log.SetOutput(openLogfileForTests(t))
 
-	runTestForAllBackends(t, func(backend string, storage storage.Storage) {
+	runTestForAllBackends(ctx, t, func(backend string, storage storage.Storage) {
 		t.Run(fmt.Sprintf("%s key is missing", backend), func(*testing.T) {
 			RunDelete(deleteCmd, []string{"this-key-does-not-exist"})
 			output := readOutputFromFile(t)
@@ -26,9 +28,9 @@ func Test__Delete(t *testing.T) {
 		})
 
 		t.Run(fmt.Sprintf("%s key is present", backend), func(*testing.T) {
-			storage.Clear()
+			storage.Clear(ctx)
 			tempFile, _ := ioutil.TempFile(os.TempDir(), "*")
-			storage.Store("abc001", tempFile.Name())
+			storage.Store(ctx, "abc001", tempFile.Name())
 
 			RunDelete(deleteCmd, []string{"abc001"})
 			output := readOutputFromFile(t)
@@ -37,7 +39,7 @@ func Test__Delete(t *testing.T) {
 		})
 
 		t.Run(fmt.Sprintf("%s normalizes key", backend), func(*testing.T) {
-			storage.Clear()
+			storage.Clear(ctx)
 			tempFile, _ := ioutil.TempFile(os.TempDir(), "*")
 			RunStore(NewStoreCommand(), []string{"abc/00/33", tempFile.Name()})
 

diff --git a/cache-cli/cmd/has_key.go b/cache-cli/cmd/has_key.go
@@ -28,12 +28,12 @@ func RunHasKey(cmd *cobra.Command, args []string) bool {
 		return true
 	}
 
-	storage, err := storage.InitStorage()
+	storage, err := storage.InitStorage(cmd.Context())
 	utils.Check(err)
 
 	rawKey := args[0]
 	key := NormalizeKey(rawKey)
-	exists, err := storage.HasKey(key)
+	exists, err := storage.HasKey(cmd.Context(), key)
 	utils.Check(err)
 
 	if exists {

diff --git a/cache-cli/cmd/has_key_test.go b/cache-cli/cmd/has_key_test.go
@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"context"
 	"fmt"
 	"io/ioutil"
 	"os"
@@ -13,11 +14,12 @@ import (
 )
 
 func Test__HasKey(t *testing.T) {
+	ctx := context.TODO()
 	log.SetFormatter(new(logging.CustomFormatter))
 	log.SetLevel(log.InfoLevel)
 	log.SetOutput(openLogfileForTests(t))
 
-	runTestForAllBackends(t, func(backend string, storage storage.Storage) {
+	runTestForAllBackends(ctx, t, func(backend string, storage storage.Storage) {
 		t.Run(fmt.Sprintf("%s key is missing", backend), func(*testing.T) {
 			RunHasKey(hasKeyCmd, []string{"this-key-does-not-exist"})
 			output := readOutputFromFile(t)
@@ -26,9 +28,9 @@ func Test__HasKey(t *testing.T) {
 		})
 
 		t.Run(fmt.Sprintf("%s key is present", backend), func(*testing.T) {
-			storage.Clear()
+			storage.Clear(ctx)
 			tempFile, _ := ioutil.TempFile(os.TempDir(), "*")
-			storage.Store("abc001", tempFile.Name())
+			storage.Store(ctx, "abc001", tempFile.Name())
 
 			RunHasKey(hasKeyCmd, []string{"abc001"})
 			output := readOutputFromFile(t)
@@ -37,7 +39,7 @@ func Test__HasKey(t *testing.T) {
 		})
 
 		t.Run(fmt.Sprintf("%s normalizes key", backend), func(*testing.T) {
-			storage.Clear()
+			storage.Clear(ctx)
 			tempFile, _ := ioutil.TempFile(os.TempDir(), "*")
 			RunStore(NewStoreCommand(), []string{"abc/00/33", tempFile.Name()})
 

diff --git a/cache-cli/cmd/is_not_empty.go b/cache-cli/cmd/is_not_empty.go
@@ -23,10 +23,10 @@ var isNotEmptyCmd = &cobra.Command{
 }
 
 func RunIsNotEmpty(cmd *cobra.Command, args []string) bool {
-	storage, err := storage.InitStorage()
+	storage, err := storage.InitStorage(cmd.Context())
 	utils.Check(err)
 
-	isNotEmpty, err := storage.IsNotEmpty()
+	isNotEmpty, err := storage.IsNotEmpty(cmd.Context())
 	utils.Check(err)
 
 	return isNotEmpty

diff --git a/cache-cli/cmd/is_not_empty_test.go b/cache-cli/cmd/is_not_empty_test.go
@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"context"
 	"fmt"
 	"io/ioutil"
 	"os"
@@ -11,16 +12,17 @@ import (
 )
 
 func Test__IsNotEmpty(t *testing.T) {
-	runTestForAllBackends(t, func(backend string, storage storage.Storage) {
+	ctx := context.TODO()
+	runTestForAllBackends(ctx, t, func(backend string, storage storage.Storage) {
 		t.Run(fmt.Sprintf("%s cache is empty", backend), func(*testing.T) {
-			storage.Clear()
+			storage.Clear(ctx)
 			assert.False(t, RunIsNotEmpty(isNotEmptyCmd, []string{}))
 		})
 
 		t.Run(fmt.Sprintf("%s cache is not empty", backend), func(*testing.T) {
-			storage.Clear()
+			storage.Clear(ctx)
 			tempFile, _ := ioutil.TempFile(os.TempDir(), "*")
-			storage.Store("abc001", tempFile.Name())
+			storage.Store(ctx, "abc001", tempFile.Name())
 
 			assert.True(t, RunIsNotEmpty(isNotEmptyCmd, []string{}))
 		})

diff --git a/cache-cli/cmd/list.go b/cache-cli/cmd/list.go
@@ -36,10 +36,10 @@ func RunList(cmd *cobra.Command, args []string) {
 	sortBy, err := cmd.Flags().GetString("sort-by")
 	utils.Check(err)
 
-	storage, err := storage.InitStorageWithConfig(storage.StorageConfig{SortKeysBy: sortBy})
+	storage, err := storage.InitStorageWithConfig(cmd.Context(), storage.StorageConfig{SortKeysBy: sortBy})
 	utils.Check(err)
 
-	keys, err := storage.List()
+	keys, err := storage.List(cmd.Context())
 	utils.Check(err)
 
 	if len(keys) == 0 {

diff --git a/cache-cli/cmd/list_test.go b/cache-cli/cmd/list_test.go
@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"context"
 	"fmt"
 	"io/ioutil"
 	"os"
@@ -13,14 +14,15 @@ import (
 )
 
 func Test__List(t *testing.T) {
+	ctx := context.TODO()
 	listCmd := NewListCommand()
 	log.SetFormatter(new(logging.CustomFormatter))
 	log.SetLevel(log.InfoLevel)
 	log.SetOutput(openLogfileForTests(t))
 
-	runTestForAllBackends(t, func(backend string, storage storage.Storage) {
+	runTestForAllBackends(ctx, t, func(backend string, storage storage.Storage) {
 		t.Run(fmt.Sprintf("%s no keys", backend), func(*testing.T) {
-			storage.Clear()
+			storage.Clear(ctx)
 
 			RunList(listCmd, []string{""})
 			output := readOutputFromFile(t)
@@ -29,11 +31,11 @@ func Test__List(t *testing.T) {
 		})
 
 		t.Run(fmt.Sprintf("%s with keys", backend), func(*testing.T) {
-			storage.Clear()
+			storage.Clear(ctx)
 			tempFile, _ := ioutil.TempFile(os.TempDir(), "*")
-			storage.Store("abc001", tempFile.Name())
-			storage.Store("abc002", tempFile.Name())
-			storage.Store("abc003", tempFile.Name())
+			storage.Store(ctx, "abc001", tempFile.Name())
+			storage.Store(ctx, "abc002", tempFile.Name())
+			storage.Store(ctx, "abc003", tempFile.Name())
 
 			RunList(listCmd, []string{})
 			output := readOutputFromFile(t)

diff --git a/cache-cli/cmd/restore.go b/cache-cli/cmd/restore.go
@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"context"
 	"fmt"
 	"io/fs"
 	"os"
@@ -34,7 +35,7 @@ func RunRestore(cmd *cobra.Command, args []string) {
 		return
 	}
 
-	storage, err := storage.InitStorage()
+	storage, err := storage.InitStorage(cmd.Context())
 	utils.Check(err)
 
 	metricsManager, err := metrics.InitMetricsManager(metrics.LocalBackend)
@@ -57,31 +58,31 @@ func RunRestore(cmd *cobra.Command, args []string) {
 			log.Infof("Detected %s.", lookupResult.DetectedFile)
 			for _, entry := range lookupResult.Entries {
 				log.Infof("Fetching '%s' directory with cache keys '%s'...", entry.Path, strings.Join(entry.Keys, ","))
-				downloadAndUnpack(storage, archiver, metricsManager, entry.Keys)
+				downloadAndUnpack(cmd.Context(), storage, archiver, metricsManager, entry.Keys)
 			}
 		}
 	} else {
 		keys := strings.Split(args[0], ",")
-		downloadAndUnpack(storage, archiver, metricsManager, keys)
+		downloadAndUnpack(cmd.Context(), storage, archiver, metricsManager, keys)
 	}
 }
 
-func downloadAndUnpack(storage storage.Storage, archiver archive.Archiver, metricsManager metrics.MetricsManager, keys []string) {
+func downloadAndUnpack(ctx context.Context, storage storage.Storage, archiver archive.Archiver, metricsManager metrics.MetricsManager, keys []string) {
 	for _, rawKey := range keys {
 		key := NormalizeKey(rawKey)
-		if ok, _ := storage.HasKey(key); ok {
+		if ok, _ := storage.HasKey(ctx, key); ok {
 			log.Infof("HIT: '%s', using key '%s'.", key, key)
-			downloadAndUnpackKey(storage, archiver, metricsManager, key)
+			downloadAndUnpackKey(ctx, storage, archiver, metricsManager, key)
 			break
 		}
 
-		availableKeys, err := storage.List()
+		availableKeys, err := storage.List(ctx)
 		utils.Check(err)
 
 		matchingKey := findMatchingKey(availableKeys, key)
 		if matchingKey != "" {
 			log.Infof("HIT: '%s', using key '%s'.", key, matchingKey)
-			downloadAndUnpackKey(storage, archiver, metricsManager, matchingKey)
+			downloadAndUnpackKey(ctx, storage, archiver, metricsManager, matchingKey)
 			break
 		} else {
 			log.Infof("MISS: '%s'.", key)
@@ -100,10 +101,10 @@ func findMatchingKey(availableKeys []storage.CacheKey, match string) string {
 	return ""
 }
 
-func downloadAndUnpackKey(storage storage.Storage, archiver archive.Archiver, metricsManager metrics.MetricsManager, key string) {
+func downloadAndUnpackKey(ctx context.Context, storage storage.Storage, archiver archive.Archiver, metricsManager metrics.MetricsManager, key string) {
 	downloadStart := time.Now()
 	log.Infof("Downloading key '%s'...", key)
-	compressed, err := storage.Restore(key)
+	compressed, err := storage.Restore(ctx, key)
 	utils.Check(err)
 
 	downloadDuration := time.Since(downloadStart)
@@ -114,7 +115,7 @@ func downloadAndUnpackKey(storage storage.Storage, archiver archive.Archiver, me
 
 	unpackStart := time.Now()
 	log.Infof("Unpacking '%s'...", compressed.Name())
-	restorationPath, err := archiver.Decompress(compressed.Name())
+	restorationPath, err := archiver.Decompress(ctx, compressed.Name())
 	utils.Check(err)
 
 	unpackDuration := time.Since(unpackStart)