chore(storage/bloom/v1): add ExtractTestableLabelMatchers (#14119)

Signed-off-by: Robert Fratto <robertfratto@gmail.com>

pkg/logql/syntax/visit.go

@@ -95,7 +95,7 @@ func (v *DepthFirstTraversal) VisitDropLabels(e *DropLabelsExpr) {
 	if e == nil {
 		return
 	}
-	if v.VisitDecolorizeFn != nil {
+	if v.VisitDropLabelsFn != nil {
 		v.VisitDropLabelsFn(v, e)
 	}
 }

pkg/storage/bloom/v1/ast_extractor.go

@@ -0,0 +1,125 @@
+package v1
+
+import (
+	"github.com/prometheus/prometheus/model/labels"
+
+	"github.com/grafana/loki/v3/pkg/logql/log"
+	"github.com/grafana/loki/v3/pkg/logql/syntax"
+)
+
+// LabelMatcher represents bloom tests for key-value pairs, mapped from
+// LabelFilterExprs from the AST.
+type LabelMatcher interface{ isLabelMatcher() }
+
+// UnsupportedLabelMatcher represents a label matcher which could not be
+// mapped. Bloom tests for UnsupportedLabelMatchers must always pass.
+type UnsupportedLabelMatcher struct{}
+
+// PlainLabelMatcher represents a direct key-value matcher. Bloom tests
+// must only pass if the key-value pair exists in the bloom.
+type PlainLabelMatcher struct{ Key, Value string }
+
+// OrLabelMatcher represents a logical OR test. Bloom tests must only pass if
+// one of the Left or Right label matcher bloom tests pass.
+type OrLabelMatcher struct{ Left, Right LabelMatcher }
+
+// AndLabelMatcher represents a logical AND test. Bloom tests must only pass
+// if both of the Left and Right label matcher bloom tests pass.
+type AndLabelMatcher struct{ Left, Right LabelMatcher }
+
+// ExtractTestableLabelMatchers extracts label matchers from the label filters
+// in an expression. The resulting label matchers can then be used for testing
+// against bloom filters. Only label matchers before the first parse stage are
+// included.
+//
+// Unsupported LabelFilterExprs map to an UnsupportedLabelMatcher, for which
+// bloom tests should always pass.
+func ExtractTestableLabelMatchers(expr syntax.Expr) []LabelMatcher {
+	var (
+		exprs           []*syntax.LabelFilterExpr
+		foundParseStage bool
+	)
+
+	visitor := &syntax.DepthFirstTraversal{
+		VisitLabelFilterFn: func(v syntax.RootVisitor, e *syntax.LabelFilterExpr) {
+			if !foundParseStage {
+				exprs = append(exprs, e)
+			}
+		},
+
+		// TODO(rfratto): Find a way to generically represent or test for an
+		// expression that modifies extracted labels (parsers, keep, drop, etc.).
+		//
+		// As the AST is now, we can't prove at compile time that the list of
+		// visitors below is complete. For example, if a new parser stage
+		// expression is added without updating this list, blooms can silently
+		// misbehave.
+
+		VisitLogfmtParserFn:           func(v syntax.RootVisitor, e *syntax.LogfmtParserExpr) { foundParseStage = true },
+		VisitLabelParserFn:            func(v syntax.RootVisitor, e *syntax.LabelParserExpr) { foundParseStage = true },
+		VisitJSONExpressionParserFn:   func(v syntax.RootVisitor, e *syntax.JSONExpressionParser) { foundParseStage = true },
+		VisitLogfmtExpressionParserFn: func(v syntax.RootVisitor, e *syntax.LogfmtExpressionParser) { foundParseStage = true },
+		VisitLabelFmtFn:               func(v syntax.RootVisitor, e *syntax.LabelFmtExpr) { foundParseStage = true },
+		VisitKeepLabelFn:              func(v syntax.RootVisitor, e *syntax.KeepLabelsExpr) { foundParseStage = true },
+		VisitDropLabelsFn:             func(v syntax.RootVisitor, e *syntax.DropLabelsExpr) { foundParseStage = true },
+	}
+	expr.Accept(visitor)
+
+	return buildLabelMatchers(exprs)
+}
+
+func buildLabelMatchers(exprs []*syntax.LabelFilterExpr) []LabelMatcher {
+	matchers := make([]LabelMatcher, 0, len(exprs))
+	for _, expr := range exprs {
+		matchers = append(matchers, buildLabelMatcher(expr.LabelFilterer))
+	}
+	return matchers
+}
+
+func buildLabelMatcher(filter log.LabelFilterer) LabelMatcher {
+	switch filter := filter.(type) {
+
+	case *log.LineFilterLabelFilter:
+		if filter.Type != labels.MatchEqual {
+			return UnsupportedLabelMatcher{}
+		}
+
+		return PlainLabelMatcher{
+			Key:   filter.Name,
+			Value: filter.Value,
+		}
+
+	case *log.StringLabelFilter:
+		if filter.Type != labels.MatchEqual {
+			return UnsupportedLabelMatcher{}
+		}
+
+		return PlainLabelMatcher{
+			Key:   filter.Name,
+			Value: filter.Value,
+		}
+
+	case *log.BinaryLabelFilter:
+		var (
+			left  = buildLabelMatcher(filter.Left)
+			right = buildLabelMatcher(filter.Right)
+		)
+
+		if filter.And {
+			return AndLabelMatcher{Left: left, Right: right}
+		}
+		return OrLabelMatcher{Left: left, Right: right}
+
+	default:
+		return UnsupportedLabelMatcher{}
+	}
+}
+
+//
+// Implement marker types:
+//
+
+func (UnsupportedLabelMatcher) isLabelMatcher() {}
+func (PlainLabelMatcher) isLabelMatcher()       {}
+func (OrLabelMatcher) isLabelMatcher()          {}
+func (AndLabelMatcher) isLabelMatcher()         {}

pkg/storage/bloom/v1/ast_extractor_test.go

@@ -0,0 +1,105 @@
+package v1_test
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/loki/v3/pkg/logql/syntax"
+	v1 "github.com/grafana/loki/v3/pkg/storage/bloom/v1"
+)
+
+func TestExtractLabelMatchers(t *testing.T) {
+	tt := []struct {
+		name   string
+		input  string
+		expect []v1.LabelMatcher
+	}{
+		{
+			name:  "basic label matcher",
+			input: `{app="foo"} | key="value"`,
+			expect: []v1.LabelMatcher{
+				v1.PlainLabelMatcher{Key: "key", Value: "value"},
+			},
+		},
+
+		{
+			name:  "or label matcher",
+			input: `{app="foo"} | key1="value1" or key2="value2"`,
+			expect: []v1.LabelMatcher{
+				v1.OrLabelMatcher{
+					Left:  v1.PlainLabelMatcher{Key: "key1", Value: "value1"},
+					Right: v1.PlainLabelMatcher{Key: "key2", Value: "value2"},
+				},
+			},
+		},
+
+		{
+			name:  "and label matcher",
+			input: `{app="foo"} | key1="value1" and key2="value2"`,
+			expect: []v1.LabelMatcher{
+				v1.AndLabelMatcher{
+					Left:  v1.PlainLabelMatcher{Key: "key1", Value: "value1"},
+					Right: v1.PlainLabelMatcher{Key: "key2", Value: "value2"},
+				},
+			},
+		},
+
+		{
+			name:  "multiple label matchers",
+			input: `{app="foo"} | key1="value1" | key2="value2"`,
+			expect: []v1.LabelMatcher{
+				v1.PlainLabelMatcher{Key: "key1", Value: "value1"},
+				v1.PlainLabelMatcher{Key: "key2", Value: "value2"},
+			},
+		},
+
+		{
+			name:  "unsupported label matchers",
+			input: `{app="foo"} | key1=~"value1"`,
+			expect: []v1.LabelMatcher{
+				v1.UnsupportedLabelMatcher{},
+			},
+		},
+	}
+
+	for _, tc := range tt {
+		t.Run(tc.name, func(t *testing.T) {
+			expr, err := syntax.ParseExpr(tc.input)
+			require.NoError(t, err)
+			require.Equal(t, tc.expect, v1.ExtractTestableLabelMatchers(expr))
+		})
+	}
+}
+
+func TestExtractLabelMatchers_IgnoreAfterParse(t *testing.T) {
+	tt := []struct {
+		name string
+		expr string
+	}{
+		{"after json parser", `json`},
+		{"after logfmt parser", `logfmt`},
+		{"after pattern parser", `pattern "<msg>"`},
+		{"after regexp parser", `regexp "(?P<message>.*)"`},
+		{"after unpack parser", `unpack`},
+		{"after label_format", `label_format foo="bar"`},
+		{"after drop labels stage", `drop foo`},
+		{"after keep labels stage", `keep foo`},
+	}
+
+	for _, tc := range tt {
+		t.Run(tc.name, func(t *testing.T) {
+			fullInput := fmt.Sprintf(`{app="foo"} | key1="value1" | %s | key2="value2"`, tc.expr)
+			expect := []v1.LabelMatcher{
+				v1.PlainLabelMatcher{Key: "key1", Value: "value1"},
+				// key2="value2" should be ignored following tc.expr
+			}
+
+			expr, err := syntax.ParseExpr(fullInput)
+			require.NoError(t, err)
+
+			require.Equal(t, expect, v1.ExtractTestableLabelMatchers(expr), "key2=value2 should be ignored with query %s", fullInput)
+		})
+	}
+}