design v2

Author: yulong-db

cmd/thanos/query.go

@@ -245,7 +245,9 @@ 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.aggregation-label-value-override", "The value override for __rollup__ label for aggregated metrics. If set to x, all queries on aggregated metrics will have a __rollup__=x matcher. Leave empty to disable this behavior. Default is empty.").Default("").String()
+	rewriteAggregationLabelName := cmd.Flag("query.aggregation-label-name", "The name for aggregation label. See query.aggregation-label-value-override for details. Default is empty.").Default("").String()
+	rewriteAggregationLabelTo := cmd.Flag("query.aggregation-label-value-override", "The value override for aggregation label. If set to x, all queries on aggregated metrics will have a `rewriteAggregationLabelName=x` matcher. Leave either flag empty to disable this behavior. Default is empty.").Default("").String()
+	rewriteAggregationInsertOnly := cmd.Flag("query.aggregation-label-insert-only", "If enabled, the above rewriter only insert missing rewriteAggregationLabelName=x and does not modify existing aggregation labels if any. Default is false.").Default("false").Bool()
 
 	var storeRateLimits store.SeriesSelectLimits
 	storeRateLimits.RegisterFlags(cmd)
@@ -386,7 +388,9 @@ func registerQuery(app *extkingpin.App) {
 			*enforceTenancy,
 			*tenantLabel,
 			*enableGroupReplicaPartialStrategy,
+			*rewriteAggregationLabelName,
 			*rewriteAggregationLabelTo,
+			*rewriteAggregationInsertOnly,
 		)
 	})
 }
@@ -471,7 +475,9 @@ func runQuery(
 	enforceTenancy bool,
 	tenantLabel string,
 	groupReplicaPartialResponseStrategy bool,
+	rewriteAggregationLabelName string,
 	rewriteAggregationLabelTo string,
+	rewriteAggregationInsertOnly bool,
 ) error {
 	comp := component.Query
 	if alertQueryURL == "" {
@@ -601,6 +607,8 @@ func runQuery(
 	opts := query.Options{
 		GroupReplicaPartialResponseStrategy: groupReplicaPartialResponseStrategy,
 		DeduplicationFunc:                   queryDeduplicationFunc,
+		RewriteAggregationInsertOnly:        rewriteAggregationInsertOnly,
+		RewriteAggregationLabelName:         rewriteAggregationLabelName,
 		RewriteAggregationLabelTo:           rewriteAggregationLabelTo,
 	}
 	level.Info(logger).Log("msg", "databricks querier features", "opts", fmt.Sprintf("%+v", opts))

pkg/query/aggregation_label_rewriter.go

@@ -14,15 +14,13 @@ import (
 	"github.com/prometheus/prometheus/model/labels"
 )
 
-const (
-	aggregationLabelName = "__agg_rule_type__"
-)
-
 type AggregationLabelRewriter struct {
 	logger  log.Logger
 	metrics *aggregationLabelRewriterMetrics
 
 	enabled           bool
+	insertOnly        bool
+	labelKey          string
 	desiredLabelValue string
 }
 
@@ -74,14 +72,16 @@ func NewNopAggregationLabelRewriter() *AggregationLabelRewriter {
 	}
 }
 
-func NewAggregationLabelRewriter(logger log.Logger, reg prometheus.Registerer, desiredLabelValue string) *AggregationLabelRewriter {
+func NewAggregationLabelRewriter(logger log.Logger, reg prometheus.Registerer, labelKey string, desiredLabelValue string, insertOnly bool) *AggregationLabelRewriter {
 	if logger == nil {
 		logger = log.NewNopLogger()
 	}
 	return &AggregationLabelRewriter{
 		enabled:           desiredLabelValue != "",
 		logger:            logger,
 		metrics:           newAggregationLabelRewriterMetrics(reg, desiredLabelValue),
+		insertOnly:        insertOnly,
+		labelKey:          labelKey,
 		desiredLabelValue: desiredLabelValue,
 	}
 }
@@ -118,16 +118,22 @@ func (a *AggregationLabelRewriter) Rewrite(ms []*labels.Matcher) []*labels.Match
 				break
 			}
 			// In any case, if we see an aggregation label, we store that for later use
-		} else if m.Name == aggregationLabelName {
+		} else if m.Name == a.labelKey {
 			aggregationLabelMatcher = m
 			aggregationLabelIndex = i
 		}
 	}
+
+	if aggregationLabelMatcher != nil && a.insertOnly {
+		needsRewrite = false
+		skipReason = "insert-only"
+	}
+
 	// After the for loop, if needsRewrite is false, no need to do anything
 	// but if it is true, we either append or modify an aggregation label
 	if needsRewrite {
 		newMatcher := &labels.Matcher{
-			Name:  aggregationLabelName,
+			Name:  a.labelKey,
 			Type:  labels.MatchRegexp,
 			Value: a.desiredLabelValue,
 		}

pkg/query/aggregation_label_rewriter_test.go

@@ -11,12 +11,17 @@ import (
 	"github.com/prometheus/prometheus/model/labels"
 )
 
+const (
+	aggregationLabelName = "__agg_rule_type__"
+)
+
 func TestAggregationLabelRewriter_Rewrite(t *testing.T) {
 	t.Parallel()
 
 	for _, tc := range []struct {
 		name               string
 		desiredLabelValue  string // Empty means disabled
+		insertOnly         bool
 		inputMatchers      []*labels.Matcher
 		expectedMatchers   []*labels.Matcher
 		expectedSkipCount  float64
@@ -104,13 +109,29 @@ func TestAggregationLabelRewriter_Rewrite(t *testing.T) {
 			},
 			expectedSkipCount: 1,
 		},
+		{
+			name:              "if insert only, should NOT rewrite existing aggregation label for aggregated metric",
+			desiredLabelValue: "5m",
+			insertOnly:        true,
+			inputMatchers: []*labels.Matcher{
+				labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
+				labels.MustNewMatcher(labels.MatchEqual, aggregationLabelName, "1h"),
+			},
+			expectedMatchers: []*labels.Matcher{
+				labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
+				labels.MustNewMatcher(labels.MatchEqual, aggregationLabelName, "1h"),
+			},
+			expectedSkipCount: 1,
+		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
 			reg := prometheus.NewRegistry()
 			rewriter := NewAggregationLabelRewriter(
 				nil,
 				reg,
+				aggregationLabelName,
 				tc.desiredLabelValue,
+				tc.insertOnly,
 			)
 
 			result := rewriter.Rewrite(tc.inputMatchers)

pkg/query/querier.go

@@ -61,7 +61,9 @@ type QueryableCreator func(
 type Options struct {
 	GroupReplicaPartialResponseStrategy bool
 	DeduplicationFunc                   string
+	RewriteAggregationLabelName         string
 	RewriteAggregationLabelTo           string
+	RewriteAggregationInsertOnly        bool
 }
 
 // NewQueryableCreator creates QueryableCreator.
@@ -94,7 +96,9 @@ func NewQueryableCreatorWithOptions(
 	aggregationLabelRewriter := NewAggregationLabelRewriter(
 		logger,
 		extprom.WrapRegistererWithPrefix("aggregation_label_rewriter_", reg),
+		opts.RewriteAggregationLabelName,
 		opts.RewriteAggregationLabelTo,
+		opts.RewriteAggregationInsertOnly,
 	)
 	return func(
 		deduplicate bool,