distribution client: return untagged tags on delete (docker#104)

* distribution client: return untagged tags on delete

Signed-off-by: Dorin Geman <dorin.geman@docker.com>

* distribution client: check expected DeleteModel response

Signed-off-by: Dorin Geman <dorin.geman@docker.com>

---------

Signed-off-by: Dorin Geman <dorin.geman@docker.com>

Author: doringeman

pkg/distribution/distribution/client.go

@@ -6,6 +6,7 @@ import (
 	"io"
 	"net/http"
 	"os"
+	"strings"
 
 	"github.com/sirupsen/logrus"
 
@@ -243,44 +244,50 @@ func (c *Client) GetModel(reference string) (types.Model, error) {
 }
 
 // DeleteModel deletes a model
-func (c *Client) DeleteModel(reference string, force bool) error {
+func (c *Client) DeleteModel(reference string, force bool) (string, error) {
 	mdl, err := c.store.Read(reference)
 	if err != nil {
-		return err
+		return "", err
 	}
 	id, err := mdl.ID()
 	if err != nil {
-		return fmt.Errorf("getting model ID: %w", err)
+		return "", fmt.Errorf("getting model ID: %w", err)
 	}
 	isTag := id != reference
 
 	if isTag {
 		c.log.Infoln("Untagging model:", reference)
 		if err := c.store.RemoveTags([]string{reference}); err != nil {
 			c.log.Errorln("Failed to untag model:", err, "tag:", reference)
-			return fmt.Errorf("untagging model: %w", err)
+			return "", fmt.Errorf("untagging model: %w", err)
 		}
 	}
 
 	if len(mdl.Tags()) > 1 {
 		if isTag {
-			return nil // we are done after untagging
+			// TODO: This should return the full matching tag.
+			return fmt.Sprintf("Untagged: %s\n", reference), nil
 		} else if !force {
 			// if the reference is not a tag and there are multiple tags, return an error unless forced
-			return fmt.Errorf(
+			return "", fmt.Errorf(
 				"unable to delete %q (must be forced) due to multiple tag references: %w",
 				reference, ErrConflict,
 			)
 		}
 	}
 
+	var untaggedInfo strings.Builder
+	for _, tag := range mdl.Tags() {
+		untaggedInfo.WriteString(fmt.Sprintf("Untagged: %s\n", tag))
+	}
+
 	c.log.Infoln("Deleting model:", id)
 	if err := c.store.Delete(id); err != nil {
 		c.log.Errorln("Failed to delete model:", err, "tag:", reference)
-		return fmt.Errorf("deleting model: %w", err)
+		return "", fmt.Errorf("deleting model: %w", err)
 	}
 	c.log.Infoln("Successfully deleted model:", reference)
-	return nil
+	return untaggedInfo.String(), nil
 }
 
 // Tag adds a tag to a model

pkg/distribution/distribution/client_test.go

@@ -246,7 +246,7 @@ func TestClientPullModel(t *testing.T) {
 		}
 
 		// Delete the local model to force a pull
-		if err := client.DeleteModel(tag, false); err != nil {
+		if _, err := client.DeleteModel(tag, false); err != nil {
 			t.Fatalf("Failed to delete model: %v", err)
 		}
 
@@ -866,7 +866,7 @@ func TestPush(t *testing.T) {
 	}
 
 	// Delete local copy (so we can test pulling)
-	if err := client.DeleteModel(tag, false); err != nil {
+	if _, err := client.DeleteModel(tag, false); err != nil {
 		t.Fatalf("Failed to delete model: %v", err)
 	}
 

pkg/distribution/distribution/delete_test.go

@@ -3,6 +3,7 @@ package distribution
 import (
 	"errors"
 	"os"
+	"slices"
 	"testing"
 
 	"github.com/docker/model-distribution/internal/gguf"
@@ -122,13 +123,28 @@ func TestDeleteModel(t *testing.T) {
 			}
 
 			// Attempt to delete the model and check for expected error
-			if err := client.DeleteModel(tc.ref, tc.force); !errors.Is(err, tc.expectedErr) {
+			out, err := client.DeleteModel(tc.ref, tc.force)
+			if !errors.Is(err, tc.expectedErr) {
 				t.Fatalf("Expected error %v, got: %v", tc.expectedErr, err)
 			}
 			if tc.expectedErr != nil {
 				return
 			}
 
+			expectedOut := ""
+			if slices.Contains(tc.tags, tc.ref) {
+				// tc.ref is a tag
+				expectedOut = "Untagged: " + tc.ref + "\n"
+			} else {
+				// tc.ref is an ID
+				for _, tag := range tc.tags {
+					expectedOut += "Untagged: " + tag + "\n"
+				}
+			}
+			if expectedOut != out {
+				t.Fatalf("Expected output %q, got: %q", expectedOut, out)
+			}
+
 			// Verify model ref unreachable by ref (untagged)
 			_, err = client.GetModel(tc.ref)
 			if !errors.Is(err, ErrModelNotFound) {

pkg/distribution/distribution/ecr_test.go

@@ -51,7 +51,7 @@ func TestECRIntegration(t *testing.T) {
 		if err := client.PushModel(context.Background(), ecrTag, nil); err != nil {
 			t.Fatalf("Failed to push model to ECR: %v", err)
 		}
-		if err := client.DeleteModel(ecrTag, false); err != nil { // cleanup
+		if _, err := client.DeleteModel(ecrTag, false); err != nil { // cleanup
 			t.Fatalf("Failed to delete model from store: %v", err)
 		}
 	})

pkg/distribution/distribution/gar_test.go

@@ -52,7 +52,7 @@ func TestGARIntegration(t *testing.T) {
 		if err := client.PushModel(context.Background(), garTag, nil); err != nil {
 			t.Fatalf("Failed to push model to ECR: %v", err)
 		}
-		if err := client.DeleteModel(garTag, false); err != nil { // cleanup
+		if _, err := client.DeleteModel(garTag, false); err != nil { // cleanup
 			t.Fatalf("Failed to delete model from store: %v", err)
 		}
 	})