Add 'kit remove' command

Add 'kit remove' to allow deleting models from local storage.

cmd/root.go

@@ -17,6 +17,7 @@ import (
 	"kitops/pkg/cmd/logout"
 	"kitops/pkg/cmd/pull"
 	"kitops/pkg/cmd/push"
+	"kitops/pkg/cmd/remove"
 	"kitops/pkg/cmd/version"
 	"kitops/pkg/lib/constants"
 	"kitops/pkg/output"
@@ -72,6 +73,7 @@ func addSubcommands(rootCmd *cobra.Command) {
 	rootCmd.AddCommand(push.PushCommand())
 	rootCmd.AddCommand(list.ListCommand())
 	rootCmd.AddCommand(export.ExportCommand())
+	rootCmd.AddCommand(remove.RemoveCommand())
 	rootCmd.AddCommand(version.VersionCommand())
 }
 

pkg/cmd/remove/cmd.go

@@ -0,0 +1,100 @@
+package remove
+
+import (
+	"context"
+	"fmt"
+	"kitops/pkg/lib/constants"
+	"kitops/pkg/lib/repo"
+	"kitops/pkg/output"
+	"strings"
+
+	"github.com/spf13/cobra"
+	"oras.land/oras-go/v2/registry"
+)
+
+const (
+	shortDesc = `Remove a modelkit from local storage`
+	longDesc  = `Remove a modelkit from local storage.
+
+Description:
+  Removes a modelkit from storage on the local disk.
+
+  The model to be removed may be specifed either by a tag or by a digest. If
+  specified by digest, that modelkit will be removed along with any tags that
+  might refer to it. If specified by tag, the modelkit will only be removed if
+  no other tags refer to it; otherwise it is only untagged.`
+
+	examples = `  kit remove my-registry.com/my-org/my-repo:my-tag
+  kit remove my-registry.com/my-org/my-repo@sha256:<digest>
+  kit remove my-registry.com/my-org/my-repo:tag1,tag2,tag3`
+)
+
+type removeOptions struct {
+	configHome string
+	modelRef   *registry.Reference
+	extraTags  []string
+}
+
+func (opts *removeOptions) complete(ctx context.Context, args []string) error {
+	configHome, ok := ctx.Value(constants.ConfigKey{}).(string)
+	if !ok {
+		return fmt.Errorf("default config path not set on command context")
+	}
+	opts.configHome = configHome
+
+	modelRef, extraTags, err := repo.ParseReference(args[0])
+	if err != nil {
+		return fmt.Errorf("failed to parse reference %s: %w", modelRef, err)
+	}
+	opts.modelRef = modelRef
+	opts.extraTags = extraTags
+
+	printConfig(opts)
+	return nil
+}
+
+func RemoveCommand() *cobra.Command {
+	opts := &removeOptions{}
+	cmd := &cobra.Command{
+		Use:     "remove [flags] <reference>",
+		Short:   shortDesc,
+		Long:    longDesc,
+		Example: examples,
+		Run:     runCommand(opts),
+	}
+	cmd.Args = cobra.ExactArgs(1)
+	return cmd
+}
+
+func runCommand(opts *removeOptions) func(*cobra.Command, []string) {
+	return func(cmd *cobra.Command, args []string) {
+		if err := opts.complete(cmd.Context(), args); err != nil {
+			output.Fatalf("Failed to process arguments: %s", err)
+		}
+		storageRoot := constants.StoragePath(opts.configHome)
+		localStore, err := repo.NewLocalStore(storageRoot, opts.modelRef)
+		if err != nil {
+			output.Fatalf("Failed to read local storage: %s", storageRoot)
+		}
+		desc, err := removeModel(cmd.Context(), localStore, opts.modelRef)
+		if err != nil {
+			output.Fatalf("Failed to remove: %s", err)
+		}
+		output.Infof("Removed %s (digest %s)", opts.modelRef.String(), desc.Digest)
+
+		for _, tag := range opts.extraTags {
+			ref := *opts.modelRef
+			ref.Reference = tag
+			desc, err := removeModel(cmd.Context(), localStore, &ref)
+			if err != nil {
+				output.Errorf("Failed to remove: %s", err)
+			}
+			output.Infof("Removed %s (digest %s)", ref.String(), desc.Digest)
+		}
+	}
+}
+
+func printConfig(opts *removeOptions) {
+	output.Debugf("Using config path: %s", opts.configHome)
+	output.Debugf("Removing %s and additional tags: [%s]", opts.modelRef.String(), strings.Join(opts.extraTags, ", "))
+}

pkg/cmd/remove/remove.go

@@ -0,0 +1,52 @@
+package remove
+
+import (
+	"context"
+	"fmt"
+	"kitops/pkg/lib/repo"
+	"kitops/pkg/output"
+	"strings"
+
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+	"oras.land/oras-go/v2"
+	"oras.land/oras-go/v2/errdef"
+	"oras.land/oras-go/v2/registry"
+)
+
+func removeModel(ctx context.Context, store repo.LocalStorage, ref *registry.Reference) (ocispec.Descriptor, error) {
+	desc, err := oras.Resolve(ctx, store, ref.Reference, oras.ResolveOptions{})
+	if err != nil {
+		if err == errdef.ErrNotFound {
+			return ocispec.DescriptorEmptyJSON, fmt.Errorf("model %s not found", ref.String())
+		}
+		return ocispec.DescriptorEmptyJSON, fmt.Errorf("error resolving model: %s", err)
+	}
+
+	// If reference passed in is a digest, remove the manifest ignoring any tags the manifest might have
+	if err := ref.ValidateReferenceAsDigest(); err == nil {
+		output.Debugf("Deleting manifest with digest %s", ref.Reference)
+		if err := store.Delete(ctx, desc); err != nil {
+			return ocispec.DescriptorEmptyJSON, fmt.Errorf("failed to delete model: %ws", err)
+		}
+		return desc, nil
+	}
+
+	tags, err := repo.GetTagsForDescriptor(ctx, store, desc)
+	if err != nil {
+		return ocispec.DescriptorEmptyJSON, err
+	}
+	if len(tags) <= 1 {
+		output.Debugf("Deleting manifest tagged %s", ref.Reference)
+		if err := store.Delete(ctx, desc); err != nil {
+			return ocispec.DescriptorEmptyJSON, fmt.Errorf("failed to delete model: %w", err)
+		}
+	} else {
+		output.Debugf("Found other tags for manifest: [%s]", strings.Join(tags, ", "))
+		output.Debugf("Untagging %s", ref.Reference)
+		if err := store.Untag(ctx, ref.Reference); err != nil {
+			return ocispec.DescriptorEmptyJSON, fmt.Errorf("failed to untag model: %w", err)
+		}
+	}
+
+	return desc, nil
+}

pkg/lib/repo/local.go

@@ -11,6 +11,7 @@ import (
 
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"oras.land/oras-go/v2"
+	"oras.land/oras-go/v2/content"
 	"oras.land/oras-go/v2/content/oci"
 	"oras.land/oras-go/v2/registry"
 )
@@ -19,6 +20,8 @@ type LocalStorage interface {
 	GetRepo() string
 	GetIndex() (*ocispec.Index, error)
 	oras.Target
+	content.Deleter
+	content.Untagger
 }
 
 type LocalStore struct {

pkg/lib/repo/repo.go

@@ -81,6 +81,20 @@ func GetConfig(ctx context.Context, store content.Storage, configDesc ocispec.De
 	return config, nil
 }
 
+func GetTagsForDescriptor(ctx context.Context, store LocalStorage, desc ocispec.Descriptor) ([]string, error) {
+	index, err := store.GetIndex()
+	if err != nil {
+		return nil, err
+	}
+	var tags []string
+	for _, manifest := range index.Manifests {
+		if manifest.Digest == desc.Digest {
+			tags = append(tags, manifest.Annotations[ocispec.AnnotationRefName])
+		}
+	}
+	return tags, nil
+}
+
 func ValidateTag(tag string) error {
 	if !validTagRegex.MatchString(tag) {
 		return fmt.Errorf("invalid tag")