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