impl aggr label rewrite capability

Author: yulong-db

cmd/thanos/query.go

@@ -245,6 +245,8 @@ func registerQuery(app *extkingpin.App) {
 	enforceTenancy := cmd.Flag("query.enforce-tenancy", "Enforce tenancy on Query APIs. Responses are returned only if the label value of the configured tenant-label-name and the value of the tenant header matches.").Default("false").Bool()
 	tenantLabel := cmd.Flag("query.tenant-label-name", "Label name to use when enforcing tenancy (if --query.enforce-tenancy is enabled).").Default(tenancy.DefaultTenantLabel).String()
 
+	rewriteAggregationLabelTo := cmd.Flag("query.rewrite-aggregation-label-to", "Rewrite the aggregation label to the provided value for metric with name ending in standard aggregation suffixes.").Default("").String()
+
 	var storeRateLimits store.SeriesSelectLimits
 	storeRateLimits.RegisterFlags(cmd)
 
@@ -384,6 +386,7 @@ func registerQuery(app *extkingpin.App) {
 			*enforceTenancy,
 			*tenantLabel,
 			*enableGroupReplicaPartialStrategy,
+			*rewriteAggregationLabelTo,
 		)
 	})
 }
@@ -468,6 +471,7 @@ func runQuery(
 	enforceTenancy bool,
 	tenantLabel string,
 	groupReplicaPartialResponseStrategy bool,
+	rewriteAggregationLabelTo string,
 ) error {
 	comp := component.Query
 	if alertQueryURL == "" {
@@ -597,6 +601,7 @@ func runQuery(
 	opts := query.Options{
 		GroupReplicaPartialResponseStrategy: groupReplicaPartialResponseStrategy,
 		DeduplicationFunc:                   queryDeduplicationFunc,
+		RewriteAggregationLabelTo:           rewriteAggregationLabelTo,
 	}
 	level.Info(logger).Log("msg", "databricks querier features", "opts", fmt.Sprintf("%+v", opts))
 	queryableCreator = query.NewQueryableCreatorWithOptions(

pkg/query/querier.go

@@ -61,6 +61,7 @@ type QueryableCreator func(
 type Options struct {
 	GroupReplicaPartialResponseStrategy bool
 	DeduplicationFunc                   string
+	RewriteAggregationLabelTo           string
 }
 
 // NewQueryableCreator creates QueryableCreator.
@@ -159,6 +160,11 @@ type querier struct {
 	selectTimeout           time.Duration
 	shardInfo               *storepb.ShardInfo
 	seriesStatsReporter     seriesStatsReporter
+
+	enablePreAggregationLabelRewrite bool
+	aggregationStandardSuffixes      []string
+	aggregationLabelName             string
+	rewriteAggregationLabelTo        string
 }
 
 // newQuerier creates implementation of storage.Querier that fetches data from the proxy
@@ -232,6 +238,11 @@ func newQuerierWithOpts(
 		skipChunks:              skipChunks,
 		shardInfo:               shardInfo,
 		seriesStatsReporter:     seriesStatsReporter,
+
+		enablePreAggregationLabelRewrite: opts.RewriteAggregationLabelTo != "",
+		aggregationStandardSuffixes:      []string{":aggr", ":count", ":sum", ":min", ":max", ":avg"},
+		aggregationLabelName:             "__rollup__",
+		rewriteAggregationLabelTo:        opts.RewriteAggregationLabelTo,
 	}
 }
 
@@ -365,6 +376,47 @@ func (q *querier) Select(ctx context.Context, _ bool, hints *storage.SelectHints
 }
 
 func (q *querier) selectFn(ctx context.Context, hints *storage.SelectHints, ms ...*labels.Matcher) (storage.SeriesSet, storepb.SeriesStatsCounter, error) {
+	if q.enablePreAggregationLabelRewrite {
+		hasNameLabel, needLabelRewrite := false, false
+		for _, m := range ms {
+			if m.Name == labels.MetricName {
+				hasNameLabel = true
+				if m.Type == labels.MatchEqual {
+					for _, suffix := range q.aggregationStandardSuffixes {
+						if strings.HasSuffix(m.Value, suffix) {
+							needLabelRewrite = true
+							break
+						}
+					}
+					break
+				} else {
+					level.Info(q.logger).Log("msg", "Skipping aggregation label rewrite attempt", "reason", "non-equal name matcher", "name_matcher", m, "labels", ms)
+				}
+			}
+		}
+		if !hasNameLabel {
+			level.Info(q.logger).Log("msg", "Skipping aggregation label rewrite", "reason", "no metric name matcher found", "labels", ms)
+		}
+		if needLabelRewrite {
+			hasAggregationLabel := false
+			for i, m := range ms {
+				if m.Name == q.aggregationLabelName {
+					hasAggregationLabel = true
+					ms[i].Type = labels.MatchEqual
+					ms[i].Value = q.rewriteAggregationLabelTo
+					break
+				}
+			}
+			if !hasAggregationLabel {
+				ms = append(ms, &labels.Matcher{
+					Name:  q.aggregationLabelName,
+					Type:  labels.MatchEqual,
+					Value: q.rewriteAggregationLabelTo,
+				})
+			}
+		}
+	}
+
 	sms, err := storepb.PromMatchersToMatchers(ms...)
 	if err != nil {
 		return nil, storepb.SeriesStatsCounter{}, errors.Wrap(err, "convert matchers")