added GetHeaders method to MetricsQueryRequest interface.

Signed-off-by: Miguel Ángel Ortuño <ortuman@gmail.com>

pkg/frontend/querymiddleware/codec.go

@@ -95,6 +95,8 @@ type MetricsQueryRequest interface {
 	GetID() int64
 	// GetPath returns the URL Path of the request
 	GetPath() string
+	// GetHeaders returns the HTTP headers in the request.
+	GetHeaders() []*PrometheusHeader
 	// GetStart returns the start timestamp of the request in milliseconds.
 	GetStart() int64
 	// GetEnd returns the end timestamp of the request in milliseconds.
@@ -163,7 +165,7 @@ type LabelsQueryRequest interface {
 type Response interface {
 	proto.Message
 	// GetHeaders returns the HTTP headers in the response.
-	GetHeaders() []*PrometheusResponseHeader
+	GetHeaders() []*PrometheusHeader
 }
 
 type prometheusCodecMetrics struct {
@@ -284,6 +286,12 @@ func (c prometheusCodec) decodeRangeQueryRequest(r *http.Request) (MetricsQueryR
 		return nil, apierror.New(apierror.TypeBadData, err.Error())
 	}
 
+	headers := make([]*PrometheusHeader, 0, len(r.Header))
+	for h, hv := range r.Header {
+		headers = append(headers, &PrometheusHeader{Name: h, Values: slices.Clone(hv)})
+	}
+	sort.Slice(headers, func(i, j int) bool { return headers[i].Name < headers[j].Name })
+
 	start, end, step, err := DecodeRangeQueryTimeParams(&reqValues)
 	if err != nil {
 		return nil, err
@@ -299,7 +307,7 @@ func (c prometheusCodec) decodeRangeQueryRequest(r *http.Request) (MetricsQueryR
 	decodeOptions(r, &options)
 
 	req := NewPrometheusRangeQueryRequest(
-		r.URL.Path, start, end, step, c.lookbackDelta, queryExpr, options, nil,
+		r.URL.Path, headers, start, end, step, c.lookbackDelta, queryExpr, options, nil,
 	)
 	return req, nil
 }
@@ -310,6 +318,12 @@ func (c prometheusCodec) decodeInstantQueryRequest(r *http.Request) (MetricsQuer
 		return nil, apierror.New(apierror.TypeBadData, err.Error())
 	}
 
+	headers := make([]*PrometheusHeader, 0, len(r.Header))
+	for h, hv := range r.Header {
+		headers = append(headers, &PrometheusHeader{Name: h, Values: slices.Clone(hv)})
+	}
+	sort.Slice(headers, func(i, j int) bool { return headers[i].Name < headers[j].Name })
+
 	time, err := DecodeInstantQueryTimeParams(&reqValues, time.Now)
 	if err != nil {
 		return nil, decorateWithParamName(err, "time")
@@ -325,7 +339,7 @@ func (c prometheusCodec) decodeInstantQueryRequest(r *http.Request) (MetricsQuer
 	decodeOptions(r, &options)
 
 	req := NewPrometheusInstantQueryRequest(
-		r.URL.Path, time, c.lookbackDelta, queryExpr, options, nil,
+		r.URL.Path, headers, time, c.lookbackDelta, queryExpr, options, nil,
 	)
 	return req, nil
 }
@@ -699,7 +713,7 @@ func (c prometheusCodec) DecodeResponse(ctx context.Context, r *http.Response, _
 	}
 
 	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/frontend/querymiddleware/codec_json_test.go

@@ -27,7 +27,7 @@ import (
 
 func TestPrometheusCodec_JSONResponse(t *testing.T) {
 	headers := http.Header{"Content-Type": []string{"application/json"}}
-	expectedRespHeaders := []*PrometheusResponseHeader{
+	expectedRespHeaders := []*PrometheusHeader{
 		{
 			Name:   "Content-Type",
 			Values: []string{"application/json"},

pkg/frontend/querymiddleware/codec_protobuf_test.go

@@ -21,7 +21,7 @@ import (
 	"github.com/grafana/mimir/pkg/mimirpb"
 )
 
-var expectedProtobufResponseHeaders = []*PrometheusResponseHeader{
+var expectedProtobufResponseHeaders = []*PrometheusHeader{
 	{
 		Name:   "Content-Type",
 		Values: []string{mimirpb.QueryResponseMimeType},

pkg/frontend/querymiddleware/codec_test.go

@@ -79,6 +79,7 @@ func TestMetricsQueryRequest(t *testing.T) {
 			url: "/api/v1/query_range?end=1536716880&query=sum+by+%28namespace%29+%28container_memory_rss%29&start=1536673680&step=120",
 			expected: NewPrometheusRangeQueryRequest(
 				"/api/v1/query_range",
+				nil,
 				1536673680*1e3,
 				1536716880*1e3,
 				(2 * time.Minute).Milliseconds(),
@@ -92,6 +93,7 @@ func TestMetricsQueryRequest(t *testing.T) {
 			url: "/api/v1/query?query=sum+by+%28namespace%29+%28container_memory_rss%29&time=1536716880",
 			expected: NewPrometheusInstantQueryRequest(
 				"/api/v1/query",
+				nil,
 				1536716880*1e3,
 				0*time.Minute,
 				parseQuery(t, "sum(container_memory_rss) by (namespace)"),
@@ -171,6 +173,7 @@ func TestMetricsQuery_MinMaxTime(t *testing.T) {
 
 	rangeRequest := NewPrometheusRangeQueryRequest(
 		"/api/v1/query_range",
+		nil,
 		startTime.UnixMilli(),
 		endTime.UnixMilli(),
 		stepDuration.Milliseconds(),
@@ -181,6 +184,7 @@ func TestMetricsQuery_MinMaxTime(t *testing.T) {
 	)
 	instantRequest := NewPrometheusInstantQueryRequest(
 		"/api/v1/query",
+		nil,
 		endTime.UnixMilli(),
 		lookbackDuration,
 		parseQuery(t, "go_goroutines{}"),
@@ -297,6 +301,7 @@ func TestMetricsQuery_WithStartEnd_TransformConsistency(t *testing.T) {
 
 	rangeRequest := NewPrometheusRangeQueryRequest(
 		"/api/v1/query_range",
+		nil,
 		startTime.UnixMilli(),
 		endTime.UnixMilli(),
 		stepDuration.Milliseconds(),
@@ -307,6 +312,7 @@ func TestMetricsQuery_WithStartEnd_TransformConsistency(t *testing.T) {
 	)
 	instantRequest := NewPrometheusInstantQueryRequest(
 		"/api/v1/query",
+		nil,
 		endTime.UnixMilli(),
 		time.Duration(0),
 		parseQuery(t, "go_goroutines{}"),
@@ -375,6 +381,7 @@ func TestMetricsQuery_WithQuery_WithExpr_TransformConsistency(t *testing.T) {
 
 	rangeRequest := NewPrometheusRangeQueryRequest(
 		"/api/v1/query_range",
+		nil,
 		startTime.UnixMilli(),
 		endTime.UnixMilli(),
 		stepDuration.Milliseconds(),
@@ -385,6 +392,7 @@ func TestMetricsQuery_WithQuery_WithExpr_TransformConsistency(t *testing.T) {
 	)
 	instantRequest := NewPrometheusInstantQueryRequest(
 		"/api/v1/query",
+		nil,
 		endTime.UnixMilli(),
 		time.Duration(0),
 		parseQuery(t, "go_goroutines{}"),