query-frontend: add option to preserve specific request headers

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## master / unreleased
 
+* [ENHANCEMENT] Querier: Retain Headers from request while decoding http requests to PrometheusRequests by prometheusCodec. #4416
 * [CHANGE] Memberlist: Expose default configuration values to the command line options. Note that setting these explicitly to zero will no longer cause the default to be used. If the default is desired, then do set the option. The following are affected: #4276
   - `-memberlist.stream-timeout`
   - `-memberlist.retransmit-factor`

pkg/querier/queryrange/query_range.go

@@ -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 headers from http request in prometheusRequest.
+	// We don't include the following header in the codec -
+	// accept-encoding - defaultroundtripper uses the response_compression_enabled setting to inflate/deflate responses
+	for h, hv := range r.Header {
+		if strings.ToLower(h) == "accept-encoding" {
+			continue
+		}
+		result.Headers = append(result.Headers, &PrometheusRequestHeader{Name: h, Values: hv})
+	}
+
 	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,16 @@ import (
 )
 
 func TestRequest(t *testing.T) {
+	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,9 +57,10 @@ 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)
 			if err != nil {