Propagate command context instead of using context.Background()

jozu-ai · Feb 23, 2024 · 601d425 · 601d425
1 parent 26d36f8
commit 601d425
Show file tree

Hide file tree

Showing 9 changed files with 32 additions and 34 deletions.
diff --git a/pkg/cmd/build/build.go b/pkg/cmd/build/build.go
@@ -1,6 +1,7 @@
 package build
 
 import (
+	"context"
 	"kitops/pkg/artifact"
 	"kitops/pkg/lib/constants"
 	"kitops/pkg/lib/filesystem"
@@ -10,7 +11,7 @@ import (
 	"path"
 )
 
-func RunBuild(options *buildOptions) error {
+func RunBuild(ctx context.Context, options *buildOptions) error {
 	// 1. Read the model file
 	modelfile, err := os.Open(options.modelFile)
 	if err != nil {
@@ -70,13 +71,13 @@ func RunBuild(options *buildOptions) error {
 		tag = options.modelRef.Reference
 	}
 	store := storage.NewLocalStore(modelStorePath, repo)
-	manifestDesc, err := store.SaveModel(model, tag)
+	manifestDesc, err := store.SaveModel(ctx, model, tag)
 	if err != nil {
 		return err
 	}
 
 	for _, tag := range options.extraRefs {
-		if err := store.TagModel(*manifestDesc, tag); err != nil {
+		if err := store.TagModel(ctx, *manifestDesc, tag); err != nil {
 			return err
 		}
 	}

diff --git a/pkg/cmd/build/cmd.go b/pkg/cmd/build/cmd.go
@@ -60,7 +60,7 @@ func runCommand(flags *buildFlags) func(cmd *cobra.Command, args []string) {
 			output.Fatalf("Failed to process configuration: %s", err)
 			return
 		}
-		err = RunBuild(opts)
+		err = RunBuild(cmd.Context(), opts)
 		if err != nil {
 			output.Fatalf("Failed to build model kit: %s", err)
 			return

diff --git a/pkg/cmd/build/cmd_test.go b/pkg/cmd/build/cmd_test.go
@@ -38,7 +38,7 @@ func TestBuildOptions_RunBuild(t *testing.T) {
 		contextDir: "/path/to/context",
 	}
 
-	err := RunBuild(options)
+	err := RunBuild(context.Background(), options)
 
 	assert.NoError(t, err)
 }
diff --git a/pkg/cmd/list/cmd.go b/pkg/cmd/list/cmd.go
@@ -76,7 +76,7 @@ func runCommand(flags *listFlags) func(*cobra.Command, []string) {
 
 		var allInfoLines []string
 		if opts.remoteRef == nil {
-			lines, err := listLocalKits(opts.storageHome)
+			lines, err := listLocalKits(cmd.Context(), opts.storageHome)
 			if err != nil {
 				output.Fatalln(err)
 			}

diff --git a/pkg/cmd/list/list.go b/pkg/cmd/list/list.go
@@ -24,7 +24,7 @@ const (
 	listTableFmt    = "%s\t%s\t%s\t%s\t%s\t%s\t"
 )
 
-func listLocalKits(storageRoot string) ([]string, error) {
+func listLocalKits(ctx context.Context, storageRoot string) ([]string, error) {
 	storeDirs, err := findRepos(storageRoot)
 	if err != nil {
 		return nil, err
@@ -34,7 +34,7 @@ func listLocalKits(storageRoot string) ([]string, error) {
 	for _, storeDir := range storeDirs {
 		store := storage.NewLocalStore(storageRoot, storeDir)
 
-		infolines, err := listKits(store)
+		infolines, err := listKits(ctx, store)
 		if err != nil {
 			return nil, err
 		}
@@ -43,22 +43,22 @@ func listLocalKits(storageRoot string) ([]string, error) {
 	return allInfoLines, nil
 }
 
-func listKits(store storage.Store) ([]string, error) {
+func listKits(ctx context.Context, store storage.Store) ([]string, error) {
 	index, err := store.ParseIndexJson()
 	if err != nil {
 		return nil, err
 	}
 
 	var infolines []string
 	for _, manifestDesc := range index.Manifests {
-		manifest, err := getManifest(store, manifestDesc)
+		manifest, err := getManifest(ctx, store, manifestDesc)
 		if err != nil {
 			return nil, err
 		}
 		if manifest.Config.MediaType != constants.ModelConfigMediaType {
 			continue
 		}
-		manifestConf, err := readManifestConfig(store, manifest)
+		manifestConf, err := readManifestConfig(ctx, store, manifest)
 		if err != nil {
 			return nil, err
 		}
@@ -69,8 +69,8 @@ func listKits(store storage.Store) ([]string, error) {
 	return infolines, nil
 }
 
-func getManifest(store storage.Store, manifestDesc ocispec.Descriptor) (*ocispec.Manifest, error) {
-	manifestBytes, err := store.Fetch(context.Background(), manifestDesc)
+func getManifest(ctx context.Context, store storage.Store, manifestDesc ocispec.Descriptor) (*ocispec.Manifest, error) {
+	manifestBytes, err := store.Fetch(ctx, manifestDesc)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read manifest %s: %w", manifestDesc.Digest, err)
 	}
@@ -81,8 +81,8 @@ func getManifest(store storage.Store, manifestDesc ocispec.Descriptor) (*ocispec
 	return manifest, nil
 }
 
-func readManifestConfig(store storage.Store, manifest *ocispec.Manifest) (*artifact.KitFile, error) {
-	configBytes, err := store.Fetch(context.Background(), manifest.Config)
+func readManifestConfig(ctx context.Context, store storage.Store, manifest *ocispec.Manifest) (*artifact.KitFile, error) {
+	configBytes, err := store.Fetch(ctx, manifest.Config)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read config: %w", err)
 	}

diff --git a/pkg/cmd/list/list_test.go b/pkg/cmd/list/list_test.go
@@ -1,6 +1,7 @@
 package list
 
 import (
+	"context"
 	"fmt"
 	"kitops/pkg/artifact"
 	"kitops/pkg/lib/constants"
@@ -155,7 +156,7 @@ func TestListKits(t *testing.T) {
 				Index:     tt.index,
 				Repo:      tt.repo,
 			}
-			summaryLines, err := listKits(testStore)
+			summaryLines, err := listKits(context.Background(), testStore)
 			if tt.expectErrRegexp != "" {
 				// Should be error
 				assert.Empty(t, summaryLines, "Should not output summary on error")

diff --git a/pkg/lib/storage/local.go b/pkg/lib/storage/local.go
@@ -43,33 +43,33 @@ func NewLocalStore(storeRoot, repo string) Store {
 	}
 }
 
-func (store *LocalStore) SaveModel(model *artifact.Model, tag string) (*ocispec.Descriptor, error) {
-	configDesc, err := store.saveConfigFile(model.Config)
+func (store *LocalStore) SaveModel(ctx context.Context, model *artifact.Model, tag string) (*ocispec.Descriptor, error) {
+	configDesc, err := store.saveConfigFile(ctx, model.Config)
 	if err != nil {
 		return nil, err
 	}
 	var layerDescs []ocispec.Descriptor
 	for _, layer := range model.Layers {
-		layerDesc, err := store.saveContentLayer(&layer)
+		layerDesc, err := store.saveContentLayer(ctx, &layer)
 		if err != nil {
 			return nil, err
 		}
 		layerDescs = append(layerDescs, layerDesc)
 	}
 
-	manifestDesc, err := store.saveModelManifest(layerDescs, configDesc, tag)
+	manifestDesc, err := store.saveModelManifest(ctx, layerDescs, configDesc, tag)
 	if err != nil {
 		return nil, err
 	}
 	return manifestDesc, nil
 }
 
-func (store *LocalStore) TagModel(manifestDesc ocispec.Descriptor, tag string) error {
+func (store *LocalStore) TagModel(ctx context.Context, manifestDesc ocispec.Descriptor, tag string) error {
 	if err := validateTag(tag); err != nil {
 		return err
 	}
 
-	if err := store.storage.Tag(context.Background(), manifestDesc, tag); err != nil {
+	if err := store.storage.Tag(ctx, manifestDesc, tag); err != nil {
 		return fmt.Errorf("failed to tag manifest: %w", err)
 	}
 
@@ -99,9 +99,7 @@ func (store *LocalStore) GetRepository() string {
 	return store.repo
 }
 
-func (store *LocalStore) saveContentLayer(layer *artifact.ModelLayer) (ocispec.Descriptor, error) {
-	ctx := context.Background()
-
+func (store *LocalStore) saveContentLayer(ctx context.Context, layer *artifact.ModelLayer) (ocispec.Descriptor, error) {
 	buf := &bytes.Buffer{}
 	err := layer.Apply(buf)
 	if err != nil {
@@ -133,8 +131,7 @@ func (store *LocalStore) saveContentLayer(layer *artifact.ModelLayer) (ocispec.D
 	return desc, nil
 }
 
-func (store *LocalStore) saveConfigFile(model *artifact.KitFile) (ocispec.Descriptor, error) {
-	ctx := context.Background()
+func (store *LocalStore) saveConfigFile(ctx context.Context, model *artifact.KitFile) (ocispec.Descriptor, error) {
 	modelBytes, err := model.MarshalToJSON()
 	if err != nil {
 		return ocispec.DescriptorEmptyJSON, err
@@ -163,8 +160,7 @@ func (store *LocalStore) saveConfigFile(model *artifact.KitFile) (ocispec.Descri
 	return desc, nil
 }
 
-func (store *LocalStore) saveModelManifest(layerDescs []ocispec.Descriptor, config ocispec.Descriptor, tag string) (*ocispec.Descriptor, error) {
-	ctx := context.Background()
+func (store *LocalStore) saveModelManifest(ctx context.Context, layerDescs []ocispec.Descriptor, config ocispec.Descriptor, tag string) (*ocispec.Descriptor, error) {
 	manifest := ocispec.Manifest{
 		Versioned:   specs.Versioned{SchemaVersion: 2},
 		Config:      config,
@@ -202,7 +198,7 @@ func (store *LocalStore) saveModelManifest(layerDescs []ocispec.Descriptor, conf
 		if err := validateTag(tag); err != nil {
 			return nil, err
 		}
-		if err := store.storage.Tag(context.Background(), desc, tag); err != nil {
+		if err := store.storage.Tag(ctx, desc, tag); err != nil {
 			return nil, fmt.Errorf("failed to tag manifest: %w", err)
 		}
 		output.Debugf("Added tag to manifest: %s", tag)

diff --git a/pkg/lib/storage/store.go b/pkg/lib/storage/store.go
@@ -11,8 +11,8 @@ import (
 )
 
 type Store interface {
-	SaveModel(model *artifact.Model, tag string) (*ocispec.Descriptor, error)
-	TagModel(manifestDesc ocispec.Descriptor, tag string) error
+	SaveModel(ctx context.Context, model *artifact.Model, tag string) (*ocispec.Descriptor, error)
+	TagModel(ctx context.Context, manifestDesc ocispec.Descriptor, tag string) error
 	GetRepository() string
 	ParseIndexJson() (*ocispec.Index, error)
 	Fetch(context.Context, ocispec.Descriptor) ([]byte, error)

diff --git a/pkg/lib/testing/testing.go b/pkg/lib/testing/testing.go
@@ -71,12 +71,12 @@ func (s *TestStore) ParseIndexJson() (*ocispec.Index, error) {
 	return nil, TestingNotFoundError
 }
 
-func (*TestStore) TagModel(ocispec.Descriptor, string) error {
+func (*TestStore) TagModel(context.Context, ocispec.Descriptor, string) error {
 	return fmt.Errorf("tag model is not implemented for testing")
 }
 
 // SaveModel is not yet implemented!
-func (*TestStore) SaveModel(*artifact.Model, string) (*ocispec.Descriptor, error) {
+func (*TestStore) SaveModel(context.Context, *artifact.Model, string) (*ocispec.Descriptor, error) {
 	return nil, fmt.Errorf("save model is not implemented for testing")
 }