Chore: Refactor quota service (#57586)

* 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
+48
View File
@@ -16,6 +16,8 @@ import (
"github.com/grafana/grafana/pkg/infra/metrics"
ac "github.com/grafana/grafana/pkg/services/accesscontrol"
"github.com/grafana/grafana/pkg/services/datasources"
"github.com/grafana/grafana/pkg/services/quota"
"github.com/grafana/grafana/pkg/services/sqlstore"
"github.com/grafana/grafana/pkg/util"
)
@@ -29,6 +31,8 @@ type Store interface {
AddDataSource(context.Context, *datasources.AddDataSourceCommand) error
UpdateDataSource(context.Context, *datasources.UpdateDataSourceCommand) error
GetAllDataSources(ctx context.Context, query *datasources.GetAllDataSourcesQuery) error
Count(context.Context, *quota.ScopeParameters) (*quota.Map, error)
}
type SqlStore struct {
@@ -171,6 +175,50 @@ func (ss *SqlStore) DeleteDataSource(ctx context.Context, cmd *datasources.Delet
})
}
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 data_source"
if _, err := sess.SQL(rawSQL).Get(&r); err != nil {
return err
}
return nil
}); err != nil {
return u, err
} else {
tag, err := quota.NewTag(datasources.QuotaTargetSrv, datasources.QuotaTarget, 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 := "SELECT COUNT(*) AS count FROM data_source 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(datasources.QuotaTargetSrv, datasources.QuotaTarget, quota.OrgScope)
if err != nil {
return u, err
}
u.Set(tag, r.Count)
}
}
return u, nil
}
func (ss *SqlStore) AddDataSource(ctx context.Context, cmd *datasources.AddDataSourceCommand) error {
return ss.db.WithTransactionalDbSession(ctx, func(sess *db.Session) error {
existing := datasources.DataSource{OrgId: cmd.OrgId, Name: cmd.Name}