-
Notifications
You must be signed in to change notification settings - Fork 548
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
add support for http header to filter down queryables (#10552)
* add support for http header to filter down queryables Signed-off-by: Mauro Stettler <[email protected]> * changelog Signed-off-by: Mauro Stettler <[email protected]> * add spdx license header to new files Signed-off-by: Mauro Stettler <[email protected]> --------- Signed-off-by: Mauro Stettler <[email protected]>
- Loading branch information
Showing
9 changed files
with
250 additions
and
5 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,59 @@ | ||
| // SPDX-License-Identifier: AGPL-3.0-only | ||
|
|
||
| package querier | ||
|
|
||
| import ( | ||
| "context" | ||
| "net/http" | ||
| "strings" | ||
|
|
||
| "github.com/grafana/dskit/middleware" | ||
| ) | ||
|
|
||
| type ( | ||
| filterQueryablesCtxKeyT int | ||
| filterQueryables map[string]struct{} | ||
| ) | ||
|
|
||
| const ( | ||
| FilterQueryablesHeader = "X-Filter-Queryables" | ||
| filterQueryablesCtxKey filterQueryablesCtxKeyT = 0 | ||
| ) | ||
|
|
||
| func newFilterQueryables(asString string) filterQueryables { | ||
| f := make(filterQueryables) | ||
| for _, name := range strings.Split(asString, ",") { | ||
| f[strings.Trim(name, " ")] = struct{}{} | ||
| } | ||
| return f | ||
| } | ||
|
|
||
| func (f filterQueryables) use(name string) bool { | ||
| _, ok := f[name] | ||
| return ok | ||
| } | ||
|
|
||
| func FilterQueryablesMiddleware() middleware.Interface { | ||
| return middleware.Func(func(next http.Handler) http.Handler { | ||
| return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { | ||
| if filterQueryables := req.Header.Get(FilterQueryablesHeader); len(filterQueryables) > 0 { | ||
| req = req.WithContext(addFilterQueryablesToContext(req.Context(), filterQueryables)) | ||
| } | ||
|
|
||
| next.ServeHTTP(w, req) | ||
| }) | ||
| }) | ||
| } | ||
|
|
||
| func addFilterQueryablesToContext(ctx context.Context, value string) context.Context { | ||
| return context.WithValue(ctx, filterQueryablesCtxKey, newFilterQueryables(value)) | ||
| } | ||
|
|
||
| func getFilterQueryablesFromContext(ctx context.Context) (filterQueryables, bool) { | ||
| value, ok := ctx.Value(filterQueryablesCtxKey).(filterQueryables) | ||
| if !ok { | ||
| return nil, false | ||
| } | ||
|
|
||
| return value, true | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,144 @@ | ||
| // SPDX-License-Identifier: AGPL-3.0-only | ||
|
|
||
| package querier | ||
|
|
||
| import ( | ||
| "context" | ||
| "testing" | ||
| "time" | ||
|
|
||
| "github.com/go-kit/log" | ||
| "github.com/grafana/dskit/flagext" | ||
| "github.com/grafana/dskit/user" | ||
| "github.com/prometheus/client_golang/prometheus" | ||
| "github.com/prometheus/prometheus/model/labels" | ||
| "github.com/prometheus/prometheus/storage" | ||
| "github.com/stretchr/testify/assert" | ||
| "github.com/stretchr/testify/mock" | ||
| "github.com/stretchr/testify/require" | ||
|
|
||
| "github.com/grafana/mimir/pkg/querier/stats" | ||
| "github.com/grafana/mimir/pkg/util/validation" | ||
| ) | ||
|
|
||
| func TestFilteringQueryablesViaHttpHeader(t *testing.T) { | ||
| type testCase struct { | ||
| transformContext func(context.Context) context.Context | ||
| expectedQuerier1Calls int | ||
| expectedQuerier2Calls int | ||
| } | ||
|
|
||
| runTestCase := func(t *testing.T, tc testCase) { | ||
| ctx := context.Background() | ||
| logger := log.NewNopLogger() | ||
| reg := prometheus.NewPedanticRegistry() | ||
| metrics := stats.NewQueryMetrics(reg) | ||
| overrides, err := validation.NewOverrides(defaultLimitsConfig(), nil) | ||
| require.NoError(t, err) | ||
|
|
||
| cfg := Config{} | ||
| flagext.DefaultValues(&cfg) | ||
|
|
||
| matcher := labels.MustNewMatcher(labels.MatchEqual, labels.MetricName, "metric") | ||
| expectedMatchers := []*labels.Matcher{matcher} | ||
| querier1 := &mockBlocksStorageQuerier{} | ||
| querier1.On("Select", mock.Anything, true, mock.Anything, expectedMatchers).Return(storage.EmptySeriesSet()) | ||
| querier2 := &mockBlocksStorageQuerier{} | ||
| querier2.On("Select", mock.Anything, true, mock.Anything, expectedMatchers).Return(storage.EmptySeriesSet()) | ||
|
|
||
| alwaysApplicable := func(_ context.Context, _ string, _ time.Time, _, _ int64, _ log.Logger, _ ...*labels.Matcher) bool { | ||
| return true | ||
| } | ||
| querierQueryables := []TimeRangeQueryable{{ | ||
| Queryable: newMockBlocksStorageQueryable(querier1), | ||
| IsApplicable: alwaysApplicable, | ||
| StorageName: "querier1", | ||
| }, { | ||
| Queryable: newMockBlocksStorageQueryable(querier2), | ||
| IsApplicable: alwaysApplicable, | ||
| StorageName: "querier2", | ||
| }} | ||
|
|
||
| queryable := newQueryable(querierQueryables, cfg, overrides, metrics, logger) | ||
| querier, err := queryable.Querier(0, 10) | ||
| require.NoError(t, err) | ||
|
|
||
| ctx = user.InjectOrgID(ctx, "0") | ||
| ctx = tc.transformContext(ctx) | ||
| series := querier.Select(ctx, false, nil, matcher) | ||
| require.NoError(t, series.Err()) | ||
|
|
||
| assert.Equal(t, tc.expectedQuerier1Calls, len(querier1.Calls)) | ||
| assert.Equal(t, tc.expectedQuerier2Calls, len(querier2.Calls)) | ||
| } | ||
|
|
||
| t.Run("do not set header", func(t *testing.T) { | ||
| runTestCase(t, testCase{ | ||
| transformContext: func(ctx context.Context) context.Context { | ||
| return ctx | ||
| }, | ||
| expectedQuerier1Calls: 1, | ||
| expectedQuerier2Calls: 1, | ||
| }) | ||
| }) | ||
|
|
||
| t.Run("set header to unknown queryable", func(t *testing.T) { | ||
| runTestCase(t, testCase{ | ||
| transformContext: func(ctx context.Context) context.Context { | ||
| return addFilterQueryablesToContext(ctx, "querier3") | ||
| }, | ||
| expectedQuerier1Calls: 0, | ||
| expectedQuerier2Calls: 0, | ||
| }) | ||
| }) | ||
|
|
||
| t.Run("set header to only querier1", func(t *testing.T) { | ||
| runTestCase(t, testCase{ | ||
| transformContext: func(ctx context.Context) context.Context { | ||
| return addFilterQueryablesToContext(ctx, "querier1") | ||
| }, | ||
| expectedQuerier1Calls: 1, | ||
| expectedQuerier2Calls: 0, | ||
| }) | ||
| }) | ||
|
|
||
| t.Run("set header to only querier2", func(t *testing.T) { | ||
| runTestCase(t, testCase{ | ||
| transformContext: func(ctx context.Context) context.Context { | ||
| return addFilterQueryablesToContext(ctx, "querier2") | ||
| }, | ||
| expectedQuerier1Calls: 0, | ||
| expectedQuerier2Calls: 1, | ||
| }) | ||
| }) | ||
|
|
||
| t.Run("set header to both querier1 and querier2", func(t *testing.T) { | ||
| runTestCase(t, testCase{ | ||
| transformContext: func(ctx context.Context) context.Context { | ||
| return addFilterQueryablesToContext(ctx, "querier2,querier1") | ||
| }, | ||
| expectedQuerier1Calls: 1, | ||
| expectedQuerier2Calls: 1, | ||
| }) | ||
| }) | ||
|
|
||
| t.Run("set header to querier1 and querier2 and unknown querier", func(t *testing.T) { | ||
| runTestCase(t, testCase{ | ||
| transformContext: func(ctx context.Context) context.Context { | ||
| return addFilterQueryablesToContext(ctx, "querier2,querier1,querier3") | ||
| }, | ||
| expectedQuerier1Calls: 1, | ||
| expectedQuerier2Calls: 1, | ||
| }) | ||
| }) | ||
|
|
||
| t.Run("set header to querier1 and unknown querier, with spaces", func(t *testing.T) { | ||
| runTestCase(t, testCase{ | ||
| transformContext: func(ctx context.Context) context.Context { | ||
| return addFilterQueryablesToContext(ctx, "querier1 , querier3") | ||
| }, | ||
| expectedQuerier1Calls: 1, | ||
| expectedQuerier2Calls: 0, | ||
| }) | ||
| }) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters