Upstream/harbor
1
0
Fork 0
mirror of https://github.com/goharbor/harbor.git synced 2024-12-24 09:38:09 +01:00
Code Issues Projects Releases Wiki Activity

refactor(quota): cleanup code for quota

Browse Source 
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 <hweiwei@vmware.com>
This commit is contained in:
He Weiwei 2020-04-12 16:14:12 +00:00
parent 93f316ccfe
commit c0349da812
40 changed files with 677 additions and 2498 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
src/common/dao/base.go
View File

@ -169,29 +169,3 @@ func Escape(str string) string {
str = strings.Replace(str, `_`, `\_`, -1) str = strings.Replace(str, `_`, `\_`, -1)
return str 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
}

50
src/common/dao/dao_test.go
View File

@ -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.True(t, IsDupRecErr(fmt.Errorf("pq: duplicate key value violates unique constraint \"properties_k_key\"")))
assert.False(t, IsDupRecErr(fmt.Errorf("other error"))) 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(&quota)
return fmt.Errorf("failed")
}
var quotaID int64
success := func(o orm.Ormer) error {
id, err := o.Insert(&quota)
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(&quota, "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(&quota, "reference", "reference_id")
assert.Nil(err)
assert.True(quota.ID != 0)
assert.Equal(quotaID, quota.ID)
GetOrmer().Delete(&models.Quota{ID: quotaID}, "id")
}
}

238
src/common/dao/quota.go
View File

@ -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(&quota)
}
// 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(&quota)
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(&quotas); 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
}

174
src/common/dao/quota_test.go
View File

@ -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(&quota1, 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)
}
})
}
}

149
src/common/dao/quota_usage.go
View File

@ -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(&quotaUsage)
}
// 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(&quotaUsage)
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(&quotaUsages); 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
}

164
src/common/dao/quota_usage_test.go
View File

@ -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(&quotaUsage1, 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)
}
})
}
}

2
src/common/models/base.go
View File

@ -38,7 +38,5 @@ func init() {
new(ProjectBlob), new(ProjectBlob),
new(ArtifactAndBlob), new(ArtifactAndBlob),
new(CVEWhitelist), new(CVEWhitelist),
new(Quota),
new(QuotaUsage),
) )
} }

85
src/common/models/quota.go
View File

@ -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"`
}

77
src/common/models/quota_usage.go
View File

@ -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
}

59
src/common/quota/driver/driver.go
View File

@ -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
}

65
src/common/quota/driver/mocks/driver.go
View File

@ -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
}

158
src/common/quota/driver/project/driver.go
View File

@ -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}
}

75
src/common/quota/driver/project/driver_test.go
View File

@ -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))
}

111
src/common/quota/errors.go
View File

@ -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}
}

287
src/common/quota/manager.go
View File

@ -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
}

359
src/common/quota/manager_test.go
View File

@ -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()
}

35
src/common/quota/quota.go
View File

@ -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)
}

45
src/common/quota/quota_test.go
View File

@ -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)
}
})
}
}

30
src/common/quota/types.go
View File

@ -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

57
src/common/quota/util.go
View File

@ -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, ",")
}

83
src/common/quota/util_test.go
View File

@ -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)
}
})
}
}

54
src/controller/quota/controller.go
View File

