From 782da6a822d7e25416ac747098affbccfb06632a Mon Sep 17 00:00:00 2001 From: "kade.lee" Date: Tue, 5 Nov 2024 17:59:44 +0900 Subject: [PATCH] Add experimental promQL funcs Signed-off-by: kade.lee --- opts.go | 15 +++++++++++---- opts_test.go | 6 ++++++ promqlsmith.go | 32 +++++++++++++++++--------------- walk.go | 21 +++++++++++++++++++++ walk_test.go | 33 ++++++++++++++++++++++++++++++++- 5 files changed, 87 insertions(+), 20 deletions(-) diff --git a/opts.go b/opts.go index 2f806fb..9e9d2b6 100644 --- a/opts.go +++ b/opts.go @@ -70,10 +70,11 @@ type options struct { enabledFuncs []*parser.Function enabledBinops []parser.ItemType - enableOffset bool - enableAtModifier bool - enableVectorMatching bool - atModifierMaxTimestamp int64 + enableOffset bool + enableAtModifier bool + enableVectorMatching bool + enableExperimentalPromQL bool + atModifierMaxTimestamp int64 enforceLabelMatchers []*labels.Matcher } @@ -135,6 +136,12 @@ func WithEnableVectorMatching(enableVectorMatching bool) Option { }) } +func WithEnableExperimentalPromQL(enableExperimentalPromQL bool) Option { + return optionFunc(func(o *options) { + o.enableExperimentalPromQL = enableExperimentalPromQL + }) +} + 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..1bcaa47 100644 --- a/opts_test.go +++ b/opts_test.go @@ -20,6 +20,12 @@ func TestWithEnableAtModifier(t *testing.T) { require.True(t, o.enableAtModifier) } +func TestWithEnableExperimentalPromQL(t *testing.T) { + o := &options{} + WithEnableExperimentalPromQL(true).apply(o) + require.True(t, o.enableExperimentalPromQL) +} + func TestWithEnabledAggrs(t *testing.T) { o := &options{} WithEnabledAggrs([]parser.ItemType{parser.SUM}).apply(o) diff --git a/promqlsmith.go b/promqlsmith.go index 436c3a5..e84391c 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.enableExperimentalPromQL, + enforceMatchers: options.enforceLabelMatchers, } ps.labelNames, ps.labelValues = labelNameAndValuesFromLabelSet(seriesSet) return ps diff --git a/walk.go b/walk.go index 930bba9..975d392 100644 --- a/walk.go +++ b/walk.go @@ -83,6 +83,10 @@ 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: + if s.enableExperimentalPromQL { + return s.Walk(parser.ValueTypeScalar) + } } return nil } @@ -256,6 +260,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 +320,20 @@ 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])) + + cnt := 0 + // randomly pick two labels + for _, name := range s.labelNames { + if cnt < 2 { + 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..92de9d0 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), WithEnableExperimentalPromQL(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,25 @@ func TestGetIncludeLabels(t *testing.T) { } } +func TestWalkSortByLabel(t *testing.T) { + rnd := rand.New(rand.NewSource(time.Now().Unix())) + opts := []Option{WithEnableOffset(true), WithEnableAtModifier(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]) + require.Equal(t, expr.Args[1].Type(), f.ArgTypes[1]) + for i := 2; 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)}