Merge pull request #37 from grafana/prepare-rules

Introduce the `rules prepare` command

Author: gotjosh

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## unreleased / master
 
+* [FEATURE] Add `rules prepare` command. It allows you add a label to PromQL aggregations and lint your expressions in rule files.
+
 ## v0.1.4 / 2020-03-10
 
 * [CHANGE] Ensure 404 deletes do not trigger an error for `rules` and `alertmanager` commands #28

README.md

@@ -60,6 +60,22 @@ This command will load each rule group in the specified files and load them into
 
     cortextool rules load ./example_rules_one.yaml ./example_rules_two.yaml  ...
 
+#### Rules Prepare
+
+This command prepares a rules file for upload to Cortex. It lints all your PromQL expressions and adds an specific label to your PromQL query aggregations in the file. Unlike, the previous command this one does not interact with your Cortex cluster.
+
+    cortextool rules prepare -i ./example_rules_one.yaml ./example_rules_two.yaml ...
+
+There are two flags of note for this command:
+- `-i` which allows you to edit in place, otherwise a a new file with a `.output` extension is created with the results of the run.
+- `-l` which allows you specify the label you want you add for your aggregations, it is `cluster` by default.
+
+At the end of the run, the command tells you whenever the operation was a success in the form of
+
+    INFO[0000] SUCESS: 194 rules found, 0 modified expressions
+
+It is important to note that a modification can be a PromQL expression lint or a label add to your aggregation.
+
 # chunktool
 
 This repo also contains the `chunktool`. A client meant to interact with chunks stored and indexed in cortex backends.

go.mod

