query-frontend: add option to preserve specific request headers

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## master / unreleased
 
+* [CHANGE] Querier: Headers from request are lost in 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

@@ -94,7 +94,7 @@ type Request interface {
 type Response interface {
 	proto.Message
 	// GetHeaders returns the HTTP headers in the response.
-	GetHeaders() []*PrometheusResponseHeader
+	GetHeaders() []*PrometheusHeader
 }
 
 type prometheusCodec struct{}
@@ -178,7 +178,7 @@ func (prometheusCodec) MergeResponse(responses ...Response) (Response, error) {
 	}
 
 	if len(resultsCacheGenNumberHeaderValues) != 0 {
-		response.Headers = []*PrometheusResponseHeader{{
+		response.Headers = []*PrometheusHeader{{
 			Name:   ResultsCacheGenNumberHeaderName,
 			Values: resultsCacheGenNumberHeaderValues,
 		}}
@@ -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, &PrometheusHeader{Name: h, Values: hv})
+	}
+
 	for _, value := range r.Header.Values(cacheControlHeader) {
 		if strings.Contains(value, noStoreValue) {
 			result.CachingOptions.Disabled = true
@@ -247,12 +257,23 @@ func (prometheusCodec) EncodeRequest(ctx context.Context, r Request) (*http.Requ
 		Path:     promReq.Path,
 		RawQuery: params.Encode(),
 	}
+	h := make(http.Header)
+
+	for _, hv := range promReq.Headers {
+		if hv == nil {
+			continue
+		}
+		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
@@ -279,7 +300,7 @@ func (prometheusCodec) DecodeResponse(ctx context.Context, r *http.Response, _ R
 	}
 
 	for h, hv := range r.Header {
-		resp.Headers = append(resp.Headers, &PrometheusResponseHeader{Name: h, Values: hv})
+		resp.Headers = append(resp.Headers, &PrometheusHeader{Name: h, Values: hv})
 	}
 	return &resp, 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 {

pkg/querier/queryrange/queryable.go

@@ -30,7 +30,7 @@ func (q *ShardedQueryable) Querier(ctx context.Context, mint, maxt int64) (stora
 	return q.sharededQuerier, nil
 }
 
-func (q *ShardedQueryable) getResponseHeaders() []*PrometheusResponseHeader {
+func (q *ShardedQueryable) getResponseHeaders() []*PrometheusHeader {
 	q.sharededQuerier.ResponseHeadersMtx.Lock()
 	defer q.sharededQuerier.ResponseHeadersMtx.Unlock()
 
@@ -117,7 +117,7 @@ func (q *ShardedQuerier) handleEmbeddedQuery(encoded string) storage.SeriesSet {
 	return NewSeriesSet(samples)
 }
 
-func (q *ShardedQuerier) setResponseHeaders(headers []*PrometheusResponseHeader) {
+func (q *ShardedQuerier) setResponseHeaders(headers []*PrometheusHeader) {
 	q.ResponseHeadersMtx.Lock()
 	defer q.ResponseHeadersMtx.Unlock()
 
@@ -145,9 +145,9 @@ func (q *ShardedQuerier) Close() error {
 	return nil
 }
 
-func headersMapToPrometheusResponseHeaders(headersMap map[string][]string) (prs []*PrometheusResponseHeader) {
+func headersMapToPrometheusResponseHeaders(headersMap map[string][]string) (prs []*PrometheusHeader) {
 	for h, v := range headersMap {
-		prs = append(prs, &PrometheusResponseHeader{Name: h, Values: v})
+		prs = append(prs, &PrometheusHeader{Name: h, Values: v})
 	}
 
 	return