Enable pull time on getting manifest (#11110)

Signed-off-by: wang yan <wangyan@vmware.com>

src/api/artifact/controller.go

@@ -33,6 +33,7 @@ import (
 	"github.com/goharbor/harbor/src/pkg/notification"
 	"github.com/goharbor/harbor/src/pkg/registry"
 	"github.com/goharbor/harbor/src/pkg/signature"
+	model_tag "github.com/goharbor/harbor/src/pkg/tag/model/tag"
 	"github.com/opencontainers/go-digest"
 	"strings"
 	"time"
@@ -454,20 +455,25 @@ func (c *controller) copyDeeply(ctx context.Context, srcRepo, reference, dstRepo
 }
 
 func (c *controller) UpdatePullTime(ctx context.Context, artifactID int64, tagID int64, time time.Time) error {
-	tag, err := c.tagCtl.Get(ctx, tagID, nil)
-	if err != nil {
-		return err
-	}
-	if tag.ArtifactID != artifactID {
-		return fmt.Errorf("tag %d isn't attached to artifact %d", tagID, artifactID)
-	}
 	if err := c.artMgr.Update(ctx, &artifact.Artifact{
 		ID:       artifactID,
 		PullTime: time,
 	}, "PullTime"); err != nil {
 		return err
 	}
-	return c.tagCtl.Update(ctx, tag, "PullTime")
+	tg, err := c.tagCtl.Get(ctx, tagID, nil)
+	if err != nil {
+		return err
+	}
+	if tg.ArtifactID != artifactID {
+		return fmt.Errorf("tag %d isn't attached to artifact %d", tagID, artifactID)
+	}
+	return c.tagCtl.Update(ctx, &tag.Tag{
+		Tag: model_tag.Tag{
+			ID:       tg.ID,
+			PullTime: time,
+		},
+	}, "PullTime")
 }
 
 func (c *controller) GetAddition(ctx context.Context, artifactID int64, addition string) (*processor.Addition, error) {

src/api/artifact/controller_test.go

@@ -461,6 +461,7 @@ func (c *controllerTestSuite) TestUpdatePullTime() {
 			ArtifactID: 2,
 		},
 	}, nil)
+	c.artMgr.On("Update").Return(nil)
 	err = c.ctl.UpdatePullTime(nil, 1, 1, time.Now())
 	c.Require().NotNil(err)
 	c.tagCtl.AssertExpectations(c.T())

src/api/event/handler/init.go

@@ -3,6 +3,7 @@ package handler
 import (
 	"github.com/goharbor/harbor/src/api/event"
 	"github.com/goharbor/harbor/src/api/event/handler/auditlog"
+	"github.com/goharbor/harbor/src/api/event/handler/internal"
 	"github.com/goharbor/harbor/src/api/event/handler/replication"
 	"github.com/goharbor/harbor/src/api/event/handler/webhook/artifact"
 	"github.com/goharbor/harbor/src/api/event/handler/webhook/chart"
@@ -40,4 +41,7 @@ func init() {
 	notifier.Subscribe(event.TopicDeleteRepository, &auditlog.Handler{})
 	notifier.Subscribe(event.TopicCreateTag, &auditlog.Handler{})
 	notifier.Subscribe(event.TopicDeleteTag, &auditlog.Handler{})
+
+	// internal
+	notifier.Subscribe(event.TopicPullArtifact, &internal.Handler{})
 }

src/api/event/handler/internal/artifact.go

@@ -0,0 +1,83 @@
+// Copyright Project Harbor Authors
+//
+// 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 internal
+
+import (
+	"context"
+	beegorm "github.com/astaxie/beego/orm"
+	"github.com/goharbor/harbor/src/api/artifact"
+	"github.com/goharbor/harbor/src/api/event"
+	"github.com/goharbor/harbor/src/api/repository"
+	"github.com/goharbor/harbor/src/api/tag"
+	"github.com/goharbor/harbor/src/common/utils/log"
+	"github.com/goharbor/harbor/src/internal/orm"
+	"github.com/goharbor/harbor/src/pkg/q"
+	"time"
+)
+
+// Handler preprocess artifact event data
+type Handler struct {
+}
+
+// Handle ...
+func (a *Handler) Handle(value interface{}) error {
+	switch v := value.(type) {
+	case *event.PullArtifactEvent:
+		return a.handle(v.ArtifactEvent)
+	default:
+		log.Errorf("Can not handler this event type! %#v", v)
+	}
+	return nil
+}
+
+// IsStateful ...
+func (a *Handler) IsStateful() bool {
+	return false
+}
+
+func (a *Handler) handle(event *event.ArtifactEvent) error {
+	ctx := orm.NewContext(context.Background(), beegorm.NewOrm())
+	go func() { a.updatePullTime(ctx, event) }()
+	go func() { a.addPullCount(ctx, event) }()
+	return nil
+}
+
+func (a *Handler) updatePullTime(ctx context.Context, event *event.ArtifactEvent) {
+	var tagID int64
+	if len(event.Tags) != 0 {
+		tags, err := tag.Ctl.List(ctx, &q.Query{
+			Keywords: map[string]interface{}{
+				"ArtifactID": event.Artifact.ID,
+				"Name":       event.Tags[0],
+			},
+		}, nil)
+		if err != nil {
+			log.Infof("failed to list tags when to update pull time, %v", err)
+		} else {
+			tagID = tags[0].ID
+		}
+	}
+	if err := artifact.Ctl.UpdatePullTime(ctx, event.Artifact.ID, tagID, time.Now()); err != nil {
+		log.Debugf("failed to update pull time form artifact %d, %v", event.Artifact.ID, err)
+	}
+	return
+}
+
+func (a *Handler) addPullCount(ctx context.Context, event *event.ArtifactEvent) {
+	if err := repository.Ctl.AddPullCount(ctx, event.Artifact.RepositoryID); err != nil {
+		log.Debugf("failed to add pull count repository %d, %v", event.Artifact.RepositoryID, err)
+	}
+	return
+}

src/api/event/handler/internal/artifact_test.go

@@ -0,0 +1 @@
+package internal

src/api/repository/controller.go

@@ -49,6 +49,8 @@ type Controller interface {
 	Delete(ctx context.Context, id int64) (err error)
 	// Update the repository. Specify the properties or all properties will be updated
 	Update(ctx context.Context, repository *models.RepoRecord, properties ...string) (err error)
+	// AddPullCount increase one pull count for the specified repository
+	AddPullCount(ctx context.Context, id int64) error
 }
 
 // NewController creates an instance of the default repository controller
@@ -154,3 +156,7 @@ func (c *controller) Delete(ctx context.Context, id int64) error {
 func (c *controller) Update(ctx context.Context, repository *models.RepoRecord, properties ...string) error {
 	return c.repoMgr.Update(ctx, repository, properties...)
 }
+
+func (c *controller) AddPullCount(ctx context.Context, id int64) error {
+	return c.repoMgr.AddPullCount(ctx, id)
+}

src/api/repository/controller_test.go

@@ -135,6 +135,12 @@ func (c *controllerTestSuite) TestUpdate() {
 	c.Require().Nil(err)
 }
 
+func (c *controllerTestSuite) TestAddPullCount() {
+	c.repoMgr.On("AddPullCount").Return(nil)
+	err := c.ctl.AddPullCount(nil, 1)
+	c.Require().Nil(err)
+}
+
 func TestControllerTestSuite(t *testing.T) {
 	suite.Run(t, &controllerTestSuite{})
 }

src/pkg/repository/dao/dao.go

@@ -16,10 +16,12 @@ package dao
 
 import (
 	"context"
+	o "github.com/astaxie/beego/orm"
 	"github.com/goharbor/harbor/src/common/models"
 	ierror "github.com/goharbor/harbor/src/internal/error"
 	"github.com/goharbor/harbor/src/internal/orm"
 	"github.com/goharbor/harbor/src/pkg/q"
+	"time"
 )
 
 // DAO is the data access object interface for repository
@@ -36,6 +38,8 @@ type DAO interface {
 	Delete(ctx context.Context, id int64) (err error)
 	// Update updates the repository. Only the properties specified by "props" will be updated if it is set
 	Update(ctx context.Context, repository *models.RepoRecord, props ...string) (err error)
+	// AddPullCount increase one pull count for the specified repository
+	AddPullCount(ctx context.Context, id int64) error
 }
 
 // New returns an instance of the default DAO
@@ -130,3 +134,23 @@ func (d *dao) Update(ctx context.Context, repository *models.RepoRecord, props .
 	}
 	return nil
 }
+
+func (d *dao) AddPullCount(ctx context.Context, id int64) error {
+	ormer, err := orm.FromContext(ctx)
+	if err != nil {
+		return err
+	}
+	num, err := ormer.QueryTable(new(models.RepoRecord)).Filter("RepositoryID", id).Update(
+		o.Params{
+			"pull_count":  o.ColValue(o.ColAdd, 1),
+			"update_time": time.Now(),
+		})
+	if err != nil {
+		return err
+	}
+	if num == 0 {
+		return ierror.New(nil).WithMessage("failed to increase repository pull count: %d", id)
+
+	}
+	return nil
+}

src/pkg/repository/dao/dao_test.go

@@ -162,6 +162,27 @@ func (d *daoTestSuite) TestUpdate() {
 	d.Equal(ierror.NotFoundCode, e.Code)
 }
 
+func (d *daoTestSuite) TestAddPullCount() {
+	repository := &models.RepoRecord{
+		Name:        "test/pullcount",
+		ProjectID:   10,
+		Description: "test pull count",
+		PullCount:   1,
+	}
+	id, err := d.dao.Create(d.ctx, repository)
+	d.Require().Nil(err)
+
+	err = d.dao.AddPullCount(d.ctx, id)
+	d.Require().Nil(err)
+
+	repository, err = d.dao.Get(d.ctx, id)
+	d.Require().Nil(err)
+	d.Require().NotNil(repository)
+	d.Equal(int64(2), repository.PullCount)
+
+	d.dao.Delete(d.ctx, id)
+}
+
 func TestDaoTestSuite(t *testing.T) {
 	suite.Run(t, &daoTestSuite{})
 }

src/pkg/repository/manager.go

@@ -41,6 +41,8 @@ type Manager interface {
 	Delete(ctx context.Context, id int64) (err error)
 	// Update updates the repository. Only the properties specified by "props" will be updated if it is set
 	Update(ctx context.Context, repository *models.RepoRecord, props ...string) (err error)
+	// AddPullCount increase one pull count for the specified repository
+	AddPullCount(ctx context.Context, id int64) error
 }
 
 // New returns a default implementation of Manager
@@ -96,3 +98,7 @@ func (m *manager) Delete(ctx context.Context, id int64) error {
 func (m *manager) Update(ctx context.Context, repository *models.RepoRecord, props ...string) error {
 	return m.dao.Update(ctx, repository, props...)
 }
+
+func (m *manager) AddPullCount(ctx context.Context, id int64) error {
+	return m.dao.AddPullCount(ctx, id)
+}

src/pkg/repository/manager_test.go

@@ -51,6 +51,10 @@ func (f *fakeDao) Update(ctx context.Context, repository *models.RepoRecord, pro
 	args := f.Called()
 	return args.Error(0)
 }
+func (f *fakeDao) AddPullCount(ctx context.Context, id int64) error {
+	args := f.Called()
+	return args.Error(0)
+}
 
 type managerTestSuite struct {
 	suite.Suite
@@ -140,6 +144,13 @@ func (m *managerTestSuite) TestUpdate() {
 	m.dao.AssertExpectations(m.T())
 }
 
+func (m *managerTestSuite) TestAddPullCount() {
+	m.dao.On("AddPullCount", mock.Anything).Return(nil)
+	err := m.mgr.AddPullCount(nil, 1)
+	m.Require().Nil(err)
+	m.dao.AssertExpectations(m.T())
+}
+
 func TestManager(t *testing.T) {
 	suite.Run(t, &managerTestSuite{})
 }

src/server/registry/manifest.go

@@ -34,7 +34,7 @@ import (
 func getManifest(w http.ResponseWriter, req *http.Request) {
 	repository := router.Param(req.Context(), ":splat")
 	reference := router.Param(req.Context(), ":reference")
-	artifact, err := artifact.Ctl.GetByReference(req.Context(), repository, reference, nil)
+	art, err := artifact.Ctl.GetByReference(req.Context(), repository, reference, nil)
 	if err != nil {
 		serror.SendError(w, err)
 		return
@@ -43,7 +43,7 @@ func getManifest(w http.ResponseWriter, req *http.Request) {
 	// the reference is tag, replace it with digest
 	if _, err = digest.Parse(reference); err != nil {
 		req = req.Clone(req.Context())
-		req.URL.Path = strings.TrimSuffix(req.URL.Path, reference) + artifact.Digest
+		req.URL.Path = strings.TrimSuffix(req.URL.Path, reference) + art.Digest
 		req.URL.RawPath = req.URL.EscapedPath()
 	}
 
@@ -56,7 +56,7 @@ func getManifest(w http.ResponseWriter, req *http.Request) {
 	}
 	e := &metadata.PullArtifactEventMetadata{
 		Ctx:      req.Context(),
-		Artifact: &artifact.Artifact,
+		Artifact: &art.Artifact,
 	}
 	// the reference is tag
 	if _, err = digest.Parse(reference); err != nil {

src/testing/api/repository/controller.go

@@ -81,3 +81,9 @@ func (f *FakeController) Update(ctx context.Context, repository *models.RepoReco
 	args := f.Called()
 	return args.Error(0)
 }
+
+// AddPullCount ...
+func (f *FakeController) AddPullCount(ctx context.Context, id int64) error {
+	args := f.Called()
+	return args.Error(0)
+}

src/testing/pkg/repository/manager.go

@@ -79,3 +79,9 @@ func (f *FakeManager) Update(ctx context.Context, repository *models.RepoRecord,
 	args := f.Called()
 	return args.Error(0)
 }
+
+// AddPullCount ...
+func (f *FakeManager) AddPullCount(ctx context.Context, id int64) error {
+	args := f.Called()
+	return args.Error(0)
+}