Returning Prometheus Warnings (#5916)

* Returning Prometheus Warnings

Signed-off-by: alanprot <[email protected]>

* changelog

Signed-off-by: alanprot <[email protected]>

---------

Signed-off-by: alanprot <[email protected]>

Author: alanprot

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## master / unreleased
 * [ENHANCEMENT] Query Frontend/Querier: Added store gateway postings touched count and touched size in Querier stats and log in Query Frontend. #5892
+* [ENHANCEMENT] Query Frontend/Querier: Returns `warnings` on prometheus query responses. #5916
 * [CHANGE] Upgrade Dockerfile Node version from 14x to 18x. #5906
 
 ## 1.17.0 2024-04-30

pkg/querier/tripperware/instantquery/instant_query.go

@@ -18,6 +18,7 @@ import (
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/model/timestamp"
 	promqlparser "github.com/prometheus/prometheus/promql/parser"
+	"github.com/thanos-io/thanos/pkg/strutil"
 	"github.com/weaveworks/common/httpgrpc"
 	"google.golang.org/grpc/status"
 
@@ -260,8 +261,12 @@ func (instantQueryCodec) MergeResponse(ctx context.Context, req tripperware.Requ
 	}
 
 	promResponses := make([]*PrometheusInstantQueryResponse, 0, len(responses))
+	warnings := make([][]string, 0, len(responses))
 	for _, resp := range responses {
 		promResponses = append(promResponses, resp.(*PrometheusInstantQueryResponse))
+		if w := resp.(*PrometheusInstantQueryResponse).Warnings; w != nil {
+			warnings = append(warnings, w)
+		}
 	}
 
 	var data PrometheusInstantQueryData
@@ -303,8 +308,9 @@ func (instantQueryCodec) MergeResponse(ctx context.Context, req tripperware.Requ
 	}
 
 	res := &PrometheusInstantQueryResponse{
-		Status: queryrange.StatusSuccess,
-		Data:   data,
+		Status:   queryrange.StatusSuccess,
+		Data:     data,
+		Warnings: strutil.MergeUnsortedSlices(warnings...),
 	}
 	return res, nil
 }

pkg/querier/tripperware/instantquery/instant_query_test.go

@@ -305,6 +305,15 @@ func TestMergeResponse(t *testing.T) {
 			},
 			expectedResp: `{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"foo"},"value":[1,"1"]},{"metric":{"__name__":"up","job":"bar"},"value":[1,"2"]}]}}`,
 		},
+		{
+			name: "merge with warnings.",
+			req:  &PrometheusRequest{Query: "topk(10, up) by(job)"},
+			resps: []string{
+				`{"status":"success","warnings":["warning1","warning2"],"data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"foo"},"value":[1,"1"]}]}}`,
+				`{"status":"success","warnings":["warning1","warning3"],"data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"bar"},"value":[1,"2"]}]}}`,
+			},
+			expectedResp: `{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"foo"},"value":[1,"1"]},{"metric":{"__name__":"up","job":"bar"},"value":[1,"2"]}]},"warnings":["warning1","warning2","warning3"]}`,
+		},
 		{
 			name: "merge two responses with stats",
 			req:  defaultReq,

pkg/querier/tripperware/instantquery/instantquery.pb.go

@@ -18,6 +18,7 @@ message PrometheusInstantQueryResponse {
   string ErrorType = 3 [(gogoproto.jsontag) = "errorType,omitempty"];
   string Error = 4 [(gogoproto.jsontag) = "error,omitempty"];
   repeated tripperware.PrometheusResponseHeader Headers = 5 [(gogoproto.jsontag) = "-"];
+  repeated string Warnings = 6 [(gogoproto.jsontag) = "warnings,omitempty"];
 }
 
 message PrometheusInstantQueryData {

pkg/querier/tripperware/instantquery/instantquery.proto

@@ -19,6 +19,7 @@ import (
 	otlog "github.com/opentracing/opentracing-go/log"
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/model/timestamp"
+	"github.com/thanos-io/thanos/pkg/strutil"
 	"github.com/weaveworks/common/httpgrpc"
 
 	"github.com/cortexproject/cortex/pkg/querier/tripperware"
@@ -135,8 +136,12 @@ func (c prometheusCodec) MergeResponse(ctx context.Context, _ tripperware.Reques
 	}
 
 	promResponses := make([]*PrometheusResponse, 0, len(responses))
+	warnings := make([][]string, 0, len(responses))
 	for _, res := range responses {
 		promResponses = append(promResponses, res.(*PrometheusResponse))
+		if w := res.(*PrometheusResponse).Warnings; w != nil {
+			warnings = append(warnings, w)
+		}
 	}
 
 	// Merge the responses.
@@ -153,6 +158,7 @@ func (c prometheusCodec) MergeResponse(ctx context.Context, _ tripperware.Reques
 			Result:     sampleStreams,
 			Stats:      statsMerge(c.sharded, promResponses),
 		},
+		Warnings: strutil.MergeUnsortedSlices(warnings...),
 	}
 
 	return &response, nil

pkg/querier/tripperware/queryrange/query_range.go

@@ -32,6 +32,8 @@ func TestRoundTrip(t *testing.T) {
 				var err error
 				if r.RequestURI == query {
 					_, err = w.Write([]byte(responseBody))
+				} else if r.RequestURI == queryWithWarnings {
+					_, err = w.Write([]byte(responseBodyWithWarnings))
 				} else {
 					_, err = w.Write([]byte("bar"))
 				}

pkg/querier/tripperware/queryrange/query_range_middlewares_test.go

@@ -91,7 +91,9 @@ func TestRequest(t *testing.T) {
 func TestResponse(t *testing.T) {
 	t.Parallel()
 	r := *parsedResponse
+	rWithWarnings := *parsedResponseWithWarnings
 	r.Headers = respHeaders
+	rWithWarnings.Headers = respHeaders
 	for i, tc := range []struct {
 		body                  string
 		expected              *PrometheusResponse
@@ -102,6 +104,10 @@ func TestResponse(t *testing.T) {
 			body:     responseBody,
 			expected: &r,
 		},
+		{
+			body:     responseBodyWithWarnings,
+			expected: &rWithWarnings,
+		},
 		{
 			body:                  responseBody,
 			cancelCtxBeforeDecode: true,
@@ -390,6 +396,28 @@ func TestMergeAPIResponses(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "Merge response with warnings.",
+			input: []tripperware.Response{
+				mustParse(t, `{"status":"success","warnings":["warning1","warning2"],"data":{"resultType":"matrix","result":[{"metric":{"a":"b","c":"d"},"values":[[1,"1"]]}]}}`),
+				mustParse(t, `{"status":"success","warnings":["warning1","warning3"],"data":{"resultType":"matrix","result":[{"metric":{"a":"b","c":"d"},"values":[[1,"1"]]}]}}`),
+			},
+			expected: &PrometheusResponse{
+				Status:   StatusSuccess,
+				Warnings: []string{"warning1", "warning2", "warning3"},
+				Data: PrometheusData{
+					ResultType: matrix,
+					Result: []tripperware.SampleStream{
+						{
+							Labels: []cortexpb.LabelAdapter{{Name: "a", Value: "b"}, {Name: "c", Value: "d"}},
+							Samples: []cortexpb.Sample{
+								{Value: 1, TimestampMs: 1000},
+							},
+						},
+					},
+				},
+			},
+		},
 		{
 			name: "Merging of samples where there is complete overlap.",
 			input: []tripperware.Response{