feat(artifact): add Walk method to artifact controller (#10881)

1. Add Walk method to artifact controller.
2. Only query references when artifact is image index.

Signed-off-by: He Weiwei <hweiwei@vmware.com>

src/api/artifact/controller.go

@@ -15,8 +15,12 @@
 package artifact
 
 import (
+	"container/list"
 	"context"
 	"fmt"
+	"strings"
+	"time"
+
 	"github.com/goharbor/harbor/src/api/artifact/abstractor"
 	"github.com/goharbor/harbor/src/api/artifact/abstractor/resolver"
 	"github.com/goharbor/harbor/src/api/artifact/descriptor"
@@ -31,8 +35,7 @@ import (
 	"github.com/goharbor/harbor/src/pkg/registry"
 	"github.com/goharbor/harbor/src/pkg/signature"
 	"github.com/opencontainers/go-digest"
-	"strings"
-	"time"
+
 	// registry image resolvers
 	_ "github.com/goharbor/harbor/src/api/artifact/abstractor/resolver/image"
 	// register chart resolver
@@ -84,6 +87,8 @@ type Controller interface {
 	AddLabel(ctx context.Context, artifactID int64, labelID int64) (err error)
 	// RemoveLabel from the specified artifact
 	RemoveLabel(ctx context.Context, artifactID int64, labelID int64) (err error)
+	// Walk walks the artifact tree rooted at root, calling walkFn for each artifact in the tree, including root.
+	Walk(ctx context.Context, root *Artifact, walkFn func(*Artifact) error, option *Option) error
 }
 
 // NewController creates an instance of the default artifact controller
@@ -313,7 +318,7 @@ func (c *controller) deleteDeeply(ctx context.Context, id int64, isRoot bool) er
 		return err
 	}
 
-	// clean associations between blob and project when when the blob is not needed by project
+	// clean associations between blob and project when the blob is not needed by project
 	if err := c.blobMgr.CleanupAssociationsForProject(ctx, art.ProjectID, blobs); err != nil {
 		return err
 	}
@@ -435,6 +440,39 @@ func (c *controller) RemoveLabel(ctx context.Context, artifactID int64, labelID
 	return c.labelMgr.RemoveFrom(ctx, labelID, artifactID)
 }
 
+func (c *controller) Walk(ctx context.Context, root *Artifact, walkFn func(*Artifact) error, option *Option) error {
+	queue := list.New()
+	queue.PushBack(root)
+
+	for queue.Len() > 0 {
+		elem := queue.Front()
+		queue.Remove(elem)
+
+		artifact := elem.Value.(*Artifact)
+		if err := walkFn(artifact); err != nil {
+			return err
+		}
+
+		if len(artifact.References) > 0 {
+			var ids []int64
+			for _, ref := range artifact.References {
+				ids = append(ids, ref.ChildID)
+			}
+
+			artifacts, err := c.List(ctx, q.New(q.KeyWords{"id__in": ids}), option)
+			if err != nil {
+				return err
+			}
+
+			for _, artifact := range artifacts {
+				queue.PushBack(artifact)
+			}
+		}
+	}
+
+	return nil
+}
+
 // assemble several part into a single artifact
 func (c *controller) assembleArtifact(ctx context.Context, art *artifact.Artifact, option *Option) *Artifact {
 	artifact := &Artifact{

src/api/artifact/controller_test.go

@@ -16,6 +16,9 @@ package artifact
 
 import (
 	"context"
+	"testing"
+	"time"
+
 	"github.com/goharbor/harbor/src/api/artifact/abstractor/resolver"
 	"github.com/goharbor/harbor/src/api/artifact/descriptor"
 	"github.com/goharbor/harbor/src/api/tag"
@@ -33,10 +36,9 @@ import (
 	"github.com/goharbor/harbor/src/testing/pkg/label"
 	"github.com/goharbor/harbor/src/testing/pkg/registry"
 	repotesting "github.com/goharbor/harbor/src/testing/pkg/repository"
+	v1 "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/suite"
-	"testing"
-	"time"
 )
 
 // TODO find another way to test artifact controller, it's hard to maintain currently
@@ -528,6 +530,41 @@ func (c *controllerTestSuite) TestRemoveFrom() {
 	c.Require().Nil(err)
 }
 
+func (c *controllerTestSuite) TestWalk() {
+	c.artMgr.On("List").Return([]*artifact.Artifact{
+		{ManifestMediaType: v1.MediaTypeImageManifest},
+		{ManifestMediaType: v1.MediaTypeImageManifest},
+	}, nil)
+
+	{
+		root := &Artifact{}
+
+		var n int
+		c.ctl.Walk(context.TODO(), root, func(a *Artifact) error {
+			n++
+			return nil
+		}, nil)
+
+		c.Equal(1, n)
+	}
+
+	{
+		root := &Artifact{}
+		root.References = []*artifact.Reference{
+			{ParentID: 1, ChildID: 2},
+			{ParentID: 1, ChildID: 3},
+		}
+
+		var n int
+		c.ctl.Walk(context.TODO(), root, func(a *Artifact) error {
+			n++
+			return nil
+		}, nil)
+
+		c.Equal(3, n)
+	}
+}
+
 func TestControllerTestSuite(t *testing.T) {
 	suite.Run(t, &controllerTestSuite{})
 }

src/pkg/artifact/manager.go

@@ -158,15 +158,15 @@ func (m *manager) assemble(ctx context.Context, art *dao.Artifact) (*Artifact, e
 	artifact := &Artifact{}
 	// convert from database object
 	artifact.From(art)
+
 	// populate the references
-	references, err := m.ListReferences(ctx, &q.Query{
-		Keywords: map[string]interface{}{
-			"ParentID": artifact.ID,
-		},
-	})
+	if artifact.HasChildren() {
+		references, err := m.ListReferences(ctx, q.New(q.KeyWords{"ParentID": artifact.ID}))
 		if err != nil {
 			return nil, err
 		}
 		artifact.References = references
+	}
+
 	return artifact, nil
 }

src/pkg/artifact/manager_test.go

@@ -16,12 +16,14 @@ package artifact
 
 import (
 	"context"
-	"github.com/goharbor/harbor/src/pkg/artifact/dao"
-	"github.com/goharbor/harbor/src/pkg/q"
-	"github.com/stretchr/testify/mock"
-	"github.com/stretchr/testify/suite"
 	"testing"
 	"time"
+
+	"github.com/goharbor/harbor/src/pkg/artifact/dao"
+	"github.com/goharbor/harbor/src/pkg/q"
+	v1 "github.com/opencontainers/image-spec/specs-go/v1"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/suite"
 )
 
 type fakeDao struct {
@@ -98,7 +100,7 @@ func (m *managerTestSuite) TestAssemble() {
 		ID:                1,
 		Type:              "IMAGE",
 		MediaType:         "application/vnd.oci.image.config.v1+json",
-		ManifestMediaType: "application/vnd.oci.image.manifest.v1+json",
+		ManifestMediaType: v1.MediaTypeImageIndex,
 		ProjectID:         1,
 		RepositoryID:      1,
 		Digest:            "sha256:418fb88ec412e340cdbef913b8ca1bbe8f9e8dc705f9617414c1f2c8db980180",
@@ -159,7 +161,7 @@ func (m *managerTestSuite) TestGet() {
 		ID:                1,
 		Type:              "IMAGE",
 		MediaType:         "application/vnd.oci.image.config.v1+json",
-		ManifestMediaType: "application/vnd.oci.image.manifest.v1+json",
+		ManifestMediaType: v1.MediaTypeImageIndex,
 		ProjectID:         1,
 		RepositoryID:      1,
 		Digest:            "sha256:418fb88ec412e340cdbef913b8ca1bbe8f9e8dc705f9617414c1f2c8db980180",

src/pkg/artifact/model.go

@@ -18,6 +18,7 @@ import (
 	"encoding/json"
 	"time"
 
+	"github.com/docker/distribution/manifest/manifestlist"
 	"github.com/goharbor/harbor/src/common/utils/log"
 	"github.com/goharbor/harbor/src/pkg/artifact/dao"
 	v1 "github.com/opencontainers/image-spec/specs-go/v1"
@@ -43,6 +44,12 @@ type Artifact struct {
 	References        []*Reference           `json:"references"` // child artifacts referenced by the parent artifact if the artifact is an index
 }
 
+// HasChildren returns true when artifact has children artifacts, most times that means the artifact is Image Index.
+func (a *Artifact) HasChildren() bool {
+	return a.ManifestMediaType == v1.MediaTypeImageIndex ||
+		a.ManifestMediaType == manifestlist.MediaTypeManifestList
+}
+
 // From converts the database level artifact to the business level object
 func (a *Artifact) From(art *dao.Artifact) {
 	a.ID = art.ID

src/pkg/artifact/model_test.go

@@ -18,7 +18,9 @@ import (
 	"testing"
 	"time"
 
+	"github.com/docker/distribution/manifest/manifestlist"
 	"github.com/goharbor/harbor/src/pkg/artifact/dao"
+	v1 "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/suite"
 )
@@ -94,6 +96,17 @@ func (m *modelTestSuite) TestArtifactTo() {
 	assert.Equal(t, `{"anno1":"value1"}`, dbArt.Annotations)
 }
 
+func (m *modelTestSuite) TestHasChildren() {
+	art1 := Artifact{ManifestMediaType: v1.MediaTypeImageIndex}
+	m.True(art1.HasChildren())
+
+	art2 := Artifact{ManifestMediaType: manifestlist.MediaTypeManifestList}
+	m.True(art2.HasChildren())
+
+	art3 := Artifact{ManifestMediaType: v1.MediaTypeImageManifest}
+	m.False(art3.HasChildren())
+}
+
 func TestModel(t *testing.T) {
 	suite.Run(t, &modelTestSuite{})
 }

src/testing/api/artifact/controller.go

@@ -16,11 +16,12 @@ package artifact
 
 import (
 	"context"
+	"time"
+
 	"github.com/goharbor/harbor/src/api/artifact"
 	"github.com/goharbor/harbor/src/api/artifact/abstractor/resolver"
 	"github.com/goharbor/harbor/src/pkg/q"
 	"github.com/stretchr/testify/mock"
-	"time"
 )
 
 // FakeController is a fake artifact controller that implement src/api/artifact.Controller interface
@@ -109,3 +110,9 @@ func (f *FakeController) RemoveLabel(ctx context.Context, artifactID int64, labe
 	args := f.Called()
 	return args.Error(0)
 }
+
+// Walk ...
+func (f *FakeController) Walk(ctx context.Context, root *artifact.Artifact, workFn func(*artifact.Artifact) error, option *artifact.Option) error {
+	args := f.Called()
+	return args.Error(0)
+}