From aa5d3a65eb212a91e0be9a6727602af100c0fe86 Mon Sep 17 00:00:00 2001
From: Philip Wernersbach <philip.wernersbach@mongodb.com>
Date: Wed, 26 Feb 2025 11:06:22 -0500
Subject: [PATCH] feat: Add hasLabelsForEveryAggregation validator

---
 pkg/validator/config.go                    |  1 +
 pkg/validator/labels.go                    | 53 +++++++++++++++++++++-
 pkg/validator/promql_expression_helpers.go | 41 +++++++++++++++++
 3 files changed, 94 insertions(+), 1 deletion(-)

diff --git a/pkg/validator/config.go b/pkg/validator/config.go
index 79eed3e..45bd410 100644
--- a/pkg/validator/config.go
+++ b/pkg/validator/config.go
@@ -14,6 +14,7 @@ type validatorCreator func(params yaml.Node) (Validator, error)
 var registeredUniversalRuleValidators = map[string]validatorCreator{
 	// Labels
 	"hasLabels":            newHasLabels,
+	"hasLabelsForEveryAggregation": newHasLabelsForEveryAggregation,
 	"doesNotHaveLabels":    newDoesNotHaveLabels,
 	"hasAnyOfLabels":       newHasAnyOfLabels,
 	"labelMatchesRegexp":   newLabelMatchesRegexp,
diff --git a/pkg/validator/labels.go b/pkg/validator/labels.go
index e8f2f8b..36e6536 100644
--- a/pkg/validator/labels.go
+++ b/pkg/validator/labels.go
@@ -12,7 +12,6 @@ import (
 	"github.com/fusakla/promruval/v3/pkg/unmarshaler"
 	"github.com/prometheus/prometheus/model/rulefmt"
 )
-
 func newHasLabels(paramsConfig yaml.Node) (Validator, error) {
 	params := struct {
 		Labels       []string `yaml:"labels"`
@@ -66,6 +65,58 @@ func (h hasLabels) Validate(_ unmarshaler.RuleGroup, rule rulefmt.Rule, _ *prome
 	return errs
 }
 
+func newHasLabelsForEveryAggregation(paramsConfig yaml.Node) (Validator, error) {
+	params := struct {
+		Labels       []string `yaml:"labels"`
+		ExcludeRuleLabels bool `yaml:"exclude_rule_labels"`
+	}{}
+	if err := paramsConfig.Decode(&params); err != nil {
+		return nil, err
+	}
+	if len(params.Labels) == 0 {
+		return nil, fmt.Errorf("missing labels")
+	}
+	return &hasLabelsForEveryAggregation{labels: params.Labels, excludeRuleLabels: params.ExcludeRuleLabels}, nil
+}
+
+type hasLabelsForEveryAggregation struct {
+	labels       []string
+	excludeRuleLabels bool
+}
+
+func (h hasLabelsForEveryAggregation) String() string {
+	return fmt.Sprintf("every aggregation has labels: `%s`", strings.Join(h.labels, "`,`"))
+}
+
+func (h hasLabelsForEveryAggregation) Validate(_ unmarshaler.RuleGroup, rule rulefmt.Rule, _ *prometheus.Client) []error {
+	var (
+		errs       []error
+		err        error
+		exprLabels map[string]struct{}
+	)
+
+	exprLabels, err = getExpressionUsedLabelsForEveryAggregation(rule.Expr)
+	if err != nil {
+		errs = append(errs, err)
+	}
+	if exprLabels == nil {
+		return errs
+	}
+
+	if h.excludeRuleLabels {
+		for label, _ := range rule.Labels {
+			exprLabels[label] = struct{}{}
+		}
+	}
+
+	for _, label := range h.labels {
+		if _, ok := exprLabels[label]; !ok {
+			errs = append(errs, fmt.Errorf("missing label `%s`", label))
+		}
+	}
+	return errs
+}
+
 func newDoesNotHaveLabels(paramsConfig yaml.Node) (Validator, error) {
 	params := struct {
 		Labels       []string `yaml:"labels"`
diff --git a/pkg/validator/promql_expression_helpers.go b/pkg/validator/promql_expression_helpers.go
index 3442195..047a34b 100644
--- a/pkg/validator/promql_expression_helpers.go
+++ b/pkg/validator/promql_expression_helpers.go
@@ -159,6 +159,47 @@ func getExpressionUsedLabels(expr string) ([]string, error) {
 	return getExpressionUsedLabelsForMetric(expr, regexp.MustCompile(".*"))
 }
 
+func getExpressionUsedLabelsForEveryAggregation(expr string) (map[string]struct{}, error) {
+	promQl, err := parser.ParseExpr(expr)
+	if err != nil {
+		return map[string]struct{}{}, fmt.Errorf("failed to parse expression `%s`: %w", expr, err)
+	}
+	var usedLabels map[string]struct{}
+
+	labelsUpInExpr := func(path []parser.Node) map[string]struct{} {
+		for _, n := range path {
+			switch v := n.(type) {
+			case *parser.AggregateExpr:
+				usedLabels := map[string]struct{}{}
+				for _, l := range v.Grouping {
+					usedLabels[l] = struct{}{}
+				}
+				return usedLabels
+			}
+		}
+
+		return nil
+	}
+
+	parser.Inspect(promQl, func(n parser.Node, path []parser.Node) error {
+		usedLabelsInExpr := labelsUpInExpr(path)
+		if usedLabels == nil {
+			usedLabels = usedLabelsInExpr
+		}
+		if usedLabelsInExpr != nil {
+			for k := range usedLabels {
+				if _, exists := usedLabelsInExpr[k]; !exists {
+					delete(usedLabels, k)
+				}
+			}
+
+		}
+		return nil
+	})
+
+	return usedLabels, nil
+}
+
 func getExpressionVectorSelectors(expr string) ([]*parser.VectorSelector, error) {
 	promQl, err := parser.ParseExpr(expr)
 	if err != nil {