Skip to content

Commit

Permalink
feat(lokicompliance): detect template args automatically
Browse files Browse the repository at this point in the history
  • Loading branch information
tdakkota committed Apr 27, 2024
1 parent cc91485 commit c5e4fb5
Show file tree
Hide file tree
Showing 2 changed files with 198 additions and 85 deletions.
243 changes: 161 additions & 82 deletions internal/lokicompliance/expand.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,98 +5,63 @@ import (
"strconv"
"strings"
"text/template"
"text/template/parse"
"time"

"github.com/go-faster/errors"
"golang.org/x/exp/maps"

"github.com/go-faster/oteldb/internal/lokiapi"
)

var testVariantArgs = map[string][]string{
"range": {"1s", "15s", "1m", "5m", "15m", "1h"},
"offset": {"1m", "5m", "10m"},
"simpleVecAggOp": {"sum", "avg", "max", "min", "count", "stddev", "stdvar"},
"simpleRangeAggOp": {
"count_over_time",
"rate",
"bytes_over_time",
"bytes_rate",
},
"unwrapRangeAggOp": {
"rate_counter",
"avg_over_time",
"sum_over_time",
"min_over_time",
"max_over_time",
"stdvar_over_time",
"stddev_over_time",
"first_over_time",
"last_over_time",
},
"topBottomVectorOp": {"topk", "bottomk"},
"unwrapExpr": {
"unwrap status",
"unwrap duration(took)",
"unwrap bytes(size)",
},

"quantile": {
"0.1",
"0.5",
"0.75",
"0.95",
"0.90",
"0.99",
"1",
"1.5",
},
"lineFilterOp": {"|=", "!=", "~=", "!~"},
"arithBinOp": {"+", "-", "*", "/", "%", "^"},
"logicBinOp": {"and", "or", "unless"},
}
var (
testVariantArgs = map[string][]string{
"range": {"1s", "15s", "1m", "5m", "15m", "1h"},
"offset": {"1m", "5m", "10m"},
"simpleVecAggOp": {"sum", "avg", "max", "min", "count", "stddev", "stdvar"},
"simpleRangeAggOp": {
"count_over_time",
"rate",
"bytes_over_time",
"bytes_rate",
},
"unwrapRangeAggOp": {
"rate_counter",
"avg_over_time",
"sum_over_time",
"min_over_time",
"max_over_time",
"stdvar_over_time",
"stddev_over_time",
"first_over_time",
"last_over_time",
},
"topBottomVectorOp": {"topk", "bottomk"},
"unwrapExpr": {
"unwrap status",
"unwrap duration(took)",
"unwrap bytes(size)",
},

func execTemplateToString(t *template.Template, data any) (string, error) {
var sb strings.Builder
if err := t.Execute(&sb, data); err != nil {
return "", err
}
return sb.String(), nil
}

func getQueries(
t *template.Template,
variantArgs []string,
args map[string]string,
add func(string),
) error {
if len(variantArgs) == 0 {
q, err := execTemplateToString(t, args)
if err != nil {
return err
}
add(q)
return nil
"quantile": {
"0.1",
"0.5",
"0.75",
"0.95",
"0.90",
"0.99",
"1",
"1.5",
},
"lineFilterOp": {"|=", "!=", "~=", "!~"},
"arithBinOp": {"+", "-", "*", "/", "%", "^"},
"logicBinOp": {"and", "or", "unless"},
}

arg := variantArgs[0]
values, ok := testVariantArgs[arg]
if !ok {
return errors.Errorf("unknown arg %q", arg)
templateFuncMap = template.FuncMap{
"quote": strconv.Quote,
}
variantArgs = variantArgs[1:]

for _, val := range values {
args[arg] = val
if err := getQueries(t, variantArgs, args, add); err != nil {
return err
}
}
return nil
}

var templateFuncMap = template.FuncMap{
"quote": strconv.Quote,
}
)

// ExpandQuery expands given test case.
func ExpandQuery(cfg *Config, start, end time.Time, step time.Duration) (r []*TestCase, _ error) {
Expand All @@ -121,7 +86,7 @@ func ExpandQuery(cfg *Config, start, end time.Time, step time.Duration) (r []*Te
}

// Sort and deduplicate args.
args := tc.VariantArgs
args := collectVariantArgs(templ)
slices.Sort(args)
args = slices.Compact(args)

Expand Down Expand Up @@ -149,3 +114,117 @@ func ExpandQuery(cfg *Config, start, end time.Time, step time.Duration) (r []*Te

return r, nil
}

func getQueries(
t *template.Template,
variantArgs []string,
args map[string]string,
add func(string),
) error {
if len(variantArgs) == 0 {
q, err := execTemplateToString(t, args)
if err != nil {
return err
}
add(q)
return nil
}

arg := variantArgs[0]
values, ok := testVariantArgs[arg]
if !ok {
return errors.Errorf("unknown arg %q", arg)
}
variantArgs = variantArgs[1:]

for _, val := range values {
args[arg] = val
if err := getQueries(t, variantArgs, args, add); err != nil {
return err
}
}
return nil
}

func execTemplateToString(t *template.Template, data any) (string, error) {
var sb strings.Builder
if err := t.Execute(&sb, data); err != nil {
return "", err
}
return sb.String(), nil
}

func collectVariantArgs(t *template.Template) []string {
n := t.Root
if n == nil {
return nil
}

args := map[string]struct{}{}
walkTemplate(t.Root, func(n parse.Node) {
field, ok := n.(*parse.FieldNode)
if !ok || len(field.Ident) != 1 {
return
}
arg := field.Ident[0]

if _, ok := testVariantArgs[arg]; !ok {
return
}
args[arg] = struct{}{}
})

if len(args) == 0 {
return nil
}
return maps.Keys(args)
}

func walkTemplate(n parse.Node, cb func(parse.Node)) {
cb(n)
switch n := n.(type) {
case *parse.ListNode:
for _, sub := range n.Nodes {
walkTemplate(sub, cb)
}
case *parse.PipeNode:
for _, sub := range n.Decl {
walkTemplate(sub, cb)
}
for _, sub := range n.Cmds {
walkTemplate(sub, cb)
}
case *parse.ActionNode:
if sub := n.Pipe; sub != nil {
walkTemplate(sub, cb)
}
case *parse.CommandNode:
for _, sub := range n.Args {
walkTemplate(sub, cb)
}
case *parse.ChainNode:
if sub := n.Node; sub != nil {
walkTemplate(sub, cb)
}
case *parse.BranchNode:
if sub := n.Pipe; sub != nil {
walkTemplate(sub, cb)
}
if sub := n.List; sub != nil {
walkTemplate(sub, cb)
}
if sub := n.ElseList; sub != nil {
walkTemplate(sub, cb)
}
case *parse.IfNode:
walkTemplate(&n.BranchNode, cb)
case *parse.RangeNode:
walkTemplate(&n.BranchNode, cb)
case *parse.WithNode:
walkTemplate(&n.BranchNode, cb)
case *parse.TemplateNode:
if sub := n.Pipe; sub != nil {
walkTemplate(sub, cb)
}
}
}
40 changes: 37 additions & 3 deletions internal/lokicompliance/expand_test.go
Original file line number Diff line number Diff line change
@@ -1,19 +1,20 @@
package lokicompliance

import (
"fmt"
"testing"
"text/template"
"time"

"github.com/stretchr/testify/require"
)

func TestExpandQuery(t *testing.T) {
arg := "offset"
arg := "quantile"
cfg := &Config{
TestCases: []*TestCasePattern{
{
Query: "{{ ." + arg + " }}",
VariantArgs: []string{arg},
Query: "{{ ." + arg + " }}",
},
},
}
Expand All @@ -26,5 +27,38 @@ func TestExpandQuery(t *testing.T) {
for _, tc := range tcs {
queries = append(queries, tc.Query)
}
require.NotEmpty(t, queries)
require.ElementsMatch(t, queries, testVariantArgs[arg])
}

func TestCollectVariantArgs(t *testing.T) {
tests := []struct {
input string
want []string
}{
{
``,
nil,
},
{
`{{ .quantile }}`,
[]string{"quantile"},
},
{
`{{ .quantile }}+{{ .quantile }}`,
[]string{"quantile"},
},
{
`{{ .unwrapRangeAggOp }}( {} [{{ .range }}] )`,
[]string{"unwrapRangeAggOp", "range"},
},
}
for i, tt := range tests {
tt := tt
t.Run(fmt.Sprintf("Test%d", i+1), func(t *testing.T) {
templ, err := template.New("templ").Parse(tt.input)
require.NoError(t, err)
require.Equal(t, tt.want, collectVariantArgs(templ))
})
}
}

0 comments on commit c5e4fb5

Please sign in to comment.