add storage consumption support (#14772)

Return the total storage consumption in the statistic API

Signed-off-by: Wang Yan <wangyan@vmware.com>

api/v2.0/swagger.yaml

@@ -8152,3 +8152,8 @@ definitions:
         format: int64
         description: The count of the total repositories, only be seen by the system admin
         x-omitempty: false
+      total_storage_consumption:
+        type: integer
+        format: int64
+        description: The total storage consumption of blobs, only be seen by the system admin
+        x-omitempty: false

src/controller/blob/controller.go

@@ -49,6 +49,9 @@ type Controller interface {
 	// CalculateTotalSizeByProject returns the sum of the blob size for the project
 	CalculateTotalSizeByProject(ctx context.Context, projectID int64, excludeForeign bool) (int64, error)
 
+	// CalculateTotalSize returns the sum of all the blobs size
+	CalculateTotalSize(ctx context.Context, excludeForeign bool) (int64, error)
+
 	// Ensure create blob when it not exist.
 	Ensure(ctx context.Context, digest string, contentType string, size int64) (int64, error)
 
@@ -144,6 +147,10 @@ func (c *controller) CalculateTotalSizeByProject(ctx context.Context, projectID
 	return c.blobMgr.CalculateTotalSizeByProject(ctx, projectID, excludeForeign)
 }
 
+func (c *controller) CalculateTotalSize(ctx context.Context, excludeForeign bool) (int64, error) {
+	return c.blobMgr.CalculateTotalSize(ctx, excludeForeign)
+}
+
 func (c *controller) Ensure(ctx context.Context, digest string, contentType string, size int64) (blobID int64, err error) {
 	blob, err := c.blobMgr.Get(ctx, digest)
 	if err == nil {

src/controller/blob/controller_test.go

@@ -112,6 +112,16 @@ func (suite *ControllerTestSuite) TestCalculateTotalSizeByProject() {
 	})
 }
 
+func (suite *ControllerTestSuite) TestCalculateTotalSize() {
+	ctx := suite.Context()
+	size1, err := Ctl.CalculateTotalSize(ctx, true)
+	Ctl.Ensure(ctx, suite.DigestString(), schema2.MediaTypeForeignLayer, 100)
+	Ctl.Ensure(ctx, suite.DigestString(), schema2.MediaTypeLayer, 100)
+	size2, err := Ctl.CalculateTotalSize(ctx, false)
+	suite.Nil(err)
+	suite.Equal(int64(200), size2-size1)
+}
+
 func (suite *ControllerTestSuite) TestEnsure() {
 	ctx := suite.Context()
 

src/pkg/blob/dao/dao.go

@@ -62,6 +62,9 @@ type DAO interface {
 	// SumBlobsSizeByProject returns sum size of blobs by project, skip foreign blobs when `excludeForeignLayer` is true
 	SumBlobsSizeByProject(ctx context.Context, projectID int64, excludeForeignLayer bool) (int64, error)
 
+	// SumBlobsSize returns sum size of all blobs skip foreign blobs when `excludeForeignLayer` is true
+	SumBlobsSize(ctx context.Context, excludeForeignLayer bool) (int64, error)
+
 	// CreateProjectBlob create ProjectBlob and ignore conflict on project id and blob id
 	CreateProjectBlob(ctx context.Context, projectID, blobID int64) (int64, error)
 
@@ -291,6 +294,31 @@ func (d *dao) SumBlobsSizeByProject(ctx context.Context, projectID int64, exclud
 	return totalSize, nil
 }
 
+// SumBlobsSize returns sum size of all blobs skip foreign blobs when `excludeForeignLayer` is true
+func (d *dao) SumBlobsSize(ctx context.Context, excludeForeignLayer bool) (int64, error) {
+	o, err := orm.FromContext(ctx)
+	if err != nil {
+		return 0, err
+	}
+
+	params := []interface{}{}
+	sql := `SELECT SUM(size) FROM blob`
+	if excludeForeignLayer {
+		foreignLayerTypes := []interface{}{
+			schema2.MediaTypeForeignLayer,
+		}
+		sql = fmt.Sprintf(`%s Where content_type NOT IN (%s)`, sql, orm.ParamPlaceholderForIn(len(foreignLayerTypes)))
+		params = append(params, foreignLayerTypes...)
+	}
+
+	var totalSize int64
+	if err := o.Raw(sql, params...).QueryRow(&totalSize); err != nil {
+		return 0, err
+	}
+
+	return totalSize, nil
+}
+
 func (d *dao) CreateProjectBlob(ctx context.Context, projectID, blobID int64) (int64, error) {
 	o, err := orm.FromContext(ctx)
 	if err != nil {

src/pkg/blob/dao/dao_test.go

@@ -245,6 +245,21 @@ func (suite *DaoTestSuite) TestListBlobsAssociatedWithArtifact() {
 
 }
 
+func (suite *DaoTestSuite) TestSumBlobsSize() {
+	ctx := suite.Context()
+
+	size1, err := suite.dao.SumBlobsSize(ctx, true)
+	suite.Nil(err)
+
+	digest1 := suite.DigestString()
+	suite.dao.CreateBlob(ctx, &models.Blob{Digest: digest1, Size: 999})
+
+	size2, err := suite.dao.SumBlobsSize(ctx, true)
+	suite.Nil(err)
+
+	suite.Equal(int64(999), size2-size1)
+}
+
 func (suite *DaoTestSuite) TestFindBlobsShouldUnassociatedWithProject() {
 	ctx := suite.Context()
 

src/pkg/blob/manager.go

@@ -41,6 +41,9 @@ type Manager interface {
 	// CalculateTotalSizeByProject returns total blob size by project, skip foreign blobs when `excludeForeignLayer` is true
 	CalculateTotalSizeByProject(ctx context.Context, projectID int64, excludeForeignLayer bool) (int64, error)
 
+	// SumBlobsSize returns sum size of all blobs skip foreign blobs when `excludeForeignLayer` is true
+	CalculateTotalSize(ctx context.Context, excludeForeignLayer bool) (int64, error)
+
 	// Create create blob
 	Create(ctx context.Context, digest string, contentType string, size int64) (int64, error)
 
@@ -139,6 +142,10 @@ func (m *manager) UselessBlobs(ctx context.Context, timeWindowHours int64) ([]*m
 	return m.dao.GetBlobsNotRefedByProjectBlob(ctx, timeWindowHours)
 }
 
+func (m *manager) CalculateTotalSize(ctx context.Context, excludeForeignLayer bool) (int64, error) {
+	return m.dao.SumBlobsSize(ctx, excludeForeignLayer)
+}
+
 // NewManager returns blob manager
 func NewManager() Manager {
 	return &manager{dao: dao.New()}

src/pkg/blob/manager_test.go

@@ -22,6 +22,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/docker/distribution/manifest/schema2"
 	"github.com/goharbor/harbor/src/pkg/blob/models"
 )
 
@@ -80,6 +81,21 @@ func (suite *ManagerTestSuite) TestAssociateWithProject() {
 	suite.True(associated)
 }
 
+func (suite *ManagerTestSuite) TestCalculateTotalSize() {
+	ctx := suite.Context()
+
+	size1, err := Mgr.CalculateTotalSize(ctx, true)
+	suite.Nil(err)
+
+	digest := suite.DigestString()
+	Mgr.Create(ctx, digest, schema2.MediaTypeLayer, 100)
+
+	size2, err := Mgr.CalculateTotalSize(ctx, true)
+	suite.Nil(err)
+
+	suite.Equal(int64(100), size2-size1)
+}
+
 func (suite *ManagerTestSuite) TestCleanupAssociationsForArtifact() {
 	ctx := suite.Context()
 

src/server/v2.0/handler/statistic.go

@@ -16,6 +16,7 @@ package handler
 
 import (
 	"context"
+	"github.com/goharbor/harbor/src/controller/blob"
 
 	"github.com/go-openapi/runtime/middleware"
 	"github.com/goharbor/harbor/src/common/security/local"
@@ -30,6 +31,7 @@ func newStatisticAPI() *statisticAPI {
 	return &statisticAPI{
 		proCtl:  project.Ctl,
 		repoCtl: repository.Ctl,
+		blobCtl: blob.Ctl,
 	}
 }
 
@@ -37,6 +39,7 @@ type statisticAPI struct {
 	BaseAPI
 	proCtl  project.Controller
 	repoCtl repository.Controller
+	blobCtl blob.Controller
 }
 
 func (s *statisticAPI) GetStatistic(ctx context.Context, params operation.GetStatisticParams) middleware.Responder {
@@ -88,6 +91,13 @@ func (s *statisticAPI) GetStatistic(ctx context.Context, params operation.GetSta
 		}
 		statistic.TotalRepoCount = n
 		statistic.PrivateRepoCount = n - statistic.PublicRepoCount
+
+		sum, err := s.blobCtl.CalculateTotalSize(ctx, true)
+		if err != nil {
+			return s.SendError(ctx, err)
+		}
+		statistic.TotalStorageConsumption = sum
+
 	} else {
 		var privProjectIDs []interface{}
 		if sc, ok := securityCtx.(*local.SecurityContext); ok && sc.IsAuthenticated() {

src/testing/controller/blob/controller.go

@@ -63,6 +63,27 @@ func (_m *Controller) AssociateWithProjectByID(ctx context.Context, blobID int64
 	return r0
 }
 
+// CalculateTotalSize provides a mock function with given fields: ctx, excludeForeign
+func (_m *Controller) CalculateTotalSize(ctx context.Context, excludeForeign bool) (int64, error) {
+	ret := _m.Called(ctx, excludeForeign)
+
+	var r0 int64
+	if rf, ok := ret.Get(0).(func(context.Context, bool) int64); ok {
+		r0 = rf(ctx, excludeForeign)
+	} else {
+		r0 = ret.Get(0).(int64)
+	}
+
+	var r1 error
+	if rf, ok := ret.Get(1).(func(context.Context, bool) error); ok {
+		r1 = rf(ctx, excludeForeign)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
 // CalculateTotalSizeByProject provides a mock function with given fields: ctx, projectID, excludeForeign
 func (_m *Controller) CalculateTotalSizeByProject(ctx context.Context, projectID int64, excludeForeign bool) (int64, error) {
 	ret := _m.Called(ctx, projectID, excludeForeign)

src/testing/pkg/blob/manager.go

@@ -58,6 +58,27 @@ func (_m *Manager) AssociateWithProject(ctx context.Context, blobID int64, proje
 	return r0, r1
 }
 
+// CalculateTotalSize provides a mock function with given fields: ctx, excludeForeignLayer
+func (_m *Manager) CalculateTotalSize(ctx context.Context, excludeForeignLayer bool) (int64, error) {
+	ret := _m.Called(ctx, excludeForeignLayer)
+
+	var r0 int64
+	if rf, ok := ret.Get(0).(func(context.Context, bool) int64); ok {
+		r0 = rf(ctx, excludeForeignLayer)
+	} else {
+		r0 = ret.Get(0).(int64)
+	}
+
+	var r1 error
+	if rf, ok := ret.Get(1).(func(context.Context, bool) error); ok {
+		r1 = rf(ctx, excludeForeignLayer)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
 // CalculateTotalSizeByProject provides a mock function with given fields: ctx, projectID, excludeForeignLayer
 func (_m *Manager) CalculateTotalSizeByProject(ctx context.Context, projectID int64, excludeForeignLayer bool) (int64, error) {
 	ret := _m.Called(ctx, projectID, excludeForeignLayer)