query-frontend: Avoid repeated query parsing

pkg/frontend/querymiddleware/cardinality.go

@@ -167,14 +167,15 @@ func isCardinalitySimilar(actualCardinality, estimatedCardinality uint64) bool {
 // estimates at the bucket boundary, an offset is added based on the hash of the
 // query string.
 func generateCardinalityEstimationCacheKey(userID string, r MetricsQueryRequest, bucketSize time.Duration) string {
+	query := r.GetQuery()
 	hasher := fnv.New64a()
-	_, _ = hasher.Write([]byte(r.GetQuery()))
+	_, _ = hasher.Write([]byte(query))
 
 	// This assumes that `bucketSize` is positive.
 	offset := hasher.Sum64() % uint64(bucketSize.Milliseconds())
 	startBucket := (r.GetStart() + int64(offset)) / bucketSize.Milliseconds()
 	rangeBucket := (r.GetEnd() - r.GetStart()) / bucketSize.Milliseconds()
 
 	// Prefix key with `QS` (short for "query statistics").
-	return fmt.Sprintf("QS:%s:%s:%d:%d", userID, cacheHashKey(r.GetQuery()), startBucket, rangeBucket)
+	return fmt.Sprintf("QS:%s:%s:%d:%d", userID, cacheHashKey(query), startBucket, rangeBucket)
 }

pkg/frontend/querymiddleware/codec.go

@@ -125,6 +125,8 @@ type MetricsQueryRequest interface {
 	GetStep() int64
 	// GetQuery returns the query of the request.
 	GetQuery() string
+	// GetQueryExpr returns the parsed query expression of the request.
+	GetQueryExpr() parser.Expr
 	// GetMinT returns the minimum timestamp in milliseconds of data to be queried,
 	// as determined from the start timestamp and any range vector or offset in the query.
 	GetMinT() int64

pkg/frontend/querymiddleware/experimental_functions.go

@@ -58,10 +58,7 @@ func (m *experimentalFunctionsMiddleware) Do(ctx context.Context, req MetricsQue
 		return m.next.Do(ctx, req)
 	}
 
-	expr, err := parser.ParseExpr(req.GetQuery())
-	if err != nil {
-		return nil, apierror.New(apierror.TypeBadData, DecorateWithParamName(err, "query").Error())
-	}
+	expr := req.GetQueryExpr()
 	funcs := containedExperimentalFunctions(expr)
 	if len(funcs) == 0 {
 		// This query does not contain any experimental functions, so we can continue to the next middleware.

pkg/frontend/querymiddleware/model_extra.go

@@ -133,6 +133,10 @@ func (r *PrometheusRangeQueryRequest) GetQuery() string {
 	return ""
 }
 
+func (r *PrometheusRangeQueryRequest) GetQueryExpr() parser.Expr {
+	return r.queryExpr
+}
+
 // GetMinT returns the minimum timestamp in milliseconds of data to be queried,
 // as determined from the start timestamp and any range vector or offset in the query.
 func (r *PrometheusRangeQueryRequest) GetMinT() int64 {
@@ -317,6 +321,10 @@ func (r *PrometheusInstantQueryRequest) GetQuery() string {
 	return ""
 }
 
+func (r *PrometheusInstantQueryRequest) GetQueryExpr() parser.Expr {
+	return r.queryExpr
+}
+
 func (r *PrometheusInstantQueryRequest) GetStart() int64 {
 	return r.GetTime()
 }

pkg/frontend/querymiddleware/prom2_range_compat.go

@@ -11,7 +11,6 @@ import (
 	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/prometheus/prometheus/promql/parser"
 
-	apierror "github.com/grafana/mimir/pkg/api/error"
 	"github.com/grafana/mimir/pkg/frontend/querymiddleware/astmapper"
 	"github.com/grafana/mimir/pkg/util/spanlogger"
 	"github.com/grafana/mimir/pkg/util/validation"
@@ -53,8 +52,10 @@ func (c *prom2RangeCompatHandler) Do(ctx context.Context, r MetricsQueryRequest)
 		return c.next.Do(ctx, r)
 	}
 
-	origQuery := r.GetQuery()
-	rewritten, err := c.rewrite(ctx, origQuery)
+	origExpr := r.GetQueryExpr()
+	origQuery := origExpr.String()
+
+	rewritten, err := c.rewrite(ctx, origExpr)
 	if err != nil {
 		return c.next.Do(ctx, r)
 	}
@@ -75,12 +76,7 @@ func (c *prom2RangeCompatHandler) Do(ctx context.Context, r MetricsQueryRequest)
 	return c.next.Do(ctx, r)
 }
 
-func (c *prom2RangeCompatHandler) rewrite(ctx context.Context, query string) (parser.Expr, error) {
-	expr, err := parser.ParseExpr(query)
-	if err != nil {
-		return nil, apierror.New(apierror.TypeBadData, DecorateWithParamName(err, "query").Error())
-	}
-
+func (c *prom2RangeCompatHandler) rewrite(ctx context.Context, expr parser.Expr) (parser.Expr, error) {
 	mapper := astmapper.NewProm2RangeCompat(ctx)
 	rewritten, err := mapper.Map(expr)
 	if err != nil {

pkg/frontend/querymiddleware/prune.go

@@ -9,7 +9,6 @@ import (
 	"github.com/go-kit/log/level"
 	"github.com/prometheus/prometheus/promql/parser"
 
-	apierror "github.com/grafana/mimir/pkg/api/error"
 	"github.com/grafana/mimir/pkg/frontend/querymiddleware/astmapper"
 	"github.com/grafana/mimir/pkg/util/spanlogger"
 )
@@ -31,44 +30,39 @@ func newPruneMiddleware(logger log.Logger) MetricsQueryMiddleware {
 }
 
 func (p *pruneMiddleware) Do(ctx context.Context, r MetricsQueryRequest) (Response, error) {
-	log := spanlogger.FromContext(ctx, p.logger)
+	spanLogger := spanlogger.FromContext(ctx, p.logger)
 
-	prunedQuery, success, err := p.pruneQuery(ctx, r.GetQuery())
+	originalExpr := r.GetQueryExpr()
+	originalQuery := originalExpr.String()
+
+	prunedExpr, success, err := p.pruneQuery(ctx, originalExpr, originalQuery)
 	if err != nil {
-		level.Warn(log).Log("msg", "failed to prune the input query, falling back to the original query", "query", r.GetQuery(), "err", err)
+		level.Warn(spanLogger).Log("msg", "failed to prune the input query, falling back to the original query", "query", originalQuery, "err", err)
 
 		return p.next.Do(ctx, r)
 	}
 
 	if !success {
-		level.Debug(log).Log("msg", "query pruning had no effect", "query", r.GetQuery())
+		spanLogger.DebugLog("msg", "query pruning had no effect", "query", originalQuery)
 		return p.next.Do(ctx, r)
 	}
 
-	level.Debug(log).Log("msg", "query has been rewritten by pruning", "original", r.GetQuery(), "rewritten", prunedQuery)
-
-	updatedReq, err := r.WithQuery(prunedQuery)
+	spanLogger.DebugLog("msg", "query has been rewritten by pruning", "original", originalQuery, "rewritten", prunedExpr)
+	updatedReq, err := r.WithExpr(prunedExpr)
 	if err != nil {
 		return nil, err
 	}
 
 	return p.next.Do(ctx, updatedReq)
 }
 
-func (p *pruneMiddleware) pruneQuery(ctx context.Context, query string) (string, bool, error) {
-	// Parse the query.
-	expr, err := parser.ParseExpr(query)
-	if err != nil {
-		return "", false, apierror.New(apierror.TypeBadData, DecorateWithParamName(err, "query").Error())
-	}
-	origQueryString := expr.String()
-
+func (p *pruneMiddleware) pruneQuery(ctx context.Context, queryExpr parser.Expr, queryString string) (parser.Expr, bool, error) {
 	mapper := astmapper.NewQueryPruner(ctx, p.logger)
-	prunedQuery, err := mapper.Map(expr)
+	prunedQuery, err := mapper.Map(queryExpr)
 	if err != nil {
-		return "", false, err
+		return nil, false, err
 	}
 	prunedQueryString := prunedQuery.String()
 
-	return prunedQueryString, origQueryString != prunedQueryString, nil
+	return prunedQuery, queryString != prunedQueryString, nil
 }

pkg/frontend/querymiddleware/remote_read.go

@@ -282,6 +282,11 @@ func (r *remoteReadQueryRequest) GetQuery() string {
 	return r.promQuery
 }
 
+func (r *remoteReadQueryRequest) GetQueryExpr() parser.Expr {
+	// TODO: Explain this
+	return nil
+}
+
 func (r *remoteReadQueryRequest) GetHeaders() []*PrometheusHeader {
 	return nil
 }