@ -24,6 +24,7 @@ import (
"github.com/goharbor/harbor/src/lib/errors" "github.com/goharbor/harbor/src/lib/errors"
"github.com/goharbor/harbor/src/lib/log" "github.com/goharbor/harbor/src/lib/log"
"github.com/goharbor/harbor/src/lib/orm" "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"
"github.com/goharbor/harbor/src/pkg/quota/driver" "github.com/goharbor/harbor/src/pkg/quota/driver"
"github.com/goharbor/harbor/src/pkg/types" "github.com/goharbor/harbor/src/pkg/types"
@ -44,6 +45,9 @@ var (
// Controller defines the operations related with quotas // Controller defines the operations related with quotas
type Controller interface { 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 ensure quota for the reference object
Create(ctx context.Context, reference, referenceID string, hardLimits types.ResourceList, used ...types.ResourceList) (int64, error) 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 returns true when quota enabled for reference object
IsEnabled(ctx context.Context, reference, referenceID string) (bool, error) 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 refresh quota for the reference object
Refresh(ctx context.Context, reference, referenceID string, options ...Option) error 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， // then runs f and refresh quota when f success，
// in the finally it releases the resources which reserved at the beginning. // 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 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 // NewController creates an instance of the default quota controller
@ -83,6 +93,10 @@ type controller struct {
quotaMgr quota.Manager 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) { 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...) 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) 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) { func (c *controller) getReservedResources(ctx context.Context, reference, referenceID string) (types.ResourceList, error) {
conn := util.DefaultPool().Get() conn := util.DefaultPool().Get()
defer conn.Close() defer conn.Close()
@ -283,6 +301,32 @@ func (c *controller) Request(ctx context.Context, reference, referenceID string,
return c.Refresh(ctx, reference, referenceID) 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 // Driver returns quota driver for the reference
func Driver(ctx context.Context, reference string) (driver.Driver, error) { func Driver(ctx context.Context, reference string) (driver.Driver, error) {
d, ok := driver.Get(reference) d, ok := driver.Get(reference)
@ -293,6 +337,16 @@ func Driver(ctx context.Context, reference string) (driver.Driver, error) {
return d, nil 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 { func reservedResourcesKey(reference, referenceID string) string {
return fmt.Sprintf("quota:%s:%s:reserved", reference, referenceID) return fmt.Sprintf("quota:%s:%s:reserved", reference, referenceID)
} }

10
src/core/api/harborapi_test.go
View File

@ -28,8 +28,6 @@ import (
"strconv" "strconv"
"strings" "strings"
"github.com/goharbor/harbor/src/server/middleware/security"
"github.com/astaxie/beego" "github.com/astaxie/beego"
"github.com/dghubble/sling" "github.com/dghubble/sling"
"github.com/goharbor/harbor/src/common/api" "github.com/goharbor/harbor/src/common/api"
@ -44,6 +42,9 @@ import (
"github.com/goharbor/harbor/src/core/config" "github.com/goharbor/harbor/src/core/config"
"github.com/goharbor/harbor/src/pkg/notification" "github.com/goharbor/harbor/src/pkg/notification"
"github.com/goharbor/harbor/src/replication/model" "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" "github.com/goharbor/harbor/src/testing/apitests/apilib"
) )
@ -217,7 +218,8 @@ func init() {
mockServer := test.NewJobServiceServer() mockServer := test.NewJobServiceServer()
defer mockServer.Close() 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) { 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 // 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 path := "/api/quotas/" + quotaID
_sling := sling.New().Put(a.basePath).Path(path).BodyJSON(req) _sling := sling.New().Put(a.basePath).Path(path).BodyJSON(req)

53
src/core/api/project.go
View File

@ -15,6 +15,7 @@
package api package api
import ( import (
"context"
"fmt" "fmt"
"net/http" "net/http"
"regexp" "regexp"
@ -26,12 +27,12 @@ import (
"github.com/goharbor/harbor/src/common/dao" "github.com/goharbor/harbor/src/common/dao"
pro "github.com/goharbor/harbor/src/common/dao/project" pro "github.com/goharbor/harbor/src/common/dao/project"
"github.com/goharbor/harbor/src/common/models" "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/rbac"
"github.com/goharbor/harbor/src/common/security/local" "github.com/goharbor/harbor/src/common/security/local"
"github.com/goharbor/harbor/src/common/utils" "github.com/goharbor/harbor/src/common/utils"
errutil "github.com/goharbor/harbor/src/common/utils/error" errutil "github.com/goharbor/harbor/src/common/utils/error"
"github.com/goharbor/harbor/src/controller/event/metadata" "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/core/config"
"github.com/goharbor/harbor/src/lib/log" "github.com/goharbor/harbor/src/lib/log"
evt "github.com/goharbor/harbor/src/pkg/notifier/event" evt "github.com/goharbor/harbor/src/pkg/notifier/event"
@ -138,7 +139,7 @@ func (p *ProjectAPI) Post() {
pro.StorageLimit = &setting.StoragePerProject pro.StorageLimit = &setting.StoragePerProject
} }
hardLimits, err = projectQuotaHardLimits(pro, setting) hardLimits, err = projectQuotaHardLimits(p.Ctx.Request.Context(), pro, setting)
if err != nil { if err != nil {
log.Errorf("Invalid project request, error: %v", err) log.Errorf("Invalid project request, error: %v", err)
p.SendBadRequestError(fmt.Errorf("invalid request: %v", err)) p.SendBadRequestError(fmt.Errorf("invalid request: %v", err))
@ -201,12 +202,9 @@ func (p *ProjectAPI) Post() {
} }
if config.QuotaPerProjectEnable() { if config.QuotaPerProjectEnable() {
quotaMgr, err := quota.NewManager("project", strconv.FormatInt(projectID, 10)) ctx := p.Ctx.Request.Context()
if err != nil { referenceID := quota.ReferenceID(projectID)
p.SendInternalServerError(fmt.Errorf("failed to get quota manager: %v", err)) if _, err := quota.Ctl.Create(ctx, quota.ProjectReference, referenceID, hardLimits); err != nil {
return
}
if _, err := quotaMgr.NewQuota(hardLimits); err != nil {
p.SendInternalServerError(fmt.Errorf("failed to create quota for project: %v", err)) p.SendInternalServerError(fmt.Errorf("failed to create quota for project: %v", err))
return return
} }
@ -285,15 +283,17 @@ func (p *ProjectAPI) Delete() {
return 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 { if err != nil {
p.SendInternalServerError(fmt.Errorf("failed to get quota manager: %v", err)) log.Warningf("failed to get quota for project %s, error: %v", p.project.Name, err)
return } else {
} if err := quota.Ctl.Delete(ctx, q.ID); err != nil {
if err := quotaMgr.DeleteQuota(); err != nil {
p.SendInternalServerError(fmt.Errorf("failed to delete quota for project: %v", err)) p.SendInternalServerError(fmt.Errorf("failed to delete quota for project: %v", err))
return return
} }
}
// fire event // fire event
evt.BuildAndPublish(&metadata.DeleteProjectEventMetadata{ evt.BuildAndPublish(&metadata.DeleteProjectEventMetadata{
@ -516,7 +516,7 @@ func (p *ProjectAPI) Summary() {
ChartCount: p.project.ChartCount, 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 { if hasPerm, _ := p.HasProjectPermission(p.project.ProjectID, rbac.ActionRead, rbac.ResourceQuota); hasPerm {
fetchSummaries = append(fetchSummaries, getProjectQuotaSummary) fetchSummaries = append(fetchSummaries, getProjectQuotaSummary)
@ -526,6 +526,8 @@ func (p *ProjectAPI) Summary() {
fetchSummaries = append(fetchSummaries, getProjectMemberSummary) fetchSummaries = append(fetchSummaries, getProjectMemberSummary)
} }
ctx := p.Ctx.Request.Context()
var wg sync.WaitGroup var wg sync.WaitGroup
for _, fn := range fetchSummaries { for _, fn := range fetchSummaries {
fn := fn fn := fn
@ -533,7 +535,7 @@ func (p *ProjectAPI) Summary() {
wg.Add(1) wg.Add(1)
go func() { go func() {
defer wg.Done() defer wg.Done()
fn(p.project.ProjectID, summary) fn(ctx, p.project.ProjectID, summary)
}() }()
} }
wg.Wait() wg.Wait()
@ -563,7 +565,7 @@ func validateProjectReq(req *models.ProjectRequest) error {
return nil 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{} hardLimits := types.ResourceList{}
if req.StorageLimit != nil { if req.StorageLimit != nil {
@ -572,38 +574,31 @@ func projectQuotaHardLimits(req *models.ProjectRequest, setting *models.QuotaSet
hardLimits[types.ResourceStorage] = setting.StoragePerProject 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 nil, err
} }
return hardLimits, nil return hardLimits, nil
} }
func getProjectQuotaSummary(projectID int64, summary *models.ProjectSummary) { func getProjectQuotaSummary(ctx context.Context, projectID int64, summary *models.ProjectSummary) {
if !config.QuotaPerProjectEnable() { if !config.QuotaPerProjectEnable() {
log.Debug("Quota per project disabled") log.Debug("Quota per project disabled")
return 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 { if err != nil {
log.Debugf("failed to get quota for project: %d", projectID) log.Debugf("failed to get quota for project: %d", projectID)
return return
} }
if len(quotas) == 0 {
log.Debugf("quota not found for project: %d", projectID)
return
}
quota := quotas[0]
summary.Quota = &models.QuotaSummary{} summary.Quota = &models.QuotaSummary{}
summary.Quota.Hard, _ = types.NewResourceList(quota.Hard) summary.Quota.Hard, _ = types.NewResourceList(q.Hard)
summary.Quota.Used, _ = types.NewResourceList(quota.Used) 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 var wg sync.WaitGroup
for _, e := range []struct { for _, e := range []struct {

73
src/core/api/quota.go
View File

@ -17,12 +17,18 @@ package api
import ( import (
"fmt" "fmt"
"github.com/goharbor/harbor/src/common/dao" "github.com/goharbor/harbor/src/controller/quota"
"github.com/goharbor/harbor/src/common/models" "github.com/goharbor/harbor/src/lib/q"
"github.com/goharbor/harbor/src/common/quota" "github.com/goharbor/harbor/src/pkg/quota/models"
"github.com/goharbor/harbor/src/pkg/types"
"github.com/pkg/errors" "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/ // QuotaAPI handles request to /api/quotas/
type QuotaAPI struct { type QuotaAPI struct {
BaseController BaseController
@ -56,14 +62,9 @@ func (qa *QuotaAPI) Prepare() {
return return
} }
quota, err := dao.GetQuota(id) quota, err := quota.Ctl.Get(qa.Ctx.Request.Context(), id)
if err != nil { if err != nil {
qa.SendInternalServerError(fmt.Errorf("failed to get quota %d, error: %v", id, err)) qa.SendError(err)
return
}
if quota == nil {
qa.SendNotFoundError(fmt.Errorf("quota %d not found", id))
return return
} }
@ -73,45 +74,27 @@ func (qa *QuotaAPI) Prepare() {
// Get returns quota by id // Get returns quota by id
func (qa *QuotaAPI) Get() { func (qa *QuotaAPI) Get() {
query := &models.QuotaQuery{ qa.Data["json"] = qa.quota
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.ServeJSON() qa.ServeJSON()
} }
// Put update the quota // Put update the quota
func (qa *QuotaAPI) Put() { func (qa *QuotaAPI) Put() {
var req *models.QuotaUpdateRequest var req *QuotaUpdateRequest
if err := qa.DecodeJSONReq(&req); err != nil { if err := qa.DecodeJSONReq(&req); err != nil {
qa.SendBadRequestError(err) qa.SendBadRequestError(err)
return 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) qa.SendBadRequestError(err)
return return
} }
mgr, err := quota.NewManager(qa.quota.Reference, qa.quota.ReferenceID) qa.quota.SetHard(req.Hard)
if err != nil {
qa.SendInternalServerError(fmt.Errorf("failed to create quota manager, error: %v", err))
return
}
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)) qa.SendInternalServerError(fmt.Errorf("failed to update hard limits of the quota, error: %v", err))
return return
} }
@ -125,25 +108,25 @@ func (qa *QuotaAPI) List() {
return return
} }
query := &models.QuotaQuery{ query := &q.Query{
Reference: qa.GetString("reference"), Keywords: q.KeyWords{
ReferenceID: qa.GetString("reference_id"), "reference": qa.GetString("reference"),
Pagination: models.Pagination{ "reference_id": qa.GetString("reference_id"),
Page: page,
Size: size,
},
Sorting: models.Sorting{
Sort: qa.GetString("sort"),
}, },
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 { if err != nil {
qa.SendInternalServerError(fmt.Errorf("failed to query database for total of quotas, error: %v", err)) qa.SendInternalServerError(fmt.Errorf("failed to query database for total of quotas, error: %v", err))
return return
} }
quotas, err := dao.ListQuotas(query) quotas, err := quota.Ctl.List(ctx, query)
if err != nil { if err != nil {
qa.SendInternalServerError(fmt.Errorf("failed to query database for quotas, error: %v", err)) qa.SendInternalServerError(fmt.Errorf("failed to query database for quotas, error: %v", err))
return return

55
src/core/api/quota_test.go
View File

@ -15,32 +15,35 @@
package api package api
import ( import (
"context"
"fmt" "fmt"
"testing" "testing"
"github.com/goharbor/harbor/src/common/models" o "github.com/astaxie/beego/orm"
"github.com/goharbor/harbor/src/common/quota" "github.com/goharbor/harbor/src/controller/quota"
"github.com/goharbor/harbor/src/common/quota/driver" "github.com/goharbor/harbor/src/lib/orm"
"github.com/goharbor/harbor/src/common/quota/driver/mocks" "github.com/goharbor/harbor/src/pkg/quota/driver"
"github.com/goharbor/harbor/src/pkg/types" "github.com/goharbor/harbor/src/pkg/types"
"github.com/goharbor/harbor/src/testing/apitests/apilib" "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/assert"
"github.com/stretchr/testify/mock"
) )
var ( var (
reference = "mock" reference = uuid.New().String()
hardLimits = types.ResourceList{types.ResourceStorage: -1} hardLimits = types.ResourceList{types.ResourceStorage: -1}
) )
func init() { func init() {
mockDriver := &mocks.Driver{} mockDriver := &drivertesting.Driver{}
mockHardLimitsFn := func() types.ResourceList { mockHardLimitsFn := func() types.ResourceList {
return hardLimits return hardLimits
} }
mockLoadFn := func(key string) driver.RefObject { mockLoadFn := func(ctx context.Context, key string) driver.RefObject {
return driver.RefObject{"id": key} return driver.RefObject{"id": key}
} }
@ -53,8 +56,8 @@ func init() {
} }
mockDriver.On("HardLimits").Return(mockHardLimitsFn) mockDriver.On("HardLimits").Return(mockHardLimitsFn)
mockDriver.On("Load", mock.AnythingOfType("string")).Return(mockLoadFn, nil) mock.OnAnything(mockDriver, "Load").Return(mockLoadFn, nil)
mockDriver.On("Validate", mock.AnythingOfType("types.ResourceList")).Return(mockValidateFn) mock.OnAnything(mockDriver, "Validate").Return(mockValidateFn)
driver.Register(reference, mockDriver) driver.Register(reference, mockDriver)
} }
@ -63,13 +66,19 @@ func TestQuotaAPIList(t *testing.T) {
assert := assert.New(t) assert := assert.New(t)
apiTest := newHarborAPI() 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 count := 10
for i := 0; i < count; i++ { for i := 0; i < count; i++ {
mgr, err := quota.NewManager(reference, fmt.Sprintf("%d", i)) quotaID, err := quota.Ctl.Create(ctx, reference, uuid.New().String(), hardLimits)
assert.Nil(err)
_, err = mgr.NewQuota(hardLimits)
assert.Nil(err) assert.Nil(err)
quotaIDs = append(quotaIDs, quotaID)
} }
code, quotas, err := apiTest.QuotasGet(&apilib.QuotaQuery{Reference: reference}, *admin) code, quotas, err := apiTest.QuotasGet(&apilib.QuotaQuery{Reference: reference}, *admin)
@ -87,11 +96,10 @@ func TestQuotaAPIGet(t *testing.T) {
assert := assert.New(t) assert := assert.New(t)
apiTest := newHarborAPI() apiTest := newHarborAPI()
mgr, err := quota.NewManager(reference, "quota-get") ctx := orm.NewContext(context.TODO(), o.NewOrm())
assert.Nil(err) quotaID, err := quota.Ctl.Create(ctx, reference, uuid.New().String(), hardLimits)
quotaID, err := mgr.NewQuota(hardLimits)
assert.Nil(err) assert.Nil(err)
defer quota.Ctl.Delete(ctx, quotaID)
code, quota, err := apiTest.QuotasGetByID(*admin, fmt.Sprintf("%d", quotaID)) code, quota, err := apiTest.QuotasGetByID(*admin, fmt.Sprintf("%d", quotaID))
assert.Nil(err) assert.Nil(err)
@ -107,22 +115,21 @@ func TestQuotaPut(t *testing.T) {
assert := assert.New(t) assert := assert.New(t)
apiTest := newHarborAPI() apiTest := newHarborAPI()
mgr, err := quota.NewManager(reference, "quota-put") ctx := orm.NewContext(context.TODO(), o.NewOrm())
assert.Nil(err) quotaID, err := quota.Ctl.Create(ctx, reference, uuid.New().String(), hardLimits)
quotaID, err := mgr.NewQuota(hardLimits)
assert.Nil(err) assert.Nil(err)
defer quota.Ctl.Delete(ctx, quotaID)
code, quota, err := apiTest.QuotasGetByID(*admin, fmt.Sprintf("%d", quotaID)) code, quota, err := apiTest.QuotasGetByID(*admin, fmt.Sprintf("%d", quotaID))
assert.Nil(err) assert.Nil(err)
assert.Equal(int(200), code) assert.Equal(int(200), code)
assert.Equal(map[string]int64{"storage": -1}, quota.Hard) 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.Nil(err, err)
assert.Equal(int(400), code) 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.Nil(err)
assert.Equal(int(200), code) assert.Equal(int(200), code)

2
src/lib/q/query.go
View File

@ -25,6 +25,8 @@ type Query struct {
PageSize int64 PageSize int64
// List of key words // List of key words
Keywords KeyWords Keywords KeyWords
// Sorting
Sorting string
} }
// New returns Query with keywords // New returns Query with keywords

8
src/lib/response_buffer.go
View File

@ -79,7 +79,8 @@ func (r *ResponseBuffer) Flush() (int, error) {
// Success checks whether the status code is >= 200 & <= 399 // Success checks whether the status code is >= 200 & <= 399
func (r *ResponseBuffer) Success() bool { 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 // Reset reset the response buffer
@ -98,5 +99,10 @@ func (r *ResponseBuffer) Reset() error {
// StatusCode returns the status code // StatusCode returns the status code
func (r *ResponseBuffer) StatusCode() int { 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 return r.code
} }

9
src/lib/response_recorder.go
View File

@ -49,5 +49,12 @@ func (r *ResponseRecorder) WriteHeader(statusCode int) {
// Success checks whether the status code is >= 200 & <= 399 // Success checks whether the status code is >= 200 & <= 399
func (r *ResponseRecorder) Success() bool { 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
} }

88
src/pkg/quota/dao/dao.go
View File

@ -16,15 +16,22 @@ package dao
import ( import (
"context" "context"
"fmt"
"time" "time"
"github.com/goharbor/harbor/src/lib/log"
"github.com/goharbor/harbor/src/lib/orm" "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/quota/models"
"github.com/goharbor/harbor/src/pkg/types" "github.com/goharbor/harbor/src/pkg/types"
) )
// DAO the dao for Quota and QuotaUsage // DAO the dao for Quota and QuotaUsage
type DAO interface { 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 create quota for reference object
Create(ctx context.Context, reference, referenceID string, hardLimits, used types.ResourceList) (int64, error) Create(ctx context.Context, reference, referenceID string, hardLimits, used types.ResourceList) (int64, error)
@ -42,6 +49,9 @@ type DAO interface {
// Update update quota // Update update quota
Update(ctx context.Context, quota *models.Quota) error 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 // New returns an instance of the default DAO
@ -51,6 +61,23 @@ func New() DAO {
type dao struct{} 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) { func (d *dao) Create(ctx context.Context, reference, referenceID string, hardLimits, used types.ResourceList) (int64, error) {
o, err := orm.FromContext(ctx) o, err := orm.FromContext(ctx)
if err != nil { if err != nil {
@ -192,6 +219,67 @@ func (d *dao) Update(ctx context.Context, quota *models.Quota) error {
return nil 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(&quotas); 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 { func toQuota(quota *Quota, usage *QuotaUsage) *models.Quota {
return &models.Quota{ return &models.Quota{
ID: quota.ID, ID: quota.ID,

92
src/pkg/quota/dao/dao_test.go
View File

@ -20,8 +20,10 @@ import (
"testing" "testing"
"github.com/goharbor/harbor/src/lib/orm" "github.com/goharbor/harbor/src/lib/orm"
"github.com/goharbor/harbor/src/lib/q"
"github.com/goharbor/harbor/src/pkg/types" "github.com/goharbor/harbor/src/pkg/types"
htesting "github.com/goharbor/harbor/src/testing" htesting "github.com/goharbor/harbor/src/testing"
"github.com/google/uuid"
"github.com/stretchr/testify/suite" "github.com/stretchr/testify/suite"
) )
@ -39,6 +41,42 @@ func (suite *DaoTestSuite) SetupSuite() {
suite.dao = New() 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() { func (suite *DaoTestSuite) TestCreate() {
hardLimits := types.ResourceList{types.ResourceStorage: 100} hardLimits := types.ResourceList{types.ResourceStorage: 100}
usage := types.ResourceList{types.ResourceStorage: 0} 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) { func TestDaoTestSuite(t *testing.T) {
suite.Run(t, &DaoTestSuite{}) suite.Run(t, &DaoTestSuite{})
} }

41
src/pkg/quota/dao/model.go
View File

@ -15,13 +15,42 @@
package dao package dao
import ( 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 // Quota model for quota
type Quota = models.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 // TableName returns table name for orm
type QuotaUsage = models.QuotaUsage 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"
}

117
src/pkg/quota/dao/util.go Normal file
View File

@ -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
}

52
src/pkg/quota/dao/util_test.go Normal file
View File

@ -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)
}
})
}
}

23
src/pkg/quota/manager.go
View File

@ -18,6 +18,7 @@ import (
"context" "context"
"github.com/goharbor/harbor/src/lib/orm" "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/dao"
"github.com/goharbor/harbor/src/pkg/quota/models" "github.com/goharbor/harbor/src/pkg/quota/models"
"github.com/goharbor/harbor/src/pkg/types" "github.com/goharbor/harbor/src/pkg/types"
@ -31,6 +32,9 @@ type Manager interface {
// Create create quota for the reference object // Create create quota for the reference object
Create(ctx context.Context, reference, referenceID string, hardLimits types.ResourceList, usages ...types.ResourceList) (int64, error) 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 delete quota by id
Delete(ctx context.Context, id int64) error Delete(ctx context.Context, id int64) error
@ -45,6 +49,9 @@ type Manager interface {
// Update update quota // Update update quota
Update(ctx context.Context, quota *Quota) error Update(ctx context.Context, quota *Quota) error
// List list quotas
List(ctx context.Context, query *q.Query) ([]*Quota, error)
} }
var ( var (
@ -75,10 +82,18 @@ func (m *manager) Create(ctx context.Context, reference, referenceID string, har
return id, err 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 { func (m *manager) Delete(ctx context.Context, id int64) error {
h := func(ctx context.Context) error {
return m.dao.Delete(ctx, id) return m.dao.Delete(ctx, id)
} }
return orm.WithTransaction(h)(ctx)
}
func (m *manager) Get(ctx context.Context, id int64) (*Quota, error) { func (m *manager) Get(ctx context.Context, id int64) (*Quota, error) {
return m.dao.Get(ctx, id) return m.dao.Get(ctx, id)
} }
@ -92,9 +107,17 @@ func (m *manager) GetByRefForUpdate(ctx context.Context, reference, referenceID
} }
func (m *manager) Update(ctx context.Context, q *Quota) error { func (m *manager) Update(ctx context.Context, q *Quota) error {
h := func(ctx context.Context) error {
return m.dao.Update(ctx, q) 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 // NewManager returns quota manager
func NewManager() Manager { func NewManager() Manager {
return &manager{dao: dao.New()} return &manager{dao: dao.New()}

40
src/testing/common/security/context.go
View File

@ -3,7 +3,6 @@
package security package security
import ( import (
models "github.com/goharbor/harbor/src/common/models"
mock "github.com/stretchr/testify/mock" mock "github.com/stretchr/testify/mock"
types "github.com/goharbor/harbor/src/pkg/permission/types" 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 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: // GetUsername provides a mock function with given fields:
func (_m *Context) GetUsername() string { func (_m *Context) GetUsername() string {
ret := _m.Called() ret := _m.Called()

1
src/testing/controller/controller.go
View File

@ -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/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/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/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/quota -name Controller -output ./quota -outpkg quota
//go:generate mockery -case snake -dir ../../controller/scan -name Controller -output ./scan -outpkg scan //go:generate mockery -case snake -dir ../../controller/scan -name Controller -output ./scan -outpkg scan

60
src/testing/controller/quota/controller.go
View File

@ -8,6 +8,8 @@ import (
models "github.com/goharbor/harbor/src/pkg/quota/models" models "github.com/goharbor/harbor/src/pkg/quota/models"
mock "github.com/stretchr/testify/mock" mock "github.com/stretchr/testify/mock"
q "github.com/goharbor/harbor/src/lib/q"
quota "github.com/goharbor/harbor/src/controller/quota" quota "github.com/goharbor/harbor/src/controller/quota"
types "github.com/goharbor/harbor/src/pkg/types" types "github.com/goharbor/harbor/src/pkg/types"
@ -18,6 +20,27 @@ type Controller struct {
mock.Mock 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 // 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) { func (_m *Controller) Create(ctx context.Context, reference string, referenceID string, hardLimits types.ResourceList, used ...types.ResourceList) (int64, error) {
_va := make([]interface{}, len(used)) _va := make([]interface{}, len(used))
@ -127,6 +150,29 @@ func (_m *Controller) IsEnabled(ctx context.Context, reference string, reference
return r0, r1 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 // 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 { func (_m *Controller) Refresh(ctx context.Context, reference string, referenceID string, options ...quota.Option) error {
_va := make([]interface{}, len(options)) _va := make([]interface{}, len(options))
@ -161,3 +207,17 @@ func (_m *Controller) Request(ctx context.Context, reference string, referenceID
return r0 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
}

46
src/testing/pkg/quota/manager.go
View File

@ -8,6 +8,8 @@ import (
models "github.com/goharbor/harbor/src/pkg/quota/models" models "github.com/goharbor/harbor/src/pkg/quota/models"
mock "github.com/stretchr/testify/mock" mock "github.com/stretchr/testify/mock"
q "github.com/goharbor/harbor/src/lib/q"
types "github.com/goharbor/harbor/src/pkg/types" types "github.com/goharbor/harbor/src/pkg/types"
) )
@ -16,6 +18,27 @@ type Manager struct {
mock.Mock 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 // 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) { func (_m *Manager) Create(ctx context.Context, reference string, referenceID string, hardLimits types.ResourceList, usages ...types.ResourceList) (int64, error) {
_va := make([]interface{}, len(usages)) _va := make([]interface{}, len(usages))
@ -127,6 +150,29 @@ func (_m *Manager) GetByRefForUpdate(ctx context.Context, reference string, refe
return r0, r1 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 // Update provides a mock function with given fields: ctx, _a1
func (_m *Manager) Update(ctx context.Context, _a1 *models.Quota) error { func (_m *Manager) Update(ctx context.Context, _a1 *models.Quota) error {
ret := _m.Called(ctx, _a1) ret := _m.Called(ctx, _a1)

14
src/testing/suite.go
View File

@ -16,12 +16,10 @@ package testing
import ( import (
"context" "context"
"fmt"
"io" "io"
"math/rand" "math/rand"
"net/http" "net/http"
"net/http/httptest" "net/http/httptest"
"strconv"
"sync" "sync"
"time" "time"
@ -31,7 +29,6 @@ import (
"github.com/goharbor/harbor/src/core/config" "github.com/goharbor/harbor/src/core/config"
"github.com/goharbor/harbor/src/lib/errors" "github.com/goharbor/harbor/src/lib/errors"
"github.com/goharbor/harbor/src/lib/orm" "github.com/goharbor/harbor/src/lib/orm"
"github.com/goharbor/harbor/src/pkg/types"
"github.com/opencontainers/go-digest" "github.com/opencontainers/go-digest"
"github.com/stretchr/testify/suite" "github.com/stretchr/testify/suite"
) )
@ -161,14 +158,3 @@ func (suite *Suite) ExecSQL(query string, args ...interface{}) {
func (suite *Suite) IsNotFoundErr(err error) bool { func (suite *Suite) IsNotFoundErr(err error) bool {
return suite.True(errors.IsNotFoundErr(err)) 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])
}