Enable usage sync when switch quota setting

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

src/common/dao/project_blob.go

@@ -18,7 +18,9 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/astaxie/beego/orm"
 	"github.com/goharbor/harbor/src/common/models"
+	"strconv"
 )
 
 // AddBlobToProject ...
@@ -103,3 +105,20 @@ func GetBlobsNotInProject(projectID int64, blobDigests ...string) ([]*models.Blo
 
 	return blobs, nil
 }
+
+// CountSizeOfProject ...
+func CountSizeOfProject(pid int64) (int64, error) {
+	var res []orm.Params
+	num, err := GetOrmer().Raw(`SELECT sum(bb.size) FROM project_blob pb LEFT JOIN blob bb ON pb.blob_id = bb.id WHERE pb.project_id = ? `, pid).Values(&res)
+	if err != nil {
+		return -1, err
+	}
+	if num > 0 {
+		size, err := strconv.ParseInt(res[0]["sum"].(string), 0, 64)
+		if err != nil {
+			return -1, err
+		}
+		return size, nil
+	}
+	return -1, err
+}

src/common/dao/project_blob_test.go

@@ -38,3 +38,31 @@ func TestHasBlobInProject(t *testing.T) {
 	require.Nil(t, err)
 	assert.True(t, has)
 }
+
+func TestCountSizeOfProject(t *testing.T) {
+	id1, err := AddBlob(&models.Blob{
+		Digest: "CountSizeOfProject_blob1",
+		Size:   101,
+	})
+	require.Nil(t, err)
+
+	id2, err := AddBlob(&models.Blob{
+		Digest: "CountSizeOfProject_blob2",
+		Size:   202,
+	})
+	require.Nil(t, err)
+
+	pid1, err := AddProject(models.Project{
+		Name:    "CountSizeOfProject_project1",
+		OwnerID: 1,
+	})
+	require.Nil(t, err)
+
+	_, err = AddBlobToProject(id1, pid1)
+	require.Nil(t, err)
+	_, err = AddBlobToProject(id2, pid1)
+	require.Nil(t, err)
+
+	pSize, err := CountSizeOfProject(pid1)
+	assert.Equal(t, pSize, int64(303))
+}

src/common/quota/manager.go

