add support for label_join and label_replace function

Signed-off-by: Ben Ye <[email protected]>

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)

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.

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])
+	}
+}