streamingpromql: use label/metric name validation scheme from limits (rather than NameValidatingEngine)

Author: juliusmh

pkg/frontend/querymiddleware/querysharding_test_utils_test.go

@@ -333,7 +333,7 @@ func newEngineForTesting(t *testing.T, engine string, opts ...engineOpt) (promql
 	case querier.PrometheusEngine:
 		return promOpts, promql.NewEngine(promOpts)
 	case querier.MimirEngine:
-		limits := streamingpromql.NewStaticQueryLimitsProvider(0)
+		limits := streamingpromql.NewStaticQueryLimitsProvider(0, model.UTF8Validation)
 		metrics := stats.NewQueryMetrics(promOpts.Reg)
 		planner := streamingpromql.NewQueryPlanner(mqeOpts)
 		logger := log.NewNopLogger()

pkg/mimir/modules.go

@@ -28,6 +28,7 @@ import (
 	"github.com/prometheus/alertmanager/matchers/compat"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/common/config"
+	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/promql"
 	"github.com/prometheus/prometheus/rules"
@@ -812,9 +813,10 @@ func (t *Mimir) initQueryFrontendTripperware() (serv services.Service, err error
 	var eng promql.QueryEngine
 	switch t.Cfg.Frontend.QueryEngine {
 	case querier.PrometheusEngine:
-		eng = promql.NewEngine(promOpts)
+		// TODO: Decide whether this is a good idea.
+		eng = streamingpromqlcompat.NameValidatingEngine(promql.NewEngine(promOpts), t.Overrides)
 	case querier.MimirEngine:
-		streamingEngine, err := streamingpromql.NewEngine(mqeOpts, streamingpromql.NewStaticQueryLimitsProvider(0), stats.NewQueryMetrics(mqeOpts.CommonOpts.Reg), t.QueryPlanner, util_log.Logger)
+		streamingEngine, err := streamingpromql.NewEngine(mqeOpts, streamingpromql.NewStaticQueryLimitsProvider(0, model.UTF8Validation), stats.NewQueryMetrics(mqeOpts.CommonOpts.Reg), t.QueryPlanner, util_log.Logger)
 		if err != nil {
 			return nil, fmt.Errorf("unable to create Mimir Query Engine: %w", err)
 		}
@@ -828,9 +830,6 @@ func (t *Mimir) initQueryFrontendTripperware() (serv services.Service, err error
 		panic(fmt.Sprintf("invalid config not caught by validation: unknown PromQL engine '%s'", t.Cfg.Querier.QueryEngine))
 	}
 
-	// TODO: Decide whether this is a good idea.
-	eng = streamingpromqlcompat.NameValidatingEngine(eng, t.Overrides)
-
 	tripperware, err := querymiddleware.NewTripperware(
 		t.Cfg.Frontend.QueryMiddleware,
 		util_log.Logger,

pkg/querier/querier.go

@@ -185,7 +185,8 @@ func New(cfg Config, limits *validation.Overrides, distributor Distributor, quer
 
 	switch cfg.QueryEngine {
 	case PrometheusEngine:
-		eng = promql.NewEngine(opts)
+		// TODO: Check whether this approach is a good idea.
+		eng = compat.NameValidatingEngine(promql.NewEngine(opts), limits)
 	case MimirEngine:
 		limitsProvider := NewTenantQueryLimitsProvider(limits)
 		streamingEngine, err := streamingpromql.NewEngine(mqeOpts, limitsProvider, queryMetrics, planner, logger)
@@ -194,7 +195,8 @@ func New(cfg Config, limits *validation.Overrides, distributor Distributor, quer
 		}
 
 		if cfg.EnableQueryEngineFallback {
-			prometheusEngine := promql.NewEngine(opts)
+			// TODO: Check whether this approach is a good idea.
+			prometheusEngine := compat.NameValidatingEngine(promql.NewEngine(opts), limits)
 			eng = compat.NewEngineWithFallback(streamingEngine, prometheusEngine, reg, logger)
 		} else {
 			eng = streamingEngine
@@ -203,8 +205,6 @@ func New(cfg Config, limits *validation.Overrides, distributor Distributor, quer
 		panic(fmt.Sprintf("invalid config not caught by validation: unknown PromQL engine '%s'", cfg.QueryEngine))
 	}
 
-	// TODO: Check whether this approach is a good idea.
-	eng = compat.NameValidatingEngine(eng, limits)
 	return NewSampleAndChunkQueryable(lazyQueryable), exemplarQueryable, eng, nil
 }
 
@@ -795,3 +795,19 @@ func (p *TenantQueryLimitsProvider) GetMaxEstimatedMemoryConsumptionPerQuery(ctx
 
 	return totalLimit, nil
 }
+
+// GetValidationScheme computes the validation scheme for tenants injected into ctx. Returns LegacyValidation if
+// at least one tenant uses LegacyValidation, UTF8Validation otherwise.
+func (p *TenantQueryLimitsProvider) GetValidationScheme(ctx context.Context) (model.ValidationScheme, error) {
+	tenantIDs, err := tenant.TenantIDs(ctx)
+	if err != nil {
+		return 0, err
+	}
+	for _, tenantID := range tenantIDs {
+		validationScheme := p.limits.ValidationScheme(tenantID)
+		if validationScheme == model.LegacyValidation {
+			return validationScheme, nil
+		}
+	}
+	return model.UTF8Validation, nil
+}

pkg/streamingpromql/benchmarks/comparison_test.go

@@ -20,6 +20,7 @@ import (
 	"github.com/grafana/dskit/services"
 	"github.com/grafana/dskit/test"
 	"github.com/grafana/dskit/user"
+	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/promql"
 	"github.com/prometheus/prometheus/storage"
@@ -43,7 +44,7 @@ func BenchmarkQuery(b *testing.B) {
 
 	opts := streamingpromql.NewTestEngineOpts()
 	prometheusEngine := promql.NewEngine(opts.CommonOpts)
-	mimirEngine, err := streamingpromql.NewEngine(opts, streamingpromql.NewStaticQueryLimitsProvider(0), stats.NewQueryMetrics(nil), streamingpromql.NewQueryPlanner(opts), log.NewNopLogger())
+	mimirEngine, err := streamingpromql.NewEngine(opts, streamingpromql.NewStaticQueryLimitsProvider(0, model.UTF8Validation), stats.NewQueryMetrics(nil), streamingpromql.NewQueryPlanner(opts), log.NewNopLogger())
 	require.NoError(b, err)
 
 	// Important: the names below must remain in sync with the names used in tools/benchmark-query-engine.
@@ -95,7 +96,7 @@ func TestBothEnginesReturnSameResultsForBenchmarkQueries(t *testing.T) {
 
 	opts := streamingpromql.NewTestEngineOpts()
 	prometheusEngine := promql.NewEngine(opts.CommonOpts)
-	limitsProvider := streamingpromql.NewStaticQueryLimitsProvider(0)
+	limitsProvider := streamingpromql.NewStaticQueryLimitsProvider(0, model.UTF8Validation)
 	queryMetrics := stats.NewQueryMetrics(nil)
 	mimirEngine, err := streamingpromql.NewEngine(opts, limitsProvider, queryMetrics, streamingpromql.NewQueryPlanner(opts), log.NewNopLogger())
 	require.NoError(t, err)
@@ -124,7 +125,7 @@ func TestBenchmarkSetup(t *testing.T) {
 	q := createBenchmarkQueryable(t, []int{1})
 
 	opts := streamingpromql.NewTestEngineOpts()
-	mimirEngine, err := streamingpromql.NewEngine(opts, streamingpromql.NewStaticQueryLimitsProvider(0), stats.NewQueryMetrics(nil), streamingpromql.NewQueryPlanner(opts), log.NewNopLogger())
+	mimirEngine, err := streamingpromql.NewEngine(opts, streamingpromql.NewStaticQueryLimitsProvider(0, model.UTF8Validation), stats.NewQueryMetrics(nil), streamingpromql.NewQueryPlanner(opts), log.NewNopLogger())
 	require.NoError(t, err)
 
 	ctx := user.InjectOrgID(context.Background(), UserID)

pkg/streamingpromql/engine.go

@@ -15,6 +15,7 @@ import (
 	"github.com/go-kit/log"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
+	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/promql"
 	"github.com/prometheus/prometheus/storage"
 	"go.opentelemetry.io/otel"
@@ -143,25 +144,33 @@ func (e *Engine) newQueryFromPlanner(ctx context.Context, q storage.Queryable, o
 type QueryLimitsProvider interface {
 	// GetMaxEstimatedMemoryConsumptionPerQuery returns the maximum estimated memory allowed to be consumed by a query in bytes, or 0 to disable the limit.
 	GetMaxEstimatedMemoryConsumptionPerQuery(ctx context.Context) (uint64, error)
+	// GetValidationScheme returns the label/metric name validation scheme to use for a query.
+	GetValidationScheme(ctx context.Context) (model.ValidationScheme, error)
 }
 
 // NewStaticQueryLimitsProvider returns a QueryLimitsProvider that always returns the provided limits.
 //
 // This should generally only be used in tests.
-func NewStaticQueryLimitsProvider(maxEstimatedMemoryConsumptionPerQuery uint64) QueryLimitsProvider {
+func NewStaticQueryLimitsProvider(maxEstimatedMemoryConsumptionPerQuery uint64, validationScheme model.ValidationScheme) QueryLimitsProvider {
 	return staticQueryLimitsProvider{
 		maxEstimatedMemoryConsumptionPerQuery: maxEstimatedMemoryConsumptionPerQuery,
+		validationScheme:                      validationScheme,
 	}
 }
 
 type staticQueryLimitsProvider struct {
 	maxEstimatedMemoryConsumptionPerQuery uint64
+	validationScheme                      model.ValidationScheme
 }
 
 func (p staticQueryLimitsProvider) GetMaxEstimatedMemoryConsumptionPerQuery(_ context.Context) (uint64, error) {
 	return p.maxEstimatedMemoryConsumptionPerQuery, nil
 }
 
+func (p staticQueryLimitsProvider) GetValidationScheme(_ context.Context) (model.ValidationScheme, error) {
+	return p.validationScheme, nil
+}
+
 type NoopQueryTracker struct{}
 
 func (n *NoopQueryTracker) GetMaxConcurrent() int {

pkg/streamingpromql/engine_concurrency_test.go

@@ -10,6 +10,7 @@ import (
 	"time"
 
 	"github.com/go-kit/log"
+	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/model/timestamp"
 	"github.com/prometheus/prometheus/promql"
 	"github.com/prometheus/prometheus/promql/promqltest"
@@ -187,7 +188,7 @@ func TestConcurrentQueries(t *testing.T) {
 	t.Cleanup(func() { require.NoError(t, storage.Close()) })
 
 	opts := NewTestEngineOpts()
-	engine, err := NewEngine(opts, NewStaticQueryLimitsProvider(0), stats.NewQueryMetrics(nil), NewQueryPlanner(opts), log.NewNopLogger())
+	engine, err := NewEngine(opts, NewStaticQueryLimitsProvider(0, model.UTF8Validation), stats.NewQueryMetrics(nil), NewQueryPlanner(opts), log.NewNopLogger())
 	require.NoError(t, err)
 
 	// Populate the expected result for each query.