Add quota setting for correlations (#65076)

* Add quota setting for correlations

* Fix linter

conf/defaults.ini

@@ -920,6 +920,9 @@ global_alert_rule = -1
 # global limit of files uploaded to the SQL DB
 global_file = 1000
 
+# global limit of correlations
+global_correlations = -1
+
 #################################### Unified Alerting ####################
 [unified_alerting]
 # Enable the Unified Alerting sub-system and interface. When enabled we'll migrate all of your alert rules and notification channels to the new system. New alert rules will be created and your notification channels will be converted into an Alertmanager configuration. Previous data is preserved to enable backwards compatibility but new data is removed when switching. When this configuration section and flag are not defined, the state is defined at runtime. See the documentation for more details.

conf/sample.ini

@@ -904,6 +904,9 @@
 # global limit of alerts
 ;global_alert_rule = -1
 
+# global limit of correlations
+; global_correlations = -1
+
 #################################### Unified Alerting ####################
 [unified_alerting]
 #Enable the Unified Alerting sub-system and interface. When enabled we'll migrate all of your alert rules and notification channels to the new system. New alert rules will be created and your notification channels will be converted into an Alertmanager configuration. Previous data is preserved to enable backwards compatibility but new data is removed.```

docs/sources/setup-grafana/configure-grafana/_index.md

@@ -1358,6 +1358,10 @@ Sets a global limit on number of users that can be logged in at one time. Defaul
 
 Sets a global limit on number of alert rules that can be created. Default is -1 (unlimited).
 
+### global_correlations
+
+Sets a global limit on number of correlations that can be created. Default is -1 (unlimited).
+
 <hr>
 
 ## [unified_alerting]

pkg/services/correlations/correlations.go

@@ -10,22 +10,43 @@ import (
 	"github.com/grafana/grafana/pkg/infra/log"
 	"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/setting"
 )
 
-func ProvideService(sqlStore db.DB, routeRegister routing.RouteRegister, ds datasources.DataSourceService, ac accesscontrol.AccessControl, bus bus.Bus) *CorrelationsService {
+var (
+	logger = log.New("correlations")
+)
+
+func ProvideService(sqlStore db.DB, routeRegister routing.RouteRegister, ds datasources.DataSourceService, ac accesscontrol.AccessControl, bus bus.Bus, qs quota.Service, cfg *setting.Cfg,
+) (*CorrelationsService, error) {
 	s := &CorrelationsService{
 		SQLStore:          sqlStore,
 		RouteRegister:     routeRegister,
-		log:               log.New("correlations"),
+		log:               logger,
 		DataSourceService: ds,
 		AccessControl:     ac,
+		QuotaService:      qs,
 	}
 
 	s.registerAPIEndpoints()
 
 	bus.AddEventListener(s.handleDatasourceDeletion)
 
-	return s
+	defaultLimits, err := readQuotaConfig(cfg)
+	if err != nil {
+		return s, err
+	}
+
+	if err := qs.RegisterQuotaReporter(&quota.NewUsageReporter{
+		TargetSrv:     QuotaTargetSrv,
+		DefaultLimits: defaultLimits,
+		Reporter:      s.Usage,
+	}); err != nil {
+		return s, err
+	}
+
+	return s, nil
 }
 
 type Service interface {
@@ -41,9 +62,19 @@ type CorrelationsService struct {
 	log               log.Logger
 	DataSourceService datasources.DataSourceService
 	AccessControl     accesscontrol.AccessControl
+	QuotaService      quota.Service
 }
 
 func (s CorrelationsService) CreateCorrelation(ctx context.Context, cmd CreateCorrelationCommand) (Correlation, error) {
+	quotaReached, err := s.QuotaService.CheckQuotaReached(ctx, QuotaTargetSrv, nil)
+	if err != nil {
+		logger.Warn("Error getting correlation quota.", "error", err)
+		return Correlation{}, ErrCorrelationsQuotaFailed
+	}
+	if quotaReached {
+		return Correlation{}, ErrCorrelationsQuotaReached
+	}
+
 	return s.createCorrelation(ctx, cmd)
 }
 
@@ -92,3 +123,23 @@ func (s CorrelationsService) handleDatasourceDeletion(ctx context.Context, event
 		return nil
 	})
 }
+
+func (s *CorrelationsService) Usage(ctx context.Context, scopeParams *quota.ScopeParameters) (*quota.Map, error) {
+	return s.CountCorrelations(ctx)
+}
+
+func readQuotaConfig(cfg *setting.Cfg) (*quota.Map, error) {
+	limits := &quota.Map{}
+
+	if cfg == nil {
+		return limits, nil
+	}
+
+	globalQuotaTag, err := quota.NewTag(QuotaTargetSrv, QuotaTarget, quota.GlobalScope)
+	if err != nil {
+		return limits, err
+	}
+
+	limits.Set(globalQuotaTag, cfg.Quota.Global.Correlations)
+	return limits, nil
+}

pkg/services/correlations/database.go

@@ -5,6 +5,7 @@ import (
 
 	"github.com/grafana/grafana/pkg/infra/db"
 	"github.com/grafana/grafana/pkg/services/datasources"
+	"github.com/grafana/grafana/pkg/services/quota"
 	"github.com/grafana/grafana/pkg/util"
 )
 
@@ -177,6 +178,31 @@ func (s CorrelationsService) getCorrelation(ctx context.Context, cmd GetCorrelat
 	return correlation, nil
 }
 
+func (s CorrelationsService) CountCorrelations(ctx context.Context) (*quota.Map, error) {
+	u := &quota.Map{}
+	var err error
+	count := int64(0)
+	err = s.SQLStore.WithDbSession(ctx, func(sess *db.Session) error {
+		q := sess.Table("correlation")
+		count, err = q.Count()
+
+		if err != nil {
+			return err
+		}
+
+		tag, err := quota.NewTag(QuotaTargetSrv, QuotaTarget, quota.GlobalScope)
+		if err != nil {
+			return err
+		}
+		u.Set(tag, count)
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+	return u, err
+}
+
 func (s CorrelationsService) getCorrelationsBySourceUID(ctx context.Context, cmd GetCorrelationsBySourceUIDQuery) ([]Correlation, error) {
 	correlations := make([]Correlation, 0)
 

pkg/services/correlations/models.go

@@ -4,6 +4,8 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+
+	"github.com/grafana/grafana/pkg/services/quota"
 )
 
 var (
@@ -17,6 +19,13 @@ var (
 	ErrInvalidTransformationType          = errors.New("invalid transformation type")
 	ErrTransformationNotNested            = errors.New("transformations must be nested under config")
 	ErrTransformationRegexReqExp          = errors.New("regex transformations require expression")
+	ErrCorrelationsQuotaFailed            = errors.New("error getting correlations quota")
+	ErrCorrelationsQuotaReached           = errors.New("correlations quota reached")
+)
+
+const (
+	QuotaTargetSrv quota.TargetSrv = "correlations"
+	QuotaTarget    quota.Target    = "correlations"
 )
 
 type CorrelationConfigType string

pkg/setting/setting_quota.go

@@ -13,14 +13,15 @@ type UserQuota struct {
 }
 
 type GlobalQuota struct {
-	Org        int64 `target:"org"`
-	User       int64 `target:"user"`
-	DataSource int64 `target:"data_source"`
-	Dashboard  int64 `target:"dashboard"`
-	ApiKey     int64 `target:"api_key"`
-	Session    int64 `target:"-"`
-	AlertRule  int64 `target:"alert_rule"`
-	File       int64 `target:"file"`
+	Org          int64 `target:"org"`
+	User         int64 `target:"user"`
+	DataSource   int64 `target:"data_source"`
+	Dashboard    int64 `target:"dashboard"`
+	ApiKey       int64 `target:"api_key"`
+	Session      int64 `target:"-"`
+	AlertRule    int64 `target:"alert_rule"`
+	File         int64 `target:"file"`
+	Correlations int64 `target:"correlations"`
 }
 
 type QuotaSettings struct {
@@ -57,13 +58,14 @@ func (cfg *Cfg) readQuotaSettings() {
 
 	// Global Limits
 	cfg.Quota.Global = GlobalQuota{
-		User:       quota.Key("global_user").MustInt64(-1),
-		Org:        quota.Key("global_org").MustInt64(-1),
-		DataSource: quota.Key("global_data_source").MustInt64(-1),
-		Dashboard:  quota.Key("global_dashboard").MustInt64(-1),
-		ApiKey:     quota.Key("global_api_key").MustInt64(-1),
-		Session:    quota.Key("global_session").MustInt64(-1),
-		File:       quota.Key("global_file").MustInt64(-1),
-		AlertRule:  alertGlobalQuota,
+		User:         quota.Key("global_user").MustInt64(-1),
+		Org:          quota.Key("global_org").MustInt64(-1),
+		DataSource:   quota.Key("global_data_source").MustInt64(-1),
+		Dashboard:    quota.Key("global_dashboard").MustInt64(-1),
+		ApiKey:       quota.Key("global_api_key").MustInt64(-1),
+		Session:      quota.Key("global_session").MustInt64(-1),
+		File:         quota.Key("global_file").MustInt64(-1),
+		AlertRule:    alertGlobalQuota,
+		Correlations: quota.Key("global_correlations").MustInt64(-1),
 	}
 }