query-frontend: add option to preserve specific request headers

To make Cortex query-frontend more flexible when used with other Prometheus engines.

Option is named -frontend.forward-headers-list.

Signed-off-by: Manish Kumar Gupta <[email protected]>

Author: srijan55

CHANGELOG.md

@@ -26,6 +26,7 @@
 * [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] Query Frontend: Add setting `-frontend.forward-headers-list` in frontend  to configure the set of headers from the requests to be forwarded to downstream requests. #4486
 * [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.
 
+- `-frontend.forward-headers-list`
+
+   Request headers  forwarded by query frontend to downstream queriers. 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

@@ -1166,6 +1166,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 forwarded by the query Frontend to downstream querier.
+# CLI flag: -frontend.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 request headers  forwarded by query frontend to downstream queriers.
+  forward_headers_list:
+    - Authorization
+
   results_cache:
     cache:
 

pkg/querier/queryrange/query_range.go

@@ -52,7 +52,7 @@ var (
 type Codec interface {
 	Merger
 	// DecodeRequest decodes a Request from an http request.
-	DecodeRequest(context.Context, *http.Request) (Request, error)
+	DecodeRequest(_ context.Context, request *http.Request, forwardHeaders []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 +187,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, forwardHeaders []string) (Request, error) {
 	var result PrometheusRequest
 	var err error
 	result.Start, err = util.ParseTime(r.FormValue("start"))
@@ -222,6 +222,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 forwardHeaders {
+		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 +257,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) {
+	// Create a Copy parsedRequest to assign the expected headers to the request without affecting other tests using the global.
+	// 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
+	parsedRequestWithHeaders := *parsedRequest
+	parsedRequestWithHeaders.Headers = reqHeaders
 	for i, tc := range []struct {
 		url         string
 		expected    Request
 		expectedErr error
 	}{
 		{
 			url:      query,
-			expected: parsedRequest,
+			expected: &parsedRequestWithHeaders,
 		},
 		{
 			url:         "api/v1/query_range?start=foo",
@@ -55,11 +59,14 @@ 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)
 
-			req, err := PrometheusCodec.DecodeRequest(ctx, r)
+			// Get a deep copy of the request with Context changed to ctx
+			r = r.Clone(ctx)
+
+			req, err := PrometheusCodec.DecodeRequest(ctx, r, []string{"Test-Header"})
 			if err != nil {
 				require.EqualValues(t, tc.expectedErr, err)
 				return