From 4f21b69a58aa2bd39331a6b838a68904e7e79315 Mon Sep 17 00:00:00 2001
From: Angel Misevski <amisevsk@gmail.com>
Date: Wed, 28 Feb 2024 14:34:18 -0500
Subject: [PATCH] Add --force flag to kit remove

Add flag --force to the remove subcommand. If specified, it will result
in modelkits being deleted even if other tags refer to them.
---
 pkg/cmd/remove/cmd.go    | 20 ++++++++++++--------
 pkg/cmd/remove/remove.go |  4 ++--
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/pkg/cmd/remove/cmd.go b/pkg/cmd/remove/cmd.go
index 1a4bf6ae..7bacb605 100644
--- a/pkg/cmd/remove/cmd.go
+++ b/pkg/cmd/remove/cmd.go
@@ -21,8 +21,9 @@ Description:
 
   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.`
+  might refer to it. If specified by tag (and the --force flag is not used),
+  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>
@@ -30,9 +31,10 @@ Description:
 )
 
 type removeOptions struct {
-	configHome string
-	modelRef   *registry.Reference
-	extraTags  []string
+	configHome  string
+	forceDelete bool
+	modelRef    *registry.Reference
+	extraTags   []string
 }
 
 func (opts *removeOptions) complete(ctx context.Context, args []string) error {
@@ -63,6 +65,7 @@ func RemoveCommand() *cobra.Command {
 		Run:     runCommand(opts),
 	}
 	cmd.Args = cobra.ExactArgs(1)
+	cmd.Flags().BoolVarP(&opts.forceDelete, "force", "f", false, "remove manifest even if other tags refer to it")
 	return cmd
 }
 
@@ -76,7 +79,7 @@ func runCommand(opts *removeOptions) func(*cobra.Command, []string) {
 		if err != nil {
 			output.Fatalf("Failed to read local storage: %s", storageRoot)
 		}
-		desc, err := removeModel(cmd.Context(), localStore, opts.modelRef)
+		desc, err := removeModel(cmd.Context(), localStore, opts.modelRef, opts.forceDelete)
 		if err != nil {
 			output.Fatalf("Failed to remove: %s", err)
 		}
@@ -85,11 +88,12 @@ func runCommand(opts *removeOptions) func(*cobra.Command, []string) {
 		for _, tag := range opts.extraTags {
 			ref := *opts.modelRef
 			ref.Reference = tag
-			desc, err := removeModel(cmd.Context(), localStore, &ref)
+			desc, err := removeModel(cmd.Context(), localStore, &ref, opts.forceDelete)
 			if err != nil {
 				output.Errorf("Failed to remove: %s", err)
+			} else {
+				output.Infof("Removed %s (digest %s)", ref.String(), desc.Digest)
 			}
-			output.Infof("Removed %s (digest %s)", ref.String(), desc.Digest)
 		}
 	}
 }
diff --git a/pkg/cmd/remove/remove.go b/pkg/cmd/remove/remove.go
index 7c1c0588..a34f8485 100644
--- a/pkg/cmd/remove/remove.go
+++ b/pkg/cmd/remove/remove.go
@@ -13,7 +13,7 @@ import (
 	"oras.land/oras-go/v2/registry"
 )
 
-func removeModel(ctx context.Context, store repo.LocalStorage, ref *registry.Reference) (ocispec.Descriptor, error) {
+func removeModel(ctx context.Context, store repo.LocalStorage, ref *registry.Reference, forceDelete bool) (ocispec.Descriptor, error) {
 	desc, err := oras.Resolve(ctx, store, ref.Reference, oras.ResolveOptions{})
 	if err != nil {
 		if err == errdef.ErrNotFound {
@@ -23,7 +23,7 @@ func removeModel(ctx context.Context, store repo.LocalStorage, ref *registry.Ref
 	}
 
 	// If reference passed in is a digest, remove the manifest ignoring any tags the manifest might have
-	if err := ref.ValidateReferenceAsDigest(); err == nil {
+	if err := ref.ValidateReferenceAsDigest(); err == nil || forceDelete {
 		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)