Merge pull request #133 from SungJin1212/Support-experimental-promQLs

Add experimental promQL funcs

opts.go

@@ -33,6 +33,11 @@ var (
 		parser.COUNT_VALUES,
 	}
 
+	experimentalPromQLAggrs = []parser.ItemType{
+		parser.LIMITK,
+		parser.LIMIT_RATIO,
+	}
+
 	defaultSupportedBinOps = []parser.ItemType{
 		parser.SUB,
 		parser.ADD,
@@ -52,14 +57,17 @@ var (
 		parser.LUNLESS,
 	}
 
-	defaultSupportedFuncs []*parser.Function
+	defaultSupportedFuncs      []*parser.Function
+	experimentalSupportedFuncs []*parser.Function
 )
 
 func init() {
 	for _, f := range parser.Functions {
 		// Ignore experimental functions for now.
 		if !f.Experimental {
 			defaultSupportedFuncs = append(defaultSupportedFuncs, f)
+		} else {
+			experimentalSupportedFuncs = append(experimentalSupportedFuncs, f)
 		}
 	}
 }
@@ -70,10 +78,11 @@ type options struct {
 	enabledFuncs  []*parser.Function
 	enabledBinops []parser.ItemType
 
-	enableOffset           bool
-	enableAtModifier       bool
-	enableVectorMatching   bool
-	atModifierMaxTimestamp int64
+	enableOffset                      bool
+	enableAtModifier                  bool
+	enableVectorMatching              bool
+	enableExperimentalPromQLFunctions bool
+	atModifierMaxTimestamp            int64
 
 	enforceLabelMatchers []*labels.Matcher
 }
@@ -95,6 +104,11 @@ func (o *options) applyDefaults() {
 		o.enabledFuncs = defaultSupportedFuncs
 	}
 
+	if o.enableExperimentalPromQLFunctions {
+		o.enabledAggrs = append(o.enabledAggrs, experimentalPromQLAggrs...)
+		o.enabledFuncs = append(o.enabledFuncs, experimentalSupportedFuncs...)
+	}
+
 	if o.atModifierMaxTimestamp == 0 {
 		o.atModifierMaxTimestamp = time.Now().UnixMilli()
 	}
@@ -135,6 +149,12 @@ func WithEnableVectorMatching(enableVectorMatching bool) Option {
 	})
 }
 
+func WithEnableExperimentalPromQLFunctions(enableExperimentalPromQLFunctions bool) Option {
+	return optionFunc(func(o *options) {
+		o.enableExperimentalPromQLFunctions = enableExperimentalPromQLFunctions
+	})
+}
+
 func WithEnabledBinOps(enabledBinops []parser.ItemType) Option {
 	return optionFunc(func(o *options) {
 		o.enabledBinops = enabledBinops

opts_test.go

@@ -20,6 +20,19 @@ func TestWithEnableAtModifier(t *testing.T) {
 	require.True(t, o.enableAtModifier)
 }
 
+func TestWithEnableExperimentalPromQL(t *testing.T) {
+	o := &options{}
+	WithEnableExperimentalPromQLFunctions(true).apply(o)
+	WithEnabledFunctions(nil).apply(o)
+	WithEnabledAggrs(nil).apply(o)
+	o.applyDefaults()
+
+	// check experimental aggrs and funcs are appended well
+	require.True(t, o.enableExperimentalPromQLFunctions)
+	require.Equal(t, len(defaultSupportedAggrs)+len(experimentalPromQLAggrs), len(o.enabledAggrs))
+	require.Equal(t, len(defaultSupportedFuncs)+len(experimentalSupportedFuncs), len(o.enabledFuncs))
+}
+
 func TestWithEnabledAggrs(t *testing.T) {
 	o := &options{}
 	WithEnabledAggrs([]parser.ItemType{parser.SUM}).apply(o)

promqlsmith.go

@@ -40,10 +40,11 @@ var (
 type PromQLSmith struct {
 	rnd *rand.Rand
 
-	enableOffset           bool
-	enableAtModifier       bool
-	enableVectorMatching   bool
-	atModifierMaxTimestamp int64
+	enableOffset             bool
+	enableAtModifier         bool
+	enableVectorMatching     bool
+	enableExperimentalPromQL bool
+	atModifierMaxTimestamp   int64
 
 	seriesSet       []labels.Labels
 	labelNames      []string
@@ -65,17 +66,18 @@ func New(rnd *rand.Rand, seriesSet []labels.Labels, opts ...Option) *PromQLSmith
 	options.applyDefaults()
 
 	ps := &PromQLSmith{
-		rnd:                    rnd,
-		seriesSet:              filterEmptySeries(seriesSet),
-		supportedExprs:         options.enabledExprs,
-		supportedAggrs:         options.enabledAggrs,
-		supportedBinops:        options.enabledBinops,
-		supportedFuncs:         options.enabledFuncs,
-		enableOffset:           options.enableOffset,
-		enableAtModifier:       options.enableAtModifier,
-		atModifierMaxTimestamp: options.atModifierMaxTimestamp,
-		enableVectorMatching:   options.enableVectorMatching,
-		enforceMatchers:        options.enforceLabelMatchers,
+		rnd:                      rnd,
+		seriesSet:                filterEmptySeries(seriesSet),
+		supportedExprs:           options.enabledExprs,
+		supportedAggrs:           options.enabledAggrs,
+		supportedBinops:          options.enabledBinops,
+		supportedFuncs:           options.enabledFuncs,
+		enableOffset:             options.enableOffset,
+		enableAtModifier:         options.enableAtModifier,
+		atModifierMaxTimestamp:   options.atModifierMaxTimestamp,
+		enableVectorMatching:     options.enableVectorMatching,
+		enableExperimentalPromQL: options.enableExperimentalPromQLFunctions,
+		enforceMatchers:          options.enforceLabelMatchers,
 	}
 	ps.labelNames, ps.labelValues = labelNameAndValuesFromLabelSet(seriesSet)
 	return ps

walk.go

@@ -83,6 +83,8 @@ func (s *PromQLSmith) walkAggregateParam(op parser.ItemType) parser.Expr {
 		return s.Walk(parser.ValueTypeScalar)
 	case parser.COUNT_VALUES:
 		return &parser.StringLiteral{Val: "value"}
+	case parser.LIMITK, parser.LIMIT_RATIO:
+		return s.Walk(parser.ValueTypeScalar)
 	}
 	return nil
 }
@@ -256,6 +258,9 @@ func (s *PromQLSmith) walkFunctions(expr *parser.Call) {
 	case "label_join":
 		s.walkLabelJoin(expr)
 		return
+	case "sort_by_label", "sort_by_label_desc":
+		s.walkSortByLabel(expr)
+		return
 	default:
 	}
 
@@ -313,6 +318,38 @@ func (s *PromQLSmith) walkLabelReplace(expr *parser.Call) {
 	expr.Args[4] = &parser.StringLiteral{Val: "(.*)"}
 }
 
+func (s *PromQLSmith) walkSortByLabel(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])
+
+	// Let's try to not sort more than 1 label 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 < 1 {
+			if s.rnd.Int()%2 == 0 {
+				expr.Args = append(expr.Args, &parser.StringLiteral{Val: name})
+				cnt++
+			}
+		}
+	}
+}
+
 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]))

