Add quota setting for correlations (#65076)

* Add quota setting for correlations

* Fix linter

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),
 	}
 }