feat: Implement tag subcommand

cmd/ormb/cmd/tag.go

@@ -0,0 +1,36 @@
+/*
+Copyright © 2020 NAME HERE <EMAIL ADDRESS>
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+package cmd
+
+import (
+	"github.com/spf13/cobra"
+)
+
+// tagCmd represents the tag command
+var tagCmd = &cobra.Command{
+	Use:     "tag",
+	Short:   "Tag the model",
+	Long:    ``,
+	PreRunE: preRunE,
+	RunE: func(cmd *cobra.Command, args []string) error {
+		// TODO(gaocegege): Validate.
+		return ormbClient.Tag(args[0], args[1])
+	},
+}
+
+func init() {
+	rootCmd.AddCommand(tagCmd)
+}

pkg/oras/cache/cache.go

@@ -248,14 +248,36 @@ func (cache *Cache) ListReferences() ([]*CacheRefSummary, error) {
 	return rr, nil
 }
 
+// TagReference tags the reference to the target.
+func (cache *Cache) TagReference(ref *oci.Reference, target *oci.Reference) error {
+	if err := cache.init(); err != nil {
+		return err
+	}
+	for _, desc := range cache.ociStore.ListReferences() {
+		if desc.Annotations[ocispec.AnnotationRefName] == ref.FullName() {
+			// We cannot use desc directly because annatations is a map,
+			// it is not copied to the new value desc.
+			new := ocispec.Descriptor{
+				MediaType: desc.MediaType,
+				Digest:    desc.Digest,
+				Size:      desc.Size,
+				URLs:      desc.URLs,
+				Platform:  desc.Platform,
+			}
+			cache.ociStore.AddReference(target.FullName(), new)
+			return cache.ociStore.SaveIndex()
+		}
+	}
+	return fmt.Errorf("Failed to find the ref %s", ref.FullName())
+}
+
 // AddManifest provides a manifest to the cache index.json.
 func (cache *Cache) AddManifest(ref *oci.Reference, manifest *ocispec.Descriptor) error {
 	if err := cache.init(); err != nil {
 		return err
 	}
 	cache.ociStore.AddReference(ref.FullName(), *manifest)
-	err := cache.ociStore.SaveIndex()
-	return err
+	return cache.ociStore.SaveIndex()
 }
 
 // Provider provides a valid containerd Provider

pkg/oras/cache/cache_test.go

@@ -20,16 +20,19 @@ var _ = Describe("Cache", func() {
 			var i Interface
 			rootPath = ".cache"
 
+			if err := os.RemoveAll(rootPath); err != nil {
+				Expect(err).To(BeNil())
+			}
 			i, err = New(CacheOptRoot(rootPath), CacheOptDebug(true), CacheOptWriter(os.Stdout))
+			Expect(err).To(BeNil())
 			c = i.(*Cache)
 		})
 
 		It("Should create the cache successfully", func() {
-			Expect(err).To(BeNil())
 			Expect(c.rootDir).To(Equal(rootPath))
 		})
 
-		Describe("with a cached artifact caicloud/test:v2", func() {
+		Describe("with a cached artifact caicloud/test:v1", func() {
 			var m *model.Model
 			var ref *oci.Reference
 
@@ -63,6 +66,13 @@ var _ = Describe("Cache", func() {
 				Expect(err).To(BeNil())
 				Expect(actual.Name).To(Equal(ref.FullName()))
 			})
+
+			It("Should tag the reference successfully", func() {
+				target, err := oci.ParseReference("test:1")
+				Expect(err).To(BeNil())
+				Expect(c.TagReference(ref, target)).To(BeNil())
+			})
 		})
+
 	})
 })

pkg/oras/cache/interface.go

@@ -15,6 +15,7 @@ type Interface interface {
 	DeleteReference(ref *oci.Reference) (*CacheRefSummary, error)
 	ListReferences() ([]*CacheRefSummary, error)
 	AddManifest(ref *oci.Reference, manifest *ocispec.Descriptor) error
+	TagReference(ref *oci.Reference, target *oci.Reference) error
 	Provider() content.Provider
 	Ingester() content.Ingester
 	ProvideIngester() orascontent.ProvideIngester

pkg/oras/client.go

@@ -130,6 +130,15 @@ func (c *Client) SaveModel(ch *model.Model, ref *oci.Reference) error {
 	return nil
 }
 
+// TagModel tags the ref to target.
+func (c *Client) TagModel(ref *oci.Reference, target *oci.Reference) error {
+	if err := c.cache.TagReference(ref, target); err != nil {
+		return err
+	}
+	fmt.Fprintf(c.out, "%s: tagged\n", ref.FullName())
+	return nil
+}
+
 // PushModel uploads a model to a registry.
 func (c *Client) PushModel(ref *oci.Reference) error {
 	r, err := c.cache.FetchReference(ref)

pkg/oras/client_test.go

@@ -265,5 +265,21 @@ var _ = Describe("OCI Client", func() {
 			).Return(returnedSummary, nil).Times(1)
 			Expect(c.LoadModel(ref)).To(Equal(m), BeNil())
 		})
+
+		It("Should tag the model successfully", func() {
+			refStr := "caicloud/resnet50:v1"
+			ref, err := oci.ParseReference(refStr)
+			Expect(err).To(BeNil())
+
+			targetStr := "caicloud/resnet50:v2"
+			target, err := oci.ParseReference(targetStr)
+			Expect(err).To(BeNil())
+
+			c.cache.(*cachemock.MockInterface).EXPECT().TagReference(
+				gomock.Eq(ref),
+				gomock.Eq(target),
+			).Return(nil).Times(1)
+			Expect(c.TagModel(ref, target)).To(BeNil())
+		})
 	})
 })

pkg/oras/interface.go

@@ -14,4 +14,5 @@ type Interface interface {
 	RemoveModel(ref *oci.Reference) error
 	PullModel(ref *oci.Reference) error
 	LoadModel(ref *oci.Reference) (*model.Model, error)
+	TagModel(ref *oci.Reference, target *oci.Reference) error
 }

pkg/ormb/ormb.go

@@ -17,6 +17,7 @@ type Interface interface {
 	Pull(refStr string) error
 	Export(refStr, dst string) error
 	Save(src, refStr string) error
+	Tag(refStr, targetStr string) error
 	Remove(refStr string) error
 }
 
@@ -59,6 +60,20 @@ func (o ORMB) Pull(refStr string) error {
 	return o.client.PullModel(ref)
 }
 
+func (o ORMB) Tag(refStr, targetStr string) error {
+	ref, err := oci.ParseReference(refStr)
+	if err != nil {
+		return err
+	}
+
+	target, err := oci.ParseReference(targetStr)
+	if err != nil {
+		return err
+	}
+
+	return o.client.TagModel(ref, target)
+}
+
 func (o ORMB) Export(refStr, dst string) error {
 	path, err := filepath.Abs(dst)
 	if err != nil {