diff --git a/src/controller/quota/controller.go b/src/controller/quota/controller.go index b63aa262b..068e355fc 100644 --- a/src/controller/quota/controller.go +++ b/src/controller/quota/controller.go @@ -55,16 +55,16 @@ type Controller interface { Delete(ctx context.Context, id int64) error // Get returns quota by id - Get(ctx context.Context, id int64) (*quota.Quota, error) + Get(ctx context.Context, id int64, options ...Option) (*quota.Quota, error) // GetByRef returns quota by reference object - GetByRef(ctx context.Context, reference, referenceID string) (*quota.Quota, error) + GetByRef(ctx context.Context, reference, referenceID string, options ...Option) (*quota.Quota, error) // IsEnabled returns true when quota enabled for reference object IsEnabled(ctx context.Context, reference, referenceID string) (bool, error) // List list quotas - List(ctx context.Context, query *q.Query) ([]*quota.Quota, error) + List(ctx context.Context, query *q.Query, options ...Option) ([]*quota.Quota, error) // Refresh refresh quota for the reference object Refresh(ctx context.Context, reference, referenceID string, options ...Option) error @@ -105,12 +105,40 @@ func (c *controller) Delete(ctx context.Context, id int64) error { return c.quotaMgr.Delete(ctx, id) } -func (c *controller) Get(ctx context.Context, id int64) (*quota.Quota, error) { - return c.quotaMgr.Get(ctx, id) +func (c *controller) Get(ctx context.Context, id int64, options ...Option) (*quota.Quota, error) { + q, err := c.quotaMgr.Get(ctx, id) + if err != nil { + return nil, err + } + + return c.assembleQuota(ctx, q, newOptions(options...)) } -func (c *controller) GetByRef(ctx context.Context, reference, referenceID string) (*quota.Quota, error) { - return c.quotaMgr.GetByRef(ctx, reference, referenceID) +func (c *controller) GetByRef(ctx context.Context, reference, referenceID string, options ...Option) (*quota.Quota, error) { + q, err := c.quotaMgr.GetByRef(ctx, reference, referenceID) + if err != nil { + return nil, err + } + + return c.assembleQuota(ctx, q, newOptions(options...)) +} + +func (c *controller) assembleQuota(ctx context.Context, q *quota.Quota, opts *Options) (*quota.Quota, error) { + if opts.WithReferenceObject { + driver, err := Driver(ctx, q.Reference) + if err != nil { + return nil, err + } + + ref, err := driver.Load(ctx, q.ReferenceID) + if err != nil { + return nil, err + } + + q.Ref = ref + } + + return q, nil } func (c *controller) IsEnabled(ctx context.Context, reference, referenceID string) (bool, error) { @@ -122,8 +150,20 @@ func (c *controller) IsEnabled(ctx context.Context, reference, referenceID strin return d.Enabled(ctx, referenceID) } -func (c *controller) List(ctx context.Context, query *q.Query) ([]*quota.Quota, error) { - return c.quotaMgr.List(ctx, query) +func (c *controller) List(ctx context.Context, query *q.Query, options ...Option) ([]*quota.Quota, error) { + quotas, err := c.quotaMgr.List(ctx, query) + if err != nil { + return nil, err + } + + opts := newOptions(options...) + for _, q := range quotas { + if _, err := c.assembleQuota(ctx, q, opts); err != nil { + return nil, err + } + } + + return quotas, nil } func (c *controller) getReservedResources(ctx context.Context, reference, referenceID string) (types.ResourceList, error) { diff --git a/src/controller/quota/driver/project/project.go b/src/controller/quota/driver/project/project.go index a546be8bd..5c577390d 100644 --- a/src/controller/quota/driver/project/project.go +++ b/src/controller/quota/driver/project/project.go @@ -120,7 +120,7 @@ func (d *driver) CalculateUsage(ctx context.Context, key string) (types.Resource func newDriver() dr.Driver { cfg := config.NewDBCfgManager() - loader := dataloader.NewBatchedLoader(getProjectsBatchFn) + loader := dataloader.NewBatchedLoader(getProjectsBatchFn, dataloader.WithClearCacheOnBatch()) return &driver{ cfg: cfg, diff --git a/src/controller/quota/options.go b/src/controller/quota/options.go index fec096859..372d9296c 100644 --- a/src/controller/quota/options.go +++ b/src/controller/quota/options.go @@ -17,9 +17,10 @@ package quota // Option option for `Refresh` method of `Controller` type Option func(*Options) -// Options options used by `Refresh` method of `Controller` +// Options options used by `Refresh`, `Get`, `List` methods of `Controller` type Options struct { - IgnoreLimitation bool + IgnoreLimitation bool + WithReferenceObject bool } // IgnoreLimitation set IgnoreLimitation for the Options @@ -29,6 +30,13 @@ func IgnoreLimitation(ignoreLimitation bool) func(*Options) { } } +// WithReferenceObject set WithReferenceObject to true for the Options +func WithReferenceObject() func(*Options) { + return func(opts *Options) { + opts.WithReferenceObject = true + } +} + func newOptions(options ...Option) *Options { opts := &Options{} for _, f := range options { diff --git a/src/core/api/quota.go b/src/core/api/quota.go index cf9a4b303..45a5858e0 100644 --- a/src/core/api/quota.go +++ b/src/core/api/quota.go @@ -126,7 +126,7 @@ func (qa *QuotaAPI) List() { return } - quotas, err := quota.Ctl.List(ctx, query) + quotas, err := quota.Ctl.List(ctx, query, quota.WithReferenceObject()) if err != nil { qa.SendInternalServerError(fmt.Errorf("failed to query database for quotas, error: %v", err)) return diff --git a/src/pkg/quota/dao/dao.go b/src/pkg/quota/dao/dao.go index 41420c2e6..ffb66f647 100644 --- a/src/pkg/quota/dao/dao.go +++ b/src/pkg/quota/dao/dao.go @@ -19,10 +19,8 @@ import ( "fmt" "time" - "github.com/goharbor/harbor/src/lib/log" "github.com/goharbor/harbor/src/lib/orm" "github.com/goharbor/harbor/src/lib/q" - "github.com/goharbor/harbor/src/pkg/quota/driver" "github.com/goharbor/harbor/src/pkg/quota/models" "github.com/goharbor/harbor/src/pkg/quota/types" ) @@ -262,21 +260,6 @@ FROM 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 } diff --git a/src/pkg/quota/models/quota.go b/src/pkg/quota/models/quota.go index 2f44e601c..4238f960a 100644 --- a/src/pkg/quota/models/quota.go +++ b/src/pkg/quota/models/quota.go @@ -26,7 +26,7 @@ import ( // Quota quota model for manager type Quota struct { ID int64 `orm:"pk;auto;column(id)" json:"id"` - Ref driver.RefObject `json:"ref"` + Ref driver.RefObject `orm:"-" json:"ref"` Reference string `orm:"column(reference)" json:"-"` ReferenceID string `orm:"column(reference_id)" json:"-"` Hard string `orm:"column(hard);type(jsonb)" json:"-"` diff --git a/src/testing/controller/quota/controller.go b/src/testing/controller/quota/controller.go index 37bf6bf71..97c3125b5 100644 --- a/src/testing/controller/quota/controller.go +++ b/src/testing/controller/quota/controller.go @@ -83,13 +83,20 @@ func (_m *Controller) Delete(ctx context.Context, id int64) error { return r0 } -// Get provides a mock function with given fields: ctx, id -func (_m *Controller) Get(ctx context.Context, id int64) (*models.Quota, error) { - ret := _m.Called(ctx, id) +// Get provides a mock function with given fields: ctx, id, options +func (_m *Controller) Get(ctx context.Context, id int64, options ...quota.Option) (*models.Quota, error) { + _va := make([]interface{}, len(options)) + for _i := range options { + _va[_i] = options[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, id) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) var r0 *models.Quota - if rf, ok := ret.Get(0).(func(context.Context, int64) *models.Quota); ok { - r0 = rf(ctx, id) + if rf, ok := ret.Get(0).(func(context.Context, int64, ...quota.Option) *models.Quota); ok { + r0 = rf(ctx, id, options...) } else { if ret.Get(0) != nil { r0 = ret.Get(0).(*models.Quota) @@ -97,8 +104,8 @@ func (_m *Controller) Get(ctx context.Context, id int64) (*models.Quota, error) } var r1 error - if rf, ok := ret.Get(1).(func(context.Context, int64) error); ok { - r1 = rf(ctx, id) + if rf, ok := ret.Get(1).(func(context.Context, int64, ...quota.Option) error); ok { + r1 = rf(ctx, id, options...) } else { r1 = ret.Error(1) } @@ -106,13 +113,20 @@ func (_m *Controller) Get(ctx context.Context, id int64) (*models.Quota, error) return r0, r1 } -// GetByRef provides a mock function with given fields: ctx, reference, referenceID -func (_m *Controller) GetByRef(ctx context.Context, reference string, referenceID string) (*models.Quota, error) { - ret := _m.Called(ctx, reference, referenceID) +// GetByRef provides a mock function with given fields: ctx, reference, referenceID, options +func (_m *Controller) GetByRef(ctx context.Context, reference string, referenceID string, options ...quota.Option) (*models.Quota, error) { + _va := make([]interface{}, len(options)) + for _i := range options { + _va[_i] = options[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, reference, referenceID) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) var r0 *models.Quota - if rf, ok := ret.Get(0).(func(context.Context, string, string) *models.Quota); ok { - r0 = rf(ctx, reference, referenceID) + if rf, ok := ret.Get(0).(func(context.Context, string, string, ...quota.Option) *models.Quota); ok { + r0 = rf(ctx, reference, referenceID, options...) } else { if ret.Get(0) != nil { r0 = ret.Get(0).(*models.Quota) @@ -120,8 +134,8 @@ func (_m *Controller) GetByRef(ctx context.Context, reference string, referenceI } var r1 error - if rf, ok := ret.Get(1).(func(context.Context, string, string) error); ok { - r1 = rf(ctx, reference, referenceID) + if rf, ok := ret.Get(1).(func(context.Context, string, string, ...quota.Option) error); ok { + r1 = rf(ctx, reference, referenceID, options...) } else { r1 = ret.Error(1) } @@ -150,13 +164,20 @@ func (_m *Controller) IsEnabled(ctx context.Context, reference string, reference return r0, r1 } -// List provides a mock function with given fields: ctx, query -func (_m *Controller) List(ctx context.Context, query *q.Query) ([]*models.Quota, error) { - ret := _m.Called(ctx, query) +// List provides a mock function with given fields: ctx, query, options +func (_m *Controller) List(ctx context.Context, query *q.Query, options ...quota.Option) ([]*models.Quota, error) { + _va := make([]interface{}, len(options)) + for _i := range options { + _va[_i] = options[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, query) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) var r0 []*models.Quota - if rf, ok := ret.Get(0).(func(context.Context, *q.Query) []*models.Quota); ok { - r0 = rf(ctx, query) + if rf, ok := ret.Get(0).(func(context.Context, *q.Query, ...quota.Option) []*models.Quota); ok { + r0 = rf(ctx, query, options...) } else { if ret.Get(0) != nil { r0 = ret.Get(0).([]*models.Quota) @@ -164,8 +185,8 @@ func (_m *Controller) List(ctx context.Context, query *q.Query) ([]*models.Quota } var r1 error - if rf, ok := ret.Get(1).(func(context.Context, *q.Query) error); ok { - r1 = rf(ctx, query) + if rf, ok := ret.Get(1).(func(context.Context, *q.Query, ...quota.Option) error); ok { + r1 = rf(ctx, query, options...) } else { r1 = ret.Error(1) }