diff --git a/pkg/api/quota.go b/pkg/api/quota.go index c78626c4ea1..9f0ff492fd1 100644 --- a/pkg/api/quota.go +++ b/pkg/api/quota.go @@ -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") diff --git a/pkg/middleware/quota_test.go b/pkg/middleware/quota_test.go index 2319c470191..c719e99b5da 100644 --- a/pkg/middleware/quota_test.go +++ b/pkg/middleware/quota_test.go @@ -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, diff --git a/pkg/services/quota/quota.go b/pkg/services/quota/quota.go index 7cb1cb115bd..7e2a923ceeb 100644 --- a/pkg/services/quota/quota.go +++ b/pkg/services/quota/quota.go @@ -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 { diff --git a/pkg/services/sqlstore/quota.go b/pkg/services/sqlstore/quota.go index 24d3bb9a802..08c85f076d4 100644 --- a/pkg/services/sqlstore/quota.go +++ b/pkg/services/sqlstore/quota.go @@ -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("a) - 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("a) + 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("as); 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("as); 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("a) - 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("a) + 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("as); 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("as); 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 + }) } diff --git a/pkg/services/sqlstore/quota_test.go b/pkg/services/sqlstore/quota_test.go index e2751f81c98..81a24316b94 100644 --- a/pkg/services/sqlstore/quota_test.go +++ b/pkg/services/sqlstore/quota_test.go @@ -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) }) diff --git a/pkg/services/sqlstore/sqlstore.go b/pkg/services/sqlstore/sqlstore.go index 122d2470a7f..03459a421f2 100644 --- a/pkg/services/sqlstore/sqlstore.go +++ b/pkg/services/sqlstore/sqlstore.go @@ -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 diff --git a/pkg/tests/api/alerting/api_alertmanager_test.go b/pkg/tests/api/alerting/api_alertmanager_test.go index 9ad29cdfbd6..2db8b6687f8 100644 --- a/pkg/tests/api/alerting/api_alertmanager_test.go +++ b/pkg/tests/api/alerting/api_alertmanager_test.go @@ -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) })