Add --label-excluded-rule-groups to cortextool rules prepare (#174)

* Add --label-excluded-rule-groups to cortextool rules prepare

Signed-off-by: Marco Pracucci <[email protected]>

* Fixed PR number in CHANGELOG

Signed-off-by: Marco Pracucci <[email protected]>

* Simplify CLI flag parsing

Signed-off-by: Marco Pracucci <[email protected]>

* Update pkg/commands/rules.go

Signed-off-by: Marco Pracucci <[email protected]>

Co-authored-by: gotjosh <[email protected]>

* Addressed review feedback

Signed-off-by: Marco Pracucci <[email protected]>

Co-authored-by: gotjosh <[email protected]>

Author: pracucci

CHANGELOG.md

@@ -4,6 +4,7 @@ Order should be `CHANGE`, `FEATURE`, `ENHANCEMENT`, and `BUGFIX`
 
 ## unreleased/master
 
+* [FEATURE] Add `--label-excluded-rule-groups` support to `cortextool rules prepare` command. #174
 * [ENHANCEMENT] Upgrade the Go version used in build images and tests to golang 1.16.3 to match upstream Cortex. #165
 
 ## v0.9.0

pkg/commands/rules.go

@@ -67,8 +67,10 @@ type RuleCommand struct {
 	ignoredNamespacesMap map[string]struct{}
 
 	// Prepare Rules Config
-	InPlaceEdit      bool
-	AggregationLabel string
+	InPlaceEdit                            bool
+	AggregationLabel                       string
+	AggregationLabelExcludedRuleGroups     string
+	aggregationLabelExcludedRuleGroupsList map[string]struct{}
 
 	// Lint Rules Config
 	LintDryRun bool
@@ -202,6 +204,7 @@ func (r *RuleCommand) Register(app *kingpin.Application) {
 		"edits the rule file in place",
 	).Short('i').BoolVar(&r.InPlaceEdit)
 	prepareCmd.Flag("label", "label to include as part of the aggregations.").Default(defaultPrepareAggregationLabel).Short('l').StringVar(&r.AggregationLabel)
+	prepareCmd.Flag("label-excluded-rule-groups", "Comma separated list of rule group names to exclude when including the configured label to aggregations.").StringVar(&r.AggregationLabelExcludedRuleGroups)
 
 	// Lint Command
 	lintCmd.Arg("rule-files", "The rule files to check.").ExistingFilesVar(&r.RuleFilesList)
@@ -271,6 +274,14 @@ func (r *RuleCommand) setupFiles() error {
 		}
 	}
 
+	// Set up rule groups excluded from label aggregation.
+	r.aggregationLabelExcludedRuleGroupsList = map[string]struct{}{}
+	for _, name := range strings.Split(r.AggregationLabelExcludedRuleGroups, ",") {
+		if name = strings.TrimSpace(name); name != "" {
+			r.aggregationLabelExcludedRuleGroupsList[name] = struct{}{}
+		}
+	}
+
 	for _, file := range strings.Split(r.RuleFiles, ",") {
 		if file != "" {
 			log.WithFields(log.Fields{
@@ -613,9 +624,15 @@ func (r *RuleCommand) prepare(k *kingpin.ParseContext) error {
 		return errors.Wrap(err, "prepare operation unsuccessful, unable to parse rules files")
 	}
 
+	// Do not apply the aggregation label to excluded rule groups.
+	applyTo := func(group rwrulefmt.RuleGroup, rule rulefmt.RuleNode) bool {
+		_, excluded := r.aggregationLabelExcludedRuleGroupsList[group.Name]
+		return !excluded
+	}
+
 	var count, mod int
 	for _, ruleNamespace := range namespaces {
-		c, m, err := ruleNamespace.AggregateBy(r.AggregationLabel)
+		c, m, err := ruleNamespace.AggregateBy(r.AggregationLabel, applyTo)
 		if err != nil {
 			return err
 		}

pkg/rules/rules.go

@@ -106,14 +106,27 @@ func (r RuleNamespace) CheckRecordingRules(strict bool) int {
 }
 
 // AggregateBy modifies the aggregation rules in groups to include a given Label.
-func (r RuleNamespace) AggregateBy(label string) (int, int, error) {
-	// `count` represents the number of rules we evalated.
+// If the applyTo function is provided, the aggregation is applied only to rules
+// for which the applyTo function returns true.
+func (r RuleNamespace) AggregateBy(label string, applyTo func(group rwrulefmt.RuleGroup, rule rulefmt.RuleNode) bool) (int, int, error) {
+	// `count` represents the number of rules we evaluated.
 	// `mod` represents the number of rules we modified - a modification can either be a lint or adding the
 	// label in the aggregation.
 	var count, mod int
 
 	for i, group := range r.Groups {
 		for j, rule := range group.Rules {
+			// Skip it if the applyTo function returns false.
+			if applyTo != nil && !applyTo(group, rule) {
+				log.WithFields(log.Fields{
+					"group": group.Name,
+					"rule":  getRuleName(rule),
+				}).Debugf("skipped")
+
+				count++
+				continue
+			}
+
 			log.WithFields(log.Fields{"rule": getRuleName(rule)}).Debugf("evaluating...")
 			exp, err := parser.ParseExpr(rule.Expr.Value)
 			if err != nil {

pkg/rules/rules_test.go

@@ -15,6 +15,7 @@ func TestAggregateBy(t *testing.T) {
 	tt := []struct {
 		name            string
 		rn              RuleNamespace
+		applyTo         func(group rwrulefmt.RuleGroup, rule rulefmt.RuleNode) bool
 		expectedExpr    []string
 		count, modified int
 		expect          error
@@ -140,20 +141,63 @@ func TestAggregateBy(t *testing.T) {
 			expectedExpr: []string{`count by(cluster, node) (sum by(node, cpu, cluster) (node_cpu_seconds_total{job="default/node-exporter"} * on(namespace, instance, cluster) group_left(node) node_namespace_pod:kube_pod_info:))`},
 			count:        1, modified: 1, expect: nil,
 		},
+		{
+			name: "with a query skipped",
+			rn: RuleNamespace{
+				Groups: []rwrulefmt.RuleGroup{
+					{
+						RuleGroup: rulefmt.RuleGroup{
+							Name: "CountAggregation",
+							Rules: []rulefmt.RuleNode{
+								{
+									Alert: yaml.Node{
+										Value: "CountAggregation",
+									},
+									Expr: yaml.Node{
+										Value: `count by(namespace) (test_series) > 1`,
+									},
+								},
+							},
+						},
+					}, {
+						RuleGroup: rulefmt.RuleGroup{
+							Name: "CountSkipped",
+							Rules: []rulefmt.RuleNode{
+								{
+									Alert: yaml.Node{
+										Value: "CountSkipped",
+									},
+									Expr: yaml.Node{
+										Value: `count by(namespace) (test_series) > 1`,
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			applyTo: func(group rwrulefmt.RuleGroup, rule rulefmt.RuleNode) bool {
+				return group.Name != "CountSkipped"
+			},
+			expectedExpr: []string{`count by(namespace, cluster) (test_series) > 1`, `count by(namespace) (test_series) > 1`},
+			count:        2, modified: 1, expect: nil,
+		},
 	}
 
 	for _, tc := range tt {
 		t.Run(tc.name, func(t *testing.T) {
-			c, m, err := tc.rn.AggregateBy("cluster")
+			c, m, err := tc.rn.AggregateBy("cluster", tc.applyTo)
 
 			require.Equal(t, tc.expect, err)
 			assert.Equal(t, tc.count, c)
 			assert.Equal(t, tc.modified, m)
 
 			// Only verify the PromQL expression if it has been modified
+			expectedIdx := 0
 			for _, g := range tc.rn.Groups {
-				for i, r := range g.Rules {
-					require.Equal(t, tc.expectedExpr[i], r.Expr.Value)
+				for _, r := range g.Rules {
+					require.Equal(t, tc.expectedExpr[expectedIdx], r.Expr.Value)
+					expectedIdx++
 				}
 			}
 		})