Prom Query Label Matcher Replacement

example commands.md

# ./pp <query> <scope> <aggregate_group_by> <func_string_replacements:funcName,argIdx,currentValue,newValue>...
# aggregate_group_by is optional

$ ./pp  'up{foo="baz",a="q"}' '{a=~"z"}'
up{a=~"z",foo="baz"}

$ ./pp 'up{foo="baz",a="q"}' '{a=~"z",c="z"}' 
up{a=~"z",c="z",foo="baz"}

$ ./pp 'up{}' '{a=~"z",c="z"}' 
up{a=~"z",c="z"}

$ ./pp 'sum(up{}) by ()' '{a=~"z",c="z"}' 'x'  
sum by (x) (up{a=~"z",c="z"})

$ ./pp 'up[5m:1m]' '{a=~"z",c="z"}'  
up{a=~"z",c="z"}[5m:1m]

$ ./pp 'label_replace(up{job="api-server",service="a:c"}, "foo", "$1", "service", "(.*):.*") + 1' '{a=~"z",c="z"}' "" "label_replace,1,foo,meow" "label_replace,3,service,cat"
label_replace(up{a=~"z",c="z",job="api-server",service="a:c"}, "meow", "$1", "cat", "(.*):.*") + 1

main.go

package main

import (
	"fmt"
	"os"
	"strconv"
	"strings"

	"github.com/prometheus/prometheus/promql/parser"
)

func main() {
	if len(os.Args) < 3 {
		fmt.Println("need at least two args, query, scope")
		return
	}
	query := os.Args[1]
	scope := os.Args[2]
	gb := map[string]struct{}{}
	funcArgReplacements := []FuncArgReplacement{}

	if len(os.Args) > 3 {
		rawGB := os.Args[3]
		gbKeys := strings.Split(rawGB, ",")
		for _, k := range gbKeys {
			gb[k] = struct{}{}
		}
	}
	if len(os.Args) > 4 {
		for _, raw := range os.Args[4:] {
			split := strings.Split(raw, ",")
			if len(split) != 4 {
				fmt.Printf("invalid arg %v", raw)
				os.Exit(1)
			}
			argN, err := strconv.Atoi(split[1])
			if err != nil {
				if len(split) != 3 {
					fmt.Printf("invalid arg %v %v", raw, err)
					os.Exit(1)
				}
			}
			funcArgReplacements = append(funcArgReplacements, FuncArgReplacement{
				FuncName:      split[0],
				ArgN:          argN,
				ExistingValue: split[2],
				NewValue:      split[3],
			})
		}
	}
	newQ, err := replace(query, scope, gb, funcArgReplacements)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	fmt.Println(newQ)

}

type FuncArgReplacement struct {
	FuncName      string
	ArgN          int // Starting at 0
	ExistingValue string
	NewValue      string
}

func replace(rawQuery string, rawMatchers string, gb map[string]struct{}, funcArgReplacements []FuncArgReplacement) (string, error) {
	scopeMatchers, err := parser.ParseMetricSelector(rawMatchers)
	if err != nil {
		return "", err
	}

	matcherNamesToIdx := make(map[string]int, len(scopeMatchers))
	for i, matcher := range scopeMatchers {
		matcherNamesToIdx[matcher.Name] = i
	}

	if err != nil {
		return "", err
	}

	expr, err := parser.ParseExpr(rawQuery)
	if err != nil {
		return "", err
	}

	parser.Inspect(expr, func(node parser.Node, nodes []parser.Node) error {
		switch v := node.(type) {
		case *parser.VectorSelector:
			found := make([]bool, len(scopeMatchers))
			for _, matcher := range v.LabelMatchers {
				if matcher == nil || matcher.Name == "__name__" { // const prob
					continue
				}
				if _, ok := matcherNamesToIdx[matcher.Name]; ok {
					found[matcherNamesToIdx[matcher.Name]] = true
					newM := scopeMatchers[matcherNamesToIdx[matcher.Name]]
					matcher.Name = newM.Name
					matcher.Type = newM.Type
					matcher.Value = newM.Value

				}
			}
			for i, f := range found {
				if f {
					continue
				}
				v.LabelMatchers = append(v.LabelMatchers, scopeMatchers[i])
			}

			return nil
		case *parser.AggregateExpr:
			found := make(map[string]bool, len(gb))
			for k := range gb {
				found[k] = false
			}
			for _, lName := range v.Grouping {
				if _, ok := gb[lName]; ok {
					found[lName] = true
				}
			}
			for k, f := range found {
				if !f {
					v.Grouping = append(v.Grouping, k)
				}
			}
			return nil
		case *parser.Call:
			for i, aT := range v.Func.ArgTypes {
				if aT != parser.ValueTypeString {
					continue
				}
				argNode := parser.Children(v)[i]
				s, ok := argNode.(*parser.StringLiteral)
				if !ok {
					return fmt.Errorf("unexpected non-string literal argument to %v", v.Func.Name)
				}

				for _, r := range funcArgReplacements {
					if v.Func.Name != r.FuncName ||
						i != r.ArgN ||
						s.Val != r.ExistingValue {
						continue
					}
					s.Val = r.NewValue
				}
			}
			return nil

		default:
			return nil
		}
	})

	if err != nil {
		return "", err
	}
	return expr.String(), nil
}