feat: support JSON/go-template output of indirect referrers for `oras…

… discover` Signed-off-by: Billy Zha <[email protected]>
qweeah · Jun 14, 2024 · bf940ff · bf940ff
1 parent 3291d32
commit bf940ff
Show file tree

Hide file tree

Showing 6 changed files with 78 additions and 47 deletions.
diff --git a/cmd/oras/internal/display/handler.go b/cmd/oras/internal/display/handler.go
@@ -110,17 +110,17 @@ func NewPullHandler(out io.Writer, format option.Format, path string, tty *os.Fi
 }
 
 // NewDiscoverHandler returns status and metadata handlers for discover command.
-func NewDiscoverHandler(out io.Writer, format option.Format, path string, rawReference string, desc ocispec.Descriptor, verbose bool) (metadata.DiscoverHandler, error) {
+func NewDiscoverHandler(out io.Writer, format option.Format, path string, rawReference string, desc ocispec.Descriptor, verbose, forceRecursive bool) (metadata.DiscoverHandler, error) {
 	var handler metadata.DiscoverHandler
 	switch format.Type {
 	case option.FormatTypeTree.Name, "":
 		handler = tree.NewDiscoverHandler(out, path, desc, verbose)
 	case option.FormatTypeTable.Name:
 		handler = table.NewDiscoverHandler(out, rawReference, desc, verbose)
 	case option.FormatTypeJSON.Name:
-		handler = json.NewDiscoverHandler(out, desc, path)
+		handler = json.NewDiscoverHandler(out, path, desc, forceRecursive)
 	case option.FormatTypeGoTemplate.Name:
-		handler = template.NewDiscoverHandler(out, desc, path, format.Template)
+		handler = template.NewDiscoverHandler(out, path, desc, forceRecursive, format.Template)
 	default:
 		return nil, errors.UnsupportedFormatTypeError(format.Type)
 	}

diff --git a/cmd/oras/internal/display/metadata/json/discover.go b/cmd/oras/internal/display/metadata/json/discover.go
@@ -16,11 +16,9 @@ limitations under the License.
 package json
 
 import (
-	"fmt"
 	"io"
 
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
-	"oras.land/oras-go/v2/content"
 	"oras.land/oras/cmd/oras/internal/display/metadata"
 	"oras.land/oras/cmd/oras/internal/display/metadata/model"
 	"oras.land/oras/cmd/oras/internal/display/utils"
@@ -29,35 +27,32 @@ import (
 // discoverHandler handles json metadata output for discover events.
 type discoverHandler struct {
 	out       io.Writer
-	root      ocispec.Descriptor
 	path      string
-	referrers []ocispec.Descriptor
+	model     model.Discover
+	recursive bool
 }
 
 // NewDiscoverHandler creates a new handler for discover events.
-func NewDiscoverHandler(out io.Writer, root ocispec.Descriptor, path string) metadata.DiscoverHandler {
+func NewDiscoverHandler(out io.Writer, path string, subject ocispec.Descriptor, recursive bool) metadata.DiscoverHandler {
 	return &discoverHandler{
-		out:  out,
-		root: root,
-		path: path,
+		out:       out,
+		path:      path,
+		recursive: recursive,
+		model:     model.NewDiscover(path, subject),
 	}
 }
 
 // MultiLevelSupported implements metadata.DiscoverHandler.
 func (h *discoverHandler) MultiLevelSupported() bool {
-	return false
+	return h.recursive
 }
 
 // OnDiscovered implements metadata.DiscoverHandler.
 func (h *discoverHandler) OnDiscovered(referrer, subject ocispec.Descriptor) error {
-	if !content.Equal(subject, h.root) {
-		return fmt.Errorf("unexpected subject descriptor: %v", subject)
-	}
-	h.referrers = append(h.referrers, referrer)
-	return nil
+	return h.model.Add(referrer, subject)
 }
 
 // OnCompleted implements metadata.DiscoverHandler.
 func (h *discoverHandler) OnCompleted() error {
-	return utils.PrintPrettyJSON(h.out, model.NewDiscover(h.path, h.referrers))
+	return utils.PrintPrettyJSON(h.out, &h.model.Root)
 }
diff --git a/cmd/oras/internal/display/metadata/model/discover.go b/cmd/oras/internal/display/metadata/model/discover.go
@@ -15,19 +15,58 @@ limitations under the License.
 
 package model
 
-import ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+import (
+	"fmt"
+	"sync"
 
-type discover struct {
-	Manifests []Descriptor `json:"manifests"`
+	"github.com/opencontainers/go-digest"
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+)
+
+// Node represents a node in the discovered reference tree.
+type Node struct {
+	Descriptor
+	Referrers []*Node `json:"manifests,omitempty"`
+}
+
+// Discover is a model for discovered referrers.
+type Discover struct {
+	lock  sync.Mutex
+	name  string
+	nodes map[digest.Digest]*Node
+	Root  *Node
+}
+
+// Add adds a reference(referrer->subject) to the discovered referrers tree.
+func (d *Discover) Add(referrer, subject ocispec.Descriptor) error {
+	d.lock.Lock()
+	defer d.lock.Unlock()
+
+	to, ok := d.nodes[subject.Digest]
+	if !ok {
+		return fmt.Errorf("unexpected subject descriptor: %v", subject)
+	}
+	from := NewNode(d.name, referrer)
+	d.nodes[from.Digest] = from
+	to.Referrers = append(to.Referrers, from)
+	return nil
 }
 
 // NewDiscover creates a new discover model.
-func NewDiscover(name string, descs []ocispec.Descriptor) discover {
-	discover := discover{
-		Manifests: make([]Descriptor, 0, len(descs)),
+func NewDiscover(name string, root ocispec.Descriptor) Discover {
+	treeRoot := NewNode(name, root)
+	return Discover{
+		name: name,
+		nodes: map[digest.Digest]*Node{
+			root.Digest: treeRoot,
+		},
+		Root: treeRoot,
 	}
-	for _, desc := range descs {
-		discover.Manifests = append(discover.Manifests, FromDescriptor(name, desc))
+}
+
+// NewNode creates a new discover model.
+func NewNode(name string, desc ocispec.Descriptor) *Node {
+	return &Node{
+		Descriptor: FromDescriptor(name, desc),
 	}
-	return discover
 }
diff --git a/cmd/oras/internal/display/metadata/template/discover.go b/cmd/oras/internal/display/metadata/template/discover.go
@@ -16,49 +16,44 @@ limitations under the License.
 package template
 
 import (
-	"fmt"
 	"io"
 
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
-	"oras.land/oras-go/v2/content"
 	"oras.land/oras/cmd/oras/internal/display/metadata"
 	"oras.land/oras/cmd/oras/internal/display/metadata/model"
 )
 
 // discoverHandler handles json metadata output for discover events.
 type discoverHandler struct {
-	referrers []ocispec.Descriptor
 	template  string
 	path      string
-	root      ocispec.Descriptor
+	model     model.Discover
+	recursive bool
 	out       io.Writer
 }
 
 // NewDiscoverHandler creates a new handler for discover events.
-func NewDiscoverHandler(out io.Writer, root ocispec.Descriptor, path string, template string) metadata.DiscoverHandler {
+func NewDiscoverHandler(out io.Writer, path string, root ocispec.Descriptor, recursive bool, template string) metadata.DiscoverHandler {
 	return &discoverHandler{
-		out:      out,
-		root:     root,
-		path:     path,
-		template: template,
+		out:       out,
+		model:     model.NewDiscover(path, root),
+		path:      path,
+		recursive: recursive,
+		template:  template,
 	}
 }
 
 // MultiLevelSupported implements metadata.DiscoverHandler.
 func (h *discoverHandler) MultiLevelSupported() bool {
-	return false
+	return h.recursive
 }
 
 // OnDiscovered implements metadata.DiscoverHandler.
 func (h *discoverHandler) OnDiscovered(referrer, subject ocispec.Descriptor) error {
-	if !content.Equal(subject, h.root) {
-		return fmt.Errorf("unexpected subject descriptor: %v", subject)
-	}
-	h.referrers = append(h.referrers, referrer)
-	return nil
+	return h.model.Add(referrer, subject)
 }
 
 // OnCompleted implements metadata.DiscoverHandler.
 func (h *discoverHandler) OnCompleted() error {
-	return parseAndWrite(h.out, model.NewDiscover(h.path, h.referrers), h.template)
+	return parseAndWrite(h.out, &h.model.Root, h.template)
 }
diff --git a/cmd/oras/internal/display/metadata/tree/discover.go b/cmd/oras/internal/display/metadata/tree/discover.go
@@ -27,7 +27,7 @@ import (
 	"oras.land/oras/internal/tree"
 )
 
-// discoverHandler handles json metadata output for discover events.
+// discoverHandler handles tree metadata output for discover events.
 type discoverHandler struct {
 	out     io.Writer
 	path    string

diff --git a/cmd/oras/root/discover.go b/cmd/oras/root/discover.go
@@ -40,6 +40,7 @@ type discoverOptions struct {
 	option.Format
 
 	artifactType string
+	recursive    bool
 }
 
 func discoverCmd() *cobra.Command {
@@ -98,13 +99,14 @@ Example - Discover referrers of the manifest tagged 'v1' in an OCI image layout
 	}
 
 	cmd.Flags().StringVarP(&opts.artifactType, "artifact-type", "", "", "artifact type")
+	cmd.Flags().BoolVarP(&opts.recursive, "recursive", "r", false, "recursively discover indirect referrers")
 	cmd.Flags().StringVarP(&opts.Format.FormatFlag, "output", "o", "tree", "[Deprecated] format in which to display referrers (table, json, or tree). tree format will also show indirect referrers")
 	opts.FormatFlag = option.FormatTypeTree.Name
 	opts.AllowedTypes = []*option.FormatType{
 		option.FormatTypeTree,
 		option.FormatTypeTable,
-		option.FormatTypeJSON.WithUsage("Get direct referrers and output in JSON format"),
-		option.FormatTypeGoTemplate.WithUsage("Print direct referrers using the given Go template"),
+		option.FormatTypeJSON.WithUsage("Get referrers and output in JSON format"),
+		option.FormatTypeGoTemplate.WithUsage("Print referrers using the given Go template"),
 	}
 	opts.EnableDistributionSpecFlag()
 	option.ApplyFlags(&opts, cmd.Flags())
@@ -129,7 +131,7 @@ func runDiscover(cmd *cobra.Command, opts *discoverOptions) error {
 		return err
 	}
 
-	handler, err := display.NewDiscoverHandler(cmd.OutOrStdout(), opts.Format, opts.Path, opts.RawReference, desc, opts.Verbose)
+	handler, err := display.NewDiscoverHandler(cmd.OutOrStdout(), opts.Format, opts.Path, opts.RawReference, desc, opts.Verbose, opts.recursive)
 	if err != nil {
 		return err
 	}