Chore: Refactor quota service (#58643)

Chore: Refactor quota service (#57586)

* Chore: refactore quota service

* Apply suggestions from code review

pkg/services/user/userimpl/store.go

@@ -11,6 +11,7 @@ import (
 	"github.com/grafana/grafana/pkg/infra/db"
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/services/accesscontrol"
+	"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"
@@ -37,6 +38,8 @@ type store interface {
 	BatchDisableUsers(context.Context, *user.BatchDisableUsersCommand) error
 	Disable(context.Context, *user.DisableUserCommand) error
 	Search(context.Context, *user.SearchUsersQuery) (*user.SearchUserQueryResult, error)
+
+	Count(ctx context.Context) (int64, error)
 }
 
 type sqlStore struct {
@@ -461,6 +464,22 @@ func (ss *sqlStore) UpdatePermissions(ctx context.Context, userID int64, isAdmin
 	})
 }
 
+func (ss *sqlStore) Count(ctx context.Context) (int64, error) {
+	type result struct {
+		Count int64
+	}
+
+	r := result{}
+	err := ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
+		rawSQL := fmt.Sprintf("SELECT COUNT(*) as count from %s WHERE is_service_account=%s", ss.db.GetDialect().Quote("user"), ss.db.GetDialect().BooleanStr(false))
+		if _, err := sess.SQL(rawSQL).Get(&r); err != nil {
+			return err
+		}
+		return nil
+	})
+	return r.Count, err
+}
+
 // validateOneAdminLeft validate that there is an admin user left
 func validateOneAdminLeft(ctx context.Context, sess *db.Session) error {
 	count, err := sess.Where("is_admin=?", true).Count(&user.User{})

pkg/services/user/userimpl/user.go

@@ -12,6 +12,7 @@ import (
 	"github.com/grafana/grafana/pkg/models/roletype"
 	ac "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/team"
 	"github.com/grafana/grafana/pkg/services/user"
 	"github.com/grafana/grafana/pkg/setting"
@@ -32,15 +33,44 @@ func ProvideService(
 	cfg *setting.Cfg,
 	teamService team.Service,
 	cacheService *localcache.CacheService,
-) user.Service {
+	quotaService quota.Service,
+) (user.Service, error) {
 	store := ProvideStore(db, cfg)
-	return &Service{
+	s := &Service{
 		store:        &store,
 		orgService:   orgService,
 		cfg:          cfg,
 		teamService:  teamService,
 		cacheService: cacheService,
 	}
+
+	defaultLimits, err := readQuotaConfig(cfg)
+	if err != nil {
+		return s, err
+	}
+
+	if err := quotaService.RegisterQuotaReporter(&quota.NewUsageReporter{
+		TargetSrv:     quota.TargetSrv(user.QuotaTargetSrv),
+		DefaultLimits: defaultLimits,
+		Reporter:      s.Usage,
+	}); err != nil {
+		return s, err
+	}
+	return s, nil
+}
+
+func (s *Service) Usage(ctx context.Context, _ *quota.ScopeParameters) (*quota.Map, error) {
+	u := &quota.Map{}
+	if used, err := s.store.Count(ctx); err != nil {
+		return u, err
+	} else {
+		tag, err := quota.NewTag(quota.TargetSrv(user.QuotaTargetSrv), quota.Target(user.QuotaTarget), quota.GlobalScope)
+		if err != nil {
+			return u, err
+		}
+		u.Set(tag, used)
+	}
+	return u, nil
 }
 
 func (s *Service) Create(ctx context.Context, cmd *user.CreateUserCommand) (*user.User, error) {
@@ -304,3 +334,19 @@ func (s *Service) GetProfile(ctx context.Context, query *user.GetUserProfileQuer
 	result, err := s.store.GetProfile(ctx, query)
 	return result, err
 }
+
+func readQuotaConfig(cfg *setting.Cfg) (*quota.Map, error) {
+	limits := &quota.Map{}
+
+	if cfg == nil {
+		return limits, nil
+	}
+
+	globalQuotaTag, err := quota.NewTag(quota.TargetSrv(user.QuotaTargetSrv), quota.Target(user.QuotaTarget), quota.GlobalScope)
+	if err != nil {
+		return limits, err
+	}
+
+	limits.Set(globalQuotaTag, cfg.Quota.Global.User)
+	return limits, nil
+}

pkg/services/user/userimpl/user_test.go

@@ -252,3 +252,7 @@ func (f *FakeUserStore) Disable(ctx context.Context, cmd *user.DisableUserComman
 func (f *FakeUserStore) Search(ctx context.Context, query *user.SearchUsersQuery) (*user.SearchUserQueryResult, error) {
 	return f.ExpectedSearchUserQueryResult, f.ExpectedError
 }
+
+func (f *FakeUserStore) Count(ctx context.Context) (int64, error) {
+	return 0, nil
+}