Implement default DFS traversal visitor for LogQL expressions (#11489)

This can be used as a default visitor for use cases where you only want to visit specific expression without the needing to implement the traversal logic of each and every expression, see implementation in test case. Signed-off-by: Christian Haudum <christian.haudum@gmail.com>
2 years ago · cd74ddf4b4
parent 37a0ed60a0
commit cd74ddf4b4
3 changed files with 333 additions and 119 deletions
--- a/pkg/logql/syntax/visit.go
+++ b/pkg/logql/syntax/visit.go
@ -0,0 +1,285 @@
+package syntax
+
+type AcceptVisitor interface {
+	Accept(RootVisitor)
+}
+
+type RootVisitor interface {
+	SampleExprVisitor
+	LogSelectorExprVisitor
+	StageExprVisitor
+
+	VisitLogRange(*LogRange)
+}
+
+type SampleExprVisitor interface {
+	VisitBinOp(*BinOpExpr)
+	VisitVectorAggregation(*VectorAggregationExpr)
+	VisitRangeAggregation(*RangeAggregationExpr)
+	VisitLabelReplace(*LabelReplaceExpr)
+	VisitLiteral(*LiteralExpr)
+	VisitVector(*VectorExpr)
+}
+
+type LogSelectorExprVisitor interface {
+	VisitMatchers(*MatchersExpr)
+	VisitPipeline(*PipelineExpr)
+	VisitLiteral(*LiteralExpr)
+	VisitVector(*VectorExpr)
+}
+
+type StageExprVisitor interface {
+	VisitDecolorize(*DecolorizeExpr)
+	VisitDropLabels(*DropLabelsExpr)
+	VisitJSONExpressionParser(*JSONExpressionParser)
+	VisitKeepLabel(*KeepLabelsExpr)
+	VisitLabelFilter(*LabelFilterExpr)
+	VisitLabelFmt(*LabelFmtExpr)
+	VisitLabelParser(*LabelParserExpr)
+	VisitLineFilter(*LineFilterExpr)
+	VisitLineFmt(*LineFmtExpr)
+	VisitLogfmtExpressionParser(*LogfmtExpressionParser)
+	VisitLogfmtParser(*LogfmtParserExpr)
+}
+
+var _ RootVisitor = &DepthFirstTraversal{}
+
+type DepthFirstTraversal struct {
+	VisitBinOpFn                  func(v RootVisitor, e *BinOpExpr)
+	VisitDecolorizeFn             func(v RootVisitor, e *DecolorizeExpr)
+	VisitDropLabelsFn             func(v RootVisitor, e *DropLabelsExpr)
+	VisitJSONExpressionParserFn   func(v RootVisitor, e *JSONExpressionParser)
+	VisitKeepLabelFn              func(v RootVisitor, e *KeepLabelsExpr)
+	VisitLabelFilterFn            func(v RootVisitor, e *LabelFilterExpr)
+	VisitLabelFmtFn               func(v RootVisitor, e *LabelFmtExpr)
+	VisitLabelParserFn            func(v RootVisitor, e *LabelParserExpr)
+	VisitLabelReplaceFn           func(v RootVisitor, e *LabelReplaceExpr)
+	VisitLineFilterFn             func(v RootVisitor, e *LineFilterExpr)
+	VisitLineFmtFn                func(v RootVisitor, e *LineFmtExpr)
+	VisitLiteralFn                func(v RootVisitor, e *LiteralExpr)
+	VisitLogRangeFn               func(v RootVisitor, e *LogRange)
+	VisitLogfmtExpressionParserFn func(v RootVisitor, e *LogfmtExpressionParser)
+	VisitLogfmtParserFn           func(v RootVisitor, e *LogfmtParserExpr)
+	VisitMatchersFn               func(v RootVisitor, e *MatchersExpr)
+	VisitPipelineFn               func(v RootVisitor, e *PipelineExpr)
+	VisitRangeAggregationFn       func(v RootVisitor, e *RangeAggregationExpr)
+	VisitVectorFn                 func(v RootVisitor, e *VectorExpr)
+	VisitVectorAggregationFn      func(v RootVisitor, e *VectorAggregationExpr)
+}
+
+// VisitBinOp implements RootVisitor.
+func (v *DepthFirstTraversal) VisitBinOp(e *BinOpExpr) {
+	if e == nil {
+		return
+	}
+	if v.VisitBinOpFn != nil {
+		v.VisitBinOpFn(v, e)
+	} else {
+		e.SampleExpr.Accept(v)
+		e.RHS.Accept(v)
+	}
+}
+
+// VisitDecolorize implements RootVisitor.
+func (v *DepthFirstTraversal) VisitDecolorize(e *DecolorizeExpr) {
+	if e == nil {
+		return
+	}
+	if v.VisitDecolorizeFn != nil {
+		v.VisitDecolorizeFn(v, e)
+	}
+}
+
+// VisitDropLabels implements RootVisitor.
+func (v *DepthFirstTraversal) VisitDropLabels(e *DropLabelsExpr) {
+	if e == nil {
+		return
+	}
+	if v.VisitDecolorizeFn != nil {
+		v.VisitDropLabelsFn(v, e)
+	}
+}
+
+// VisitJSONExpressionParser implements RootVisitor.
+func (v *DepthFirstTraversal) VisitJSONExpressionParser(e *JSONExpressionParser) {
+	if e == nil {
+		return
+	}
+	if v.VisitJSONExpressionParserFn != nil {
+		v.VisitJSONExpressionParserFn(v, e)
+	}
+}
+
+// VisitKeepLabel implements RootVisitor.
+func (v *DepthFirstTraversal) VisitKeepLabel(e *KeepLabelsExpr) {
+	if e == nil {
+		return
+	}
+	if v.VisitKeepLabelFn != nil {
+		v.VisitKeepLabelFn(v, e)
+	}
+}
+
+// VisitLabelFilter implements RootVisitor.
+func (v *DepthFirstTraversal) VisitLabelFilter(e *LabelFilterExpr) {
+	if e == nil {
+		return
+	}
+	if v.VisitLabelFilterFn != nil {
+		v.VisitLabelFilterFn(v, e)
+	}
+}
+
+// VisitLabelFmt implements RootVisitor.
+func (v *DepthFirstTraversal) VisitLabelFmt(e *LabelFmtExpr) {
+	if e == nil {
+		return
+	}
+	if v.VisitLabelFmtFn != nil {
+		v.VisitLabelFmtFn(v, e)
+	}
+}
+
+// VisitLabelParser implements RootVisitor.
+func (v *DepthFirstTraversal) VisitLabelParser(e *LabelParserExpr) {
+	if e == nil {
+		return
+	}
+	if v.VisitLabelParserFn != nil {
+		v.VisitLabelParserFn(v, e)
+	}
+}
+
+// VisitLabelReplace implements RootVisitor.
+func (v *DepthFirstTraversal) VisitLabelReplace(e *LabelReplaceExpr) {
+	if e == nil {
+		return
+	}
+	if v.VisitLabelReplaceFn != nil {
+		v.VisitLabelReplaceFn(v, e)
+	}
+}
+
+// VisitLineFilter implements RootVisitor.
+func (v *DepthFirstTraversal) VisitLineFilter(e *LineFilterExpr) {
+	if e == nil {
+		return
+	}
+	if v.VisitLineFilterFn != nil {
+		v.VisitLineFilterFn(v, e)
+	} else {
+		e.Left.Accept(v)
+		e.Or.Accept(v)
+	}
+}
+
+// VisitLineFmt implements RootVisitor.
+func (v *DepthFirstTraversal) VisitLineFmt(e *LineFmtExpr) {
+	if e == nil {
+		return
+	}
+	if v.VisitLineFmtFn != nil {
+		v.VisitLineFmtFn(v, e)
+	}
+}
+
+// VisitLiteral implements RootVisitor.
+func (v *DepthFirstTraversal) VisitLiteral(e *LiteralExpr) {
+	if e == nil {
+		return
+	}
+	if v.VisitLiteralFn != nil {
+		v.VisitLiteralFn(v, e)
+	}
+}
+
+// VisitLogRange implements RootVisitor.
+func (v *DepthFirstTraversal) VisitLogRange(e *LogRange) {
+	if e == nil {
+		return
+	}
+	if v.VisitLogRangeFn != nil {
+		v.VisitLogRangeFn(v, e)
+	} else {
+		e.Left.Accept(v)
+	}
+}
+
+// VisitLogfmtExpressionParser implements RootVisitor.
+func (v *DepthFirstTraversal) VisitLogfmtExpressionParser(e *LogfmtExpressionParser) {
+	if e == nil {
+		return
+	}
+	if v.VisitLogfmtExpressionParserFn != nil {
+		v.VisitLogfmtExpressionParserFn(v, e)
+	}
+}
+
+// VisitLogfmtParser implements RootVisitor.
+func (v *DepthFirstTraversal) VisitLogfmtParser(e *LogfmtParserExpr) {
+	if e == nil {
+		return
+	}
+	if v.VisitLogfmtParserFn != nil {
+		v.VisitLogfmtParserFn(v, e)
+	}
+}
+
+// VisitMatchers implements RootVisitor.
+func (v *DepthFirstTraversal) VisitMatchers(e *MatchersExpr) {
+	if e == nil {
+		return
+	}
+	if v.VisitMatchersFn != nil {
+		v.VisitMatchersFn(v, e)
+	}
+}
+
+// VisitPipeline implements RootVisitor.
+func (v *DepthFirstTraversal) VisitPipeline(e *PipelineExpr) {
+	if e == nil {
+		return
+	}
+	if v.VisitPipelineFn != nil {
+		v.VisitPipelineFn(v, e)
+	} else {
+		e.Left.Accept(v)
+		for i := range e.MultiStages {
+			e.MultiStages[i].Accept(v)
+		}
+	}
+}
+
+// VisitRangeAggregation implements RootVisitor.
+func (v *DepthFirstTraversal) VisitRangeAggregation(e *RangeAggregationExpr) {
+	if e == nil {
+		return
+	}
+	if v.VisitRangeAggregationFn != nil {
+		v.VisitRangeAggregationFn(v, e)
+	} else {
+		e.Left.Accept(v)
+	}
+}
+
+// VisitVector implements RootVisitor.
+func (v *DepthFirstTraversal) VisitVector(e *VectorExpr) {
+	if e == nil {
+		return
+	}
+	if v.VisitVectorFn != nil {
+		v.VisitVectorFn(v, e)
+	}
+}
+
+// VisitVectorAggregation implements RootVisitor.
+func (v *DepthFirstTraversal) VisitVectorAggregation(e *VectorAggregationExpr) {
+	if e == nil {
+		return
+	}
+	if v.VisitVectorAggregationFn != nil {
+		v.VisitVectorAggregationFn(v, e)
+	} else {
+		e.Left.Accept(v)
+	}
+}
--- a/pkg/logql/syntax/visit_test.go
+++ b/pkg/logql/syntax/visit_test.go
@ -0,0 +1,48 @@
+package syntax
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestDepthFirstTraversalVisitor(t *testing.T) {
+
+	visited := [][2]string{}
+
+	visitor := &DepthFirstTraversal{
+		VisitLabelParserFn: func(v RootVisitor, e *LabelParserExpr) {
+			visited = append(visited, [2]string{fmt.Sprintf("%T", e), e.String()})
+		},
+		VisitLineFilterFn: func(v RootVisitor, e *LineFilterExpr) {
+			visited = append(visited, [2]string{fmt.Sprintf("%T", e), e.String()})
+		},
+		VisitLogfmtParserFn: func(v RootVisitor, e *LogfmtParserExpr) {
+			visited = append(visited, [2]string{fmt.Sprintf("%T", e), e.String()})
+		},
+		VisitMatchersFn: func(v RootVisitor, e *MatchersExpr) {
+			visited = append(visited, [2]string{fmt.Sprintf("%T", e), e.String()})
+		},
+	}
+
+	// Only expressions that have a Visit function defined are added to the list
+	expected := [][2]string{
+		{"*syntax.MatchersExpr", `{env="prod"}`},
+		{"*syntax.LineFilterExpr", `|= "foo" or "bar"`},
+		{"*syntax.LogfmtParserExpr", `| logfmt`},
+		{"*syntax.MatchersExpr", `{env="dev"}`},
+		{"*syntax.LineFilterExpr", `|~ "(foo|bar)"`},
+		{"*syntax.LabelParserExpr", `| json`},
+	}
+
+	query := `
+	sum by (container) (min_over_time({env="prod"} |= "foo" or "bar" | logfmt | unwrap duration [1m]))
+	/
+	sum by (container) (max_over_time({env="dev"} |~ "(foo|bar)" | json | unwrap duration [1m]))
+	`
+	expr, err := ParseExpr(query)
+	require.NoError(t, err)
+	expr.Accept(visitor)
+	require.Equal(t, expected, visited)
+}
--- a/pkg/logql/syntax/walk.go
+++ b/pkg/logql/syntax/walk.go
@ -1,7 +1,5 @@
 package syntax

-import "fmt"
-
 type WalkFn = func(e Expr)

 func walkAll(f WalkFn, xs ...Walkable) {
@ -13,120 +11,3 @@ func walkAll(f WalkFn, xs ...Walkable) {
 type Walkable interface {
 	Walk(f WalkFn)
 }
-
-type AcceptVisitor interface {
-	Accept(RootVisitor)
-}
-
-type RootVisitor interface {
-	SampleExprVisitor
-	LogSelectorExprVisitor
-	StageExprVisitor
-
-	VisitLogRange(*LogRange)
-}
-
-type SampleExprVisitor interface {
-	VisitBinOp(*BinOpExpr)
-	VisitVectorAggregation(*VectorAggregationExpr)
-	VisitRangeAggregation(*RangeAggregationExpr)
-	VisitLabelReplace(*LabelReplaceExpr)
-	VisitLiteral(*LiteralExpr)
-	VisitVector(*VectorExpr)
-}
-
-type LogSelectorExprVisitor interface {
-	VisitMatchers(*MatchersExpr)
-	VisitPipeline(*PipelineExpr)
-	VisitLiteral(*LiteralExpr)
-	VisitVector(*VectorExpr)
-}
-
-type StageExprVisitor interface {
-	VisitDecolorize(*DecolorizeExpr)
-	VisitDropLabels(*DropLabelsExpr)
-	VisitJSONExpressionParser(*JSONExpressionParser)
-	VisitKeepLabel(*KeepLabelsExpr)
-	VisitLabelFilter(*LabelFilterExpr)
-	VisitLabelFmt(*LabelFmtExpr)
-	VisitLabelParser(*LabelParserExpr)
-	VisitLineFilter(*LineFilterExpr)
-	VisitLineFmt(*LineFmtExpr)
-	VisitLogfmtExpressionParser(*LogfmtExpressionParser)
-	VisitLogfmtParser(*LogfmtParserExpr)
-}
-
-func Dispatch(root Expr, v RootVisitor) error {
-	switch e := root.(type) {
-	case SampleExpr:
-		DispatchSampleExpr(e, v)
-	case LogSelectorExpr:
-		DispatchLogSelectorExpr(e, v)
-	case StageExpr:
-		DispatchStageExpr(e, v)
-	case *LogRange:
-		v.VisitLogRange(e)
-	default:
-		return fmt.Errorf("unpexpected root expression type: got (%T)", e)
-	}
-
-	return nil
-}
-
-func DispatchSampleExpr(expr SampleExpr, v SampleExprVisitor) {
-	switch e := expr.(type) {
-	case *BinOpExpr:
-		v.VisitBinOp(e)
-	case *VectorAggregationExpr:
-		v.VisitVectorAggregation(e)
-	case *RangeAggregationExpr:
-		v.VisitRangeAggregation(e)
-	case *LabelReplaceExpr:
-		v.VisitLabelReplace(e)
-	case *LiteralExpr:
-		v.VisitLiteral(e)
-	case *VectorExpr:
-		v.VisitVector(e)
-	}
-}
-
-func DispatchLogSelectorExpr(expr LogSelectorExpr, v LogSelectorExprVisitor) {
-	switch e := expr.(type) {
-	case *PipelineExpr:
-		v.VisitPipeline(e)
-	case *MatchersExpr:
-		v.VisitMatchers(e)
-	case *VectorExpr:
-		v.VisitVector(e)
-	case *LiteralExpr:
-		v.VisitLiteral(e)
-	}
-}
-
-func DispatchStageExpr(expr StageExpr, v StageExprVisitor) {
-	switch e := expr.(type) {
-	case *DecolorizeExpr:
-		v.VisitDecolorize(e)
-	case *DropLabelsExpr:
-		v.VisitDropLabels(e)
-	case *JSONExpressionParser:
-		v.VisitJSONExpressionParser(e)
-	case *KeepLabelsExpr:
-		v.VisitKeepLabel(e)
-	case *LabelFilterExpr:
-		v.VisitLabelFilter(e)
-	case *LabelFmtExpr:
-		v.VisitLabelFmt(e)
-	case *LabelParserExpr:
-		v.VisitLabelParser(e)
-	case *LineFilterExpr:
-		v.VisitLineFilter(e)
-	case *LineFmtExpr:
-		v.VisitLineFmt(e)
-	case *LogfmtExpressionParser:
-		v.VisitLogfmtExpressionParser(e)
-	case *LogfmtParserExpr:
-		v.VisitLogfmtParser(e)
-	}
-
-}