@@ -176,16 +176,63 @@ func (m *Manager) DeleteQuota() error {
 
 // UpdateQuota update the quota resource spec
 func (m *Manager) UpdateQuota(hardLimits types.ResourceList) error {
+	o := dao.GetOrmer()
 	if err := m.driver.Validate(hardLimits); err != nil {
 		return err
 	}
 
 	sql := `UPDATE quota SET hard = ? WHERE reference = ? AND reference_id = ?`
-	_, err := dao.GetOrmer().Raw(sql, hardLimits.String(), m.reference, m.referenceID).Exec()
+	_, err := o.Raw(sql, hardLimits.String(), m.reference, m.referenceID).Exec()
 
 	return err
 }
 
+// EnsureQuota ensures the reference has quota and usage,
+// if non-existent, will create new quota and usage.
+// if existent, update the quota and usage.
+func (m *Manager) EnsureQuota(usages types.ResourceList) error {
+	query := &models.QuotaQuery{
+		Reference:   m.reference,
+		ReferenceID: m.referenceID,
+	}
+	quotas, err := dao.ListQuotas(query)
+	if err != nil {
+		return err
+	}
+
+	// non-existent: create quota and usage
+	defaultHardLimit := m.driver.HardLimits()
+	if len(quotas) == 0 {
+		_, err := m.NewQuota(defaultHardLimit, usages)
+		if err != nil {
+			return err
+		}
+		return nil
+	}
+
+	// existent
+	used := usages
+	quotaUsed, err := types.NewResourceList(quotas[0].Used)
+	if types.Equals(quotaUsed, used) {
+		return nil
+	}
+	dao.WithTransaction(func(o orm.Ormer) error {
+		usage, err := m.getUsageForUpdate(o)
+		if err != nil {
+			return err
+		}
+		usage.Used = used.String()
+		usage.UpdateTime = time.Now()
+		_, err = o.Update(usage)
+		if err != nil {
+			return err
+		}
+		return nil
+	})
+
+	return nil
+}
+
 // AddResources add resources to usage
 func (m *Manager) AddResources(resources types.ResourceList) error {
 	return dao.WithTransaction(func(o orm.Ormer) error {

src/common/quota/manager_test.go

@@ -21,6 +21,7 @@ import (
 	"testing"
 
 	"github.com/goharbor/harbor/src/common/dao"
+	"github.com/goharbor/harbor/src/common/models"
 	"github.com/goharbor/harbor/src/common/quota/driver"
 	"github.com/goharbor/harbor/src/common/quota/driver/mocks"
 	"github.com/goharbor/harbor/src/pkg/types"
@@ -131,6 +132,48 @@ func (suite *ManagerSuite) TestUpdateQuota() {
 	}
 }
 
+func (suite *ManagerSuite) TestEnsureQuota() {
+	// non-existent
+	nonExistRefID := "3"
+	mgr := suite.quotaManager(nonExistRefID)
+	infinite := types.ResourceList{types.ResourceCount: -1, types.ResourceStorage: -1}
+	usage := types.ResourceList{types.ResourceCount: 10, types.ResourceStorage: 10}
+	err := mgr.EnsureQuota(usage)
+	suite.Nil(err)
+	query := &models.QuotaQuery{
+		Reference:   reference,
+		ReferenceID: nonExistRefID,
+	}
+	quotas, err := dao.ListQuotas(query)
+	suite.Nil(err)
+	suite.Equal(usage, mustResourceList(quotas[0].Used))
+	suite.Equal(infinite, mustResourceList(quotas[0].Hard))
+
+	// existent
+	existRefID := "4"
+	mgr = suite.quotaManager(existRefID)
+	used := types.ResourceList{types.ResourceCount: 11, types.ResourceStorage: 11}
+	if id, err := mgr.NewQuota(hardLimits, used); suite.Nil(err) {
+		quota, _ := dao.GetQuota(id)
+		suite.Equal(hardLimits, mustResourceList(quota.Hard))
+
+		usage, _ := dao.GetQuotaUsage(id)
+		suite.Equal(used, mustResourceList(usage.Used))
+	}
+
+	usage2 := types.ResourceList{types.ResourceCount: 12, types.ResourceStorage: 12}
+	err = mgr.EnsureQuota(usage2)
+	suite.Nil(err)
+	query2 := &models.QuotaQuery{
+		Reference:   reference,
+		ReferenceID: existRefID,
+	}
+	quotas2, err := dao.ListQuotas(query2)
+	suite.Equal(usage2, mustResourceList(quotas2[0].Used))
+	suite.Equal(hardLimits, mustResourceList(quotas2[0].Hard))
+
+}
+
 func (suite *ManagerSuite) TestQuotaAutoCreation() {
 	for i := 0; i < 10; i++ {
 		mgr := suite.quotaManager(fmt.Sprintf("%d", i))

src/core/api/harborapi_test.go

@@ -202,6 +202,8 @@ func init() {
 	beego.Router("/api/quotas", quotaAPIType, "get:List")
 	beego.Router("/api/quotas/:id([0-9]+)", quotaAPIType, "get:Get;put:Put")
 
+	beego.Router("/api/internal/switchquota", &InternalAPI{}, "put:SwitchQuota")
+
 	// syncRegistry
 	if err := SyncRegistry(config.GlobalProjectMgr); err != nil {
 		log.Fatalf("failed to sync repositories from registry: %v", err)

src/core/api/internal.go

@@ -15,12 +15,19 @@
 package api
 
 import (
-	"errors"
-
+	"fmt"
+	"github.com/goharbor/harbor/src/chartserver"
 	"github.com/goharbor/harbor/src/common"
 	"github.com/goharbor/harbor/src/common/dao"
 	"github.com/goharbor/harbor/src/common/models"
+	"github.com/goharbor/harbor/src/common/quota"
 	"github.com/goharbor/harbor/src/common/utils/log"
+	"github.com/goharbor/harbor/src/core/config"
+	"github.com/goharbor/harbor/src/jobservice/logger"
+	"github.com/pkg/errors"
+	"net/url"
+	"strconv"
+	"strings"
 )
 
 // InternalAPI handles request of harbor admin...
@@ -69,3 +76,96 @@ func (ia *InternalAPI) RenameAdmin() {
 	log.Debugf("The super user has been renamed to: %s", newName)
 	ia.DestroySession()
 }
+
+type QuotaSwitcher struct {
+	Disabled bool
+}
+
+// SwitchQuota ...
+func (ia *InternalAPI) SwitchQuota() {
+	var req QuotaSwitcher
+	if err := ia.DecodeJSONReq(&req); err != nil {
+		ia.SendBadRequestError(err)
+		return
+	}
+	// From disable to enable, update the quota usage bases on the DB records.
+	if config.QuotaPerProjectEnable() == false && req.Disabled == true {
+		if err := ensureQuota(); err != nil {
+			ia.SendBadRequestError(err)
+			return
+		}
+	}
+	defer config.GetCfgManager().Set(common.QuotaPerProjectEnable, req.Disabled)
+	return
+}
+
+func ensureQuota() error {
+	projects, err := dao.GetProjects(nil)
+	if err != nil {
+		return err
+	}
+	for _, project := range projects {
+		pSize, err := dao.CountSizeOfProject(project.ProjectID)
+		if err != nil {
+			logger.Warningf("error happen on counting size of project:%d , error:%v, just skip it.", project.ProjectID, err)
+			continue
+		}
+		afQuery := &models.ArtifactQuery{
+			PID: project.ProjectID,
+		}
+		afs, err := dao.ListArtifacts(afQuery)
+		if err != nil {
+			logger.Warningf("error happen on counting number of project:%d , error:%v, just skip it.", project.ProjectID, err)
+			continue
+		}
+		pCount := int64(len(afs))
+
+		// it needs to append the chart count
+		if config.GetCfgManager().Get(common.WithChartMuseum).GetBool() {
+			chartEndpoint := strings.TrimSpace(config.GetCfgManager().Get(common.ChartRepoURL).GetString())
+			if len(chartEndpoint) == 0 {
+				return errors.New("empty chartmuseum endpoint")
+			}
+			chartCtr, err := getChartCtr(chartEndpoint)
+			if err != nil {
+				return err
+			}
+			count, err := chartCtr.GetCountOfCharts([]string{project.Name})
+			if err != nil {
+				err = errors.Wrap(err, fmt.Sprintf("get chart count of project %d failed", project.ProjectID))
+				return err
+			}
+			pCount = pCount + int64(count)
+		}
+
+		quotaMgr, err := quota.NewManager("project", strconv.FormatInt(project.ProjectID, 10))
+		if err != nil {
+			logger.Errorf("Error occurred when to new quota manager %v, just skip it.", err)
+			continue
+		}
+		used := quota.ResourceList{
+			quota.ResourceStorage: pSize,
+			quota.ResourceCount:   pCount,
+		}
+		if err := quotaMgr.EnsureQuota(used); err != nil {
+			logger.Errorf("cannot ensure quota for the project: %d, err: %v, just skip it.", project.ProjectID, err)
+			continue
+		}
+	}
+	return nil
+}
+
+func getChartCtr(chartEndpoint string) (*chartserver.Controller, error) {
+	chartEndpoint = strings.TrimSuffix(chartEndpoint, "/")
+	url, err := url.Parse(chartEndpoint)
+	if err != nil {
+		err = errors.New("endpoint URL of chart storage server is malformed")
+		return nil, err
+	}
+	ctr, err := chartserver.NewController(url)
+	if err != nil {
+		err = errors.New("failed to initialize chart API controller")
+		return nil, err
+	}
+	return ctr, nil
+}

src/core/api/internal_test.go

@@ -0,0 +1,56 @@
+// Copyright 2018 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 api
+
+import (
+	"net/http"
+	"testing"
+)
+
+// cannot verify the real scenario here
+func TestSwitchQuota(t *testing.T) {
+	cases := []*codeCheckingCase{
+		// 401
+		{
+			request: &testingRequest{
+				method: http.MethodPut,
+				url:    "/api/internal/switchquota",
+			},
+			code: http.StatusUnauthorized,
+		},
+		// 200
+		{
+			request: &testingRequest{
+				method:     http.MethodPut,
+				url:        "/api/internal/switchquota",
+				credential: sysAdmin,
+				bodyJSON: &QuotaSwitcher{
+					Disabled: true,
+				},
+			},
+			code: http.StatusOK,
+		},
+		// 403
+		{
+			request: &testingRequest{
+				url:        "/api/internal/switchquota",
+				method:     http.MethodPut,
+				credential: nonSysAdmin,
+			},
+			code: http.StatusForbidden,
+		},
+	}
+	runCodeCheckingCases(t, cases...)
+}

src/core/router.go

@@ -134,6 +134,7 @@ func initRouters() {
 
 	beego.Router("/api/internal/syncregistry", &api.InternalAPI{}, "post:SyncRegistry")
 	beego.Router("/api/internal/renameadmin", &api.InternalAPI{}, "post:RenameAdmin")
+	beego.Router("/api/internal/switchquota", &api.InternalAPI{}, "put:SwitchQuota")
 
 	// external service that hosted on harbor process:
 	beego.Router("/service/notifications", &registry.NotificationHandler{})