query-frontend: add option to preserve specific request headers

CHANGELOG.md

@@ -23,6 +23,8 @@
 * [FEATURE] Query Frontend: Add `cortex_query_fetched_series_total` and `cortex_query_fetched_chunks_bytes_total` per-user counters to expose the number of series and bytes fetched as part of queries. These metrics can be enabled with the `-frontend.query-stats-enabled` flag (or its respective YAML config option `query_stats_enabled`). #4343
 * [FEATURE] AlertManager: Add support for SNS Receiver. #4382
 * [FEATURE] Distributor: Add label `status` to metric `cortex_distributor_ingester_append_failures_total` #4442
+* [ENHANCEMENT] Querier: Add setting in query_range middleware to configure the set of headers from the requests to be forwared to downstream requests.The following setting is added. #4486
+* `-querier.forward-headers-list`
 * [FEATURE] Queries: Added `present_over_time` PromQL function, also some TSDB optimisations. #4505
 * [ENHANCEMENT] Add timeout for waiting on compactor to become ACTIVE in the ring. #4262
 * [ENHANCEMENT] Reduce memory used by streaming queries, particularly in ruler. #4341

docs/configuration/arguments.md

@@ -116,6 +116,10 @@ The ingester query API was improved over time, but defaults to the old behaviour
 
    If set to true, will cause the querier to cache query results.  The cache will be used to answer future, overlapping queries.  The query frontend calculates extra queries required to fill gaps in the cache.
 
+- `-querier.forward-headers-list`
+
+   Request headers to be forwarded to downstream via query_range middleware. Multiple headers may be specified. Defaults to empty.
+
 - `-frontend.max-cache-freshness`
 
    When caching query results, it is desirable to prevent the caching of very recent results that might still be in flux.  Use this parameter to configure the age of results that should be excluded.

docs/configuration/config-file-reference.md

@@ -1150,6 +1150,10 @@ results_cache:
 # query ASTs. This feature is supported only by the chunks storage engine.
 # CLI flag: -querier.parallelise-shardable-queries
 [parallelise_shardable_queries: <boolean> | default = false]
+
+# List of headers from requests to be forwarded via query_range middleware.
+# CLI flag: -querier.forward-headers-list
+[forward_headers_list: <list of string> | default = []]
 ```
 
 ### `ruler_config`

docs/configuration/prometheus-frontend.md

@@ -28,6 +28,10 @@ query_range:
   align_queries_with_step: true
   cache_results: true
 
+  # list of headers from requests to be forwarded via query_range middleware.
+  forward_headers_list:
+    - Authorization
+
   results_cache:
     cache:
 

pkg/querier/queryrange/query_range.go

@@ -52,7 +52,8 @@ var (
 type Codec interface {
 	Merger
 	// DecodeRequest decodes a Request from an http request.
-	DecodeRequest(context.Context, *http.Request) (Request, error)
+	// The last parameter specifies a list of headers on the http request that should be included in the decoded request
+	DecodeRequest(context.Context, *http.Request, []string) (Request, error)
 	// DecodeResponse decodes a Response from an http response.
 	// The original request is also passed as a parameter this is useful for implementation that needs the request
 	// to merge result or build the result correctly.
@@ -187,7 +188,7 @@ func (prometheusCodec) MergeResponse(responses ...Response) (Response, error) {
 	return &response, nil
 }
 
-func (prometheusCodec) DecodeRequest(_ context.Context, r *http.Request) (Request, error) {
+func (prometheusCodec) DecodeRequest(_ context.Context, r *http.Request, headers []string) (Request, error) {
 	var result PrometheusRequest
 	var err error
 	result.Start, err = util.ParseTime(r.FormValue("start"))
@@ -222,6 +223,16 @@ func (prometheusCodec) DecodeRequest(_ context.Context, r *http.Request) (Reques
 	result.Query = r.FormValue("query")
 	result.Path = r.URL.Path
 
+	// Include the specified headers from http request in prometheusRequest.
+	for _, header := range headers {
+		for h, hv := range r.Header {
+			if strings.EqualFold(h, header) {
+				result.Headers = append(result.Headers, &PrometheusRequestHeader{Name: h, Values: hv})
+				break
+			}
+		}
+	}
+
 	for _, value := range r.Header.Values(cacheControlHeader) {
 		if strings.Contains(value, noStoreValue) {
 			result.CachingOptions.Disabled = true
@@ -247,12 +258,20 @@ func (prometheusCodec) EncodeRequest(ctx context.Context, r Request) (*http.Requ
 		Path:     promReq.Path,
 		RawQuery: params.Encode(),
 	}
+	var h = http.Header{}
+
+	for _, hv := range promReq.Headers {
+		for _, v := range hv.Values {
+			h.Add(hv.Name, v)
+		}
+	}
+
 	req := &http.Request{
 		Method:     "GET",
 		RequestURI: u.String(), // This is what the httpgrpc code looks at.
 		URL:        u,
 		Body:       http.NoBody,
-		Header:     http.Header{},
+		Header:     h,
 	}
 
 	return req.WithContext(ctx), nil

pkg/querier/queryrange/query_range_test.go

@@ -18,14 +18,18 @@ import (
 )
 
 func TestRequest(t *testing.T) {
+	// Copy parsedRequest to assign the expected headers to the request.
+	// The test below adds a Test-Header header to the request and expects it back once the encode/decode of request is done via PrometheusCodec
+	r := *parsedRequest
+	r.Headers = reqHeaders
 	for i, tc := range []struct {
 		url         string
 		expected    Request
 		expectedErr error
 	}{
 		{
 			url:      query,
-			expected: parsedRequest,
+			expected: &r,
 		},
 		{
 			url:         "api/v1/query_range?start=foo",
@@ -55,11 +59,12 @@ func TestRequest(t *testing.T) {
 		t.Run(strconv.Itoa(i), func(t *testing.T) {
 			r, err := http.NewRequest("GET", tc.url, nil)
 			require.NoError(t, err)
+			r.Header.Add("Test-Header", "test")
 
 			ctx := user.InjectOrgID(context.Background(), "1")
-			r = r.WithContext(ctx)
+			r = r.Clone(ctx)
 
-			req, err := PrometheusCodec.DecodeRequest(ctx, r)
+			req, err := PrometheusCodec.DecodeRequest(ctx, r, []string{"Test-Header"})
 			if err != nil {
 				require.EqualValues(t, tc.expectedErr, err)
 				return