From 228f5f43593408930d2785c950a53303b7e2347e Mon Sep 17 00:00:00 2001 From: "kade.lee" Date: Wed, 6 Nov 2024 10:46:57 +0900 Subject: [PATCH] Add experimental promQL funcs Signed-off-by: kade.lee --- opts.go | 30 +++++++++++++++++++++++++----- opts_test.go | 13 +++++++++++++ promqlsmith.go | 32 +++++++++++++++++--------------- walk.go | 37 +++++++++++++++++++++++++++++++++++++ walk_test.go | 32 +++++++++++++++++++++++++++++++- 5 files changed, 123 insertions(+), 21 deletions(-) diff --git a/opts.go b/opts.go index 2f806fb..1ab2efa 100644 --- a/opts.go +++ b/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,7 +57,8 @@ var ( parser.LUNLESS, } - defaultSupportedFuncs []*parser.Function + defaultSupportedFuncs []*parser.Function + experimentalSupportedFuncs []*parser.Function ) func init() { @@ -60,6 +66,8 @@ func init() { // 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 diff --git a/opts_test.go b/opts_test.go index af92dd5..339423b 100644 --- a/opts_test.go +++ b/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) diff --git a/promqlsmith.go b/promqlsmith.go index 436c3a5..184d939 100644 --- a/promqlsmith.go +++ b/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 diff --git a/walk.go b/walk.go index 930bba9..5e50a31 100644 --- a/walk.go +++ b/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])) diff --git a/walk_test.go b/walk_test.go index 7515c84..f7c6a81 100644 --- a/walk_test.go +++ b/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)}