Add PromQL format and label matcher set/delete commands to promtool

Signed-off-by: Rob Skillington <rob@chronosphere.io>
Signed-off-by: Julien Pivotto <roidelapluie@o11y.eu>

cmd/promtool/main.go

@@ -58,6 +58,7 @@ import (
 	"github.com/prometheus/prometheus/notifier"
 	_ "github.com/prometheus/prometheus/plugins" // Register plugins.
 	"github.com/prometheus/prometheus/promql"
+	"github.com/prometheus/prometheus/promql/parser"
 	"github.com/prometheus/prometheus/scrape"
 	"github.com/prometheus/prometheus/util/documentcli"
 )
@@ -245,6 +246,22 @@ func main() {
 		"A list of one or more files containing recording rules to be backfilled. All recording rules listed in the files will be backfilled. Alerting rules are not evaluated.",
 	).Required().ExistingFiles()
 
+	promQLCmd := app.Command("promql", "PromQL formatting and editing.")
+
+	promQLFormatCmd := promQLCmd.Command("format", "Format PromQL query to pretty printed form.")
+	promQLFormatQuery := promQLFormatCmd.Arg("query", "PromQL query.").Required().String()
+
+	promQLLabelsCmd := promQLCmd.Command("label-matchers", "Edit label matchers contained within an existing PromQL query.")
+	promQLLabelsSetCmd := promQLLabelsCmd.Command("set", "Set a label matcher in the query.")
+	promQLLabelsSetType := promQLLabelsSetCmd.Flag("type", "Type of the label matcher to set.").Short('t').Default("=").Enum("=", "!=", "=~", "!~")
+	promQLLabelsSetQuery := promQLLabelsSetCmd.Arg("query", "PromQL query.").Required().String()
+	promQLLabelsSetName := promQLLabelsSetCmd.Arg("name", "Name of the label matcher to set.").Required().String()
+	promQLLabelsSetValue := promQLLabelsSetCmd.Arg("value", "Value of the label matcher to set.").Required().String()
+
+	promQLLabelsDeleteCmd := promQLLabelsCmd.Command("delete", "Delete a label from the query.")
+	promQLLabelsDeleteQuery := promQLLabelsDeleteCmd.Arg("query", "PromQL query.").Required().String()
+	promQLLabelsDeleteName := promQLLabelsDeleteCmd.Arg("name", "Name of the label to delete.").Required().String()
+
 	featureList := app.Flag("enable-feature", "Comma separated feature names to enable (only PromQL related and no-default-scrape-port). See https://prometheus.io/docs/prometheus/latest/feature_flags/ for the options and more details.").Default("").Strings()
 
 	documentationCmd := app.Command("write-documentation", "Generate command line documentation. Internal use.").Hidden()
@@ -364,8 +381,18 @@ func main() {
 
 	case importRulesCmd.FullCommand():
 		os.Exit(checkErr(importRules(serverURL, httpRoundTripper, *importRulesStart, *importRulesEnd, *importRulesOutputDir, *importRulesEvalInterval, *maxBlockDuration, *importRulesFiles...)))
+
 	case documentationCmd.FullCommand():
 		os.Exit(checkErr(documentcli.GenerateMarkdown(app.Model(), os.Stdout)))
+
+	case promQLFormatCmd.FullCommand():
+		os.Exit(checkErr(formatPromQL(*promQLFormatQuery)))
+
+	case promQLLabelsSetCmd.FullCommand():
+		os.Exit(checkErr(labelsSetPromQL(*promQLLabelsSetQuery, *promQLLabelsSetType, *promQLLabelsSetName, *promQLLabelsSetValue)))
+
+	case promQLLabelsDeleteCmd.FullCommand():
+		os.Exit(checkErr(labelsDeletePromQL(*promQLLabelsDeleteQuery, *promQLLabelsDeleteName)))
 	}
 }
 
@@ -1375,3 +1402,79 @@ func checkTargetGroupsForScrapeConfig(targetGroups []*targetgroup.Group, scfg *c
 
 	return nil
 }
+
+func formatPromQL(query string) error {
+	expr, err := parser.ParseExpr(query)
+	if err != nil {
+		return err
+	}
+
+	fmt.Println(expr.Pretty(0))
+	return nil
+}
+
+func labelsSetPromQL(query, labelMatchType, name, value string) error {
+	expr, err := parser.ParseExpr(query)
+	if err != nil {
+		return err
+	}
+
+	var matchType labels.MatchType
+	switch labelMatchType {
+	case parser.ItemType(parser.EQL).String():
+		matchType = labels.MatchEqual
+	case parser.ItemType(parser.NEQ).String():
+		matchType = labels.MatchNotEqual
+	case parser.ItemType(parser.EQL_REGEX).String():
+		matchType = labels.MatchRegexp
+	case parser.ItemType(parser.NEQ_REGEX).String():
+		matchType = labels.MatchNotRegexp
+	default:
+		return fmt.Errorf("invalid label match type: %s", labelMatchType)
+	}
+
+	parser.Inspect(expr, func(node parser.Node, path []parser.Node) error {
+		if n, ok := node.(*parser.VectorSelector); ok {
+			var found bool
+			for i, l := range n.LabelMatchers {
+				if l.Name == name {
+					n.LabelMatchers[i].Type = matchType
+					n.LabelMatchers[i].Value = value
+					found = true
+				}
+			}
+			if !found {
+				n.LabelMatchers = append(n.LabelMatchers, &labels.Matcher{
+					Type:  matchType,
+					Name:  name,
+					Value: value,
+				})
+			}
+		}
+		return nil
+	})
+
+	fmt.Println(expr.Pretty(0))
+	return nil
+}
+
+func labelsDeletePromQL(query, name string) error {
+	expr, err := parser.ParseExpr(query)
+	if err != nil {
+		return err
+	}
+
+	parser.Inspect(expr, func(node parser.Node, path []parser.Node) error {
+		if n, ok := node.(*parser.VectorSelector); ok {
+			for i, l := range n.LabelMatchers {
+				if l.Name == name {
+					n.LabelMatchers = append(n.LabelMatchers[:i], n.LabelMatchers[i+1:]...)
+				}
+			}
+		}
+		return nil
+	})
+
+	fmt.Println(expr.Pretty(0))
+	return nil
+}