[Alerting]: Extend quota service to optionally set limits on alerts (#33283)

* Quota: Extend service to set limit on alerts

* Add test for applying quota to alert rules

* Apply suggestions from code review

Co-authored-by: Diana Payton <52059945+oddlittlebird@users.noreply.github.com>

* Get used alert quota only if naglert is enabled

* Set alert limit to zero if nglalert is not enabled
Co-authored-by: Diana Payton <52059945+oddlittlebird@users.noreply.github.com>

pkg/services/ngalert/api/api.go

@@ -3,6 +3,8 @@ package api
 import (
 	"time"
 
+	"github.com/grafana/grafana/pkg/services/quota"
+
 	"github.com/grafana/grafana/pkg/services/ngalert/metrics"
 	"github.com/grafana/grafana/pkg/services/ngalert/state"
 
@@ -41,6 +43,7 @@ type API struct {
 	DatasourceCache datasources.CacheService
 	RouteRegister   routing.RouteRegister
 	DataService     *tsdb.Service
+	QuotaService    *quota.QuotaService
 	Schedule        schedule.ScheduleService
 	RuleStore       store.RuleStore
 	InstanceStore   store.InstanceStore
@@ -73,7 +76,7 @@ func (api *API) RegisterAPIEndpoints(m *metrics.Metrics) {
 	api.RegisterRulerApiEndpoints(NewForkedRuler(
 		api.DatasourceCache,
 		NewLotexRuler(proxy, logger),
-		RulerSrv{DatasourceCache: api.DatasourceCache, manager: api.StateManager, store: api.RuleStore, log: logger},
+		RulerSrv{DatasourceCache: api.DatasourceCache, QuotaService: api.QuotaService, manager: api.StateManager, store: api.RuleStore, log: logger},
 	), m)
 	api.RegisterTestingApiEndpoints(TestingApiSrv{
 		AlertingProxy:   proxy,

pkg/services/ngalert/api/api_ruler.go

@@ -9,6 +9,7 @@ import (
 	"github.com/grafana/grafana/pkg/services/datasources"
 	"github.com/grafana/grafana/pkg/services/ngalert/state"
 	"github.com/grafana/grafana/pkg/services/ngalert/store"
+	"github.com/grafana/grafana/pkg/services/quota"
 
 	coreapi "github.com/grafana/grafana/pkg/api"
 	"github.com/grafana/grafana/pkg/api/response"
@@ -23,6 +24,7 @@ import (
 type RulerSrv struct {
 	store           store.RuleStore
 	DatasourceCache datasources.CacheService
+	QuotaService    *quota.QuotaService
 	manager         *state.Manager
 	log             log.Logger
 }
@@ -209,8 +211,18 @@ func (srv RulerSrv) RoutePostNameRulesConfig(c *models.ReqContext, ruleGroupConf
 		return toNamespaceErrorResponse(err)
 	}
 
-	// TODO check permissions
-	// TODO check quota
+	// quotas are checked in advanced
+	// that is acceptable under the assumption that there will be only one alert rule under the rule group
+	// alternatively we should check the quotas after the rule group update
+	// and rollback the transaction in case of violation
+	limitReached, err := srv.QuotaService.QuotaReached(c, "alert_rule")
+	if err != nil {
+		return response.Error(http.StatusInternalServerError, "failed to get quota", err)
+	}
+	if limitReached {
+		return response.Error(http.StatusForbidden, "quota reached", nil)
+	}
+
 	// TODO validate UID uniqueness in the payload
 
 	//TODO: Should this belong in alerting-api?

pkg/services/ngalert/ngalert.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"time"
 
+	"github.com/grafana/grafana/pkg/services/quota"
+
 	"github.com/grafana/grafana/pkg/services/ngalert/metrics"
 	"github.com/grafana/grafana/pkg/services/ngalert/state"
 
@@ -46,6 +48,7 @@ type AlertNG struct {
 	DataService     *tsdb.Service                           `inject:""`
 	Alertmanager    *notifier.Alertmanager                  `inject:""`
 	DataProxy       *datasourceproxy.DatasourceProxyService `inject:""`
+	QuotaService    *quota.QuotaService                     `inject:""`
 	Metrics         *metrics.Metrics                        `inject:""`
 	Log             log.Logger
 	schedule        schedule.ScheduleService
@@ -83,6 +86,7 @@ func (ng *AlertNG) Init() error {
 		DataService:     ng.DataService,
 		Schedule:        ng.schedule,
 		DataProxy:       ng.DataProxy,
+		QuotaService:    ng.QuotaService,
 		InstanceStore:   store,
 		RuleStore:       store,
 		AlertingStore:   store,

pkg/services/quota/quota.go

@@ -64,7 +64,7 @@ func (qs *QuotaService) QuotaReached(c *models.ReqContext, target string) (bool,
 				}
 				continue
 			}
-			query := models.GetGlobalQuotaByTargetQuery{Target: scope.Target}
+			query := models.GetGlobalQuotaByTargetQuery{Target: scope.Target, IsNgAlertEnabled: qs.Cfg.IsNgAlertEnabled()}
 			if err := bus.Dispatch(&query); err != nil {
 				return true, err
 			}
@@ -75,7 +75,7 @@ func (qs *QuotaService) QuotaReached(c *models.ReqContext, target string) (bool,
 			if !c.IsSignedIn {
 				continue
 			}
-			query := models.GetOrgQuotaByTargetQuery{OrgId: c.OrgId, Target: scope.Target, Default: scope.DefaultLimit}
+			query := models.GetOrgQuotaByTargetQuery{OrgId: c.OrgId, Target: scope.Target, Default: scope.DefaultLimit, IsNgAlertEnabled: qs.Cfg.IsNgAlertEnabled()}
 			if err := bus.Dispatch(&query); err != nil {
 				return true, err
 			}
@@ -93,7 +93,7 @@ func (qs *QuotaService) QuotaReached(c *models.ReqContext, target string) (bool,
 			if !c.IsSignedIn || c.UserId == 0 {
 				continue
 			}
-			query := models.GetUserQuotaByTargetQuery{UserId: c.UserId, Target: scope.Target, Default: scope.DefaultLimit}
+			query := models.GetUserQuotaByTargetQuery{UserId: c.UserId, Target: scope.Target, Default: scope.DefaultLimit, IsNgAlertEnabled: qs.Cfg.IsNgAlertEnabled()}
 			if err := bus.Dispatch(&query); err != nil {
 				return true, err
 			}
@@ -151,6 +151,12 @@ func (qs *QuotaService) getQuotaScopes(target string) ([]models.QuotaScope, erro
 			models.QuotaScope{Name: "global", Target: target, DefaultLimit: qs.Cfg.Quota.Global.Session},
 		)
 		return scopes, nil
+	case "alert_rule": // target need to match the respective database name
+		scopes = append(scopes,
+			models.QuotaScope{Name: "global", Target: target, DefaultLimit: qs.Cfg.Quota.Global.AlertRule},
+			models.QuotaScope{Name: "org", Target: target, DefaultLimit: qs.Cfg.Quota.Org.AlertRule},
+		)
+		return scopes, nil
 	default:
 		return scopes, ErrInvalidQuotaTarget
 	}

pkg/services/sqlstore/quota.go

@@ -9,6 +9,8 @@ import (
 	"github.com/grafana/grafana/pkg/setting"
 )
 
+const ALERT_RULE_TARGET = "alert_rule"
+
 func init() {
 	bus.AddHandler("sql", GetOrgQuotaByTarget)
 	bus.AddHandler("sql", GetOrgQuotas)
@@ -35,18 +37,22 @@ func GetOrgQuotaByTarget(query *models.GetOrgQuotaByTargetQuery) error {
 		quota.Limit = query.Default
 	}
 
-	// get quota used.
-	rawSQL := fmt.Sprintf("SELECT COUNT(*) as count from %s where org_id=?", dialect.Quote(query.Target))
-	resp := make([]*targetCount, 0)
-	if err := x.SQL(rawSQL, query.OrgId).Find(&resp); err != nil {
-		return err
+	var used int64
+	if query.Target != ALERT_RULE_TARGET || query.IsNgAlertEnabled {
+		// get quota used.
+		rawSQL := fmt.Sprintf("SELECT COUNT(*) as count from %s where org_id=?", dialect.Quote(query.Target))
+		resp := make([]*targetCount, 0)
+		if err := x.SQL(rawSQL, query.OrgId).Find(&resp); err != nil {
+			return err
+		}
+		used = resp[0].Count
 	}
 
 	query.Result = &models.OrgQuotaDTO{
 		Target: query.Target,
 		Limit:  quota.Limit,
 		OrgId:  query.OrgId,
-		Used:   resp[0].Count,
+		Used:   used,
 	}
 
 	return nil
@@ -78,17 +84,21 @@ func GetOrgQuotas(query *models.GetOrgQuotasQuery) error {
 
 	result := make([]*models.OrgQuotaDTO, len(quotas))
 	for i, q := range quotas {
-		// 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 := x.SQL(rawSQL, q.OrgId).Find(&resp); err != nil {
-			return err
+		var used int64
+		if q.Target != ALERT_RULE_TARGET || 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 := x.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:   resp[0].Count,
+			Used:   used,
 		}
 	}
 	query.Result = result
@@ -138,18 +148,22 @@ func GetUserQuotaByTarget(query *models.GetUserQuotaByTargetQuery) error {
 		quota.Limit = query.Default
 	}
 
-	// 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 {
-		return err
+	var used int64
+	if query.Target != ALERT_RULE_TARGET || 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 {
+			return err
+		}
+		used = resp[0].Count
 	}
 
 	query.Result = &models.UserQuotaDTO{
 		Target: query.Target,
 		Limit:  quota.Limit,
 		UserId: query.UserId,
-		Used:   resp[0].Count,
+		Used:   used,
 	}
 
 	return nil
@@ -181,17 +195,21 @@ func GetUserQuotas(query *models.GetUserQuotasQuery) error {
 
 	result := make([]*models.UserQuotaDTO, len(quotas))
 	for i, q := range quotas {
-		// 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 := x.SQL(rawSQL, q.UserId).Find(&resp); err != nil {
-			return err
+		var used int64
+		if q.Target != ALERT_RULE_TARGET || 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 := x.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:   resp[0].Count,
+			Used:   used,
 		}
 	}
 	query.Result = result
@@ -230,17 +248,22 @@ func UpdateUserQuota(cmd *models.UpdateUserQuotaCmd) error {
 }
 
 func GetGlobalQuotaByTarget(query *models.GetGlobalQuotaByTargetQuery) error {
-	// get quota used.
-	rawSQL := fmt.Sprintf("SELECT COUNT(*) as count from %s", dialect.Quote(query.Target))
-	resp := make([]*targetCount, 0)
-	if err := x.SQL(rawSQL).Find(&resp); err != nil {
-		return err
+	var used int64
+	if query.Target != ALERT_RULE_TARGET || query.IsNgAlertEnabled {
+		// get quota used.
+
+		rawSQL := fmt.Sprintf("SELECT COUNT(*) as count from %s", dialect.Quote(query.Target))
+		resp := make([]*targetCount, 0)
+		if err := x.SQL(rawSQL).Find(&resp); err != nil {
+			return err
+		}
+		used = resp[0].Count
 	}
 
 	query.Result = &models.GlobalQuotaDTO{
 		Target: query.Target,
 		Limit:  query.Default,
-		Used:   resp[0].Count,
+		Used:   used,
 	}
 
 	return nil

pkg/services/sqlstore/quota_test.go

@@ -24,6 +24,7 @@ func TestQuotaCommandsAndQueries(t *testing.T) {
 				Dashboard:  5,
 				DataSource: 5,
 				ApiKey:     5,
+				AlertRule:  5,
 			},
 			User: &setting.UserQuota{
 				Org: 5,
@@ -35,6 +36,7 @@ func TestQuotaCommandsAndQueries(t *testing.T) {
 				DataSource: 5,
 				ApiKey:     5,
 				Session:    5,
+				AlertRule:  5,
 			},
 		}
 
@@ -87,12 +89,19 @@ func TestQuotaCommandsAndQueries(t *testing.T) {
 				So(err, ShouldBeNil)
 				So(query.Result.Used, ShouldEqual, 0)
 			})
+			Convey("Should be able to get zero used org alert quota when table does not exist (ngalert is not enabled - default case)", func() {
+				query := models.GetOrgQuotaByTargetQuery{OrgId: 2, Target: "alert", Default: 11}
+				err = GetOrgQuotaByTarget(&query)
+
+				So(err, ShouldBeNil)
+				So(query.Result.Used, ShouldEqual, 0)
+			})
 			Convey("Should be able to quota list for org", func() {
 				query := models.GetOrgQuotasQuery{OrgId: orgId}
 				err = GetOrgQuotas(&query)
 
 				So(err, ShouldBeNil)
-				So(len(query.Result), ShouldEqual, 4)
+				So(len(query.Result), ShouldEqual, 5)
 				for _, res := range query.Result {
 					limit := 5 // default quota limit
 					used := 0
@@ -169,6 +178,14 @@ func TestQuotaCommandsAndQueries(t *testing.T) {
 			So(query.Result.Limit, ShouldEqual, 5)
 			So(query.Result.Used, ShouldEqual, 1)
 		})
+		Convey("Should be able to get zero used global alert quota when table does not exist (ngalert is not enabled - default case)", func() {
+			query := models.GetGlobalQuotaByTargetQuery{Target: "alert_rule", Default: 5}
+			err = GetGlobalQuotaByTarget(&query)
+			So(err, ShouldBeNil)
+
+			So(query.Result.Limit, ShouldEqual, 5)
+			So(query.Result.Used, ShouldEqual, 0)
+		})
 
 		// related: https://github.com/grafana/grafana/issues/14342
 		Convey("Should org quota updating is successful even if it called multiple time", func() {

pkg/services/sqlstore/sqlstore.go

@@ -158,7 +158,6 @@ func (ss *SQLStore) ensureMainOrgAndAdminUser() error {
 
 		// ensure admin user
 		if !ss.Cfg.DisableInitAdminCreation {
-			ss.log.Debug("Creating default admin user")
 			ss.log.Debug("Creating default admin user")
 			if _, err := ss.createUser(ctx, sess, userCreationArgs{
 				Login:    ss.Cfg.AdminUser,