IntenseWebs/grafana
3
0
Fork 0
mirror of https://github.com/grafana/grafana.git synced 2025-02-25 18:55:37 -06:00
Code Issues Packages Projects Releases Wiki Activity

Chore: Refactor quota service (#57586)

Browse Source 
* Chore: refactore quota service

* Apply suggestions from code review
This commit is contained in:
Sofia Papagiannaki
2022-11-08 10:25:34 +02:00
committed by GitHub
parent faa0fda6eb
commit 326ea86a57
99 changed files with 2595 additions and 1397 deletions
Download Patch File Download Diff File Expand all files Collapse all files

54
pkg/services/apikey/apikeyimpl/apikey.go
View File

@@ -6,6 +6,7 @@ import (
"github.com/grafana/grafana/pkg/infra/db"
"github.com/grafana/grafana/pkg/services/apikey"
"github.com/grafana/grafana/pkg/services/featuremgmt"
"github.com/grafana/grafana/pkg/services/quota"
"github.com/grafana/grafana/pkg/setting"
)
@@ -13,16 +14,34 @@ type Service struct {
store store
}
func ProvideService(db db.DB, cfg *setting.Cfg) apikey.Service {
func ProvideService(db db.DB, cfg *setting.Cfg, quotaService quota.Service) (apikey.Service, error) {
s := &Service{}
if cfg.IsFeatureToggleEnabled(featuremgmt.FlagNewDBLibrary) {
return &Service{
store: &sqlxStore{
sess: db.GetSqlxSession(),
cfg: cfg,
},
s.store = &sqlxStore{
sess: db.GetSqlxSession(),
cfg: cfg,
}
}
return &Service{store: &sqlStore{db: db, cfg: cfg}}
s.store = &sqlStore{db: db, cfg: cfg}
defaultLimits, err := readQuotaConfig(cfg)
if err != nil {
return s, err
}
if err := quotaService.RegisterQuotaReporter(&quota.NewUsageReporter{
TargetSrv: apikey.QuotaTargetSrv,
DefaultLimits: defaultLimits,
Reporter: s.Usage,
}); err != nil {
return s, err
}
return s, nil
}
func (s *Service) Usage(ctx context.Context, scopeParams *quota.ScopeParameters) (*quota.Map, error) {
return s.store.Count(ctx, scopeParams)
}
func (s *Service) GetAPIKeys(ctx context.Context, query *apikey.GetApiKeysQuery) error {
@@ -49,3 +68,24 @@ func (s *Service) AddAPIKey(ctx context.Context, cmd *apikey.AddCommand) error {
func (s *Service) UpdateAPIKeyLastUsedDate(ctx context.Context, tokenID int64) error {
return s.store.UpdateAPIKeyLastUsedDate(ctx, tokenID)
}
func readQuotaConfig(cfg *setting.Cfg) (*quota.Map, error) {
limits := &quota.Map{}
if cfg == nil {
return limits, nil
}
globalQuotaTag, err := quota.NewTag(apikey.QuotaTargetSrv, apikey.QuotaTarget, quota.GlobalScope)
if err != nil {
return limits, err
}
orgQuotaTag, err := quota.NewTag(apikey.QuotaTargetSrv, apikey.QuotaTarget, quota.OrgScope)
if err != nil {
return limits, err
}
limits.Set(globalQuotaTag, cfg.Quota.Global.ApiKey)
limits.Set(orgQuotaTag, cfg.Quota.Org.ApiKey)
return limits, nil
}

33
pkg/services/apikey/apikeyimpl/sqlx_store.go
View File

@@ -10,6 +10,7 @@ import (
"github.com/grafana/grafana/pkg/services/accesscontrol"
"github.com/grafana/grafana/pkg/services/apikey"
"github.com/grafana/grafana/pkg/services/quota"
"github.com/grafana/grafana/pkg/services/sqlstore/session"
"github.com/grafana/grafana/pkg/setting"
)
@@ -142,3 +143,35 @@ func (ss *sqlxStore) UpdateAPIKeyLastUsedDate(ctx context.Context, tokenID int64
_, err := ss.sess.Exec(ctx, `UPDATE api_key SET last_used_at=? WHERE id=?`, &now, tokenID)
return err
}
func (ss *sqlxStore) Count(ctx context.Context, scopeParams *quota.ScopeParameters) (*quota.Map, error) {
u := &quota.Map{}
type result struct {
Count int64
}
r := result{}
if err := ss.sess.Get(ctx, &r, `SELECT COUNT(*) AS count FROM api_key`); err != nil {
return u, err
} else {
tag, err := quota.NewTag(apikey.QuotaTargetSrv, apikey.QuotaTarget, quota.GlobalScope)
if err != nil {
return nil, err
}
u.Set(tag, r.Count)
}
if scopeParams.OrgID != 0 {
if err := ss.sess.Get(ctx, &r, `SELECT COUNT(*) AS count FROM api_key WHERE org_id = ?`, scopeParams.OrgID); err != nil {
return u, err
} else {
tag, err := quota.NewTag(apikey.QuotaTargetSrv, apikey.QuotaTarget, quota.OrgScope)
if err != nil {
return nil, err
}
u.Set(tag, r.Count)
}
}
return u, nil
}

3
pkg/services/apikey/apikeyimpl/store.go
View File

@@ -4,6 +4,7 @@ import (
"context"
"github.com/grafana/grafana/pkg/services/apikey"
"github.com/grafana/grafana/pkg/services/quota"
)
type store interface {
@@ -15,4 +16,6 @@ type store interface {
GetApiKeyByName(ctx context.Context, query *apikey.GetByNameQuery) error
GetAPIKeyByHash(ctx context.Context, hash string) (*apikey.APIKey, error)
UpdateAPIKeyLastUsedDate(ctx context.Context, tokenID int64) error
Count(context.Context, *quota.ScopeParameters) (*quota.Map, error)
}

46
pkg/services/apikey/apikeyimpl/xorm_store.go
View File

@@ -11,6 +11,8 @@ import (
"github.com/grafana/grafana/pkg/infra/db"
"github.com/grafana/grafana/pkg/services/accesscontrol"
"github.com/grafana/grafana/pkg/services/apikey"
"github.com/grafana/grafana/pkg/services/quota"
"github.com/grafana/grafana/pkg/services/sqlstore"
"github.com/grafana/grafana/pkg/setting"
)
@@ -174,3 +176,47 @@ func (ss *sqlStore) UpdateAPIKeyLastUsedDate(ctx context.Context, tokenID int64)
return nil
})
}
func (ss *sqlStore) Count(ctx context.Context, scopeParams *quota.ScopeParameters) (*quota.Map, error) {
u := &quota.Map{}
type result struct {
Count int64
}
r := result{}
if err := ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
rawSQL := "SELECT COUNT(*) AS count FROM api_key"
if _, err := sess.SQL(rawSQL).Get(&r); err != nil {
return err
}
return nil
}); err != nil {
return u, err
} else {
tag, err := quota.NewTag(apikey.QuotaTargetSrv, apikey.QuotaTarget, quota.GlobalScope)
if err != nil {
return nil, err
}
u.Set(tag, r.Count)
}
if scopeParams.OrgID != 0 {
if err := ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
rawSQL := "SELECT COUNT(*) AS count FROM api_key WHERE org_id = ?"
if _, err := sess.SQL(rawSQL, scopeParams.OrgID).Get(&r); err != nil {
return err
}
return nil
}); err != nil {
return u, err
} else {
tag, err := quota.NewTag(apikey.QuotaTargetSrv, apikey.QuotaTarget, quota.OrgScope)
if err != nil {
return nil, err
}
u.Set(tag, r.Count)
}
}
return u, nil
}