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: Propagate context for quotas (#39205)

Browse Source 
Propagate context for quotas and queries/commands.

Ref #36734
This commit is contained in:
Marcus Efraimsson 2021-09-20 17:05:30 +02:00 committed by GitHub
parent fef8e646ea
commit 1e4a660ea7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
7 changed files with 222 additions and 211 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
pkg/api/quota.go
View File

@ -14,7 +14,7 @@ func GetOrgQuotas(c *models.ReqContext) response.Response {
}
query := models.GetOrgQuotasQuery{OrgId: c.ParamsInt64(":orgId")}
if err := bus.Dispatch(&query); err != nil {
if err := bus.DispatchCtx(c.Req.Context(), &query); err != nil {
return response.Error(500, "Failed to get org quotas", err)
}
@ -32,7 +32,7 @@ func UpdateOrgQuota(c *models.ReqContext, cmd models.UpdateOrgQuotaCmd) response
return response.Error(404, "Invalid quota target", nil)
}
if err := bus.Dispatch(&cmd); err != nil {
if err := bus.DispatchCtx(c.Req.Context(), &cmd); err != nil {
return response.Error(500, "Failed to update org quotas", err)
}
return response.Success("Organization quota updated")
@ -44,7 +44,7 @@ func GetUserQuotas(c *models.ReqContext) response.Response {
}
query := models.GetUserQuotasQuery{UserId: c.ParamsInt64(":id")}
if err := bus.Dispatch(&query); err != nil {
if err := bus.DispatchCtx(c.Req.Context(), &query); err != nil {
return response.Error(500, "Failed to get org quotas", err)
}
@ -62,7 +62,7 @@ func UpdateUserQuota(c *models.ReqContext, cmd models.UpdateUserQuotaCmd) respon
return response.Error(404, "Invalid quota target", nil)
}
if err := bus.Dispatch(&cmd); err != nil {
if err := bus.DispatchCtx(c.Req.Context(), &cmd); err != nil {
return response.Error(500, "Failed to update org quotas", err)
}
return response.Success("Organization quota updated")

8
pkg/middleware/quota_test.go
View File

@ -16,7 +16,7 @@ import (
func TestMiddlewareQuota(t *testing.T) {
t.Run("With user not logged in", func(t *testing.T) {
middlewareScenario(t, "and global quota not reached", func(t *testing.T, sc *scenarioContext) {
bus.AddHandler("globalQuota", func(query *models.GetGlobalQuotaByTargetQuery) error {
bus.AddHandlerCtx("globalQuota", func(_ context.Context, query *models.GetGlobalQuotaByTargetQuery) error {
query.Result = &models.GlobalQuotaDTO{
Target: query.Target,
Limit: query.Default,
@ -33,7 +33,7 @@ func TestMiddlewareQuota(t *testing.T) {
}, configure)
middlewareScenario(t, "and global quota reached", func(t *testing.T, sc *scenarioContext) {
bus.AddHandler("globalQuota", func(query *models.GetGlobalQuotaByTargetQuery) error {
bus.AddHandlerCtx("globalQuota", func(_ context.Context, query *models.GetGlobalQuotaByTargetQuery) error {
query.Result = &models.GlobalQuotaDTO{
Target: query.Target,
Limit: query.Default,
@ -53,7 +53,7 @@ func TestMiddlewareQuota(t *testing.T) {
})
middlewareScenario(t, "and global session quota not reached", func(t *testing.T, sc *scenarioContext) {
bus.AddHandler("globalQuota", func(query *models.GetGlobalQuotaByTargetQuery) error {
bus.AddHandlerCtx("globalQuota", func(_ context.Context, query *models.GetGlobalQuotaByTargetQuery) error {
query.Result = &models.GlobalQuotaDTO{
Target: query.Target,
Limit: query.Default,
@ -101,7 +101,7 @@ func TestMiddlewareQuota(t *testing.T) {
}, nil
}
bus.AddHandler("globalQuota", func(query *models.GetGlobalQuotaByTargetQuery) error {
bus.AddHandlerCtx("globalQuota", func(_ context.Context, query *models.GetGlobalQuotaByTargetQuery) error {
query.Result = &models.GlobalQuotaDTO{
Target: query.Target,
Limit: query.Default,

6
pkg/services/quota/quota.go
View File

@ -63,7 +63,7 @@ func (qs *QuotaService) QuotaReached(c *models.ReqContext, target string) (bool,
continue
}
query := models.GetGlobalQuotaByTargetQuery{Target: scope.Target, IsNgAlertEnabled: qs.Cfg.IsNgAlertEnabled()}
if err := bus.Dispatch(&query); err != nil {
if err := bus.DispatchCtx(c.Req.Context(), &query); err != nil {
return true, err
}
if query.Result.Used >= scope.DefaultLimit {
@ -79,7 +79,7 @@ func (qs *QuotaService) QuotaReached(c *models.ReqContext, target string) (bool,
Default: scope.DefaultLimit,
IsNgAlertEnabled: qs.Cfg.IsNgAlertEnabled(),
}
if err := bus.Dispatch(&query); err != nil {
if err := bus.DispatchCtx(c.Req.Context(), &query); err != nil {
return true, err
}
if query.Result.Limit < 0 {
@ -97,7 +97,7 @@ func (qs *QuotaService) QuotaReached(c *models.ReqContext, target string) (bool,
continue
}
query := models.GetUserQuotaByTargetQuery{UserId: c.UserId, Target: scope.Target, Default: scope.DefaultLimit, IsNgAlertEnabled: qs.Cfg.IsNgAlertEnabled()}
if err := bus.Dispatch(&query); err != nil {
if err := bus.DispatchCtx(c.Req.Context(), &query); err != nil {
return true, err
}
if query.Result.Limit < 0 {

347
pkg/services/sqlstore/quota.go
View File

@ -1,6 +1,7 @@
package sqlstore
import (
"context"
"fmt"
"time"
@ -14,108 +15,111 @@ const (
dashboardTarget = "dashboard"
)
func init() {
bus.AddHandler("sql", GetOrgQuotaByTarget)
bus.AddHandler("sql", GetOrgQuotas)
bus.AddHandler("sql", UpdateOrgQuota)
bus.AddHandler("sql", GetUserQuotaByTarget)
bus.AddHandler("sql", GetUserQuotas)
bus.AddHandler("sql", UpdateUserQuota)
bus.AddHandler("sql", GetGlobalQuotaByTarget)
func (ss *SQLStore) addQuotaQueryAndCommandHandlers() {
bus.AddHandlerCtx("sql", ss.GetOrgQuotaByTarget)
bus.AddHandlerCtx("sql", ss.GetOrgQuotas)
bus.AddHandlerCtx("sql", ss.UpdateOrgQuota)
bus.AddHandlerCtx("sql", ss.GetUserQuotaByTarget)
bus.AddHandlerCtx("sql", ss.GetUserQuotas)
bus.AddHandlerCtx("sql", ss.UpdateUserQuota)
bus.AddHandlerCtx("sql", ss.GetGlobalQuotaByTarget)
}
type targetCount struct {
Count int64
}
func GetOrgQuotaByTarget(query *models.GetOrgQuotaByTargetQuery) error {
quota := models.Quota{
Target: query.Target,
OrgId: query.OrgId,
}
has, err := x.Get(&quota)
if err != nil {
return err
} else if !has {
quota.Limit = query.Default
}
var used int64
if query.Target != alertRuleTarget || query.IsNgAlertEnabled {
// get quota used.
rawSQL := fmt.Sprintf("SELECT COUNT(*) AS count FROM %s WHERE org_id=?",
dialect.Quote(query.Target))
if query.Target == dashboardTarget {
rawSQL += fmt.Sprintf(" AND is_folder=%s", dialect.BooleanStr(false))
func (ss *SQLStore) GetOrgQuotaByTarget(ctx context.Context, query *models.GetOrgQuotaByTargetQuery) error {
return ss.WithDbSession(ctx, func(sess *DBSession) error {
quota := models.Quota{
Target: query.Target,
OrgId: query.OrgId,
}
resp := make([]*targetCount, 0)
if err := x.SQL(rawSQL, query.OrgId).Find(&resp); err != nil {
has, err := sess.Get(&quota)
if err != nil {
return err
} else if !has {
quota.Limit = query.Default
}
used = resp[0].Count
}
query.Result = &models.OrgQuotaDTO{
Target: query.Target,
Limit: quota.Limit,
OrgId: query.OrgId,
Used: used,
}
return nil
}
func GetOrgQuotas(query *models.GetOrgQuotasQuery) error {
quotas := make([]*models.Quota, 0)
sess := x.Table("quota")
if err := sess.Where("org_id=? AND user_id=0", query.OrgId).Find(&quotas); err != nil {
return err
}
defaultQuotas := setting.Quota.Org.ToMap()
seenTargets := make(map[string]bool)
for _, q := range quotas {
seenTargets[q.Target] = true
}
for t, v := range defaultQuotas {
if _, ok := seenTargets[t]; !ok {
quotas = append(quotas, &models.Quota{
OrgId: query.OrgId,
Target: t,
Limit: v,
})
}
}
result := make([]*models.OrgQuotaDTO, len(quotas))
for i, q := range quotas {
var used int64
if q.Target != alertRuleTarget || query.IsNgAlertEnabled {
if query.Target != alertRuleTarget || query.IsNgAlertEnabled {
// get quota used.
rawSQL := fmt.Sprintf("SELECT COUNT(*) as count from %s where org_id=?", dialect.Quote(q.Target))
rawSQL := fmt.Sprintf("SELECT COUNT(*) AS count FROM %s WHERE org_id=?",
dialect.Quote(query.Target))
if query.Target == dashboardTarget {
rawSQL += fmt.Sprintf(" AND is_folder=%s", dialect.BooleanStr(false))
}
resp := make([]*targetCount, 0)
if err := x.SQL(rawSQL, q.OrgId).Find(&resp); err != nil {
if err := sess.SQL(rawSQL, query.OrgId).Find(&resp); err != nil {
return err
}
used = resp[0].Count
}
result[i] = &models.OrgQuotaDTO{
Target: q.Target,
Limit: q.Limit,
OrgId: q.OrgId,
query.Result = &models.OrgQuotaDTO{
Target: query.Target,
Limit: quota.Limit,
OrgId: query.OrgId,
Used: used,
}
}
query.Result = result
return nil
return nil
})
}
func UpdateOrgQuota(cmd *models.UpdateOrgQuotaCmd) error {
return inTransaction(func(sess *DBSession) error {
func (ss *SQLStore) GetOrgQuotas(ctx context.Context, query *models.GetOrgQuotasQuery) error {
return ss.WithDbSession(ctx, func(sess *DBSession) error {
quotas := make([]*models.Quota, 0)
if err := sess.Table("quota").Where("org_id=? AND user_id=0", query.OrgId).Find(&quotas); err != nil {
return err
}
defaultQuotas := setting.Quota.Org.ToMap()
seenTargets := make(map[string]bool)
for _, q := range quotas {
seenTargets[q.Target] = true
}
for t, v := range defaultQuotas {
if _, ok := seenTargets[t]; !ok {
quotas = append(quotas, &models.Quota{
OrgId: query.OrgId,
Target: t,
Limit: v,
})
}
}
result := make([]*models.OrgQuotaDTO, len(quotas))
for i, q := range quotas {
var used int64
if q.Target != alertRuleTarget || query.IsNgAlertEnabled {
// get quota used.
rawSQL := fmt.Sprintf("SELECT COUNT(*) as count from %s where org_id=?", dialect.Quote(q.Target))
resp := make([]*targetCount, 0)
if err := sess.SQL(rawSQL, q.OrgId).Find(&resp); err != nil {
return err
}
used = resp[0].Count
}
result[i] = &models.OrgQuotaDTO{
Target: q.Target,
Limit: q.Limit,
OrgId: q.OrgId,
Used: used,
}
}
query.Result = result
return nil
})
}
func (ss *SQLStore) UpdateOrgQuota(ctx context.Context, cmd *models.UpdateOrgQuotaCmd) error {
return ss.WithTransactionalDbSession(ctx, func(sess *DBSession) error {
// Check if quota is already defined in the DB
quota := models.Quota{
Target: cmd.Target,
@ -145,88 +149,91 @@ func UpdateOrgQuota(cmd *models.UpdateOrgQuotaCmd) error {
})
}
func GetUserQuotaByTarget(query *models.GetUserQuotaByTargetQuery) error {
quota := models.Quota{
Target: query.Target,
UserId: query.UserId,
}
has, err := x.Get(&quota)
if err != nil {
return err
} else if !has {
quota.Limit = query.Default
}
var used int64
if query.Target != alertRuleTarget || query.IsNgAlertEnabled {
// get quota used.
rawSQL := fmt.Sprintf("SELECT COUNT(*) as count from %s where user_id=?", dialect.Quote(query.Target))
resp := make([]*targetCount, 0)
if err := x.SQL(rawSQL, query.UserId).Find(&resp); err != nil {
func (ss *SQLStore) GetUserQuotaByTarget(ctx context.Context, query *models.GetUserQuotaByTargetQuery) error {
return ss.WithDbSession(ctx, func(sess *DBSession) error {
quota := models.Quota{
Target: query.Target,
UserId: query.UserId,
}
has, err := sess.Get(&quota)
if err != nil {
return err
} else if !has {
quota.Limit = query.Default
}
used = resp[0].Count
}
query.Result = &models.UserQuotaDTO{
Target: query.Target,
Limit: quota.Limit,
UserId: query.UserId,
Used: used,
}
return nil
}
func GetUserQuotas(query *models.GetUserQuotasQuery) error {
quotas := make([]*models.Quota, 0)
sess := x.Table("quota")
if err := sess.Where("user_id=? AND org_id=0", query.UserId).Find(&quotas); err != nil {
return err
}
defaultQuotas := setting.Quota.User.ToMap()
seenTargets := make(map[string]bool)
for _, q := range quotas {
seenTargets[q.Target] = true
}
for t, v := range defaultQuotas {
if _, ok := seenTargets[t]; !ok {
quotas = append(quotas, &models.Quota{
UserId: query.UserId,
Target: t,
Limit: v,
})
}
}
result := make([]*models.UserQuotaDTO, len(quotas))
for i, q := range quotas {
var used int64
if q.Target != alertRuleTarget || query.IsNgAlertEnabled {
if query.Target != alertRuleTarget || query.IsNgAlertEnabled {
// get quota used.
rawSQL := fmt.Sprintf("SELECT COUNT(*) as count from %s where user_id=?", dialect.Quote(q.Target))
rawSQL := fmt.Sprintf("SELECT COUNT(*) as count from %s where user_id=?", dialect.Quote(query.Target))
resp := make([]*targetCount, 0)
if err := x.SQL(rawSQL, q.UserId).Find(&resp); err != nil {
if err := sess.SQL(rawSQL, query.UserId).Find(&resp); err != nil {
return err
}
used = resp[0].Count
}
result[i] = &models.UserQuotaDTO{
Target: q.Target,
Limit: q.Limit,
UserId: q.UserId,
query.Result = &models.UserQuotaDTO{
Target: query.Target,
Limit: quota.Limit,
UserId: query.UserId,
Used: used,
}
}
query.Result = result
return nil
return nil
})
}
func UpdateUserQuota(cmd *models.UpdateUserQuotaCmd) error {
return inTransaction(func(sess *DBSession) error {
func (ss *SQLStore) GetUserQuotas(ctx context.Context, query *models.GetUserQuotasQuery) error {
return ss.WithDbSession(ctx, func(sess *DBSession) error {
quotas := make([]*models.Quota, 0)
if err := sess.Table("quota").Where("user_id=? AND org_id=0", query.UserId).Find(&quotas); err != nil {
return err
}
defaultQuotas := setting.Quota.User.ToMap()
seenTargets := make(map[string]bool)
for _, q := range quotas {
seenTargets[q.Target] = true
}
for t, v := range defaultQuotas {
if _, ok := seenTargets[t]; !ok {
quotas = append(quotas, &models.Quota{
UserId: query.UserId,
Target: t,
Limit: v,
})
}
}
result := make([]*models.UserQuotaDTO, len(quotas))
for i, q := range quotas {
var used int64
if q.Target != alertRuleTarget || query.IsNgAlertEnabled {
// get quota used.
rawSQL := fmt.Sprintf("SELECT COUNT(*) as count from %s where user_id=?", dialect.Quote(q.Target))
resp := make([]*targetCount, 0)
if err := sess.SQL(rawSQL, q.UserId).Find(&resp); err != nil {
return err
}
used = resp[0].Count
}
result[i] = &models.UserQuotaDTO{
Target: q.Target,
Limit: q.Limit,
UserId: q.UserId,
Used: used,
}
}
query.Result = result
return nil
})
}
func (ss *SQLStore) UpdateUserQuota(ctx context.Context, cmd *models.UpdateUserQuotaCmd) error {
return ss.WithTransactionalDbSession(ctx, func(sess *DBSession) error {
// Check if quota is already defined in the DB
quota := models.Quota{
Target: cmd.Target,
@ -256,29 +263,31 @@ func UpdateUserQuota(cmd *models.UpdateUserQuotaCmd) error {
})
}
func GetGlobalQuotaByTarget(query *models.GetGlobalQuotaByTargetQuery) error {
var used int64
if query.Target != alertRuleTarget || query.IsNgAlertEnabled {
// get quota used.
rawSQL := fmt.Sprintf("SELECT COUNT(*) AS count FROM %s",
dialect.Quote(query.Target))
func (ss *SQLStore) GetGlobalQuotaByTarget(ctx context.Context, query *models.GetGlobalQuotaByTargetQuery) error {
return ss.WithDbSession(ctx, func(sess *DBSession) error {
var used int64
if query.Target != alertRuleTarget || query.IsNgAlertEnabled {
// get quota used.
rawSQL := fmt.Sprintf("SELECT COUNT(*) AS count FROM %s",
dialect.Quote(query.Target))
if query.Target == dashboardTarget {
rawSQL += fmt.Sprintf(" WHERE is_folder=%s", dialect.BooleanStr(false))
if query.Target == dashboardTarget {
rawSQL += fmt.Sprintf(" WHERE is_folder=%s", dialect.BooleanStr(false))
}
resp := make([]*targetCount, 0)
if err := sess.SQL(rawSQL).Find(&resp); err != nil {
return err
}
used = resp[0].Count
}
resp := make([]*targetCount, 0)
if err := x.SQL(rawSQL).Find(&resp); err != nil {
return err
query.Result = &models.GlobalQuotaDTO{
Target: query.Target,
Limit: query.Default,
Used: used,
}
used = resp[0].Count
}
query.Result = &models.GlobalQuotaDTO{
Target: query.Target,
Limit: query.Default,
Used: used,
}
return nil
return nil
})
}

57
pkg/services/sqlstore/quota_test.go
View File

@ -4,6 +4,7 @@
package sqlstore
import (
"context"
"testing"
"time"
@ -13,7 +14,7 @@ import (
)
func TestQuotaCommandsAndQueries(t *testing.T) {
InitTestDB(t)
sqlStore := InitTestDB(t)
userId := int64(1)
orgId := int64(0)
@ -58,12 +59,12 @@ func TestQuotaCommandsAndQueries(t *testing.T) {
Target: "org_user",
Limit: 10,
}
err := UpdateOrgQuota(&orgCmd)
err := sqlStore.UpdateOrgQuota(context.Background(), &orgCmd)
require.NoError(t, err)
t.Run("Should be able to get saved quota by org id and target", func(t *testing.T) {
query := models.GetOrgQuotaByTargetQuery{OrgId: orgId, Target: "org_user", Default: 1}
err = GetOrgQuotaByTarget(&query)
err = sqlStore.GetOrgQuotaByTarget(context.Background(), &query)
require.NoError(t, err)
require.Equal(t, int64(10), query.Result.Limit)
@ -71,7 +72,7 @@ func TestQuotaCommandsAndQueries(t *testing.T) {
t.Run("Should be able to get default quota by org id and target", func(t *testing.T) {
query := models.GetOrgQuotaByTargetQuery{OrgId: 123, Target: "org_user", Default: 11}
err = GetOrgQuotaByTarget(&query)
err = sqlStore.GetOrgQuotaByTarget(context.Background(), &query)
require.NoError(t, err)
require.Equal(t, int64(11), query.Result.Limit)
@ -79,7 +80,7 @@ func TestQuotaCommandsAndQueries(t *testing.T) {
t.Run("Should be able to get used org quota when rows exist", func(t *testing.T) {
query := models.GetOrgQuotaByTargetQuery{OrgId: orgId, Target: "org_user", Default: 11}
err = GetOrgQuotaByTarget(&query)
err = sqlStore.GetOrgQuotaByTarget(context.Background(), &query)
require.NoError(t, err)
require.Equal(t, int64(1), query.Result.Used)
@ -87,7 +88,7 @@ func TestQuotaCommandsAndQueries(t *testing.T) {
t.Run("Should be able to get used org quota when no rows exist", func(t *testing.T) {
query := models.GetOrgQuotaByTargetQuery{OrgId: 2, Target: "org_user", Default: 11}
err = GetOrgQuotaByTarget(&query)
err = sqlStore.GetOrgQuotaByTarget(context.Background(), &query)
require.NoError(t, err)
require.Equal(t, int64(0), query.Result.Used)
@ -95,7 +96,7 @@ func TestQuotaCommandsAndQueries(t *testing.T) {
t.Run("Should be able to get zero used org alert quota when table does not exist (ngalert is not enabled - default case)", func(t *testing.T) {
query := models.GetOrgQuotaByTargetQuery{OrgId: 2, Target: "alert", Default: 11}
err = GetOrgQuotaByTarget(&query)
err = sqlStore.GetOrgQuotaByTarget(context.Background(), &query)
require.NoError(t, err)
require.Equal(t, int64(0), query.Result.Used)
@ -103,7 +104,7 @@ func TestQuotaCommandsAndQueries(t *testing.T) {
t.Run("Should be able to quota list for org", func(t *testing.T) {
query := models.GetOrgQuotasQuery{OrgId: orgId}
err = GetOrgQuotas(&query)
err = sqlStore.GetOrgQuotas(context.Background(), &query)
require.NoError(t, err)
require.Len(t, query.Result, 5)
@ -126,12 +127,12 @@ func TestQuotaCommandsAndQueries(t *testing.T) {
Target: dashboardTarget,
Limit: 10,
}
err := UpdateOrgQuota(&orgCmd)
err := sqlStore.UpdateOrgQuota(context.Background(), &orgCmd)
require.NoError(t, err)
t.Run("Should be able to get saved quota by org id and target", func(t *testing.T) {
query := models.GetOrgQuotaByTargetQuery{OrgId: orgId, Target: dashboardTarget, Default: 1}
err = GetOrgQuotaByTarget(&query)
err = sqlStore.GetOrgQuotaByTarget(context.Background(), &query)
require.NoError(t, err)
require.Equal(t, int64(10), query.Result.Limit)
@ -145,12 +146,12 @@ func TestQuotaCommandsAndQueries(t *testing.T) {
Target: "org_user",
Limit: 10,
}
err := UpdateUserQuota(&userQuotaCmd)
err := sqlStore.UpdateUserQuota(context.Background(), &userQuotaCmd)
require.NoError(t, err)
t.Run("Should be able to get saved quota by user id and target", func(t *testing.T) {
query := models.GetUserQuotaByTargetQuery{UserId: userId, Target: "org_user", Default: 1}
err = GetUserQuotaByTarget(&query)
err = sqlStore.GetUserQuotaByTarget(context.Background(), &query)
require.NoError(t, err)
require.Equal(t, int64(10), query.Result.Limit)
@ -158,7 +159,7 @@ func TestQuotaCommandsAndQueries(t *testing.T) {
t.Run("Should be able to get default quota by user id and target", func(t *testing.T) {
query := models.GetUserQuotaByTargetQuery{UserId: 9, Target: "org_user", Default: 11}
err = GetUserQuotaByTarget(&query)
err = sqlStore.GetUserQuotaByTarget(context.Background(), &query)
require.NoError(t, err)
require.Equal(t, int64(11), query.Result.Limit)
@ -166,7 +167,7 @@ func TestQuotaCommandsAndQueries(t *testing.T) {
t.Run("Should be able to get used user quota when rows exist", func(t *testing.T) {
query := models.GetUserQuotaByTargetQuery{UserId: userId, Target: "org_user", Default: 11}
err = GetUserQuotaByTarget(&query)
err = sqlStore.GetUserQuotaByTarget(context.Background(), &query)
require.NoError(t, err)
require.Equal(t, int64(1), query.Result.Used)
@ -174,7 +175,7 @@ func TestQuotaCommandsAndQueries(t *testing.T) {
t.Run("Should be able to get used user quota when no rows exist", func(t *testing.T) {
query := models.GetUserQuotaByTargetQuery{UserId: 2, Target: "org_user", Default: 11}
err = GetUserQuotaByTarget(&query)
err = sqlStore.GetUserQuotaByTarget(context.Background(), &query)
require.NoError(t, err)
require.Equal(t, int64(0), query.Result.Used)
@ -182,7 +183,7 @@ func TestQuotaCommandsAndQueries(t *testing.T) {
t.Run("Should be able to quota list for user", func(t *testing.T) {
query := models.GetUserQuotasQuery{UserId: userId}
err = GetUserQuotas(&query)
err = sqlStore.GetUserQuotas(context.Background(), &query)
require.NoError(t, err)
require.Len(t, query.Result, 1)
@ -193,7 +194,7 @@ func TestQuotaCommandsAndQueries(t *testing.T) {
t.Run("Should be able to global user quota", func(t *testing.T) {
query := models.GetGlobalQuotaByTargetQuery{Target: "user", Default: 5}
err = GetGlobalQuotaByTarget(&query)
err = sqlStore.GetGlobalQuotaByTarget(context.Background(), &query)
require.NoError(t, err)
require.Equal(t, int64(5), query.Result.Limit)
@ -202,7 +203,7 @@ func TestQuotaCommandsAndQueries(t *testing.T) {
t.Run("Should be able to global org quota", func(t *testing.T) {
query := models.GetGlobalQuotaByTargetQuery{Target: "org", Default: 5}
err = GetGlobalQuotaByTarget(&query)
err = sqlStore.GetGlobalQuotaByTarget(context.Background(), &query)
require.NoError(t, err)
require.Equal(t, int64(5), query.Result.Limit)
@ -211,7 +212,7 @@ func TestQuotaCommandsAndQueries(t *testing.T) {
t.Run("Should be able to get zero used global alert quota when table does not exist (ngalert is not enabled - default case)", func(t *testing.T) {
query := models.GetGlobalQuotaByTargetQuery{Target: "alert_rule", Default: 5}
err = GetGlobalQuotaByTarget(&query)
err = sqlStore.GetGlobalQuotaByTarget(context.Background(), &query)
require.NoError(t, err)
require.Equal(t, int64(5), query.Result.Limit)
@ -220,7 +221,7 @@ func TestQuotaCommandsAndQueries(t *testing.T) {
t.Run("Should be able to global dashboard quota", func(t *testing.T) {
query := models.GetGlobalQuotaByTargetQuery{Target: dashboardTarget, Default: 5}
err = GetGlobalQuotaByTarget(&query)
err = sqlStore.GetGlobalQuotaByTarget(context.Background(), &query)
require.NoError(t, err)
require.Equal(t, int64(5), query.Result.Limit)
@ -234,11 +235,11 @@ func TestQuotaCommandsAndQueries(t *testing.T) {
Target: "org_user",
Limit: 5,
}
err := UpdateOrgQuota(&orgCmd)
err := sqlStore.UpdateOrgQuota(context.Background(), &orgCmd)
require.NoError(t, err)
query := models.GetOrgQuotaByTargetQuery{OrgId: orgId, Target: "org_user", Default: 1}
err = GetOrgQuotaByTarget(&query)
err = sqlStore.GetOrgQuotaByTarget(context.Background(), &query)
require.NoError(t, err)
require.Equal(t, int64(5), query.Result.Limit)
@ -250,11 +251,11 @@ func TestQuotaCommandsAndQueries(t *testing.T) {
Target: "org_user",
Limit: 10,
}
err = UpdateOrgQuota(&orgCmd)
err = sqlStore.UpdateOrgQuota(context.Background(), &orgCmd)
require.NoError(t, err)
query = models.GetOrgQuotaByTargetQuery{OrgId: orgId, Target: "org_user", Default: 1}
err = GetOrgQuotaByTarget(&query)
err = sqlStore.GetOrgQuotaByTarget(context.Background(), &query)
require.NoError(t, err)
require.Equal(t, int64(10), query.Result.Limit)
})
@ -266,11 +267,11 @@ func TestQuotaCommandsAndQueries(t *testing.T) {
Target: "org_user",
Limit: 5,
}
err := UpdateUserQuota(&userQuotaCmd)
err := sqlStore.UpdateUserQuota(context.Background(), &userQuotaCmd)
require.NoError(t, err)
query := models.GetUserQuotaByTargetQuery{UserId: userId, Target: "org_user", Default: 1}
err = GetUserQuotaByTarget(&query)
err = sqlStore.GetUserQuotaByTarget(context.Background(), &query)
require.NoError(t, err)
require.Equal(t, int64(5), query.Result.Limit)
@ -282,11 +283,11 @@ func TestQuotaCommandsAndQueries(t *testing.T) {
Target: "org_user",
Limit: 10,
}
err = UpdateUserQuota(&userQuotaCmd)
err = sqlStore.UpdateUserQuota(context.Background(), &userQuotaCmd)
require.NoError(t, err)
query = models.GetUserQuotaByTargetQuery{UserId: userId, Target: "org_user", Default: 1}
err = GetUserQuotaByTarget(&query)
err = sqlStore.GetUserQuotaByTarget(context.Background(), &query)
require.NoError(t, err)
require.Equal(t, int64(10), query.Result.Limit)
})

1
pkg/services/sqlstore/sqlstore.go
View File

@ -113,6 +113,7 @@ func newSQLStore(cfg *setting.Cfg, cacheService *localcache.CacheService, bus bu
ss.addAlertNotificationUidByIdHandler()
ss.addPreferencesQueryAndCommandHandlers()
ss.addDashboardQueryAndCommandHandlers()
ss.addQuotaQueryAndCommandHandlers()
// if err := ss.Reset(); err != nil {
// return nil, err

6
pkg/tests/api/alerting/api_alertmanager_test.go
View File

@ -2018,7 +2018,7 @@ func TestQuota(t *testing.T) {
t.Run("when quota limit exceed creating new rule should fail", func(t *testing.T) {
// get existing org quota
query := models.GetOrgQuotaByTargetQuery{OrgId: 1, Target: "alert_rule"}
err = sqlstore.GetOrgQuotaByTarget(&query)
err = store.GetOrgQuotaByTarget(context.Background(), &query)
require.NoError(t, err)
used := query.Result.Used
limit := query.Result.Limit
@ -2029,7 +2029,7 @@ func TestQuota(t *testing.T) {
Target: "alert_rule",
Limit: used,
}
err := sqlstore.UpdateOrgQuota(&orgCmd)
err := store.UpdateOrgQuota(context.Background(), &orgCmd)
require.NoError(t, err)
t.Cleanup(func() {
@ -2039,7 +2039,7 @@ func TestQuota(t *testing.T) {
Target: "alert_rule",
Limit: limit,
}
err := sqlstore.UpdateOrgQuota(&orgCmd)
err := store.UpdateOrgQuota(context.Background(), &orgCmd)
require.NoError(t, err)
})