From 29c8eef2073ef9db1e90227feffccade85fe6fe3 Mon Sep 17 00:00:00 2001 From: Ben Ye Date: Fri, 1 Nov 2024 17:35:54 -0700 Subject: [PATCH] add support for label_join and label_replace function Signed-off-by: Ben Ye --- opts.go | 4 --- walk.go | 77 ++++++++++++++++++++++++++++++++++++++++++++++++++++ walk_test.go | 38 ++++++++++++++++++++++++++ 3 files changed, 115 insertions(+), 4 deletions(-) diff --git a/opts.go b/opts.go index c4b080c..2f806fb 100644 --- a/opts.go +++ b/opts.go @@ -5,7 +5,6 @@ import ( "github.com/prometheus/prometheus/model/labels" "github.com/prometheus/prometheus/promql/parser" - "golang.org/x/exp/slices" ) var ( @@ -58,9 +57,6 @@ var ( func init() { for _, f := range parser.Functions { - if slices.Contains(f.ArgTypes, parser.ValueTypeString) { - continue - } // Ignore experimental functions for now. if !f.Experimental { defaultSupportedFuncs = append(defaultSupportedFuncs, f) diff --git a/walk.go b/walk.go index f98dc15..930bba9 100644 --- a/walk.go +++ b/walk.go @@ -16,6 +16,9 @@ import ( const ( // max number of grouping labels in either by or without clause. maxGroupingLabels = 5 + + // Destination label used in functions like label_replace and label_join. + destinationLabel = "__promqlsmith_dst_label__" ) // walkExpr generates the given expression type with one of the required value type. @@ -249,10 +252,20 @@ func (s *PromQLSmith) walkCall(valueTypes ...parser.ValueType) parser.Expr { } func (s *PromQLSmith) walkFunctions(expr *parser.Call) { + switch expr.Func.Name { + case "label_join": + s.walkLabelJoin(expr) + return + default: + } + expr.Args = make([]parser.Expr, len(expr.Func.ArgTypes)) if expr.Func.Name == "holt_winters" { s.walkHoltWinters(expr) return + } else if expr.Func.Name == "label_replace" { + s.walkLabelReplace(expr) + return } if expr.Func.Variadic != 0 { s.walkVariadicFunctions(expr) @@ -269,6 +282,70 @@ func (s *PromQLSmith) walkHoltWinters(expr *parser.Call) { expr.Args[2] = &parser.NumberLiteral{Val: getNonZeroFloat64(s.rnd)} } +func (s *PromQLSmith) walkLabelReplace(expr *parser.Call) { + expr.Args[0] = s.Walk(expr.Func.ArgTypes[0]) + expr.Args[1] = &parser.StringLiteral{Val: destinationLabel} + expr.Args[2] = &parser.StringLiteral{Val: "$1"} + seriesSet, _ := getOutputSeries(expr.Args[0]) + + var srcLabel string + if len(seriesSet) > 0 { + lbls := seriesSet[0] + if lbls.Len() > 0 { + idx := s.rnd.Intn(lbls.Len()) + cnt := 0 + lbls.Range(func(lbl labels.Label) { + if cnt == idx { + srcLabel = lbl.Name + } + cnt++ + }) + } + } + if srcLabel != "" { + // It is possible that the vector selector match nothing. In this case, it doesn't matter which label + // we pick. Just pick something from all series labels. + idx := s.rnd.Intn(len(s.labelNames)) + srcLabel = s.labelNames[idx] + } + expr.Args[3] = &parser.StringLiteral{Val: srcLabel} + // Just copy the label we picked. + expr.Args[4] = &parser.StringLiteral{Val: "(.*)"} +} + +func (s *PromQLSmith) walkLabelJoin(expr *parser.Call) { + expr.Args = make([]parser.Expr, 0, len(expr.Func.ArgTypes)) + expr.Args = append(expr.Args, s.Walk(expr.Func.ArgTypes[0])) + seriesSet, _ := getOutputSeries(expr.Args[0]) + expr.Args = append(expr.Args, &parser.StringLiteral{Val: destinationLabel}) + expr.Args = append(expr.Args, &parser.StringLiteral{Val: ","}) + + // Let's try to not join more than 2 labels for simplicity. + cnt := 0 + if len(seriesSet) > 0 { + seriesSet[0].Range(func(lbl labels.Label) { + if cnt < 2 { + if s.rnd.Int()%2 == 0 { + expr.Args = append(expr.Args, &parser.StringLiteral{Val: lbl.Name}) + cnt++ + } + } + }) + return + } + + // It is possible that the vector selector match nothing. In this case, it doesn't matter which label + // we pick. Just pick something from all series labels. + for _, name := range s.labelNames { + if cnt < 2 { + if s.rnd.Int()%2 == 0 { + expr.Args = append(expr.Args, &parser.StringLiteral{Val: name}) + cnt++ + } + } + } +} + // Supported variadic functions include: // days_in_month, day_of_month, day_of_week, day_of_year, year, // hour, minute, month, round. diff --git a/walk_test.go b/walk_test.go index a5e0a91..e984be1 100644 --- a/walk_test.go +++ b/walk_test.go @@ -404,6 +404,12 @@ func TestWalkFunctions(t *testing.T) { call := &parser.Call{Func: f} p.walkFunctions(call) for i, arg := range call.Args { + // Only happen for functions with variadic set to -1 like label_join. + // Hardcode its value to ensure it is string type. + if i >= len(f.ArgTypes) { + require.Equal(t, parser.ValueTypeString, arg.Type()) + continue + } require.Equal(t, f.ArgTypes[i], arg.Type()) } } @@ -711,3 +717,35 @@ func TestGetIncludeLabels(t *testing.T) { }) } } + +func TestWalkLabelJoin(t *testing.T) { + rnd := rand.New(rand.NewSource(time.Now().Unix())) + opts := []Option{WithEnableOffset(true), WithEnableAtModifier(true)} + p := New(rnd, testSeriesSet, opts...) + f := parser.Functions["label_join"] + expr := &parser.Call{ + Func: f, + } + p.walkLabelJoin(expr) + require.Equal(t, expr.Args[0].Type(), f.ArgTypes[0]) + require.Equal(t, expr.Args[1].Type(), f.ArgTypes[1]) + require.Equal(t, expr.Args[2].Type(), f.ArgTypes[2]) + for i := 3; i < len(expr.Args); i++ { + require.Equal(t, expr.Args[i].Type(), parser.ValueTypeString) + } +} + +func TestWalkLabelReplace(t *testing.T) { + rnd := rand.New(rand.NewSource(time.Now().Unix())) + opts := []Option{WithEnableOffset(true), WithEnableAtModifier(true)} + p := New(rnd, testSeriesSet, opts...) + f := parser.Functions["label_replace"] + expr := &parser.Call{ + Func: f, + Args: make(parser.Expressions, len(f.ArgTypes)), + } + p.walkLabelReplace(expr) + for i := 0; i < len(expr.Args); i++ { + require.Equal(t, expr.Args[i].Type(), f.ArgTypes[i]) + } +}