walk_test.go

@@ -190,7 +190,7 @@ func TestWalkVectorMatching(t *testing.T) {
 
 func TestWalkAggregateParam(t *testing.T) {
 	rnd := rand.New(rand.NewSource(time.Now().Unix()))
-	opts := []Option{WithEnableOffset(true), WithEnableAtModifier(true)}
+	opts := []Option{WithEnableOffset(true), WithEnableAtModifier(true), WithEnableExperimentalPromQLFunctions(true)}
 	p := New(rnd, testSeriesSet, opts...)
 	for i, tc := range []struct {
 		op           parser.ItemType
@@ -222,6 +222,18 @@ func TestWalkAggregateParam(t *testing.T) {
 				require.Equal(t, e.Val, "value")
 			},
 		},
+		{
+			op: parser.LIMITK,
+			expectedFunc: func(t *testing.T, expr parser.Expr) {
+				require.Equal(t, parser.ValueTypeScalar, expr.Type())
+			},
+		},
+		{
+			op: parser.LIMIT_RATIO,
+			expectedFunc: func(t *testing.T, expr parser.Expr) {
+				require.Equal(t, parser.ValueTypeScalar, expr.Type())
+			},
+		},
 	} {
 		t.Run(fmt.Sprintf("test_case_%d", i), func(t *testing.T) {
 			expr := p.walkAggregateParam(tc.op)
@@ -724,6 +736,24 @@ func TestGetIncludeLabels(t *testing.T) {
 	}
 }
 
+func TestWalkSortByLabel(t *testing.T) {
+	rnd := rand.New(rand.NewSource(time.Now().Unix()))
+	opts := []Option{WithEnableOffset(true), WithEnableAtModifier(true), WithEnableExperimentalPromQLFunctions(true)}
+	p := New(rnd, testSeriesSet, opts...)
+
+	for _, name := range []string{"sort_by_label", "sort_by_label_desc"} {
+		f := parser.Functions[name]
+		expr := &parser.Call{
+			Func: f,
+		}
+		p.walkSortByLabel(expr)
+		require.Equal(t, expr.Args[0].Type(), f.ArgTypes[0])
+		for i := 1; i < len(expr.Args); i++ {
+			require.Equal(t, expr.Args[i].Type(), parser.ValueTypeString)
+		}
+	}
+}
+
 func TestWalkLabelJoin(t *testing.T) {
 	rnd := rand.New(rand.NewSource(time.Now().Unix()))
 	opts := []Option{WithEnableOffset(true), WithEnableAtModifier(true)}