@@ -20,6 +20,7 @@ require (
 	github.com/prometheus/common v0.7.0
 	github.com/prometheus/prometheus v1.8.2-0.20190918104050-8744afdd1ea0
 	github.com/sirupsen/logrus v1.4.2
+	github.com/stretchr/testify v1.4.0
 	google.golang.org/api v0.8.0
 	gopkg.in/alecthomas/kingpin.v2 v2.2.6
 	gopkg.in/yaml.v2 v2.2.2

pkg/commands/rules.go

@@ -3,6 +3,7 @@ package commands
 import (
 	"context"
 	"fmt"
+	"io/ioutil"
 	"os"
 	"path/filepath"
 	"strings"
@@ -16,6 +17,11 @@ import (
 	"github.com/grafana/cortextool/pkg/rules"
 	log "github.com/sirupsen/logrus"
 	"gopkg.in/alecthomas/kingpin.v2"
+	"gopkg.in/yaml.v2"
+)
+
+const (
+	defaultPrepareAggregationLabel = "cluster"
 )
 
 var (
@@ -31,7 +37,7 @@ var (
 	})
 )
 
-// RuleCommand configures and executes rule related cortex api operations
+// RuleCommand configures and executes rule related cortex operations
 type RuleCommand struct {
 	ClientConfig client.Config
 
@@ -41,7 +47,7 @@ type RuleCommand struct {
 	Namespace string
 	RuleGroup string
 
-	// Load Rules Configs
+	// Load Rules Config
 	RuleFilesList []string
 	RuleFiles     string
 	RuleFilesPath string
@@ -50,38 +56,73 @@ type RuleCommand struct {
 	IgnoredNamespaces    string
 	ignoredNamespacesMap map[string]struct{}
 
+	// Prepare Rules Config
+	InPlaceEdit      bool
+	AggregationLabel string
+
 	DisableColor bool
 }
 
 // Register rule related commands and flags with the kingpin application
 func (r *RuleCommand) Register(app *kingpin.Application) {
 	rulesCmd := app.Command("rules", "View & edit rules stored in cortex.").PreAction(r.setup)
-	rulesCmd.Flag("address", "Address of the cortex cluster, alternatively set CORTEX_ADDRESS.").Envar("CORTEX_ADDRESS").Required().StringVar(&r.ClientConfig.Address)
-	rulesCmd.Flag("id", "Cortex tenant id, alternatively set CORTEX_TENANT_ID.").Envar("CORTEX_TENANT_ID").Required().StringVar(&r.ClientConfig.ID)
 	rulesCmd.Flag("key", "Api key to use when contacting cortex, alternatively set $CORTEX_API_KEY.").Default("").Envar("CORTEX_API_KEY").StringVar(&r.ClientConfig.Key)
 
-	// List Rules Command
-	rulesCmd.Command("list", "List the rules currently in the cortex ruler.").Action(r.listRules)
+	// Register rule commands
+	listCmd := rulesCmd.
+		Command("list", "List the rules currently in the cortex ruler.").
+		Action(r.listRules)
+	printRulesCmd := rulesCmd.
+		Command("print", "Print the rules currently in the cortex ruler.").
+		Action(r.printRules)
+	getRuleGroupCmd := rulesCmd.
+		Command("get", "Retreive a rulegroup from the ruler.").
+		Action(r.getRuleGroup)
+	deleteRuleGroupCmd := rulesCmd.
+		Command("delete", "Delete a rulegroup from the ruler.").
+		Action(r.deleteRuleGroup)
+	loadRulesCmd := rulesCmd.
+		Command("load", "load a set of rules to a designated cortex endpoint").
+		Action(r.loadRules)
+	diffRulesCmd := rulesCmd.
+		Command("diff", "diff a set of rules to a designated cortex endpoint").
+		Action(r.diffRules)
+	syncRulesCmd := rulesCmd.
+		Command("sync", "sync a set of rules to a designated cortex endpoint").
+		Action(r.syncRules)
+	prepareCmd := rulesCmd.
+		Command("prepare", "modifies a set of rules by including an specific label in aggregations.").
+		Action(r.prepare)
+
+	// Require Cortex cluster address and tentant ID on all these commands
+	for _, c := range []*kingpin.CmdClause{listCmd, printRulesCmd, getRuleGroupCmd, deleteRuleGroupCmd, loadRulesCmd, diffRulesCmd, syncRulesCmd} {
+		c.Flag("address", "Address of the cortex cluster, alternatively set CORTEX_ADDRESS.").
+			Envar("CORTEX_ADDRESS").
+			Required().
+			StringVar(&r.ClientConfig.Address)
+
+		c.Flag("id", "Cortex tenant id, alternatively set CORTEX_TENANT_ID.").
+			Envar("CORTEX_TENANT_ID").
+			Required().
+			StringVar(&r.ClientConfig.ID)
+	}
 
 	// Print Rules Command
-	printRulesCmd := rulesCmd.Command("print", "Print the rules currently in the cortex ruler.").Action(r.printRules)
 	printRulesCmd.Flag("disable-color", "disable colored output").BoolVar(&r.DisableColor)
 
 	// Get RuleGroup Command
-	getRuleGroupCmd := rulesCmd.Command("get", "Retreive a rulegroup from the ruler.").Action(r.getRuleGroup)
 	getRuleGroupCmd.Arg("namespace", "Namespace of the rulegroup to retrieve.").Required().StringVar(&r.Namespace)
 	getRuleGroupCmd.Arg("group", "Name of the rulegroup ot retrieve.").Required().StringVar(&r.RuleGroup)
 	getRuleGroupCmd.Flag("disable-color", "disable colored output").BoolVar(&r.DisableColor)
 
 	// Delete RuleGroup Command
-	deleteRuleGroupCmd := rulesCmd.Command("delete", "Delete a rulegroup from the ruler.").Action(r.deleteRuleGroup)
 	deleteRuleGroupCmd.Arg("namespace", "Namespace of the rulegroup to delete.").Required().StringVar(&r.Namespace)
 	deleteRuleGroupCmd.Arg("group", "Name of the rulegroup ot delete.").Required().StringVar(&r.RuleGroup)
 
-	loadRulesCmd := rulesCmd.Command("load", "load a set of rules to a designated cortex endpoint").Action(r.loadRules)
+	// Load Rules Command
 	loadRulesCmd.Arg("rule-files", "The rule files to check.").Required().ExistingFilesVar(&r.RuleFilesList)
 
-	diffRulesCmd := rulesCmd.Command("diff", "diff a set of rules to a designated cortex endpoint").Action(r.diffRules)
+	// Diff Command
 	diffRulesCmd.Flag("ignored-namespaces", "comma-separated list of namespaces to ignore during a diff.").StringVar(&r.IgnoredNamespaces)
 	diffRulesCmd.Flag("rule-files", "The rule files to check. Flag can be reused to load multiple files.").StringVar(&r.RuleFiles)
 	diffRulesCmd.Flag(
@@ -90,13 +131,26 @@ func (r *RuleCommand) Register(app *kingpin.Application) {
 	).StringVar(&r.RuleFilesPath)
 	diffRulesCmd.Flag("disable-color", "disable colored output").BoolVar(&r.DisableColor)
 
-	syncRulesCmd := rulesCmd.Command("sync", "sync a set of rules to a designated cortex endpoint").Action(r.syncRules)
+	// Sync Command
 	syncRulesCmd.Flag("ignored-namespaces", "comma-separated list of namespaces to ignore during a sync.").StringVar(&r.IgnoredNamespaces)
 	syncRulesCmd.Flag("rule-files", "The rule files to check. Flag can be reused to load multiple files.").StringVar(&r.RuleFiles)
 	syncRulesCmd.Flag(
 		"rule-dirs",
 		"Comma seperated list of paths to directories containing rules yaml files. Each file in a directory with a .yml or .yaml suffix will be parsed.",
 	).StringVar(&r.RuleFilesPath)
+
+	// Prepare Command
+	prepareCmd.Arg("rule-files", "The rule files to check.").Required().ExistingFilesVar(&r.RuleFilesList)
+	prepareCmd.Flag("rule-files", "The rule files to check. Flag can be reused to load multiple files.").StringVar(&r.RuleFiles)
+	prepareCmd.Flag(
+		"rule-dirs",
+		"Comma seperated list of paths to directories containing rules yaml files. Each file in a directory with a .yml or .yaml suffix will be parsed.",
+	).StringVar(&r.RuleFilesPath)
+	prepareCmd.Flag(
+		"in-place",
+		"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)
 }
 
 func (r *RuleCommand) setup(k *kingpin.ParseContext) error {
@@ -420,3 +474,48 @@ func (r *RuleCommand) executeChanges(ctx context.Context, changes []rules.Namesp
 	fmt.Printf("Sync Summary: %v Groups Created, %v Groups Updated, %v Groups Deleted\n", created, updated, deleted)
 	return nil
 }
+
+func (r *RuleCommand) prepare(k *kingpin.ParseContext) error {
+	err := r.setupFiles()
+	if err != nil {
+		return errors.Wrap(err, "prepare operation unsuccessful, unable to load rules files")
+	}
+
+	namespaces, err := rules.ParseFiles(r.RuleFilesList)
+	if err != nil {
+		return errors.Wrap(err, "prepare operation unsuccessful, unable to parse rules files")
+	}
+
+	var count, mod int
+	for _, ruleNamespace := range namespaces {
+		c, m, err := ruleNamespace.AggregateBy(r.AggregationLabel)
+		if err != nil {
+			return err
+		}
+
+		count += c
+		mod += m
+	}
+
+	// now, save all the files
+	for _, ns := range namespaces {
+		payload, err := yaml.Marshal(ns)
+		if err != nil {
+			return err
+		}
+
+		filepath := ns.Filepath
+		if !r.InPlaceEdit {
+			filepath = filepath + ".result"
+		}
+
+		err = ioutil.WriteFile(filepath, payload, 0644)
+		if err != nil {
+			return err
+		}
+	}
+
+	log.Infof("SUCESS: %d rules found, %d modified expressions", count, mod)
+
+	return nil
+}

pkg/rules/parser.go

@@ -30,6 +30,8 @@ func ParseFiles(files []string) (map[string]RuleNamespace, error) {
 			return nil, errFileReadError
 		}
 
+		ns.Filepath = f
+
 		// Determine if the namespace is explicitly set. If not
 		// the file name without the extension is used.
 		namespace := ns.Namespace

pkg/rules/rules.go

@@ -2,8 +2,11 @@ package rules
 
 import (
 	"fmt"
+	"strings"
 
 	"github.com/prometheus/prometheus/pkg/rulefmt"
+	"github.com/prometheus/prometheus/promql"
+	log "github.com/sirupsen/logrus"
 )
 
 // RuleNamespace is used to parse a slightly modified prometheus
@@ -12,10 +15,74 @@ import (
 type RuleNamespace struct {
 	// Namespace field only exists for setting namespace in namespace body instead of file name
 	Namespace string `yaml:"namespace,omitempty"`
+	Filepath  string `yaml:"-"`
 
 	Groups []rulefmt.RuleGroup `yaml:"groups"`
 }
 
+// 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.
+	// `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 {
+			log.WithFields(log.Fields{"rule": getRuleName(rule)}).Debugf("evaluating...")
+			exp, err := promql.ParseExpr(rule.Expr)
+			if err != nil {
+				return count, mod, err
+			}
+
+			count++
+			// Given inspect will help us traverse every node in the AST, Let's create the
+			// function that will modify the labels.
+			f := exprNodeInspectorFunc(rule, label)
+			promql.Inspect(exp, f)
+
+			// Only modify the ones that actually changed.
+			if r.Groups[i].Rules[j].Expr != exp.String() {
+				mod++
+				r.Groups[i].Rules[j].Expr = exp.String()
+			}
+		}
+	}
+
+	return count, mod, nil
+}
+
+// exprNodeInspectorFunc returns a PromQL inspector.
+// It modifies most PromQL aggregations to include a given label.
+func exprNodeInspectorFunc(rule rulefmt.Rule, label string) func(node promql.Node, path []promql.Node) error {
+	return func(node promql.Node, path []promql.Node) error {
+		aggregation, ok := node.(*promql.AggregateExpr)
+		if !ok {
+			return nil
+		}
+
+		// If the aggregation is about dropping labels (e.g. without), we don't want to modify
+		// this expression. Omission as long as it is not the cluster label will include it.
+		// TODO: We probably want to check whenever the label we're trying to include is included in the omission.
+		if aggregation.Without {
+			return nil
+		}
+
+		for _, lbl := range aggregation.Grouping {
+			if lbl == label {
+				return nil
+			}
+		}
+
+		log.WithFields(
+			log.Fields{"rule": getRuleName(rule), "lbls": strings.Join(aggregation.Grouping, ", ")},
+		).Debugf("aggregation without '%s' label, adding.", label)
+
+		aggregation.Grouping = append(aggregation.Grouping, label)
+		return nil
+	}
+}
+
 // Validate each rule in the rule namespace is valid
 func (r RuleNamespace) Validate() []error {
 	set := map[string]struct{}{}
@@ -63,3 +130,11 @@ func ValidateRuleGroup(g rulefmt.RuleGroup) []error {
 
 	return errs
 }
+
+func getRuleName(r rulefmt.Rule) string {
+	if r.Record != "" {
+		return r.Record
+	}
+
+	return r.Alert
+}