From c0349da8125c3b6d9e9c263bd4cd540d2a57deca Mon Sep 17 00:00:00 2001 From: He Weiwei Date: Sun, 12 Apr 2020 16:14:12 +0000 Subject: [PATCH] refactor(quota): cleanup code for quota 1. Remove `common/quota` package. 2. Remove functions about quota in `common/dao` package. 3. Move `Quota` and `QuotaUsage` models from `common/models` to `pkg/quota/dao`. 4. Add `Count` and `List` methods to `quota.Controller`. 5. Use `quota.Controller` to implement quota APIs. Signed-off-by: He Weiwei --- src/common/dao/base.go | 26 -- src/common/dao/dao_test.go | 50 --- src/common/dao/quota.go | 238 ------------ src/common/dao/quota_test.go | 174 --------- src/common/dao/quota_usage.go | 149 -------- src/common/dao/quota_usage_test.go | 164 -------- src/common/models/base.go | 2 - src/common/models/quota.go | 85 ----- src/common/models/quota_usage.go | 77 ---- src/common/quota/driver/driver.go | 59 --- src/common/quota/driver/mocks/driver.go | 65 ---- src/common/quota/driver/project/driver.go | 158 -------- .../quota/driver/project/driver_test.go | 75 ---- src/common/quota/errors.go | 111 ------ src/common/quota/manager.go | 287 -------------- src/common/quota/manager_test.go | 359 ------------------ src/common/quota/quota.go | 35 -- src/common/quota/quota_test.go | 45 --- src/common/quota/types.go | 30 -- src/common/quota/util.go | 57 --- src/common/quota/util_test.go | 83 ---- src/controller/quota/controller.go | 54 +++ src/core/api/harborapi_test.go | 10 +- src/core/api/project.go | 57 ++- src/core/api/quota.go | 73 ++-- src/core/api/quota_test.go | 55 +-- src/lib/q/query.go | 2 + src/lib/response_buffer.go | 8 +- src/lib/response_recorder.go | 9 +- src/pkg/quota/dao/dao.go | 88 +++++ src/pkg/quota/dao/dao_test.go | 92 +++++ src/pkg/quota/dao/model.go | 41 +- src/pkg/quota/dao/util.go | 117 ++++++ src/pkg/quota/dao/util_test.go | 52 +++ src/pkg/quota/manager.go | 27 +- src/testing/common/security/context.go | 40 -- src/testing/controller/controller.go | 1 - src/testing/controller/quota/controller.go | 60 +++ src/testing/pkg/quota/manager.go | 46 +++ src/testing/suite.go | 14 - 40 files changed, 677 insertions(+), 2498 deletions(-) delete mode 100644 src/common/dao/quota.go delete mode 100644 src/common/dao/quota_test.go delete mode 100644 src/common/dao/quota_usage.go delete mode 100644 src/common/dao/quota_usage_test.go delete mode 100644 src/common/models/quota.go delete mode 100644 src/common/models/quota_usage.go delete mode 100644 src/common/quota/driver/driver.go delete mode 100644 src/common/quota/driver/mocks/driver.go delete mode 100644 src/common/quota/driver/project/driver.go delete mode 100644 src/common/quota/driver/project/driver_test.go delete mode 100644 src/common/quota/errors.go delete mode 100644 src/common/quota/manager.go delete mode 100644 src/common/quota/manager_test.go delete mode 100644 src/common/quota/quota.go delete mode 100644 src/common/quota/quota_test.go delete mode 100644 src/common/quota/types.go delete mode 100644 src/common/quota/util.go delete mode 100644 src/common/quota/util_test.go create mode 100644 src/pkg/quota/dao/util.go create mode 100644 src/pkg/quota/dao/util_test.go diff --git a/src/common/dao/base.go b/src/common/dao/base.go index 92fe40265..d9f2632cb 100644 --- a/src/common/dao/base.go +++ b/src/common/dao/base.go @@ -169,29 +169,3 @@ func Escape(str string) string { str = strings.Replace(str, `_`, `\_`, -1) return str } - -// WithTransaction helper for transaction -func WithTransaction(handler func(o orm.Ormer) error) error { - o := orm.NewOrm() - - if err := o.Begin(); err != nil { - log.Errorf("begin transaction failed: %v", err) - return err - } - - if err := handler(o); err != nil { - if e := o.Rollback(); e != nil { - log.Errorf("rollback transaction failed: %v", e) - return e - } - - return err - } - - if err := o.Commit(); err != nil { - log.Errorf("commit transaction failed: %v", err) - return err - } - - return nil -} diff --git a/src/common/dao/dao_test.go b/src/common/dao/dao_test.go index 03d4dde95..c08df01f8 100644 --- a/src/common/dao/dao_test.go +++ b/src/common/dao/dao_test.go @@ -670,53 +670,3 @@ func TestIsDupRecError(t *testing.T) { assert.True(t, IsDupRecErr(fmt.Errorf("pq: duplicate key value violates unique constraint \"properties_k_key\""))) assert.False(t, IsDupRecErr(fmt.Errorf("other error"))) } - -func TestWithTransaction(t *testing.T) { - reference := "transaction" - - quota := models.Quota{ - Reference: reference, - ReferenceID: "1", - Hard: "{}", - } - - failed := func(o orm.Ormer) error { - o.Insert("a) - - return fmt.Errorf("failed") - } - - var quotaID int64 - success := func(o orm.Ormer) error { - id, err := o.Insert("a) - if err != nil { - return err - } - - quotaID = id - return nil - } - - assert := assert.New(t) - - if assert.Error(WithTransaction(failed)) { - var quota models.Quota - quota.Reference = reference - quota.ReferenceID = "1" - err := GetOrmer().Read("a, "reference", "reference_id") - assert.Error(err) - assert.False(quota.ID != 0) - } - - if assert.Nil(WithTransaction(success)) { - var quota models.Quota - quota.Reference = reference - quota.ReferenceID = "1" - err := GetOrmer().Read("a, "reference", "reference_id") - assert.Nil(err) - assert.True(quota.ID != 0) - assert.Equal(quotaID, quota.ID) - - GetOrmer().Delete(&models.Quota{ID: quotaID}, "id") - } -} diff --git a/src/common/dao/quota.go b/src/common/dao/quota.go deleted file mode 100644 index 31faa8723..000000000 --- a/src/common/dao/quota.go +++ /dev/null @@ -1,238 +0,0 @@ -// 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 dao - -import ( - "encoding/json" - "fmt" - "strings" - "time" - - "github.com/astaxie/beego/orm" - "github.com/goharbor/harbor/src/common/models" - "github.com/goharbor/harbor/src/common/quota/driver" - "github.com/goharbor/harbor/src/lib/log" - "github.com/goharbor/harbor/src/pkg/types" -) - -var ( - quotaOrderMap = map[string]string{ - "creation_time": "b.creation_time asc", - "+creation_time": "b.creation_time asc", - "-creation_time": "b.creation_time desc", - "update_time": "b.update_time asc", - "+update_time": "b.update_time asc", - "-update_time": "b.update_time desc", - } -) - -// AddQuota add quota to the database. -func AddQuota(quota models.Quota) (int64, error) { - now := time.Now() - quota.CreationTime = now - quota.UpdateTime = now - return GetOrmer().Insert("a) -} - -// GetQuota returns quota by id. -func GetQuota(id int64) (*models.Quota, error) { - q := models.Quota{ID: id} - err := GetOrmer().Read(&q, "ID") - if err == orm.ErrNoRows { - return nil, nil - } - return &q, err -} - -// UpdateQuota update the quota. -func UpdateQuota(quota models.Quota) error { - quota.UpdateTime = time.Now() - _, err := GetOrmer().Update("a) - return err -} - -// Quota quota mode for api -type Quota struct { - ID int64 `orm:"pk;auto;column(id)" json:"id"` - Ref driver.RefObject `json:"ref"` - Reference string `orm:"column(reference)" json:"-"` - ReferenceID string `orm:"column(reference_id)" json:"-"` - Hard string `orm:"column(hard);type(jsonb)" json:"-"` - Used string `orm:"column(used);type(jsonb)" json:"-"` - CreationTime time.Time `orm:"column(creation_time);auto_now_add" json:"creation_time"` - UpdateTime time.Time `orm:"column(update_time);auto_now" json:"update_time"` -} - -// MarshalJSON ... -func (q *Quota) MarshalJSON() ([]byte, error) { - hard, err := types.NewResourceList(q.Hard) - if err != nil { - return nil, err - } - - used, err := types.NewResourceList(q.Used) - if err != nil { - return nil, err - } - - type Alias Quota - return json.Marshal(&struct { - *Alias - Hard types.ResourceList `json:"hard"` - Used types.ResourceList `json:"used"` - }{ - Alias: (*Alias)(q), - Hard: hard, - Used: used, - }) -} - -// ListQuotas returns quotas by query. -func ListQuotas(query ...*models.QuotaQuery) ([]*Quota, error) { - condition, params := quotaQueryConditions(query...) - - sql := fmt.Sprintf(` -SELECT - a.id, - a.reference, - a.reference_id, - a.hard, - b.used, - b.creation_time, - b.update_time -FROM - quota AS a - JOIN quota_usage AS b ON a.id = b.id %s`, condition) - - orderBy := quotaOrderBy(query...) - if orderBy != "" { - sql += ` order by ` + orderBy - } - - if len(query) > 0 && query[0] != nil { - page, size := query[0].Page, query[0].Size - if size > 0 { - sql += ` limit ?` - params = append(params, size) - if page > 0 { - sql += ` offset ?` - params = append(params, size*(page-1)) - } - } - } - - var quotas []*Quota - if _, err := GetOrmer().Raw(sql, params).QueryRows("as); err != nil { - return nil, err - } - - for _, quota := range quotas { - d, ok := driver.Get(quota.Reference) - if !ok { - continue - } - - ref, err := d.Load(quota.ReferenceID) - if err != nil { - log.Warning(fmt.Sprintf("Load quota reference object (%s, %s) failed: %v", quota.Reference, quota.ReferenceID, err)) - continue - } - - quota.Ref = ref - } - - return quotas, nil -} - -// GetTotalOfQuotas returns total of quotas -func GetTotalOfQuotas(query ...*models.QuotaQuery) (int64, error) { - condition, params := quotaQueryConditions(query...) - sql := fmt.Sprintf("SELECT COUNT(1) FROM quota AS a JOIN quota_usage AS b ON a.id = b.id %s", condition) - - var count int64 - if err := GetOrmer().Raw(sql, params).QueryRow(&count); err != nil { - return 0, err - } - - return count, nil -} - -func quotaQueryConditions(query ...*models.QuotaQuery) (string, []interface{}) { - params := []interface{}{} - sql := "" - if len(query) == 0 || query[0] == nil { - return sql, params - } - - sql += `WHERE 1=1 ` - - q := query[0] - if q.ID != 0 { - sql += `AND a.id = ? ` - params = append(params, q.ID) - } - if q.Reference != "" { - sql += `AND a.reference = ? ` - params = append(params, q.Reference) - } - if q.ReferenceID != "" { - sql += `AND a.reference_id = ? ` - params = append(params, q.ReferenceID) - } - - if len(q.ReferenceIDs) != 0 { - sql += fmt.Sprintf(`AND a.reference_id IN (%s) `, ParamPlaceholderForIn(len(q.ReferenceIDs))) - params = append(params, q.ReferenceIDs) - } - - return sql, params -} - -func castQuantity(field string) string { - // cast -1 to max int64 when order by field - return fmt.Sprintf("CAST( (CASE WHEN (%[1]s) IS NULL THEN '0' WHEN (%[1]s) = '-1' THEN '9223372036854775807' ELSE (%[1]s) END) AS BIGINT )", field) -} - -func quotaOrderBy(query ...*models.QuotaQuery) string { - orderBy := "b.creation_time DESC" - - if len(query) > 0 && query[0] != nil && query[0].Sort != "" { - if val, ok := quotaOrderMap[query[0].Sort]; ok { - orderBy = val - } else { - sort := query[0].Sort - - order := "ASC" - if sort[0] == '-' { - order = "DESC" - sort = sort[1:] - } - - prefixes := []string{"hard.", "used."} - for _, prefix := range prefixes { - if strings.HasPrefix(sort, prefix) { - resource := strings.TrimPrefix(sort, prefix) - if types.IsValidResource(types.ResourceName(resource)) { - field := fmt.Sprintf("%s->>'%s'", strings.TrimSuffix(prefix, "."), resource) - orderBy = fmt.Sprintf("(%s) %s", castQuantity(field), order) - break - } - } - } - } - } - - return orderBy -} diff --git a/src/common/dao/quota_test.go b/src/common/dao/quota_test.go deleted file mode 100644 index 3d99f142a..000000000 --- a/src/common/dao/quota_test.go +++ /dev/null @@ -1,174 +0,0 @@ -// 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 dao - -import ( - "testing" - "time" - - "github.com/goharbor/harbor/src/common/models" - "github.com/stretchr/testify/suite" -) - -var ( - quotaReference = "dao" - quotaUserReference = "user" - quotaHard = models.QuotaHard{"storage": 1024} - quotaHardLarger = models.QuotaHard{"storage": 2048} -) - -type QuotaDaoSuite struct { - suite.Suite -} - -func (suite *QuotaDaoSuite) equalHard(quota1 *models.Quota, quota2 *models.Quota) { - hard1, err := quota1.GetHard() - suite.Nil(err, "hard1 invalid") - - hard2, err := quota2.GetHard() - suite.Nil(err, "hard2 invalid") - - suite.Equal(hard1, hard2) -} - -func (suite *QuotaDaoSuite) TearDownTest() { - ClearTable("quota") - ClearTable("quota_usage") -} - -func (suite *QuotaDaoSuite) TestAddQuota() { - _, err1 := AddQuota(models.Quota{Reference: quotaReference, ReferenceID: "1", Hard: quotaHard.String()}) - suite.Nil(err1) - - // Will failed for reference and reference_id should unique in db - _, err2 := AddQuota(models.Quota{Reference: quotaReference, ReferenceID: "1", Hard: quotaHard.String()}) - suite.Error(err2) - - _, err3 := AddQuota(models.Quota{Reference: quotaUserReference, ReferenceID: "1", Hard: quotaHard.String()}) - suite.Nil(err3) -} - -func (suite *QuotaDaoSuite) TestGetQuota() { - quota1 := models.Quota{Reference: quotaReference, ReferenceID: "1", Hard: quotaHard.String()} - id, err := AddQuota(quota1) - suite.Nil(err) - - // Get the new added quota - quota2, err := GetQuota(id) - suite.Nil(err) - suite.NotNil(quota2) - - // Get the quota which id is 10000 not found - quota3, err := GetQuota(10000) - suite.Nil(err) - suite.Nil(quota3) -} - -func (suite *QuotaDaoSuite) TestUpdateQuota() { - quota1 := models.Quota{Reference: quotaReference, ReferenceID: "1", Hard: quotaHard.String()} - id, err := AddQuota(quota1) - suite.Nil(err) - - // Get the new added quota - quota2, err := GetQuota(id) - suite.Nil(err) - suite.equalHard("a1, quota2) - - // Update the quota - quota2.SetHard(quotaHardLarger) - time.Sleep(time.Millisecond * 10) // Ensure that UpdateTime changed - suite.Nil(UpdateQuota(*quota2)) - - // Get the updated quota - quota3, err := GetQuota(id) - suite.Nil(err) - suite.equalHard(quota2, quota3) - suite.NotEqual(quota2.UpdateTime, quota3.UpdateTime) -} - -func (suite *QuotaDaoSuite) TestListQuotas() { - id1, _ := AddQuota(models.Quota{Reference: quotaReference, ReferenceID: "1", Hard: quotaHard.String()}) - AddQuotaUsage(models.QuotaUsage{ID: id1, Reference: quotaReference, ReferenceID: "1", Used: "{}"}) - - id2, _ := AddQuota(models.Quota{Reference: quotaReference, ReferenceID: "2", Hard: quotaHard.String()}) - AddQuotaUsage(models.QuotaUsage{ID: id2, Reference: quotaReference, ReferenceID: "2", Used: "{}"}) - - id3, _ := AddQuota(models.Quota{Reference: quotaUserReference, ReferenceID: "1", Hard: quotaHardLarger.String()}) - AddQuotaUsage(models.QuotaUsage{ID: id3, Reference: quotaUserReference, ReferenceID: "1", Used: "{}"}) - - id4, _ := AddQuota(models.Quota{Reference: quotaReference, ReferenceID: "3", Hard: quotaHard.String()}) - AddQuotaUsage(models.QuotaUsage{ID: id4, Reference: quotaReference, ReferenceID: "3", Used: "{}"}) - - // List all the quotas - quotas, err := ListQuotas() - suite.Nil(err) - suite.Equal(4, len(quotas)) - suite.Equal(quotaReference, quotas[0].Reference) - - // List quotas filter by reference - quotas, err = ListQuotas(&models.QuotaQuery{Reference: quotaReference}) - suite.Nil(err) - suite.Equal(3, len(quotas)) - - // List quotas filter by reference ids - quotas, err = ListQuotas(&models.QuotaQuery{Reference: quotaReference, ReferenceIDs: []string{"1", "2"}}) - suite.Nil(err) - suite.Equal(2, len(quotas)) - - // List quotas by pagination - quotas, err = ListQuotas(&models.QuotaQuery{Pagination: models.Pagination{Size: 2}}) - suite.Nil(err) - suite.Equal(2, len(quotas)) - - // List quotas by sorting - quotas, err = ListQuotas(&models.QuotaQuery{Sorting: models.Sorting{Sort: "-hard.storage"}}) - suite.Nil(err) - suite.Equal(quotaUserReference, quotas[0].Reference) -} - -func TestRunQuotaDaoSuite(t *testing.T) { - suite.Run(t, new(QuotaDaoSuite)) -} - -func Test_quotaOrderBy(t *testing.T) { - query := func(sort string) []*models.QuotaQuery { - return []*models.QuotaQuery{ - {Sorting: models.Sorting{Sort: sort}}, - } - } - - type args struct { - query []*models.QuotaQuery - } - tests := []struct { - name string - args args - want string - }{ - {"no query", args{nil}, "b.creation_time DESC"}, - {"order by unsupport field", args{query("unknow")}, "b.creation_time DESC"}, - {"order by storage of hard", args{query("hard.storage")}, "(CAST( (CASE WHEN (hard->>'storage') IS NULL THEN '0' WHEN (hard->>'storage') = '-1' THEN '9223372036854775807' ELSE (hard->>'storage') END) AS BIGINT )) ASC"}, - {"order by unsupport hard resource", args{query("hard.unknow")}, "b.creation_time DESC"}, - {"order by storage of used", args{query("used.storage")}, "(CAST( (CASE WHEN (used->>'storage') IS NULL THEN '0' WHEN (used->>'storage') = '-1' THEN '9223372036854775807' ELSE (used->>'storage') END) AS BIGINT )) ASC"}, - {"order by unsupport used resource", args{query("used.unknow")}, "b.creation_time DESC"}, - } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - if got := quotaOrderBy(tt.args.query...); got != tt.want { - t.Errorf("quotaOrderBy() = %v, want %v", got, tt.want) - } - }) - } -} diff --git a/src/common/dao/quota_usage.go b/src/common/dao/quota_usage.go deleted file mode 100644 index 2879ef609..000000000 --- a/src/common/dao/quota_usage.go +++ /dev/null @@ -1,149 +0,0 @@ -// 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 dao - -import ( - "fmt" - "strings" - "time" - - "github.com/astaxie/beego/orm" - "github.com/goharbor/harbor/src/common/models" - "github.com/goharbor/harbor/src/pkg/types" -) - -var ( - quotaUsageOrderMap = map[string]string{ - "id": "id asc", - "+id": "id asc", - "-id": "id desc", - "creation_time": "creation_time asc", - "+creation_time": "creation_time asc", - "-creation_time": "creation_time desc", - "update_time": "update_time asc", - "+update_time": "update_time asc", - "-update_time": "update_time desc", - } -) - -// AddQuotaUsage add quota usage to the database. -func AddQuotaUsage(quotaUsage models.QuotaUsage) (int64, error) { - now := time.Now() - quotaUsage.CreationTime = now - quotaUsage.UpdateTime = now - return GetOrmer().Insert("aUsage) -} - -// GetQuotaUsage returns quota usage by id. -func GetQuotaUsage(id int64) (*models.QuotaUsage, error) { - q := models.QuotaUsage{ID: id} - err := GetOrmer().Read(&q, "ID") - if err == orm.ErrNoRows { - return nil, nil - } - return &q, err -} - -// UpdateQuotaUsage update the quota usage. -func UpdateQuotaUsage(quotaUsage models.QuotaUsage) error { - quotaUsage.UpdateTime = time.Now() - _, err := GetOrmer().Update("aUsage) - return err -} - -// ListQuotaUsages returns quota usages by query. -func ListQuotaUsages(query ...*models.QuotaUsageQuery) ([]*models.QuotaUsage, error) { - condition, params := quotaUsageQueryConditions(query...) - sql := fmt.Sprintf(`select * %s`, condition) - - orderBy := quotaUsageOrderBy(query...) - if orderBy != "" { - sql += ` order by ` + orderBy - } - - if len(query) > 0 && query[0] != nil { - page, size := query[0].Page, query[0].Size - if size > 0 { - sql += ` limit ?` - params = append(params, size) - if page > 0 { - sql += ` offset ?` - params = append(params, size*(page-1)) - } - } - } - - var quotaUsages []*models.QuotaUsage - if _, err := GetOrmer().Raw(sql, params).QueryRows("aUsages); err != nil { - return nil, err - } - - return quotaUsages, nil -} - -func quotaUsageQueryConditions(query ...*models.QuotaUsageQuery) (string, []interface{}) { - params := []interface{}{} - sql := `from quota_usage ` - if len(query) == 0 || query[0] == nil { - return sql, params - } - - sql += `where 1=1 ` - - q := query[0] - if q.Reference != "" { - sql += `and reference = ? ` - params = append(params, q.Reference) - } - if q.ReferenceID != "" { - sql += `and reference_id = ? ` - params = append(params, q.ReferenceID) - } - if len(q.ReferenceIDs) != 0 { - sql += fmt.Sprintf(`and reference_id in (%s) `, ParamPlaceholderForIn(len(q.ReferenceIDs))) - params = append(params, q.ReferenceIDs) - } - - return sql, params -} - -func quotaUsageOrderBy(query ...*models.QuotaUsageQuery) string { - orderBy := "" - - if len(query) > 0 && query[0] != nil && query[0].Sort != "" { - if val, ok := quotaUsageOrderMap[query[0].Sort]; ok { - orderBy = val - } else { - sort := query[0].Sort - - order := "ASC" - if sort[0] == '-' { - order = "DESC" - sort = sort[1:] - } - - prefix := "used." - if strings.HasPrefix(sort, prefix) { - resource := strings.TrimPrefix(sort, prefix) - if types.IsValidResource(types.ResourceName(resource)) { - field := fmt.Sprintf("%s->>'%s'", strings.TrimSuffix(prefix, "."), resource) - orderBy = fmt.Sprintf("(%s) %s", castQuantity(field), order) - } - } - } - } - - return orderBy -} diff --git a/src/common/dao/quota_usage_test.go b/src/common/dao/quota_usage_test.go deleted file mode 100644 index 6d4f8349c..000000000 --- a/src/common/dao/quota_usage_test.go +++ /dev/null @@ -1,164 +0,0 @@ -// 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 dao - -import ( - "testing" - "time" - - "github.com/goharbor/harbor/src/common/models" - "github.com/stretchr/testify/suite" -) - -var ( - quotaUsageReference = "project" - quotaUsageUserReference = "user" - quotaUsageUsed = models.QuotaUsed{"storage": 1024} - quotaUsageUsedLarger = models.QuotaUsed{"storage": 2048} -) - -type QuotaUsageDaoSuite struct { - suite.Suite -} - -func (suite *QuotaUsageDaoSuite) equalUsed(usage1 *models.QuotaUsage, usage2 *models.QuotaUsage) { - used1, err := usage1.GetUsed() - suite.Nil(err, "used1 invalid") - - used2, err := usage2.GetUsed() - suite.Nil(err, "used2 invalid") - - suite.Equal(used1, used2) -} - -func (suite *QuotaUsageDaoSuite) TearDownTest() { - ClearTable("quota_usage") -} - -func (suite *QuotaUsageDaoSuite) TestAddQuotaUsage() { - _, err1 := AddQuotaUsage(models.QuotaUsage{Reference: quotaUsageReference, ReferenceID: "1", Used: quotaUsageUsed.String()}) - suite.Nil(err1) - - // Will failed for reference and reference_id should unique in db - _, err2 := AddQuotaUsage(models.QuotaUsage{Reference: quotaUsageReference, ReferenceID: "1", Used: quotaUsageUsed.String()}) - suite.Error(err2) - - _, err3 := AddQuotaUsage(models.QuotaUsage{Reference: quotaUsageUserReference, ReferenceID: "1", Used: quotaUsageUsed.String()}) - suite.Nil(err3) -} - -func (suite *QuotaUsageDaoSuite) TestGetQuotaUsage() { - quotaUsage1 := models.QuotaUsage{Reference: quotaUsageReference, ReferenceID: "1", Used: quotaUsageUsed.String()} - id, err := AddQuotaUsage(quotaUsage1) - suite.Nil(err) - - // Get the new added quotaUsage - quotaUsage2, err := GetQuotaUsage(id) - suite.Nil(err) - suite.NotNil(quotaUsage2) - - // Get the quotaUsage which id is 10000 not found - quotaUsage3, err := GetQuotaUsage(10000) - suite.Nil(err) - suite.Nil(quotaUsage3) -} - -func (suite *QuotaUsageDaoSuite) TestUpdateQuotaUsage() { - quotaUsage1 := models.QuotaUsage{Reference: quotaUsageReference, ReferenceID: "1", Used: quotaUsageUsed.String()} - id, err := AddQuotaUsage(quotaUsage1) - suite.Nil(err) - - // Get the new added quotaUsage - quotaUsage2, err := GetQuotaUsage(id) - suite.Nil(err) - suite.equalUsed("aUsage1, quotaUsage2) - - // Update the quotaUsage - quotaUsage2.SetUsed(quotaUsageUsedLarger) - time.Sleep(time.Millisecond * 10) // Ensure that UpdateTime changed - suite.Nil(UpdateQuotaUsage(*quotaUsage2)) - - // Get the updated quotaUsage - quotaUsage3, err := GetQuotaUsage(id) - suite.Nil(err) - suite.equalUsed(quotaUsage2, quotaUsage3) - suite.NotEqual(quotaUsage2.UpdateTime, quotaUsage3.UpdateTime) -} - -func (suite *QuotaUsageDaoSuite) TestListQuotaUsages() { - AddQuotaUsage(models.QuotaUsage{Reference: quotaUsageReference, ReferenceID: "1", Used: quotaUsageUsed.String()}) - AddQuotaUsage(models.QuotaUsage{Reference: quotaUsageReference, ReferenceID: "2", Used: quotaUsageUsed.String()}) - AddQuotaUsage(models.QuotaUsage{Reference: quotaUsageReference, ReferenceID: "3", Used: quotaUsageUsed.String()}) - AddQuotaUsage(models.QuotaUsage{Reference: quotaUsageUserReference, ReferenceID: "1", Used: quotaUsageUsedLarger.String()}) - - // List all the quotaUsages - quotaUsages, err := ListQuotaUsages() - suite.Nil(err) - suite.Equal(4, len(quotaUsages)) - suite.Equal(quotaUsageReference, quotaUsages[0].Reference) - - // List quotaUsages filter by reference - quotaUsages, err = ListQuotaUsages(&models.QuotaUsageQuery{Reference: quotaUsageReference}) - suite.Nil(err) - suite.Equal(3, len(quotaUsages)) - - // List quotaUsages filter by reference ids - quotaUsages, err = ListQuotaUsages(&models.QuotaUsageQuery{Reference: quotaUsageReference, ReferenceIDs: []string{"1", "2"}}) - suite.Nil(err) - suite.Equal(2, len(quotaUsages)) - - // List quotaUsages by pagination - quotaUsages, err = ListQuotaUsages(&models.QuotaUsageQuery{Pagination: models.Pagination{Size: 2}}) - suite.Nil(err) - suite.Equal(2, len(quotaUsages)) - - // List quotaUsages by sorting - quotaUsages, err = ListQuotaUsages(&models.QuotaUsageQuery{Sorting: models.Sorting{Sort: "-used.storage"}}) - suite.Nil(err) - suite.Equal(quotaUsageUserReference, quotaUsages[0].Reference) -} - -func TestRunQuotaUsageDaoSuite(t *testing.T) { - suite.Run(t, new(QuotaUsageDaoSuite)) -} - -func Test_quotaUsageOrderBy(t *testing.T) { - query := func(sort string) []*models.QuotaUsageQuery { - return []*models.QuotaUsageQuery{ - {Sorting: models.Sorting{Sort: sort}}, - } - } - - type args struct { - query []*models.QuotaUsageQuery - } - tests := []struct { - name string - args args - want string - }{ - {"no query", args{nil}, ""}, - {"order by unsupport field", args{query("unknow")}, ""}, - {"order by storage of used", args{query("used.storage")}, "(CAST( (CASE WHEN (used->>'storage') IS NULL THEN '0' WHEN (used->>'storage') = '-1' THEN '9223372036854775807' ELSE (used->>'storage') END) AS BIGINT )) ASC"}, - {"order by unsupport used resource", args{query("used.unknow")}, ""}, - } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - if got := quotaUsageOrderBy(tt.args.query...); got != tt.want { - t.Errorf("quotaUsageOrderBy() = %v, want %v", got, tt.want) - } - }) - } -} diff --git a/src/common/models/base.go b/src/common/models/base.go index 460ce9823..4ca2b678d 100644 --- a/src/common/models/base.go +++ b/src/common/models/base.go @@ -38,7 +38,5 @@ func init() { new(ProjectBlob), new(ArtifactAndBlob), new(CVEWhitelist), - new(Quota), - new(QuotaUsage), ) } diff --git a/src/common/models/quota.go b/src/common/models/quota.go deleted file mode 100644 index e7d8ade6e..000000000 --- a/src/common/models/quota.go +++ /dev/null @@ -1,85 +0,0 @@ -// 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 models - -import ( - "encoding/json" - "time" - - "github.com/goharbor/harbor/src/pkg/types" -) - -// QuotaHard a map for the quota hard -type QuotaHard map[string]int64 - -func (h QuotaHard) String() string { - bytes, _ := json.Marshal(h) - return string(bytes) -} - -// Copy returns copied quota hard -func (h QuotaHard) Copy() QuotaHard { - hard := QuotaHard{} - for key, value := range h { - hard[key] = value - } - - return hard -} - -// Quota model for quota -type Quota struct { - ID int64 `orm:"pk;auto;column(id)" json:"id"` - Reference string `orm:"column(reference)" json:"reference"` // The reference type for quota, eg: project, user - ReferenceID string `orm:"column(reference_id)" json:"reference_id"` - Hard string `orm:"column(hard);type(jsonb)" json:"-"` - CreationTime time.Time `orm:"column(creation_time);auto_now_add" json:"creation_time"` - UpdateTime time.Time `orm:"column(update_time);auto_now" json:"update_time"` -} - -// TableName returns table name for orm -func (q *Quota) TableName() string { - return "quota" -} - -// GetHard returns quota hard -func (q *Quota) GetHard() (QuotaHard, error) { - var hard QuotaHard - if err := json.Unmarshal([]byte(q.Hard), &hard); err != nil { - return nil, err - } - - return hard, nil -} - -// SetHard set new quota hard -func (q *Quota) SetHard(hard QuotaHard) { - q.Hard = hard.String() -} - -// QuotaQuery query parameters for quota -type QuotaQuery struct { - ID int64 - Reference string - ReferenceID string - ReferenceIDs []string - Pagination - Sorting -} - -// QuotaUpdateRequest the request for quota update -type QuotaUpdateRequest struct { - Hard types.ResourceList `json:"hard"` -} diff --git a/src/common/models/quota_usage.go b/src/common/models/quota_usage.go deleted file mode 100644 index c5c24eeb3..000000000 --- a/src/common/models/quota_usage.go +++ /dev/null @@ -1,77 +0,0 @@ -// 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 models - -import ( - "encoding/json" - "time" -) - -// QuotaUsed a map for the quota used -type QuotaUsed map[string]int64 - -func (u QuotaUsed) String() string { - bytes, _ := json.Marshal(u) - return string(bytes) -} - -// Copy returns copied quota used -func (u QuotaUsed) Copy() QuotaUsed { - used := QuotaUsed{} - for key, value := range u { - used[key] = value - } - - return used -} - -// QuotaUsage model for quota usage -type QuotaUsage struct { - ID int64 `orm:"pk;auto;column(id)" json:"id"` - Reference string `orm:"column(reference)" json:"reference"` // The reference type for quota usage, eg: project, user - ReferenceID string `orm:"column(reference_id)" json:"reference_id"` - Used string `orm:"column(used);type(jsonb)" json:"-"` - CreationTime time.Time `orm:"column(creation_time);auto_now_add" json:"creation_time"` - UpdateTime time.Time `orm:"column(update_time);auto_now" json:"update_time"` -} - -// TableName returns table name for orm -func (qu *QuotaUsage) TableName() string { - return "quota_usage" -} - -// GetUsed returns quota used -func (qu *QuotaUsage) GetUsed() (QuotaUsed, error) { - var used QuotaUsed - if err := json.Unmarshal([]byte(qu.Used), &used); err != nil { - return nil, err - } - - return used, nil -} - -// SetUsed set quota used -func (qu *QuotaUsage) SetUsed(used QuotaUsed) { - qu.Used = used.String() -} - -// QuotaUsageQuery query parameters for quota -type QuotaUsageQuery struct { - Reference string - ReferenceID string - ReferenceIDs []string - Pagination - Sorting -} diff --git a/src/common/quota/driver/driver.go b/src/common/quota/driver/driver.go deleted file mode 100644 index fbd339e37..000000000 --- a/src/common/quota/driver/driver.go +++ /dev/null @@ -1,59 +0,0 @@ -// 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 driver - -import ( - "sync" - - "github.com/goharbor/harbor/src/pkg/types" -) - -var ( - driversMu sync.RWMutex - drivers = map[string]Driver{} -) - -// RefObject type for quota ref object -type RefObject map[string]interface{} - -// Driver the driver for quota -type Driver interface { - // HardLimits returns default resource list - HardLimits() types.ResourceList - // Load returns quota ref object by key - Load(key string) (RefObject, error) - // Validate validate the hard limits - Validate(hardLimits types.ResourceList) error -} - -// Register register quota driver -func Register(name string, driver Driver) { - driversMu.Lock() - defer driversMu.Unlock() - if driver == nil { - panic("quota: Register driver is nil") - } - - drivers[name] = driver -} - -// Get returns quota driver by name -func Get(name string) (Driver, bool) { - driversMu.Lock() - defer driversMu.Unlock() - - driver, ok := drivers[name] - return driver, ok -} diff --git a/src/common/quota/driver/mocks/driver.go b/src/common/quota/driver/mocks/driver.go deleted file mode 100644 index 8f8c1ac82..000000000 --- a/src/common/quota/driver/mocks/driver.go +++ /dev/null @@ -1,65 +0,0 @@ -// Code generated by mockery v1.0.0. DO NOT EDIT. - -package mocks - -import driver "github.com/goharbor/harbor/src/common/quota/driver" -import mock "github.com/stretchr/testify/mock" -import types "github.com/goharbor/harbor/src/pkg/types" - -// Driver is an autogenerated mock type for the Driver type -type Driver struct { - mock.Mock -} - -// HardLimits provides a mock function with given fields: -func (_m *Driver) HardLimits() types.ResourceList { - ret := _m.Called() - - var r0 types.ResourceList - if rf, ok := ret.Get(0).(func() types.ResourceList); ok { - r0 = rf() - } else { - if ret.Get(0) != nil { - r0 = ret.Get(0).(types.ResourceList) - } - } - - return r0 -} - -// Load provides a mock function with given fields: key -func (_m *Driver) Load(key string) (driver.RefObject, error) { - ret := _m.Called(key) - - var r0 driver.RefObject - if rf, ok := ret.Get(0).(func(string) driver.RefObject); ok { - r0 = rf(key) - } else { - if ret.Get(0) != nil { - r0 = ret.Get(0).(driver.RefObject) - } - } - - var r1 error - if rf, ok := ret.Get(1).(func(string) error); ok { - r1 = rf(key) - } else { - r1 = ret.Error(1) - } - - return r0, r1 -} - -// Validate provides a mock function with given fields: resources -func (_m *Driver) Validate(resources types.ResourceList) error { - ret := _m.Called(resources) - - var r0 error - if rf, ok := ret.Get(0).(func(types.ResourceList) error); ok { - r0 = rf(resources) - } else { - r0 = ret.Error(0) - } - - return r0 -} diff --git a/src/common/quota/driver/project/driver.go b/src/common/quota/driver/project/driver.go deleted file mode 100644 index 60d095a10..000000000 --- a/src/common/quota/driver/project/driver.go +++ /dev/null @@ -1,158 +0,0 @@ -// 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 project - -import ( - "context" - "fmt" - "strconv" - - "github.com/goharbor/harbor/src/common" - "github.com/goharbor/harbor/src/common/config" - "github.com/goharbor/harbor/src/common/dao" - "github.com/goharbor/harbor/src/common/models" - dr "github.com/goharbor/harbor/src/common/quota/driver" - "github.com/goharbor/harbor/src/pkg/types" - "github.com/graph-gophers/dataloader" -) - -func init() { - dr.Register("project", newDriver()) -} - -func getProjectsBatchFn(ctx context.Context, keys dataloader.Keys) []*dataloader.Result { - handleError := func(err error) []*dataloader.Result { - var results []*dataloader.Result - var result dataloader.Result - result.Error = err - results = append(results, &result) - return results - } - - var projectIDs []int64 - for _, key := range keys { - id, err := strconv.ParseInt(key.String(), 10, 64) - if err != nil { - return handleError(err) - } - projectIDs = append(projectIDs, id) - } - - projects, err := dao.GetProjects(&models.ProjectQueryParam{}) - if err != nil { - return handleError(err) - } - - var ownerIDs []int - var projectsMap = make(map[int64]*models.Project, len(projectIDs)) - for _, project := range projects { - ownerIDs = append(ownerIDs, project.OwnerID) - projectsMap[project.ProjectID] = project - } - - owners, err := dao.ListUsers(&models.UserQuery{UserIDs: ownerIDs}) - if err != nil { - return handleError(err) - } - - var ownersMap = make(map[int]*models.User, len(owners)) - for i, owner := range owners { - ownersMap[owner.UserID] = &owners[i] - } - - var results []*dataloader.Result - for _, projectID := range projectIDs { - project, ok := projectsMap[projectID] - if !ok { - return handleError(fmt.Errorf("project not found, "+"project_id: %d", projectID)) - } - - owner, ok := ownersMap[project.OwnerID] - if ok { - project.OwnerName = owner.Username - } - - result := dataloader.Result{ - Data: project, - Error: nil, - } - results = append(results, &result) - } - - return results -} - -type driver struct { - cfg *config.CfgManager - loader *dataloader.Loader -} - -func (d *driver) HardLimits() types.ResourceList { - return types.ResourceList{ - types.ResourceStorage: d.cfg.Get(common.StoragePerProject).GetInt64(), - } -} - -func (d *driver) Load(key string) (dr.RefObject, error) { - thunk := d.loader.Load(context.TODO(), dataloader.StringKey(key)) - - result, err := thunk() - if err != nil { - return nil, err - } - - project, ok := result.(*models.Project) - if !ok { - return nil, fmt.Errorf("bad result for project: %s", key) - } - - return dr.RefObject{ - "id": project.ProjectID, - "name": project.Name, - "owner_name": project.OwnerName, - }, nil -} - -func (d *driver) Validate(hardLimits types.ResourceList) error { - resources := map[types.ResourceName]bool{ - types.ResourceStorage: true, - } - - for resource, value := range hardLimits { - if !resources[resource] { - return fmt.Errorf("resource %s not support", resource) - } - - if value <= 0 && value != types.UNLIMITED { - return fmt.Errorf("invalid value for resource %s", resource) - } - } - - for resource := range resources { - if _, found := hardLimits[resource]; !found { - return fmt.Errorf("resource %s not found", resource) - } - } - - return nil -} - -func newDriver() dr.Driver { - cfg := config.NewDBCfgManager() - - loader := dataloader.NewBatchedLoader(getProjectsBatchFn) - - return &driver{cfg: cfg, loader: loader} -} diff --git a/src/common/quota/driver/project/driver_test.go b/src/common/quota/driver/project/driver_test.go deleted file mode 100644 index 4e0232f49..000000000 --- a/src/common/quota/driver/project/driver_test.go +++ /dev/null @@ -1,75 +0,0 @@ -// 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 project - -import ( - "os" - "testing" - - "github.com/goharbor/harbor/src/common/dao" - dr "github.com/goharbor/harbor/src/common/quota/driver" - "github.com/goharbor/harbor/src/pkg/types" - "github.com/stretchr/testify/suite" -) - -type DriverSuite struct { - suite.Suite -} - -func (suite *DriverSuite) TestHardLimits() { - driver := newDriver() - - suite.Equal(types.ResourceList{types.ResourceStorage: -1}, driver.HardLimits()) -} - -func (suite *DriverSuite) TestLoad() { - driver := newDriver() - - if ref, err := driver.Load("1"); suite.Nil(err) { - obj := dr.RefObject{ - "id": int64(1), - "name": "library", - "owner_name": "admin", - } - - suite.Equal(obj, ref) - } - - if ref, err := driver.Load("100000"); suite.Error(err) { - suite.Empty(ref) - } - - if ref, err := driver.Load("library"); suite.Error(err) { - suite.Empty(ref) - } -} - -func (suite *DriverSuite) TestValidate() { - driver := newDriver() - - suite.Error(driver.Validate(types.ResourceList{})) - suite.Error(driver.Validate(types.ResourceList{types.ResourceStorage: 0})) - suite.Error(driver.Validate(types.ResourceList{types.ResourceName("foo"): 1})) -} - -func TestMain(m *testing.M) { - dao.PrepareTestForPostgresSQL() - - os.Exit(m.Run()) -} - -func TestRunDriverSuite(t *testing.T) { - suite.Run(t, new(DriverSuite)) -} diff --git a/src/common/quota/errors.go b/src/common/quota/errors.go deleted file mode 100644 index 162cf536f..000000000 --- a/src/common/quota/errors.go +++ /dev/null @@ -1,111 +0,0 @@ -// 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 quota - -import ( - "fmt" - "strings" - - "github.com/goharbor/harbor/src/pkg/types" -) - -// Errors contains all happened errors -type Errors []error - -// GetErrors gets all errors that have occurred and returns a slice of errors (Error type) -func (errs Errors) GetErrors() []error { - return errs -} - -// Add adds an error to a given slice of errors -func (errs Errors) Add(newErrors ...error) Errors { - for _, err := range newErrors { - if err == nil { - continue - } - - if errors, ok := err.(Errors); ok { - errs = errs.Add(errors...) - } else { - ok = true - for _, e := range errs { - if err == e { - ok = false - } - } - if ok { - errs = append(errs, err) - } - } - } - - return errs -} - -// Error takes a slice of all errors that have occurred and returns it as a formatted string -func (errs Errors) Error() string { - var errors = []string{} - for _, e := range errs { - errors = append(errors, e.Error()) - } - return strings.Join(errors, "; ") -} - -// ResourceOverflow ... -type ResourceOverflow struct { - Resource types.ResourceName - HardLimit int64 - CurrentUsed int64 - NewUsed int64 -} - -func (e *ResourceOverflow) Error() string { - resource := e.Resource - var ( - op string - delta int64 - ) - - if e.NewUsed > e.CurrentUsed { - op = "adding" - delta = e.NewUsed - e.CurrentUsed - } else { - op = "subtracting" - delta = e.CurrentUsed - e.NewUsed - } - - return fmt.Sprintf("%s %s of %s resource, which when updated to current usage of %s will exceed the configured upper limit of %s.", - op, resource.FormatValue(delta), resource, - resource.FormatValue(e.CurrentUsed), resource.FormatValue(e.HardLimit)) -} - -// NewResourceOverflowError ... -func NewResourceOverflowError(resource types.ResourceName, hardLimit, currentUsed, newUsed int64) error { - return &ResourceOverflow{Resource: resource, HardLimit: hardLimit, CurrentUsed: currentUsed, NewUsed: newUsed} -} - -// ResourceNotFound ... -type ResourceNotFound struct { - Resource types.ResourceName -} - -func (e *ResourceNotFound) Error() string { - return fmt.Sprintf("resource %s not found", e.Resource) -} - -// NewResourceNotFoundError ... -func NewResourceNotFoundError(resource types.ResourceName) error { - return &ResourceNotFound{Resource: resource} -} diff --git a/src/common/quota/manager.go b/src/common/quota/manager.go deleted file mode 100644 index be180abe6..000000000 --- a/src/common/quota/manager.go +++ /dev/null @@ -1,287 +0,0 @@ -// 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 quota - -import ( - "fmt" - "time" - - "github.com/astaxie/beego/orm" - "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/lib/log" - "github.com/goharbor/harbor/src/pkg/types" -) - -// Manager manager for quota -type Manager struct { - driver driver.Driver - reference string - referenceID string -} - -func (m *Manager) addQuota(o orm.Ormer, hardLimits types.ResourceList, now time.Time) (int64, error) { - quota := &models.Quota{ - Reference: m.reference, - ReferenceID: m.referenceID, - Hard: hardLimits.String(), - CreationTime: now, - UpdateTime: now, - } - - return o.Insert(quota) -} - -func (m *Manager) addUsage(o orm.Ormer, used types.ResourceList, now time.Time, ids ...int64) (int64, error) { - usage := &models.QuotaUsage{ - Reference: m.reference, - ReferenceID: m.referenceID, - Used: used.String(), - CreationTime: now, - UpdateTime: now, - } - - if len(ids) > 0 { - usage.ID = ids[0] - } - - return o.Insert(usage) -} - -func (m *Manager) newQuota(o orm.Ormer, hardLimits types.ResourceList, usages ...types.ResourceList) (int64, error) { - now := time.Now() - - id, err := m.addQuota(o, hardLimits, now) - if err != nil { - return 0, err - } - - var used types.ResourceList - if len(usages) > 0 { - used = usages[0] - } else { - used = types.Zero(hardLimits) - } - - if _, err := m.addUsage(o, used, now, id); err != nil { - return 0, err - } - - return id, nil -} - -func (m *Manager) getQuotaForUpdate(o orm.Ormer) (*models.Quota, error) { - quota := &models.Quota{Reference: m.reference, ReferenceID: m.referenceID} - if err := o.ReadForUpdate(quota, "reference", "reference_id"); err != nil { - if err == orm.ErrNoRows { - if _, err := m.newQuota(o, m.driver.HardLimits()); err != nil { - return nil, err - } - - return m.getQuotaForUpdate(o) - } - - return nil, err - } - - return quota, nil -} - -func (m *Manager) getUsageForUpdate(o orm.Ormer) (*models.QuotaUsage, error) { - usage := &models.QuotaUsage{Reference: m.reference, ReferenceID: m.referenceID} - if err := o.ReadForUpdate(usage, "reference", "reference_id"); err != nil { - return nil, err - } - - return usage, nil -} - -func (m *Manager) updateUsage(o orm.Ormer, resources types.ResourceList, - calculate func(types.ResourceList, types.ResourceList) types.ResourceList, - skipOverflow bool) error { - - quota, err := m.getQuotaForUpdate(o) - if err != nil { - return err - } - hardLimits, err := types.NewResourceList(quota.Hard) - if err != nil { - return err - } - - usage, err := m.getUsageForUpdate(o) - if err != nil { - return err - } - used, err := types.NewResourceList(usage.Used) - if err != nil { - return err - } - - newUsed := calculate(used, resources) - - // ensure that new used is never negative - if negativeUsed := types.IsNegative(newUsed); len(negativeUsed) > 0 { - return fmt.Errorf("quota usage is negative for resource(s): %s", prettyPrintResourceNames(negativeUsed)) - } - - if err := isSafe(hardLimits, used, newUsed, skipOverflow); err != nil { - return err - } - - usage.Used = newUsed.String() - usage.UpdateTime = time.Now() - - _, err = o.Update(usage) - return err -} - -// NewQuota create new quota for (reference, reference id) -func (m *Manager) NewQuota(hardLimit types.ResourceList, usages ...types.ResourceList) (int64, error) { - var id int64 - err := dao.WithTransaction(func(o orm.Ormer) (err error) { - id, err = m.newQuota(o, hardLimit, usages...) - return err - }) - - if err != nil { - return 0, err - } - - return id, nil -} - -// DeleteQuota delete the quota -func (m *Manager) DeleteQuota() error { - return dao.WithTransaction(func(o orm.Ormer) error { - quota := &models.Quota{Reference: m.reference, ReferenceID: m.referenceID} - if _, err := o.Delete(quota, "reference", "reference_id"); err != nil { - return err - } - - usage := &models.QuotaUsage{Reference: m.reference, ReferenceID: m.referenceID} - if _, err := o.Delete(usage, "reference", "reference_id"); err != nil { - return err - } - - return nil - }) -} - -// 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 := o.Raw(sql, hardLimits.String(), m.reference, m.referenceID).Exec() - - return err -} - -// SetResourceUsage sets the usage per resource name -func (m *Manager) SetResourceUsage(resource types.ResourceName, value int64) error { - o := dao.GetOrmer() - - sql := fmt.Sprintf("UPDATE quota_usage SET used = jsonb_set(used, '{%s}', to_jsonb(%d::bigint), true) WHERE reference = ? AND reference_id = ?", resource, value) - _, err := o.Raw(sql, 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 err != nil { - return err - } - 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 { - return m.updateUsage(o, resources, types.Add, false) - }) -} - -// SubtractResources subtract resources from usage -func (m *Manager) SubtractResources(resources types.ResourceList) error { - return dao.WithTransaction(func(o orm.Ormer) error { - return m.updateUsage(o, resources, types.Subtract, true) - }) -} - -// NewManager returns quota manager -func NewManager(reference string, referenceID string) (*Manager, error) { - d, ok := driver.Get(reference) - if !ok { - return nil, fmt.Errorf("quota not support for %s", reference) - } - - if _, err := d.Load(referenceID); err != nil { - log.Warning(fmt.Sprintf("Load quota reference object (%s, %s) failed: %v", reference, referenceID, err)) - return nil, err - } - - return &Manager{ - driver: d, - reference: reference, - referenceID: referenceID, - }, nil -} diff --git a/src/common/quota/manager_test.go b/src/common/quota/manager_test.go deleted file mode 100644 index e55ddacf7..000000000 --- a/src/common/quota/manager_test.go +++ /dev/null @@ -1,359 +0,0 @@ -// 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 quota - -import ( - "fmt" - "os" - "sync" - "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" - "github.com/stretchr/testify/mock" - "github.com/stretchr/testify/suite" -) - -var ( - hardLimits = types.ResourceList{types.ResourceStorage: 1000} - reference = "mock" -) - -func init() { - mockDriver := &mocks.Driver{} - - mockHardLimitsFn := func() types.ResourceList { - return types.ResourceList{ - types.ResourceStorage: -1, - } - } - - mockLoadFn := func(key string) driver.RefObject { - return driver.RefObject{"id": key} - } - - mockDriver.On("HardLimits").Return(mockHardLimitsFn) - mockDriver.On("Load", mock.AnythingOfType("string")).Return(mockLoadFn, nil) - mockDriver.On("Validate", mock.AnythingOfType("types.ResourceList")).Return(nil) - - driver.Register(reference, mockDriver) -} - -func mustResourceList(s string) types.ResourceList { - resources, _ := types.NewResourceList(s) - return resources -} - -type ManagerSuite struct { - suite.Suite -} - -func (suite *ManagerSuite) SetupTest() { - _, ok := driver.Get(reference) - if !ok { - suite.Fail("driver not found for %s", reference) - } -} - -func (suite *ManagerSuite) quotaManager(referenceIDs ...string) *Manager { - referenceID := "1" - if len(referenceIDs) > 0 { - referenceID = referenceIDs[0] - } - - mgr, _ := NewManager(reference, referenceID) - return mgr -} - -func (suite *ManagerSuite) TearDownTest() { - dao.ClearTable("quota") - dao.ClearTable("quota_usage") -} - -func (suite *ManagerSuite) TestNewQuota() { - mgr := suite.quotaManager() - - if id, err := mgr.NewQuota(hardLimits); suite.Nil(err) { - quota, _ := dao.GetQuota(id) - suite.Equal(hardLimits, mustResourceList(quota.Hard)) - } - - mgr = suite.quotaManager("2") - used := types.ResourceList{types.ResourceStorage: 100} - 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)) - } -} - -func (suite *ManagerSuite) TestDeleteQuota() { - mgr := suite.quotaManager() - - id, err := mgr.NewQuota(hardLimits) - if suite.Nil(err) { - quota, _ := dao.GetQuota(id) - suite.Equal(hardLimits, mustResourceList(quota.Hard)) - } - - if err := mgr.DeleteQuota(); suite.Nil(err) { - quota, _ := dao.GetQuota(id) - suite.Nil(quota) - } -} - -func (suite *ManagerSuite) TestUpdateQuota() { - mgr := suite.quotaManager() - - id, _ := mgr.NewQuota(hardLimits) - largeHardLimits := types.ResourceList{types.ResourceStorage: 1000000} - - if err := mgr.UpdateQuota(largeHardLimits); suite.Nil(err) { - quota, _ := dao.GetQuota(id) - suite.Equal(largeHardLimits, mustResourceList(quota.Hard)) - } -} - -func (suite *ManagerSuite) TestSetResourceUsage() { - mgr := suite.quotaManager() - id, _ := mgr.NewQuota(hardLimits) - - if err := mgr.SetResourceUsage(types.ResourceStorage, 999999999999999999); suite.Nil(err) { - quota, _ := dao.GetQuota(id) - suite.Equal(hardLimits, mustResourceList(quota.Hard)) - - usage, _ := dao.GetQuotaUsage(id) - suite.Equal(types.ResourceList{types.ResourceStorage: 999999999999999999}, mustResourceList(usage.Used)) - } - - if err := mgr.SetResourceUsage(types.ResourceStorage, 234); suite.Nil(err) { - usage, _ := dao.GetQuotaUsage(id) - suite.Equal(types.ResourceList{types.ResourceStorage: 234}, mustResourceList(usage.Used)) - } -} - -func (suite *ManagerSuite) TestEnsureQuota() { - // non-existent - nonExistRefID := "3" - mgr := suite.quotaManager(nonExistRefID) - infinite := types.ResourceList{types.ResourceStorage: -1} - usage := types.ResourceList{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.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.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)) - resource := types.ResourceList{types.ResourceStorage: 100} - - suite.Nil(mgr.AddResources(resource)) - } -} - -func (suite *ManagerSuite) TestAddResources() { - mgr := suite.quotaManager() - id, _ := mgr.NewQuota(hardLimits) - - resource := types.ResourceList{types.ResourceStorage: 100} - - if suite.Nil(mgr.AddResources(resource)) { - usage, _ := dao.GetQuotaUsage(id) - suite.Equal(resource, mustResourceList(usage.Used)) - } - - if suite.Nil(mgr.AddResources(resource)) { - usage, _ := dao.GetQuotaUsage(id) - suite.Equal(types.ResourceList{types.ResourceStorage: 200}, mustResourceList(usage.Used)) - } - - if err := mgr.AddResources(types.ResourceList{types.ResourceStorage: 10000}); suite.Error(err) { - if errs, ok := err.(Errors); suite.True(ok) { - for _, err := range errs { - suite.IsType(&ResourceOverflow{}, err) - } - } - } -} - -func (suite *ManagerSuite) TestSubtractResources() { - mgr := suite.quotaManager() - id, _ := mgr.NewQuota(hardLimits) - - resource := types.ResourceList{types.ResourceStorage: 100} - - if suite.Nil(mgr.AddResources(resource)) { - usage, _ := dao.GetQuotaUsage(id) - suite.Equal(resource, mustResourceList(usage.Used)) - } - - if suite.Nil(mgr.SubtractResources(resource)) { - usage, _ := dao.GetQuotaUsage(id) - suite.Equal(types.ResourceList{types.ResourceStorage: 0}, mustResourceList(usage.Used)) - } -} - -func (suite *ManagerSuite) TestRaceAddResources() { - mgr := suite.quotaManager() - mgr.NewQuota(hardLimits) - - resources := types.ResourceList{ - types.ResourceStorage: 100, - } - - var wg sync.WaitGroup - - results := make([]bool, 100) - for i := 0; i < 100; i++ { - wg.Add(1) - go func(i int) { - defer wg.Done() - - results[i] = mgr.AddResources(resources) == nil - }(i) - } - wg.Wait() - - var success int - for _, result := range results { - if result { - success++ - } - } - - suite.Equal(10, success) -} - -func (suite *ManagerSuite) TestRaceSubtractResources() { - mgr := suite.quotaManager() - mgr.NewQuota(hardLimits, types.ResourceList{types.ResourceStorage: 1000}) - - resources := types.ResourceList{ - types.ResourceStorage: 100, - } - - var wg sync.WaitGroup - - results := make([]bool, 100) - for i := 0; i < 100; i++ { - wg.Add(1) - go func(i int) { - defer wg.Done() - - results[i] = mgr.SubtractResources(resources) == nil - }(i) - } - wg.Wait() - - var success int - for _, result := range results { - if result { - success++ - } - } - - suite.Equal(10, success) -} - -func TestMain(m *testing.M) { - dao.PrepareTestForPostgresSQL() - - if result := m.Run(); result != 0 { - os.Exit(result) - } -} - -func TestRunManagerSuite(t *testing.T) { - suite.Run(t, new(ManagerSuite)) -} - -func BenchmarkAddResources(b *testing.B) { - defer func() { - dao.ClearTable("quota") - dao.ClearTable("quota_usage") - }() - - mgr, _ := NewManager(reference, "1") - mgr.NewQuota(types.ResourceList{types.ResourceStorage: int64(b.N)}) - - resource := types.ResourceList{ - types.ResourceStorage: 1, - } - - b.ResetTimer() - for i := 0; i < b.N; i++ { - mgr.AddResources(resource) - } - b.StopTimer() -} - -func BenchmarkAddResourcesParallel(b *testing.B) { - defer func() { - dao.ClearTable("quota") - dao.ClearTable("quota_usage") - }() - - mgr, _ := NewManager(reference, "1") - mgr.NewQuota(types.ResourceList{}) - - resource := types.ResourceList{ - types.ResourceStorage: 1, - } - - b.ResetTimer() - b.RunParallel(func(b *testing.PB) { - for b.Next() { - mgr.AddResources(resource) - } - }) - b.StopTimer() -} diff --git a/src/common/quota/quota.go b/src/common/quota/quota.go deleted file mode 100644 index 4446d61eb..000000000 --- a/src/common/quota/quota.go +++ /dev/null @@ -1,35 +0,0 @@ -// 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 quota - -import ( - "fmt" - - "github.com/goharbor/harbor/src/common/quota/driver" - "github.com/goharbor/harbor/src/pkg/types" - - // project driver for quota - _ "github.com/goharbor/harbor/src/common/quota/driver/project" -) - -// Validate validate hard limits -func Validate(reference string, hardLimits types.ResourceList) error { - d, ok := driver.Get(reference) - if !ok { - return fmt.Errorf("quota not support for %s", reference) - } - - return d.Validate(hardLimits) -} diff --git a/src/common/quota/quota_test.go b/src/common/quota/quota_test.go deleted file mode 100644 index da88d5427..000000000 --- a/src/common/quota/quota_test.go +++ /dev/null @@ -1,45 +0,0 @@ -// 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 quota - -import ( - "testing" - - _ "github.com/goharbor/harbor/src/common/quota/driver/project" - "github.com/goharbor/harbor/src/pkg/types" -) - -func TestValidate(t *testing.T) { - type args struct { - reference string - hardLimits types.ResourceList - } - tests := []struct { - name string - args args - wantErr bool - }{ - {"valid", args{"project", types.ResourceList{types.ResourceStorage: 1}}, false}, - {"invalid", args{"project", types.ResourceList{types.ResourceStorage: 0}}, true}, - {"not support", args{"not support", types.ResourceList{types.ResourceStorage: 1}}, true}, - } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - if err := Validate(tt.args.reference, tt.args.hardLimits); (err != nil) != tt.wantErr { - t.Errorf("Validate() error = %v, wantErr %v", err, tt.wantErr) - } - }) - } -} diff --git a/src/common/quota/types.go b/src/common/quota/types.go deleted file mode 100644 index 385ffc1da..000000000 --- a/src/common/quota/types.go +++ /dev/null @@ -1,30 +0,0 @@ -// 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 quota - -import ( - "github.com/goharbor/harbor/src/pkg/types" -) - -var ( - // ResourceStorage alias types.ResourceStorage - ResourceStorage = types.ResourceStorage -) - -// ResourceName alias types.ResourceName -type ResourceName = types.ResourceName - -// ResourceList alias types.ResourceList -type ResourceList = types.ResourceList diff --git a/src/common/quota/util.go b/src/common/quota/util.go deleted file mode 100644 index e57e05170..000000000 --- a/src/common/quota/util.go +++ /dev/null @@ -1,57 +0,0 @@ -// 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 quota - -import ( - "sort" - "strings" - - "github.com/goharbor/harbor/src/pkg/types" -) - -func isSafe(hardLimits types.ResourceList, currentUsed types.ResourceList, newUsed types.ResourceList, skipOverflow bool) error { - var errs Errors - - for resource, value := range newUsed { - hardLimit, found := hardLimits[resource] - if !found { - errs = errs.Add(NewResourceNotFoundError(resource)) - continue - } - - if hardLimit == types.UNLIMITED || value == currentUsed[resource] { - continue - } - - if value > hardLimit && !skipOverflow { - errs = errs.Add(NewResourceOverflowError(resource, hardLimit, currentUsed[resource], value)) - } - } - - if len(errs) > 0 { - return errs - } - - return nil -} - -func prettyPrintResourceNames(a []types.ResourceName) string { - values := []string{} - for _, value := range a { - values = append(values, string(value)) - } - sort.Strings(values) - return strings.Join(values, ",") -} diff --git a/src/common/quota/util_test.go b/src/common/quota/util_test.go deleted file mode 100644 index 7fdacb9e1..000000000 --- a/src/common/quota/util_test.go +++ /dev/null @@ -1,83 +0,0 @@ -// 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 quota - -import ( - "testing" - - "github.com/goharbor/harbor/src/pkg/types" -) - -func Test_isSafe(t *testing.T) { - type args struct { - hardLimits types.ResourceList - currentUsed types.ResourceList - newUsed types.ResourceList - skipOverflow bool - } - tests := []struct { - name string - args args - wantErr bool - }{ - { - "unlimited", - args{ - types.ResourceList{types.ResourceStorage: types.UNLIMITED}, - types.ResourceList{types.ResourceStorage: 1000}, - types.ResourceList{types.ResourceStorage: 1000}, - false, - }, - false, - }, - { - "ok", - args{ - types.ResourceList{types.ResourceStorage: 100}, - types.ResourceList{types.ResourceStorage: 10}, - types.ResourceList{types.ResourceStorage: 1}, - false, - }, - false, - }, - { - "over the hard limit", - args{ - types.ResourceList{types.ResourceStorage: 100}, - types.ResourceList{types.ResourceStorage: 0}, - types.ResourceList{types.ResourceStorage: 200}, - false, - }, - true, - }, - { - "skip overflow", - args{ - types.ResourceList{types.ResourceStorage: 100}, - types.ResourceList{types.ResourceStorage: 0}, - types.ResourceList{types.ResourceStorage: 200}, - true, - }, - false, - }, - } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - if err := isSafe(tt.args.hardLimits, tt.args.currentUsed, tt.args.newUsed, tt.args.skipOverflow); (err != nil) != tt.wantErr { - t.Errorf("isSafe() error = %v, wantErr %v", err, tt.wantErr) - } - }) - } -} diff --git a/src/controller/quota/controller.go b/src/controller/quota/controller.go index a38c67c51..489850922 100644 --- a/src/controller/quota/controller.go +++ b/src/controller/quota/controller.go @@ -24,6 +24,7 @@ import ( "github.com/goharbor/harbor/src/lib/errors" "github.com/goharbor/harbor/src/lib/log" "github.com/goharbor/harbor/src/lib/orm" + "github.com/goharbor/harbor/src/lib/q" "github.com/goharbor/harbor/src/pkg/quota" "github.com/goharbor/harbor/src/pkg/quota/driver" "github.com/goharbor/harbor/src/pkg/types" @@ -44,6 +45,9 @@ var ( // Controller defines the operations related with quotas type Controller interface { + // Count returns the total count of quotas according to the query. + Count(ctx context.Context, query *q.Query) (int64, error) + // Create ensure quota for the reference object Create(ctx context.Context, reference, referenceID string, hardLimits types.ResourceList, used ...types.ResourceList) (int64, error) @@ -59,6 +63,9 @@ type Controller interface { // IsEnabled returns true when quota enabled for reference object IsEnabled(ctx context.Context, reference, referenceID string) (bool, error) + // List list quotas + List(ctx context.Context, query *q.Query) ([]*quota.Quota, error) + // Refresh refresh quota for the reference object Refresh(ctx context.Context, reference, referenceID string, options ...Option) error @@ -67,6 +74,9 @@ type Controller interface { // then runs f and refresh quota when f success， // in the finally it releases the resources which reserved at the beginning. Request(ctx context.Context, reference, referenceID string, resources types.ResourceList, f func() error) error + + // Update update quota + Update(ctx context.Context, q *quota.Quota) error } // NewController creates an instance of the default quota controller @@ -83,6 +93,10 @@ type controller struct { quotaMgr quota.Manager } +func (c *controller) Count(ctx context.Context, query *q.Query) (int64, error) { + return c.quotaMgr.Count(ctx, query) +} + func (c *controller) Create(ctx context.Context, reference, referenceID string, hardLimits types.ResourceList, used ...types.ResourceList) (int64, error) { return c.quotaMgr.Create(ctx, reference, referenceID, hardLimits, used...) } @@ -108,6 +122,10 @@ func (c *controller) IsEnabled(ctx context.Context, reference, referenceID strin return d.Enabled(ctx, referenceID) } +func (c *controller) List(ctx context.Context, query *q.Query) ([]*quota.Quota, error) { + return c.quotaMgr.List(ctx, query) +} + func (c *controller) getReservedResources(ctx context.Context, reference, referenceID string) (types.ResourceList, error) { conn := util.DefaultPool().Get() defer conn.Close() @@ -283,6 +301,32 @@ func (c *controller) Request(ctx context.Context, reference, referenceID string, return c.Refresh(ctx, reference, referenceID) } +func (c *controller) Update(ctx context.Context, u *quota.Quota) error { + update := func(ctx context.Context) error { + q, err := c.quotaMgr.GetByRefForUpdate(ctx, u.Reference, u.ReferenceID) + if err != nil { + return err + } + + if q.Hard != u.Hard { + if hard, err := u.GetHard(); err == nil { + q.SetHard(hard) + } + } + + if q.Used != u.Used { + if used, err := u.GetUsed(); err == nil { + q.SetUsed(used) + } + } + + q.UpdateTime = time.Now() + return c.quotaMgr.Update(ctx, q) + } + + return orm.WithTransaction(update)(ctx) +} + // Driver returns quota driver for the reference func Driver(ctx context.Context, reference string) (driver.Driver, error) { d, ok := driver.Get(reference) @@ -293,6 +337,16 @@ func Driver(ctx context.Context, reference string) (driver.Driver, error) { return d, nil } +// Validate validate hard limits +func Validate(ctx context.Context, reference string, hardLimits types.ResourceList) error { + d, err := Driver(ctx, reference) + if err != nil { + return err + } + + return d.Validate(hardLimits) +} + func reservedResourcesKey(reference, referenceID string) string { return fmt.Sprintf("quota:%s:%s:reserved", reference, referenceID) } diff --git a/src/core/api/harborapi_test.go b/src/core/api/harborapi_test.go index 314a1df11..df1f7b66b 100644 --- a/src/core/api/harborapi_test.go +++ b/src/core/api/harborapi_test.go @@ -28,8 +28,6 @@ import ( "strconv" "strings" - "github.com/goharbor/harbor/src/server/middleware/security" - "github.com/astaxie/beego" "github.com/dghubble/sling" "github.com/goharbor/harbor/src/common/api" @@ -44,6 +42,9 @@ import ( "github.com/goharbor/harbor/src/core/config" "github.com/goharbor/harbor/src/pkg/notification" "github.com/goharbor/harbor/src/replication/model" + "github.com/goharbor/harbor/src/server/middleware" + "github.com/goharbor/harbor/src/server/middleware/orm" + "github.com/goharbor/harbor/src/server/middleware/security" "github.com/goharbor/harbor/src/testing/apitests/apilib" ) @@ -217,7 +218,8 @@ func init() { mockServer := test.NewJobServiceServer() defer mockServer.Close() - handler = security.Middleware()(beego.BeeApp.Handlers) + chain := middleware.Chain(orm.Middleware(), security.Middleware()) + handler = chain(beego.BeeApp.Handlers) } func request0(_sling *sling.Sling, acceptHeader string, authInfo ...usrInfo) (int, http.Header, []byte, error) { @@ -1064,7 +1066,7 @@ func (a testapi) QuotasGetByID(authInfo usrInfo, quotaID string) (int, apilib.Qu } // Update spec for the quota -func (a testapi) QuotasPut(authInfo usrInfo, quotaID string, req models.QuotaUpdateRequest) (int, error) { +func (a testapi) QuotasPut(authInfo usrInfo, quotaID string, req QuotaUpdateRequest) (int, error) { path := "/api/quotas/" + quotaID _sling := sling.New().Put(a.basePath).Path(path).BodyJSON(req) diff --git a/src/core/api/project.go b/src/core/api/project.go index 1790ebc1d..3f2f562d9 100644 --- a/src/core/api/project.go +++ b/src/core/api/project.go @@ -15,6 +15,7 @@ package api import ( + "context" "fmt" "net/http" "regexp" @@ -26,12 +27,12 @@ import ( "github.com/goharbor/harbor/src/common/dao" pro "github.com/goharbor/harbor/src/common/dao/project" "github.com/goharbor/harbor/src/common/models" - "github.com/goharbor/harbor/src/common/quota" "github.com/goharbor/harbor/src/common/rbac" "github.com/goharbor/harbor/src/common/security/local" "github.com/goharbor/harbor/src/common/utils" errutil "github.com/goharbor/harbor/src/common/utils/error" "github.com/goharbor/harbor/src/controller/event/metadata" + "github.com/goharbor/harbor/src/controller/quota" "github.com/goharbor/harbor/src/core/config" "github.com/goharbor/harbor/src/lib/log" evt "github.com/goharbor/harbor/src/pkg/notifier/event" @@ -138,7 +139,7 @@ func (p *ProjectAPI) Post() { pro.StorageLimit = &setting.StoragePerProject } - hardLimits, err = projectQuotaHardLimits(pro, setting) + hardLimits, err = projectQuotaHardLimits(p.Ctx.Request.Context(), pro, setting) if err != nil { log.Errorf("Invalid project request, error: %v", err) p.SendBadRequestError(fmt.Errorf("invalid request: %v", err)) @@ -201,12 +202,9 @@ func (p *ProjectAPI) Post() { } if config.QuotaPerProjectEnable() { - quotaMgr, err := quota.NewManager("project", strconv.FormatInt(projectID, 10)) - if err != nil { - p.SendInternalServerError(fmt.Errorf("failed to get quota manager: %v", err)) - return - } - if _, err := quotaMgr.NewQuota(hardLimits); err != nil { + ctx := p.Ctx.Request.Context() + referenceID := quota.ReferenceID(projectID) + if _, err := quota.Ctl.Create(ctx, quota.ProjectReference, referenceID, hardLimits); err != nil { p.SendInternalServerError(fmt.Errorf("failed to create quota for project: %v", err)) return } @@ -285,14 +283,16 @@ func (p *ProjectAPI) Delete() { return } - quotaMgr, err := quota.NewManager("project", strconv.FormatInt(p.project.ProjectID, 10)) + ctx := p.Ctx.Request.Context() + referenceID := quota.ReferenceID(p.project.ProjectID) + q, err := quota.Ctl.GetByRef(ctx, quota.ProjectReference, referenceID) if err != nil { - p.SendInternalServerError(fmt.Errorf("failed to get quota manager: %v", err)) - return - } - if err := quotaMgr.DeleteQuota(); err != nil { - p.SendInternalServerError(fmt.Errorf("failed to delete quota for project: %v", err)) - return + log.Warningf("failed to get quota for project %s, error: %v", p.project.Name, err) + } else { + if err := quota.Ctl.Delete(ctx, q.ID); err != nil { + p.SendInternalServerError(fmt.Errorf("failed to delete quota for project: %v", err)) + return + } } // fire event @@ -516,7 +516,7 @@ func (p *ProjectAPI) Summary() { ChartCount: p.project.ChartCount, } - var fetchSummaries []func(int64, *models.ProjectSummary) + var fetchSummaries []func(context.Context, int64, *models.ProjectSummary) if hasPerm, _ := p.HasProjectPermission(p.project.ProjectID, rbac.ActionRead, rbac.ResourceQuota); hasPerm { fetchSummaries = append(fetchSummaries, getProjectQuotaSummary) @@ -526,6 +526,8 @@ func (p *ProjectAPI) Summary() { fetchSummaries = append(fetchSummaries, getProjectMemberSummary) } + ctx := p.Ctx.Request.Context() + var wg sync.WaitGroup for _, fn := range fetchSummaries { fn := fn @@ -533,7 +535,7 @@ func (p *ProjectAPI) Summary() { wg.Add(1) go func() { defer wg.Done() - fn(p.project.ProjectID, summary) + fn(ctx, p.project.ProjectID, summary) }() } wg.Wait() @@ -563,7 +565,7 @@ func validateProjectReq(req *models.ProjectRequest) error { return nil } -func projectQuotaHardLimits(req *models.ProjectRequest, setting *models.QuotaSetting) (types.ResourceList, error) { +func projectQuotaHardLimits(ctx context.Context, req *models.ProjectRequest, setting *models.QuotaSetting) (types.ResourceList, error) { hardLimits := types.ResourceList{} if req.StorageLimit != nil { @@ -572,38 +574,31 @@ func projectQuotaHardLimits(req *models.ProjectRequest, setting *models.QuotaSet hardLimits[types.ResourceStorage] = setting.StoragePerProject } - if err := quota.Validate("project", hardLimits); err != nil { + if err := quota.Validate(ctx, quota.ProjectReference, hardLimits); err != nil { return nil, err } return hardLimits, nil } -func getProjectQuotaSummary(projectID int64, summary *models.ProjectSummary) { +func getProjectQuotaSummary(ctx context.Context, projectID int64, summary *models.ProjectSummary) { if !config.QuotaPerProjectEnable() { log.Debug("Quota per project disabled") return } - quotas, err := dao.ListQuotas(&models.QuotaQuery{Reference: "project", ReferenceID: strconv.FormatInt(projectID, 10)}) + q, err := quota.Ctl.GetByRef(ctx, quota.ProjectReference, quota.ReferenceID(projectID)) if err != nil { log.Debugf("failed to get quota for project: %d", projectID) return } - if len(quotas) == 0 { - log.Debugf("quota not found for project: %d", projectID) - return - } - - quota := quotas[0] - summary.Quota = &models.QuotaSummary{} - summary.Quota.Hard, _ = types.NewResourceList(quota.Hard) - summary.Quota.Used, _ = types.NewResourceList(quota.Used) + summary.Quota.Hard, _ = types.NewResourceList(q.Hard) + summary.Quota.Used, _ = types.NewResourceList(q.Used) } -func getProjectMemberSummary(projectID int64, summary *models.ProjectSummary) { +func getProjectMemberSummary(ctx context.Context, projectID int64, summary *models.ProjectSummary) { var wg sync.WaitGroup for _, e := range []struct { diff --git a/src/core/api/quota.go b/src/core/api/quota.go index eb55a6df3..82db5c96b 100644 --- a/src/core/api/quota.go +++ b/src/core/api/quota.go @@ -17,12 +17,18 @@ package api import ( "fmt" - "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/controller/quota" + "github.com/goharbor/harbor/src/lib/q" + "github.com/goharbor/harbor/src/pkg/quota/models" + "github.com/goharbor/harbor/src/pkg/types" "github.com/pkg/errors" ) +// QuotaUpdateRequest struct for the body of put quota API +type QuotaUpdateRequest struct { + Hard types.ResourceList `json:"hard"` +} + // QuotaAPI handles request to /api/quotas/ type QuotaAPI struct { BaseController @@ -56,14 +62,9 @@ func (qa *QuotaAPI) Prepare() { return } - quota, err := dao.GetQuota(id) + quota, err := quota.Ctl.Get(qa.Ctx.Request.Context(), id) if err != nil { - qa.SendInternalServerError(fmt.Errorf("failed to get quota %d, error: %v", id, err)) - return - } - - if quota == nil { - qa.SendNotFoundError(fmt.Errorf("quota %d not found", id)) + qa.SendError(err) return } @@ -73,45 +74,27 @@ func (qa *QuotaAPI) Prepare() { // Get returns quota by id func (qa *QuotaAPI) Get() { - query := &models.QuotaQuery{ - ID: qa.quota.ID, - } - - quotas, err := dao.ListQuotas(query) - if err != nil { - qa.SendInternalServerError(fmt.Errorf("failed to get quota %d, error: %v", qa.quota.ID, err)) - return - } - - if len(quotas) == 0 { - qa.SendNotFoundError(fmt.Errorf("quota %d not found", qa.quota.ID)) - return - } - - qa.Data["json"] = quotas[0] + qa.Data["json"] = qa.quota qa.ServeJSON() } // Put update the quota func (qa *QuotaAPI) Put() { - var req *models.QuotaUpdateRequest + var req *QuotaUpdateRequest if err := qa.DecodeJSONReq(&req); err != nil { qa.SendBadRequestError(err) return } - if err := quota.Validate(qa.quota.Reference, req.Hard); err != nil { + ctx := qa.Ctx.Request.Context() + if err := quota.Validate(ctx, qa.quota.Reference, req.Hard); err != nil { qa.SendBadRequestError(err) return } - mgr, err := quota.NewManager(qa.quota.Reference, qa.quota.ReferenceID) - if err != nil { - qa.SendInternalServerError(fmt.Errorf("failed to create quota manager, error: %v", err)) - return - } + qa.quota.SetHard(req.Hard) - if err := mgr.UpdateQuota(req.Hard); err != nil { + if err := quota.Ctl.Update(ctx, qa.quota); err != nil { qa.SendInternalServerError(fmt.Errorf("failed to update hard limits of the quota, error: %v", err)) return } @@ -125,25 +108,25 @@ func (qa *QuotaAPI) List() { return } - query := &models.QuotaQuery{ - Reference: qa.GetString("reference"), - ReferenceID: qa.GetString("reference_id"), - Pagination: models.Pagination{ - Page: page, - Size: size, - }, - Sorting: models.Sorting{ - Sort: qa.GetString("sort"), + query := &q.Query{ + Keywords: q.KeyWords{ + "reference": qa.GetString("reference"), + "reference_id": qa.GetString("reference_id"), }, + PageNumber: page, + PageSize: size, + Sorting: qa.GetString("sort"), } - total, err := dao.GetTotalOfQuotas(query) + ctx := qa.Ctx.Request.Context() + + total, err := quota.Ctl.Count(ctx, query) if err != nil { qa.SendInternalServerError(fmt.Errorf("failed to query database for total of quotas, error: %v", err)) return } - quotas, err := dao.ListQuotas(query) + quotas, err := quota.Ctl.List(ctx, query) if err != nil { qa.SendInternalServerError(fmt.Errorf("failed to query database for quotas, error: %v", err)) return diff --git a/src/core/api/quota_test.go b/src/core/api/quota_test.go index d96fc86a5..630f5583f 100644 --- a/src/core/api/quota_test.go +++ b/src/core/api/quota_test.go @@ -15,32 +15,35 @@ package api import ( + "context" "fmt" "testing" - "github.com/goharbor/harbor/src/common/models" - "github.com/goharbor/harbor/src/common/quota" - "github.com/goharbor/harbor/src/common/quota/driver" - "github.com/goharbor/harbor/src/common/quota/driver/mocks" + o "github.com/astaxie/beego/orm" + "github.com/goharbor/harbor/src/controller/quota" + "github.com/goharbor/harbor/src/lib/orm" + "github.com/goharbor/harbor/src/pkg/quota/driver" "github.com/goharbor/harbor/src/pkg/types" "github.com/goharbor/harbor/src/testing/apitests/apilib" + "github.com/goharbor/harbor/src/testing/mock" + drivertesting "github.com/goharbor/harbor/src/testing/pkg/quota/driver" + "github.com/google/uuid" "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/mock" ) var ( - reference = "mock" + reference = uuid.New().String() hardLimits = types.ResourceList{types.ResourceStorage: -1} ) func init() { - mockDriver := &mocks.Driver{} + mockDriver := &drivertesting.Driver{} mockHardLimitsFn := func() types.ResourceList { return hardLimits } - mockLoadFn := func(key string) driver.RefObject { + mockLoadFn := func(ctx context.Context, key string) driver.RefObject { return driver.RefObject{"id": key} } @@ -53,8 +56,8 @@ func init() { } mockDriver.On("HardLimits").Return(mockHardLimitsFn) - mockDriver.On("Load", mock.AnythingOfType("string")).Return(mockLoadFn, nil) - mockDriver.On("Validate", mock.AnythingOfType("types.ResourceList")).Return(mockValidateFn) + mock.OnAnything(mockDriver, "Load").Return(mockLoadFn, nil) + mock.OnAnything(mockDriver, "Validate").Return(mockValidateFn) driver.Register(reference, mockDriver) } @@ -63,13 +66,19 @@ func TestQuotaAPIList(t *testing.T) { assert := assert.New(t) apiTest := newHarborAPI() + ctx := orm.NewContext(context.TODO(), o.NewOrm()) + var quotaIDs []int64 + defer func() { + for _, quotaID := range quotaIDs { + quota.Ctl.Delete(ctx, quotaID) + } + }() + count := 10 for i := 0; i < count; i++ { - mgr, err := quota.NewManager(reference, fmt.Sprintf("%d", i)) - assert.Nil(err) - - _, err = mgr.NewQuota(hardLimits) + quotaID, err := quota.Ctl.Create(ctx, reference, uuid.New().String(), hardLimits) assert.Nil(err) + quotaIDs = append(quotaIDs, quotaID) } code, quotas, err := apiTest.QuotasGet(&apilib.QuotaQuery{Reference: reference}, *admin) @@ -87,11 +96,10 @@ func TestQuotaAPIGet(t *testing.T) { assert := assert.New(t) apiTest := newHarborAPI() - mgr, err := quota.NewManager(reference, "quota-get") - assert.Nil(err) - - quotaID, err := mgr.NewQuota(hardLimits) + ctx := orm.NewContext(context.TODO(), o.NewOrm()) + quotaID, err := quota.Ctl.Create(ctx, reference, uuid.New().String(), hardLimits) assert.Nil(err) + defer quota.Ctl.Delete(ctx, quotaID) code, quota, err := apiTest.QuotasGetByID(*admin, fmt.Sprintf("%d", quotaID)) assert.Nil(err) @@ -107,22 +115,21 @@ func TestQuotaPut(t *testing.T) { assert := assert.New(t) apiTest := newHarborAPI() - mgr, err := quota.NewManager(reference, "quota-put") - assert.Nil(err) - - quotaID, err := mgr.NewQuota(hardLimits) + ctx := orm.NewContext(context.TODO(), o.NewOrm()) + quotaID, err := quota.Ctl.Create(ctx, reference, uuid.New().String(), hardLimits) assert.Nil(err) + defer quota.Ctl.Delete(ctx, quotaID) code, quota, err := apiTest.QuotasGetByID(*admin, fmt.Sprintf("%d", quotaID)) assert.Nil(err) assert.Equal(int(200), code) assert.Equal(map[string]int64{"storage": -1}, quota.Hard) - code, err = apiTest.QuotasPut(*admin, fmt.Sprintf("%d", quotaID), models.QuotaUpdateRequest{}) + code, err = apiTest.QuotasPut(*admin, fmt.Sprintf("%d", quotaID), QuotaUpdateRequest{}) assert.Nil(err, err) assert.Equal(int(400), code) - code, err = apiTest.QuotasPut(*admin, fmt.Sprintf("%d", quotaID), models.QuotaUpdateRequest{Hard: types.ResourceList{types.ResourceStorage: 100}}) + code, err = apiTest.QuotasPut(*admin, fmt.Sprintf("%d", quotaID), QuotaUpdateRequest{Hard: types.ResourceList{types.ResourceStorage: 100}}) assert.Nil(err) assert.Equal(int(200), code) diff --git a/src/lib/q/query.go b/src/lib/q/query.go index 9add324d0..d30920230 100644 --- a/src/lib/q/query.go +++ b/src/lib/q/query.go @@ -25,6 +25,8 @@ type Query struct { PageSize int64 // List of key words Keywords KeyWords + // Sorting + Sorting string } // New returns Query with keywords diff --git a/src/lib/response_buffer.go b/src/lib/response_buffer.go index edfbd4fe8..70a397f34 100644 --- a/src/lib/response_buffer.go +++ b/src/lib/response_buffer.go @@ -79,7 +79,8 @@ func (r *ResponseBuffer) Flush() (int, error) { // Success checks whether the status code is >= 200 & <= 399 func (r *ResponseBuffer) Success() bool { - return r.code >= http.StatusOK && r.code < http.StatusBadRequest + code := r.StatusCode() + return code >= http.StatusOK && code < http.StatusBadRequest } // Reset reset the response buffer @@ -98,5 +99,10 @@ func (r *ResponseBuffer) Reset() error { // StatusCode returns the status code func (r *ResponseBuffer) StatusCode() int { + if r.code == 0 { + // NOTE: r.code is zero means that `WriteHeader` not called by the http handler, + // so process it as http.StatusOK + return http.StatusOK + } return r.code } diff --git a/src/lib/response_recorder.go b/src/lib/response_recorder.go index 2a71de758..396fe7bf2 100644 --- a/src/lib/response_recorder.go +++ b/src/lib/response_recorder.go @@ -49,5 +49,12 @@ func (r *ResponseRecorder) WriteHeader(statusCode int) { // Success checks whether the status code is >= 200 & <= 399 func (r *ResponseRecorder) Success() bool { - return r.StatusCode >= http.StatusOK && r.StatusCode < http.StatusBadRequest + statusCode := r.StatusCode + if statusCode == 0 { + // NOTE: r.code is zero means that `WriteHeader` not called by the http handler, + // so process it as http.StatusOK + statusCode = http.StatusOK + } + + return statusCode >= http.StatusOK && statusCode < http.StatusBadRequest } diff --git a/src/pkg/quota/dao/dao.go b/src/pkg/quota/dao/dao.go index 76bea61f7..d0c2d89d2 100644 --- a/src/pkg/quota/dao/dao.go +++ b/src/pkg/quota/dao/dao.go @@ -16,15 +16,22 @@ package dao import ( "context" + "fmt" "time" + "github.com/goharbor/harbor/src/lib/log" "github.com/goharbor/harbor/src/lib/orm" + "github.com/goharbor/harbor/src/lib/q" + "github.com/goharbor/harbor/src/pkg/quota/driver" "github.com/goharbor/harbor/src/pkg/quota/models" "github.com/goharbor/harbor/src/pkg/types" ) // DAO the dao for Quota and QuotaUsage type DAO interface { + // Count returns the total count of quotas according to the query. + Count(ctx context.Context, query *q.Query) (int64, error) + // Create create quota for reference object Create(ctx context.Context, reference, referenceID string, hardLimits, used types.ResourceList) (int64, error) @@ -42,6 +49,9 @@ type DAO interface { // Update update quota Update(ctx context.Context, quota *models.Quota) error + + // List list quotas + List(ctx context.Context, query *q.Query) ([]*models.Quota, error) } // New returns an instance of the default DAO @@ -51,6 +61,23 @@ func New() DAO { type dao struct{} +func (d *dao) Count(ctx context.Context, query *q.Query) (int64, error) { + o, err := orm.FromContext(ctx) + if err != nil { + return 0, err + } + + condition, params := listConditions(query) + sql := fmt.Sprintf("SELECT COUNT(1) FROM quota AS a JOIN quota_usage AS b ON a.id = b.id %s", condition) + + var count int64 + if err := o.Raw(sql, params).QueryRow(&count); err != nil { + return 0, err + } + + return count, nil +} + func (d *dao) Create(ctx context.Context, reference, referenceID string, hardLimits, used types.ResourceList) (int64, error) { o, err := orm.FromContext(ctx) if err != nil { @@ -192,6 +219,67 @@ func (d *dao) Update(ctx context.Context, quota *models.Quota) error { return nil } +func (d *dao) List(ctx context.Context, query *q.Query) ([]*models.Quota, error) { + o, err := orm.FromContext(ctx) + if err != nil { + return nil, err + } + + condition, params := listConditions(query) + + sql := fmt.Sprintf(` +SELECT + a.id, + a.reference, + a.reference_id, + a.hard, + b.used, + b.creation_time, + b.update_time +FROM + quota AS a + JOIN quota_usage AS b ON a.id = b.id %s`, condition) + + orderBy := listOrderBy(query) + if orderBy != "" { + sql += ` order by ` + orderBy + } + + if query != nil { + page, size := query.PageNumber, query.PageSize + if size > 0 { + sql += ` limit ?` + params = append(params, size) + if page > 0 { + sql += ` offset ?` + params = append(params, size*(page-1)) + } + } + } + + var quotas []*models.Quota + if _, err := o.Raw(sql, params).QueryRows("as); err != nil { + return nil, err + } + + for _, quota := range quotas { + d, ok := driver.Get(quota.Reference) + if !ok { + continue + } + + ref, err := d.Load(ctx, quota.ReferenceID) + if err != nil { + log.Warning(fmt.Sprintf("Load quota reference object (%s, %s) failed: %v", quota.Reference, quota.ReferenceID, err)) + continue + } + + quota.Ref = ref + } + + return quotas, nil +} + func toQuota(quota *Quota, usage *QuotaUsage) *models.Quota { return &models.Quota{ ID: quota.ID, diff --git a/src/pkg/quota/dao/dao_test.go b/src/pkg/quota/dao/dao_test.go index 2d6475e1b..55b49c406 100644 --- a/src/pkg/quota/dao/dao_test.go +++ b/src/pkg/quota/dao/dao_test.go @@ -20,8 +20,10 @@ import ( "testing" "github.com/goharbor/harbor/src/lib/orm" + "github.com/goharbor/harbor/src/lib/q" "github.com/goharbor/harbor/src/pkg/types" htesting "github.com/goharbor/harbor/src/testing" + "github.com/google/uuid" "github.com/stretchr/testify/suite" ) @@ -39,6 +41,42 @@ func (suite *DaoTestSuite) SetupSuite() { suite.dao = New() } +func (suite *DaoTestSuite) TestCount() { + suite.Suite.TearDownSuite() // Clean other quotas + + reference := uuid.New().String() + hardLimits := types.ResourceList{types.ResourceStorage: 100} + usage := types.ResourceList{types.ResourceStorage: 0} + + ctx := suite.Context() + + suite.dao.Create(ctx, reference, "1", types.ResourceList{types.ResourceStorage: 200}, usage) + suite.dao.Create(ctx, reference, "2", hardLimits, usage) + suite.dao.Create(ctx, reference, "3", hardLimits, usage) + suite.dao.Create(ctx, uuid.New().String(), "4", types.ResourceList{types.ResourceStorage: 10}, usage) + + { + // Count all the quotas + count, err := suite.dao.Count(ctx, nil) + suite.Nil(err) + suite.Equal(int64(5), count) // 4 + library project quota + } + + { + // Count quotas filter by reference + count, err := suite.dao.Count(ctx, q.New(q.KeyWords{"reference": reference})) + suite.Nil(err) + suite.Equal(int64(3), count) + } + + { + // Count quotas filter by reference ids + count, err := suite.dao.Count(ctx, q.New(q.KeyWords{"reference": reference, "reference_ids": []string{"1", "2"}})) + suite.Nil(err) + suite.Equal(int64(2), count) + } +} + func (suite *DaoTestSuite) TestCreate() { hardLimits := types.ResourceList{types.ResourceStorage: 100} usage := types.ResourceList{types.ResourceStorage: 0} @@ -171,6 +209,60 @@ func (suite *DaoTestSuite) TestUpdate() { } } +func (suite *DaoTestSuite) TestList() { + suite.Suite.TearDownSuite() // Clean other quotas + + reference := uuid.New().String() + hardLimits := types.ResourceList{types.ResourceStorage: 100} + usage := types.ResourceList{types.ResourceStorage: 0} + + ctx := suite.Context() + + suite.dao.Create(ctx, reference, "1", types.ResourceList{types.ResourceStorage: 200}, usage) + suite.dao.Create(ctx, reference, "2", hardLimits, usage) + suite.dao.Create(ctx, reference, "3", hardLimits, usage) + suite.dao.Create(ctx, uuid.New().String(), "4", types.ResourceList{types.ResourceStorage: 10}, usage) + + { + // List all the quotas + quotas, err := suite.dao.List(ctx, nil) + suite.Nil(err) + suite.Equal(5, len(quotas)) // 4 + library project quota + suite.NotEqual(reference, quotas[0].Reference) + suite.Equal("4", quotas[0].ReferenceID) + } + + { + // List quotas filter by reference + quotas, err := suite.dao.List(ctx, q.New(q.KeyWords{"reference": reference})) + suite.Nil(err) + suite.Equal(3, len(quotas)) + } + + { + // List quotas filter by reference ids + quotas, err := suite.dao.List(ctx, q.New(q.KeyWords{"reference": reference, "reference_ids": []string{"1", "2"}})) + suite.Nil(err) + suite.Equal(2, len(quotas)) + } + + { + // List quotas by pagination + quotas, err := suite.dao.List(ctx, &q.Query{PageSize: 2}) + suite.Nil(err) + suite.Equal(2, len(quotas)) + } + + { + // List quotas by sorting + quotas, err := suite.dao.List(ctx, &q.Query{Keywords: q.KeyWords{"reference": reference}, Sorting: "-hard.storage"}) + suite.Nil(err) + suite.Equal(reference, quotas[0].Reference) + suite.Equal("1", quotas[0].ReferenceID) + } + +} + func TestDaoTestSuite(t *testing.T) { suite.Run(t, &DaoTestSuite{}) } diff --git a/src/pkg/quota/dao/model.go b/src/pkg/quota/dao/model.go index 11965ca84..a9b83a739 100644 --- a/src/pkg/quota/dao/model.go +++ b/src/pkg/quota/dao/model.go @@ -15,13 +15,42 @@ package dao import ( - "github.com/goharbor/harbor/src/common/models" + "time" + + "github.com/astaxie/beego/orm" ) -// TODO: move Quota and QuotaUsage models to here +func init() { + orm.RegisterModel(&Quota{}) + orm.RegisterModel(&QuotaUsage{}) +} -// Quota quota model alias from models -type Quota = models.Quota +// Quota model for quota +type Quota struct { + ID int64 `orm:"pk;auto;column(id)" json:"id"` + Reference string `orm:"column(reference)" json:"reference"` // The reference type for quota, eg: project, user + ReferenceID string `orm:"column(reference_id)" json:"reference_id"` + Hard string `orm:"column(hard);type(jsonb)" json:"-"` + CreationTime time.Time `orm:"column(creation_time);auto_now_add" json:"creation_time"` + UpdateTime time.Time `orm:"column(update_time);auto_now" json:"update_time"` +} -// QuotaUsage quota usage model alias from models -type QuotaUsage = models.QuotaUsage +// TableName returns table name for orm +func (q *Quota) TableName() string { + return "quota" +} + +// QuotaUsage model for quota usage +type QuotaUsage struct { + ID int64 `orm:"pk;auto;column(id)" json:"id"` + Reference string `orm:"column(reference)" json:"reference"` // The reference type for quota usage, eg: project, user + ReferenceID string `orm:"column(reference_id)" json:"reference_id"` + Used string `orm:"column(used);type(jsonb)" json:"-"` + CreationTime time.Time `orm:"column(creation_time);auto_now_add" json:"creation_time"` + UpdateTime time.Time `orm:"column(update_time);auto_now" json:"update_time"` +} + +// TableName returns table name for orm +func (qu *QuotaUsage) TableName() string { + return "quota_usage" +} diff --git a/src/pkg/quota/dao/util.go b/src/pkg/quota/dao/util.go new file mode 100644 index 000000000..7f0c2448d --- /dev/null +++ b/src/pkg/quota/dao/util.go @@ -0,0 +1,117 @@ +// 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 dao + +import ( + "encoding/json" + "fmt" + "strings" + + "github.com/goharbor/harbor/src/lib/orm" + "github.com/goharbor/harbor/src/lib/q" + "github.com/goharbor/harbor/src/pkg/types" +) + +var ( + quotaOrderMap = map[string]string{ + "creation_time": "b.creation_time asc", + "+creation_time": "b.creation_time asc", + "-creation_time": "b.creation_time desc", + "update_time": "b.update_time asc", + "+update_time": "b.update_time asc", + "-update_time": "b.update_time desc", + } +) + +type listQuery struct { + ID int64 `json:"id"` + Reference string `json:"reference"` + ReferenceID string `json:"reference_id"` + ReferenceIDs []string `json:"reference_ids"` +} + +func listConditions(query *q.Query) (string, []interface{}) { + params := []interface{}{} + sql := "" + if query == nil { + return sql, params + } + + sql += `WHERE 1=1 ` + + var q listQuery + + bytes, err := json.Marshal(query.Keywords) + if err == nil { + json.Unmarshal(bytes, &q) + } + + if q.ID != 0 { + sql += `AND a.id = ? ` + params = append(params, q.ID) + } + if q.Reference != "" { + sql += `AND a.reference = ? ` + params = append(params, q.Reference) + } + if q.ReferenceID != "" { + sql += `AND a.reference_id = ? ` + params = append(params, q.ReferenceID) + } + + if len(q.ReferenceIDs) != 0 { + sql += fmt.Sprintf(`AND a.reference_id IN (%s) `, orm.ParamPlaceholderForIn(len(q.ReferenceIDs))) + params = append(params, q.ReferenceIDs) + } + + return sql, params +} + +func castQuantity(field string) string { + // cast -1 to max int64 when order by field + return fmt.Sprintf("CAST( (CASE WHEN (%[1]s) IS NULL THEN '0' WHEN (%[1]s) = '-1' THEN '9223372036854775807' ELSE (%[1]s) END) AS BIGINT )", field) +} + +func listOrderBy(query *q.Query) string { + orderBy := "b.creation_time DESC" + + if query != nil && query.Sorting != "" { + if val, ok := quotaOrderMap[query.Sorting]; ok { + orderBy = val + } else { + sort := query.Sorting + + order := "ASC" + if sort[0] == '-' { + order = "DESC" + sort = sort[1:] + } + + prefixes := []string{"hard.", "used."} + for _, prefix := range prefixes { + if strings.HasPrefix(sort, prefix) { + resource := strings.TrimPrefix(sort, prefix) + if types.IsValidResource(types.ResourceName(resource)) { + field := fmt.Sprintf("%s->>'%s'", strings.TrimSuffix(prefix, "."), resource) + orderBy = fmt.Sprintf("(%s) %s", castQuantity(field), order) + break + } + } + } + } + } + + return orderBy +} diff --git a/src/pkg/quota/dao/util_test.go b/src/pkg/quota/dao/util_test.go new file mode 100644 index 000000000..7a26df194 --- /dev/null +++ b/src/pkg/quota/dao/util_test.go @@ -0,0 +1,52 @@ +// 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 dao + +import ( + "testing" + + "github.com/goharbor/harbor/src/lib/q" +) + +func Test_listOrderBy(t *testing.T) { + query := func(sort string) *q.Query { + return &q.Query{ + Sorting: sort, + } + } + + type args struct { + query *q.Query + } + tests := []struct { + name string + args args + want string + }{ + {"no query", args{nil}, "b.creation_time DESC"}, + {"order by unsupported field", args{query("unknown")}, "b.creation_time DESC"}, + {"order by storage of hard", args{query("hard.storage")}, "(CAST( (CASE WHEN (hard->>'storage') IS NULL THEN '0' WHEN (hard->>'storage') = '-1' THEN '9223372036854775807' ELSE (hard->>'storage') END) AS BIGINT )) ASC"}, + {"order by unsupported hard resource", args{query("hard.unknown")}, "b.creation_time DESC"}, + {"order by storage of used", args{query("used.storage")}, "(CAST( (CASE WHEN (used->>'storage') IS NULL THEN '0' WHEN (used->>'storage') = '-1' THEN '9223372036854775807' ELSE (used->>'storage') END) AS BIGINT )) ASC"}, + {"order by unsupported used resource", args{query("used.unknown")}, "b.creation_time DESC"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := listOrderBy(tt.args.query); got != tt.want { + t.Errorf("listOrderBy() = %v, want %v", got, tt.want) + } + }) + } +} diff --git a/src/pkg/quota/manager.go b/src/pkg/quota/manager.go index 03cea9f03..b8d5f66a5 100644 --- a/src/pkg/quota/manager.go +++ b/src/pkg/quota/manager.go @@ -18,6 +18,7 @@ import ( "context" "github.com/goharbor/harbor/src/lib/orm" + "github.com/goharbor/harbor/src/lib/q" "github.com/goharbor/harbor/src/pkg/quota/dao" "github.com/goharbor/harbor/src/pkg/quota/models" "github.com/goharbor/harbor/src/pkg/types" @@ -31,6 +32,9 @@ type Manager interface { // Create create quota for the reference object Create(ctx context.Context, reference, referenceID string, hardLimits types.ResourceList, usages ...types.ResourceList) (int64, error) + // Count returns the total count of quotas according to the query. + Count(ctx context.Context, query *q.Query) (int64, error) + // Delete delete quota by id Delete(ctx context.Context, id int64) error @@ -45,6 +49,9 @@ type Manager interface { // Update update quota Update(ctx context.Context, quota *Quota) error + + // List list quotas + List(ctx context.Context, query *q.Query) ([]*Quota, error) } var ( @@ -75,8 +82,16 @@ func (m *manager) Create(ctx context.Context, reference, referenceID string, har return id, err } +func (m *manager) Count(ctx context.Context, query *q.Query) (int64, error) { + return m.dao.Count(ctx, query) +} + func (m *manager) Delete(ctx context.Context, id int64) error { - return m.dao.Delete(ctx, id) + h := func(ctx context.Context) error { + return m.dao.Delete(ctx, id) + } + + return orm.WithTransaction(h)(ctx) } func (m *manager) Get(ctx context.Context, id int64) (*Quota, error) { @@ -92,7 +107,15 @@ func (m *manager) GetByRefForUpdate(ctx context.Context, reference, referenceID } func (m *manager) Update(ctx context.Context, q *Quota) error { - return m.dao.Update(ctx, q) + h := func(ctx context.Context) error { + return m.dao.Update(ctx, q) + } + + return orm.WithTransaction(h)(ctx) +} + +func (m *manager) List(ctx context.Context, query *q.Query) ([]*Quota, error) { + return m.dao.List(ctx, query) } // NewManager returns quota manager diff --git a/src/testing/common/security/context.go b/src/testing/common/security/context.go index 4290409e3..575965cc2 100644 --- a/src/testing/common/security/context.go +++ b/src/testing/common/security/context.go @@ -3,7 +3,6 @@ package security import ( - models "github.com/goharbor/harbor/src/common/models" mock "github.com/stretchr/testify/mock" types "github.com/goharbor/harbor/src/pkg/permission/types" @@ -28,45 +27,6 @@ func (_m *Context) Can(action types.Action, resource types.Resource) bool { return r0 } -// GetMyProjects provides a mock function with given fields: -func (_m *Context) GetMyProjects() ([]*models.Project, error) { - ret := _m.Called() - - var r0 []*models.Project - if rf, ok := ret.Get(0).(func() []*models.Project); ok { - r0 = rf() - } else { - if ret.Get(0) != nil { - r0 = ret.Get(0).([]*models.Project) - } - } - - var r1 error - if rf, ok := ret.Get(1).(func() error); ok { - r1 = rf() - } else { - r1 = ret.Error(1) - } - - return r0, r1 -} - -// GetProjectRoles provides a mock function with given fields: projectIDOrName -func (_m *Context) GetProjectRoles(projectIDOrName interface{}) []int { - ret := _m.Called(projectIDOrName) - - var r0 []int - if rf, ok := ret.Get(0).(func(interface{}) []int); ok { - r0 = rf(projectIDOrName) - } else { - if ret.Get(0) != nil { - r0 = ret.Get(0).([]int) - } - } - - return r0 -} - // GetUsername provides a mock function with given fields: func (_m *Context) GetUsername() string { ret := _m.Called() diff --git a/src/testing/controller/controller.go b/src/testing/controller/controller.go index 0a062997a..cb965b427 100644 --- a/src/testing/controller/controller.go +++ b/src/testing/controller/controller.go @@ -16,7 +16,6 @@ package controller //go:generate mockery -case snake -dir ../../controller/artifact -name Controller -output ./artifact -outpkg artifact //go:generate mockery -case snake -dir ../../controller/blob -name Controller -output ./blob -outpkg blob -//go:generate mockery -case snake -dir ../../controller/chartmuseum -name Controller -output ./chartmuseum -outpkg chartmuseum //go:generate mockery -case snake -dir ../../controller/project -name Controller -output ./project -outpkg project //go:generate mockery -case snake -dir ../../controller/quota -name Controller -output ./quota -outpkg quota //go:generate mockery -case snake -dir ../../controller/scan -name Controller -output ./scan -outpkg scan diff --git a/src/testing/controller/quota/controller.go b/src/testing/controller/quota/controller.go index 95331467b..f20593a3c 100644 --- a/src/testing/controller/quota/controller.go +++ b/src/testing/controller/quota/controller.go @@ -8,6 +8,8 @@ import ( models "github.com/goharbor/harbor/src/pkg/quota/models" mock "github.com/stretchr/testify/mock" + q "github.com/goharbor/harbor/src/lib/q" + quota "github.com/goharbor/harbor/src/controller/quota" types "github.com/goharbor/harbor/src/pkg/types" @@ -18,6 +20,27 @@ type Controller struct { mock.Mock } +// Count provides a mock function with given fields: ctx, query +func (_m *Controller) Count(ctx context.Context, query *q.Query) (int64, error) { + ret := _m.Called(ctx, query) + + var r0 int64 + if rf, ok := ret.Get(0).(func(context.Context, *q.Query) int64); ok { + r0 = rf(ctx, query) + } else { + r0 = ret.Get(0).(int64) + } + + var r1 error + if rf, ok := ret.Get(1).(func(context.Context, *q.Query) error); ok { + r1 = rf(ctx, query) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + // Create provides a mock function with given fields: ctx, reference, referenceID, hardLimits, used func (_m *Controller) Create(ctx context.Context, reference string, referenceID string, hardLimits types.ResourceList, used ...types.ResourceList) (int64, error) { _va := make([]interface{}, len(used)) @@ -127,6 +150,29 @@ func (_m *Controller) IsEnabled(ctx context.Context, reference string, reference return r0, r1 } +// List provides a mock function with given fields: ctx, query +func (_m *Controller) List(ctx context.Context, query *q.Query) ([]*models.Quota, error) { + ret := _m.Called(ctx, query) + + var r0 []*models.Quota + if rf, ok := ret.Get(0).(func(context.Context, *q.Query) []*models.Quota); ok { + r0 = rf(ctx, query) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*models.Quota) + } + } + + var r1 error + if rf, ok := ret.Get(1).(func(context.Context, *q.Query) error); ok { + r1 = rf(ctx, query) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + // Refresh provides a mock function with given fields: ctx, reference, referenceID, options func (_m *Controller) Refresh(ctx context.Context, reference string, referenceID string, options ...quota.Option) error { _va := make([]interface{}, len(options)) @@ -161,3 +207,17 @@ func (_m *Controller) Request(ctx context.Context, reference string, referenceID return r0 } + +// Update provides a mock function with given fields: ctx, _a1 +func (_m *Controller) Update(ctx context.Context, _a1 *models.Quota) error { + ret := _m.Called(ctx, _a1) + + var r0 error + if rf, ok := ret.Get(0).(func(context.Context, *models.Quota) error); ok { + r0 = rf(ctx, _a1) + } else { + r0 = ret.Error(0) + } + + return r0 +} diff --git a/src/testing/pkg/quota/manager.go b/src/testing/pkg/quota/manager.go index 1b654fd75..155820087 100644 --- a/src/testing/pkg/quota/manager.go +++ b/src/testing/pkg/quota/manager.go @@ -8,6 +8,8 @@ import ( models "github.com/goharbor/harbor/src/pkg/quota/models" mock "github.com/stretchr/testify/mock" + q "github.com/goharbor/harbor/src/lib/q" + types "github.com/goharbor/harbor/src/pkg/types" ) @@ -16,6 +18,27 @@ type Manager struct { mock.Mock } +// Count provides a mock function with given fields: ctx, query +func (_m *Manager) Count(ctx context.Context, query *q.Query) (int64, error) { + ret := _m.Called(ctx, query) + + var r0 int64 + if rf, ok := ret.Get(0).(func(context.Context, *q.Query) int64); ok { + r0 = rf(ctx, query) + } else { + r0 = ret.Get(0).(int64) + } + + var r1 error + if rf, ok := ret.Get(1).(func(context.Context, *q.Query) error); ok { + r1 = rf(ctx, query) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + // Create provides a mock function with given fields: ctx, reference, referenceID, hardLimits, usages func (_m *Manager) Create(ctx context.Context, reference string, referenceID string, hardLimits types.ResourceList, usages ...types.ResourceList) (int64, error) { _va := make([]interface{}, len(usages)) @@ -127,6 +150,29 @@ func (_m *Manager) GetByRefForUpdate(ctx context.Context, reference string, refe return r0, r1 } +// List provides a mock function with given fields: ctx, query +func (_m *Manager) List(ctx context.Context, query *q.Query) ([]*models.Quota, error) { + ret := _m.Called(ctx, query) + + var r0 []*models.Quota + if rf, ok := ret.Get(0).(func(context.Context, *q.Query) []*models.Quota); ok { + r0 = rf(ctx, query) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*models.Quota) + } + } + + var r1 error + if rf, ok := ret.Get(1).(func(context.Context, *q.Query) error); ok { + r1 = rf(ctx, query) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + // Update provides a mock function with given fields: ctx, _a1 func (_m *Manager) Update(ctx context.Context, _a1 *models.Quota) error { ret := _m.Called(ctx, _a1) diff --git a/src/testing/suite.go b/src/testing/suite.go index b418a6a23..f6bd2bb13 100644 --- a/src/testing/suite.go +++ b/src/testing/suite.go @@ -16,12 +16,10 @@ package testing import ( "context" - "fmt" "io" "math/rand" "net/http" "net/http/httptest" - "strconv" "sync" "time" @@ -31,7 +29,6 @@ import ( "github.com/goharbor/harbor/src/core/config" "github.com/goharbor/harbor/src/lib/errors" "github.com/goharbor/harbor/src/lib/orm" - "github.com/goharbor/harbor/src/pkg/types" "github.com/opencontainers/go-digest" "github.com/stretchr/testify/suite" ) @@ -161,14 +158,3 @@ func (suite *Suite) ExecSQL(query string, args ...interface{}) { func (suite *Suite) IsNotFoundErr(err error) bool { return suite.True(errors.IsNotFoundErr(err)) } - -// AssertResourceUsage ... -func (suite *Suite) AssertResourceUsage(expected int64, resource types.ResourceName, projectID int64) { - usage := models.QuotaUsage{Reference: "project", ReferenceID: strconv.FormatInt(projectID, 10)} - err := dao.GetOrmer().Read(&usage, "reference", "reference_id") - suite.Nil(err, fmt.Sprintf("Failed to get resource %s usage of project %d, error: %v", resource, projectID, err)) - - used, err := types.NewResourceList(usage.Used) - suite.Nil(err, "Bad resource usage of project %d", projectID) - suite.Equal(expected, used[resource]) -}