Skip to content

Commit 08b1526

Browse files
yulong-dbyulong-db
committed
address feedback
1 parent 55a240a commit 08b1526

File tree

4 files changed

+57
-36
lines changed

4 files changed

+57
-36
lines changed

‎cmd/thanos/query.go‎

Lines changed: 5 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -245,9 +245,8 @@ func registerQuery(app *extkingpin.App) {
245245
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()
246246
tenantLabel := cmd.Flag("query.tenant-label-name", "Label name to use when enforcing tenancy (if --query.enforce-tenancy is enabled).").Default(tenancy.DefaultTenantLabel).String()
247247

248-
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()
249-
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()
250-
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()
248+
rewriteAggregationLabelStrategy := cmd.Flag("query.aggregation-label-strategy", "The strategy to use when rewriting aggregation labels. Used during aggregator migration only.").Default(string(query.NoopLabelRewriter)).Hidden().Enum(string(query.NoopLabelRewriter), string(query.UpsertLabelRewriter), string(query.InsertOnlyLabelRewriter))
249+
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 `__agg_rule_type__=x` matcher. If empty, this behavior is disabled. Default is empty.").Hidden().Default("").String()
251250

252251
var storeRateLimits store.SeriesSelectLimits
253252
storeRateLimits.RegisterFlags(cmd)
@@ -388,9 +387,8 @@ func registerQuery(app *extkingpin.App) {
388387
*enforceTenancy,
389388
*tenantLabel,
390389
*enableGroupReplicaPartialStrategy,
391-
*rewriteAggregationLabelName,
390+
*rewriteAggregationLabelStrategy,
392391
*rewriteAggregationLabelTo,
393-
*rewriteAggregationInsertOnly,
394392
)
395393
})
396394
}
@@ -475,9 +473,8 @@ func runQuery(
475473
enforceTenancy bool,
476474
tenantLabel string,
477475
groupReplicaPartialResponseStrategy bool,
478-
rewriteAggregationLabelName string,
476+
rewriteAggregationLabelStrategy string,
479477
rewriteAggregationLabelTo string,
480-
rewriteAggregationInsertOnly bool,
481478
) error {
482479
comp := component.Query
483480
if alertQueryURL == "" {
@@ -607,8 +604,7 @@ func runQuery(
607604
opts := query.Options{
608605
GroupReplicaPartialResponseStrategy: groupReplicaPartialResponseStrategy,
609606
DeduplicationFunc: queryDeduplicationFunc,
610-
RewriteAggregationInsertOnly: rewriteAggregationInsertOnly,
611-
RewriteAggregationLabelName: rewriteAggregationLabelName,
607+
RewriteAggregationLabelStrategy: rewriteAggregationLabelStrategy,
612608
RewriteAggregationLabelTo: rewriteAggregationLabelTo,
613609
}
614610
level.Info(logger).Log("msg", "databricks querier features", "opts", fmt.Sprintf("%+v", opts))

‎pkg/query/aggregation_label_rewriter.go‎

Lines changed: 29 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -14,13 +14,26 @@ import (
1414
"github.com/prometheus/prometheus/model/labels"
1515
)
1616

17-
type AggregationLabelRewriter struct {
18-
logger log.Logger
19-
metrics *aggregationLabelRewriterMetrics
17+
// RewriterStrategy defines the strategy used by the AggregationLabelRewriter.
18+
type RewriterStrategy string
19+
20+
const (
21+
// NoopLabelRewriter is a no-op strategy that basically disables the rewriter.
22+
NoopLabelRewriter RewriterStrategy = "noop"
23+
// UpsertLabelRewriter is a strategy that upserts the aggregation label.
24+
UpsertLabelRewriter RewriterStrategy = "upsert"
25+
// InsertOnlyLabelRewriter is a strategy that only inserts the aggregation label if it does not exist.
26+
InsertOnlyLabelRewriter RewriterStrategy = "insert-only"
27+
)
28+
29+
const (
30+
aggregationLabelName = "__agg_rule_type__"
31+
)
2032

21-
enabled bool
22-
insertOnly bool
23-
labelKey string
33+
type AggregationLabelRewriter struct {
34+
logger log.Logger
35+
metrics *aggregationLabelRewriterMetrics
36+
strategy RewriterStrategy
2437
desiredLabelValue string
2538
}
2639

@@ -68,26 +81,27 @@ func newAggregationLabelRewriterMetrics(reg prometheus.Registerer, desiredLabelV
6881

6982
func NewNopAggregationLabelRewriter() *AggregationLabelRewriter {
7083
return &AggregationLabelRewriter{
71-
enabled: false,
84+
strategy: NoopLabelRewriter,
7285
}
7386
}
7487

75-
func NewAggregationLabelRewriter(logger log.Logger, reg prometheus.Registerer, labelKey string, desiredLabelValue string, insertOnly bool) *AggregationLabelRewriter {
88+
func NewAggregationLabelRewriter(logger log.Logger, reg prometheus.Registerer, strategy RewriterStrategy, desiredLabelValue string) *AggregationLabelRewriter {
7689
if logger == nil {
7790
logger = log.NewNopLogger()
7891
}
92+
if desiredLabelValue == "" {
93+
strategy = NoopLabelRewriter
94+
}
7995
return &AggregationLabelRewriter{
80-
enabled: desiredLabelValue != "" && labelKey != "",
96+
strategy: strategy,
8197
logger: logger,
8298
metrics: newAggregationLabelRewriterMetrics(reg, desiredLabelValue),
83-
insertOnly: insertOnly,
84-
labelKey: labelKey,
8599
desiredLabelValue: desiredLabelValue,
86100
}
87101
}
88102

89103
func (a *AggregationLabelRewriter) Rewrite(ms []*labels.Matcher) []*labels.Matcher {
90-
if !a.enabled {
104+
if a.strategy == NoopLabelRewriter {
91105
return ms
92106
}
93107

@@ -118,13 +132,13 @@ func (a *AggregationLabelRewriter) Rewrite(ms []*labels.Matcher) []*labels.Match
118132
break
119133
}
120134
// In any case, if we see an aggregation label, we store that for later use
121-
} else if m.Name == a.labelKey {
135+
} else if m.Name == aggregationLabelName {
122136
aggregationLabelMatcher = m
123137
aggregationLabelIndex = i
124138
}
125139
}
126140

127-
if aggregationLabelMatcher != nil && a.insertOnly {
141+
if aggregationLabelMatcher != nil && a.strategy == InsertOnlyLabelRewriter {
128142
needsRewrite = false
129143
skipReason = "insert-only"
130144
}
@@ -133,7 +147,7 @@ func (a *AggregationLabelRewriter) Rewrite(ms []*labels.Matcher) []*labels.Match
133147
// but if it is true, we either append or modify an aggregation label
134148
if needsRewrite {
135149
newMatcher := &labels.Matcher{
136-
Name: a.labelKey,
150+
Name: aggregationLabelName,
137151
Type: labels.MatchRegexp,
138152
Value: a.desiredLabelValue,
139153
}

‎pkg/query/aggregation_label_rewriter_test.go‎

Lines changed: 21 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -11,17 +11,13 @@ import (
1111
"github.com/prometheus/prometheus/model/labels"
1212
)
1313

14-
const (
15-
aggregationLabelName = "__agg_rule_type__"
16-
)
17-
1814
func TestAggregationLabelRewriter_Rewrite(t *testing.T) {
1915
t.Parallel()
2016

2117
for _, tc := range []struct {
2218
name string
2319
desiredLabelValue string // Empty means disabled
24-
insertOnly bool
20+
strategy RewriterStrategy
2521
inputMatchers []*labels.Matcher
2622
expectedMatchers []*labels.Matcher
2723
expectedSkipCount float64
@@ -30,7 +26,19 @@ func TestAggregationLabelRewriter_Rewrite(t *testing.T) {
3026
}{
3127
{
3228
name: "disabled rewriter should not modify label matchers",
29+
desiredLabelValue: "v1",
30+
strategy: NoopLabelRewriter,
31+
inputMatchers: []*labels.Matcher{
32+
labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
33+
},
34+
expectedMatchers: []*labels.Matcher{
35+
labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
36+
},
37+
},
38+
{
39+
name: "no desired label value makes a disabled rewriter and should not modify label matchers",
3340
desiredLabelValue: "",
41+
strategy: UpsertLabelRewriter,
3442
inputMatchers: []*labels.Matcher{
3543
labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
3644
},
@@ -41,6 +49,7 @@ func TestAggregationLabelRewriter_Rewrite(t *testing.T) {
4149
{
4250
name: "should add label for aggregated metric if no existing aggregation label",
4351
desiredLabelValue: "5m",
52+
strategy: UpsertLabelRewriter,
4453
inputMatchers: []*labels.Matcher{
4554
labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
4655
},
@@ -53,6 +62,7 @@ func TestAggregationLabelRewriter_Rewrite(t *testing.T) {
5362
{
5463
name: "should rewrite existing equal aggregation label for aggregated metric",
5564
desiredLabelValue: "5m",
65+
strategy: UpsertLabelRewriter,
5666
inputMatchers: []*labels.Matcher{
5767
labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
5868
labels.MustNewMatcher(labels.MatchEqual, aggregationLabelName, "1h"),
@@ -66,6 +76,7 @@ func TestAggregationLabelRewriter_Rewrite(t *testing.T) {
6676
{
6777
name: "should rewrite existing regex aggregation label for aggregated metric",
6878
desiredLabelValue: "5m",
79+
strategy: UpsertLabelRewriter,
6980
inputMatchers: []*labels.Matcher{
7081
labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
7182
labels.MustNewMatcher(labels.MatchRegexp, aggregationLabelName, "1h"),
@@ -79,6 +90,7 @@ func TestAggregationLabelRewriter_Rewrite(t *testing.T) {
7990
{
8091
name: "should skip non-aggregated metric",
8192
desiredLabelValue: "5m",
93+
strategy: UpsertLabelRewriter,
8294
inputMatchers: []*labels.Matcher{
8395
labels.MustNewMatcher(labels.MatchEqual, "__name__", "test_metric"),
8496
},
@@ -90,6 +102,7 @@ func TestAggregationLabelRewriter_Rewrite(t *testing.T) {
90102
{
91103
name: "should skip non-equal name matcher",
92104
desiredLabelValue: "5m",
105+
strategy: UpsertLabelRewriter,
93106
inputMatchers: []*labels.Matcher{
94107
labels.MustNewMatcher(labels.MatchRegexp, "__name__", "test:sum"),
95108
},
@@ -101,6 +114,7 @@ func TestAggregationLabelRewriter_Rewrite(t *testing.T) {
101114
{
102115
name: "should skip when no name matcher",
103116
desiredLabelValue: "5m",
117+
strategy: UpsertLabelRewriter,
104118
inputMatchers: []*labels.Matcher{
105119
labels.MustNewMatcher(labels.MatchEqual, "job", "prometheus"),
106120
},
@@ -112,7 +126,7 @@ func TestAggregationLabelRewriter_Rewrite(t *testing.T) {
112126
{
113127
name: "if insert only, should NOT rewrite existing aggregation label for aggregated metric",
114128
desiredLabelValue: "5m",
115-
insertOnly: true,
129+
strategy: InsertOnlyLabelRewriter,
116130
inputMatchers: []*labels.Matcher{
117131
labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
118132
labels.MustNewMatcher(labels.MatchEqual, aggregationLabelName, "1h"),
@@ -129,9 +143,8 @@ func TestAggregationLabelRewriter_Rewrite(t *testing.T) {
129143
rewriter := NewAggregationLabelRewriter(
130144
nil,
131145
reg,
132-
aggregationLabelName,
146+
tc.strategy,
133147
tc.desiredLabelValue,
134-
tc.insertOnly,
135148
)
136149

137150
result := rewriter.Rewrite(tc.inputMatchers)

‎pkg/query/querier.go‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -61,9 +61,8 @@ type QueryableCreator func(
6161
type Options struct {
6262
GroupReplicaPartialResponseStrategy bool
6363
DeduplicationFunc string
64-
RewriteAggregationLabelName string
64+
RewriteAggregationLabelStrategy string
6565
RewriteAggregationLabelTo string
66-
RewriteAggregationInsertOnly bool
6766
}
6867

6968
// NewQueryableCreator creates QueryableCreator.
@@ -96,9 +95,8 @@ func NewQueryableCreatorWithOptions(
9695
aggregationLabelRewriter := NewAggregationLabelRewriter(
9796
logger,
9897
extprom.WrapRegistererWithPrefix("aggregation_label_rewriter_", reg),
99-
opts.RewriteAggregationLabelName,
98+
RewriterStrategy(opts.RewriteAggregationLabelStrategy),
10099
opts.RewriteAggregationLabelTo,
101-
opts.RewriteAggregationInsertOnly,
102100
)
103101
return func(
104102
deduplicate bool,

0 commit comments

Comments
 (0)