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 (#58643)

Browse Source 
Chore: Refactor quota service (#57586)

* Chore: refactore quota service

* Apply suggestions from code review
This commit is contained in:
Sofia Papagiannaki
2022-11-14 21:08:10 +02:00
committed by GitHub
parent dd0d034796
commit 9855e74b92
99 changed files with 2596 additions and 1398 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
pkg/services/org/model.go
View File

@@ -204,3 +204,9 @@ func (o ByOrgName) Less(i, j int) bool {
return o[i].Name < o[j].Name
}
const (
QuotaTargetSrv string = "org"
OrgQuotaTarget string = "org"
OrgUserQuotaTarget string = "org_user"
)

51
pkg/services/org/orgimpl/org.go
View File

@@ -8,6 +8,7 @@ import (
"github.com/grafana/grafana/pkg/infra/db"
"github.com/grafana/grafana/pkg/infra/log"
"github.com/grafana/grafana/pkg/services/org"
"github.com/grafana/grafana/pkg/services/quota"
"github.com/grafana/grafana/pkg/setting"
"github.com/grafana/grafana/pkg/util"
)
@@ -18,9 +19,9 @@ type Service struct {
log log.Logger
}
func ProvideService(db db.DB, cfg *setting.Cfg) org.Service {
func ProvideService(db db.DB, cfg *setting.Cfg, quotaService quota.Service) (org.Service, error) {
log := log.New("org service")
return &Service{
s := &Service{
store: &sqlStore{
db: db,
dialect: db.GetDialect(),
@@ -30,6 +31,24 @@ func ProvideService(db db.DB, cfg *setting.Cfg) org.Service {
cfg: cfg,
log: log,
}
defaultLimits, err := readQuotaConfig(cfg)
if err != nil {
return s, err
}
if err := quotaService.RegisterQuotaReporter(&quota.NewUsageReporter{
TargetSrv: quota.TargetSrv(org.QuotaTargetSrv),
DefaultLimits: defaultLimits,
Reporter: s.Usage,
}); err != nil {
return s, nil
}
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) GetIDForNewUser(ctx context.Context, cmd org.GetOrgIDForNewUserCommand) (int64, error) {
@@ -179,3 +198,31 @@ func (s *Service) GetOrgUsers(ctx context.Context, query *org.GetOrgUsersQuery)
func (s *Service) SearchOrgUsers(ctx context.Context, query *org.SearchOrgUsersQuery) (*org.SearchOrgUsersQueryResult, error) {
return s.store.SearchOrgUsers(ctx, query)
}
func readQuotaConfig(cfg *setting.Cfg) (*quota.Map, error) {
limits := &quota.Map{}
if cfg == nil {
return limits, nil
}
globalQuotaTag, err := quota.NewTag(quota.TargetSrv(org.QuotaTargetSrv), quota.Target(org.OrgQuotaTarget), quota.GlobalScope)
if err != nil {
return limits, err
}
orgQuotaTag, err := quota.NewTag(quota.TargetSrv(org.QuotaTargetSrv), quota.Target(org.OrgUserQuotaTarget), quota.OrgScope)
if err != nil {
return limits, err
}
userTag, err := quota.NewTag(quota.TargetSrv(org.QuotaTargetSrv), quota.Target(org.OrgUserQuotaTarget), quota.UserScope)
if err != nil {
return limits, err
}
limits.Set(globalQuotaTag, cfg.Quota.Global.Org)
// users per org
limits.Set(orgQuotaTag, cfg.Quota.Org.User)
// orgs per user
limits.Set(userTag, cfg.Quota.User.Org)
return limits, nil
}

5
pkg/services/org/orgimpl/org_test.go
View File

@@ -5,6 +5,7 @@ import (
"testing"
"github.com/grafana/grafana/pkg/services/org"
"github.com/grafana/grafana/pkg/services/quota"
"github.com/grafana/grafana/pkg/setting"
"github.com/stretchr/testify/require"
)
@@ -135,3 +136,7 @@ func (f *FakeOrgStore) SearchOrgUsers(ctx context.Context, query *org.SearchOrgU
func (f *FakeOrgStore) RemoveOrgUser(ctx context.Context, cmd *org.RemoveOrgUserCommand) error {
return f.ExpectedError
}
func (f *FakeOrgStore) Count(ctx context.Context, _ *quota.ScopeParameters) (*quota.Map, error) {
return nil, nil
}

70
pkg/services/org/orgimpl/store.go
View File

@@ -14,6 +14,8 @@ import (
"github.com/grafana/grafana/pkg/models"
"github.com/grafana/grafana/pkg/services/accesscontrol"
"github.com/grafana/grafana/pkg/services/org"
"github.com/grafana/grafana/pkg/services/quota"
"github.com/grafana/grafana/pkg/services/sqlstore"
"github.com/grafana/grafana/pkg/services/sqlstore/migrator"
"github.com/grafana/grafana/pkg/services/user"
"github.com/grafana/grafana/pkg/setting"
@@ -42,6 +44,8 @@ type store interface {
GetByName(context.Context, *org.GetOrgByNameQuery) (*org.Org, error)
SearchOrgUsers(context.Context, *org.SearchOrgUsersQuery) (*org.SearchOrgUsersQueryResult, error)
RemoveOrgUser(context.Context, *org.RemoveOrgUserCommand) error
Count(context.Context, *quota.ScopeParameters) (*quota.Map, error)
}
type sqlStore struct {
@@ -395,6 +399,72 @@ func (ss *sqlStore) AddOrgUser(ctx context.Context, cmd *org.AddOrgUserCommand)
})
}
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 org"
if _, err := sess.SQL(rawSQL).Get(&r); err != nil {
return err
}
return nil
}); err != nil {
return u, err
} else {
tag, err := quota.NewTag(quota.TargetSrv(org.QuotaTargetSrv), quota.Target(org.OrgQuotaTarget), quota.GlobalScope)
if err != nil {
return u, err
}
u.Set(tag, r.Count)
}
if scopeParams.OrgID != 0 {
if err := ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
rawSQL := fmt.Sprintf("SELECT COUNT(*) AS count FROM (SELECT user_id FROM org_user WHERE org_id=? AND user_id IN (SELECT id AS user_id FROM %s WHERE is_service_account=%s)) as subq",
ss.db.GetDialect().Quote("user"),
ss.db.GetDialect().BooleanStr(false),
)
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(quota.TargetSrv(org.QuotaTargetSrv), quota.Target(org.OrgUserQuotaTarget), quota.OrgScope)
if err != nil {
return u, err
}
u.Set(tag, r.Count)
}
}
if scopeParams.UserID != 0 {
if err := ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
// should we exclude service accounts?
rawSQL := "SELECT COUNT(*) AS count FROM org_user WHERE user_id=?"
if _, err := sess.SQL(rawSQL, scopeParams.UserID).Get(&r); err != nil {
return err
}
return nil
}); err != nil {
return u, err
} else {
tag, err := quota.NewTag(quota.TargetSrv(org.QuotaTargetSrv), quota.Target(org.OrgUserQuotaTarget), quota.UserScope)
if err != nil {
return u, err
}
u.Set(tag, r.Count)
}
}
return u, nil
}
func setUsingOrgInTransaction(sess *db.Session, userID int64, orgID int64) error {
user := user.User{
ID: userID,