From a3741d10686b250b346da509c0ef7aa6d5102b10 Mon Sep 17 00:00:00 2001
From: Karsten Jeschkies <karsten.jeschkies@grafana.com>
Date: Thu, 10 Aug 2023 18:27:22 +0200
Subject: [PATCH] Include response type in merging error message and merge
 already merged views. (#10218)

---
 pkg/querier/queryrange/codec.go      |  8 +++++-
 pkg/querier/queryrange/codec_test.go | 41 +++++++++++++++++-----------
 pkg/querier/queryrange/views_test.go |  1 -
 3 files changed, 32 insertions(+), 18 deletions(-)

diff --git a/pkg/querier/queryrange/codec.go b/pkg/querier/queryrange/codec.go
index 269ea7cff501..7fc9149cb086 100644
--- a/pkg/querier/queryrange/codec.go
+++ b/pkg/querier/queryrange/codec.go
@@ -892,6 +892,12 @@ func (Codec) MergeResponse(responses ...queryrangebase.Response) (queryrangebase
 			v.responses = append(v.responses, r.(*LokiSeriesResponseView))
 		}
 		return v, nil
+	case *MergedSeriesResponseView:
+		v := &MergedSeriesResponseView{}
+		for _, r := range responses {
+			v.responses = append(v.responses, r.(*MergedSeriesResponseView).responses...)
+		}
+		return v, nil
 	case *LokiLabelNamesResponse:
 		labelNameRes := responses[0].(*LokiLabelNamesResponse)
 		uniqueNames := make(map[string]struct{})
@@ -942,7 +948,7 @@ func (Codec) MergeResponse(responses ...queryrangebase.Response) (queryrangebase
 			Headers:  headers,
 		}, nil
 	default:
-		return nil, errors.New("unknown response in merging responses")
+		return nil, fmt.Errorf("unknown response type (%T) in merging responses", responses[0])
 	}
 }
 
diff --git a/pkg/querier/queryrange/codec_test.go b/pkg/querier/queryrange/codec_test.go
index f4ccd0ab3b74..2d124e0a7663 100644
--- a/pkg/querier/queryrange/codec_test.go
+++ b/pkg/querier/queryrange/codec_test.go
@@ -847,13 +847,23 @@ func Test_codec_EncodeResponse(t *testing.T) {
 
 func Test_codec_MergeResponse(t *testing.T) {
 	tests := []struct {
-		name      string
-		responses []queryrangebase.Response
-		want      queryrangebase.Response
-		wantErr   bool
+		name         string
+		responses    []queryrangebase.Response
+		want         queryrangebase.Response
+		errorMessage string
 	}{
-		{"empty", []queryrangebase.Response{}, nil, true},
-		{"unknown response", []queryrangebase.Response{&badResponse{}}, nil, true},
+		{
+			"empty",
+			[]queryrangebase.Response{},
+			nil,
+			"merging responses requires at least one response",
+		},
+		{
+			"unknown response",
+			[]queryrangebase.Response{&badResponse{}},
+			nil,
+			"unknown response type (*queryrange.badResponse) in merging responses",
+		},
 		{
 			"prom",
 			[]queryrangebase.Response{
@@ -877,7 +887,7 @@ func Test_codec_MergeResponse(t *testing.T) {
 					},
 				},
 			},
-			false,
+			"",
 		},
 		{
 			"loki backward",
@@ -965,7 +975,7 @@ func Test_codec_MergeResponse(t *testing.T) {
 					},
 				},
 			},
-			false,
+			"",
 		},
 		{
 			"loki backward limited",
@@ -1050,7 +1060,7 @@ func Test_codec_MergeResponse(t *testing.T) {
 					},
 				},
 			},
-			false,
+			"",
 		},
 		{
 			"loki forward",
@@ -1138,7 +1148,7 @@ func Test_codec_MergeResponse(t *testing.T) {
 					},
 				},
 			},
-			false,
+			"",
 		},
 		{
 			"loki forward limited",
@@ -1222,7 +1232,7 @@ func Test_codec_MergeResponse(t *testing.T) {
 					},
 				},
 			},
-			false,
+			"",
 		},
 		{
 			"loki series",
@@ -1268,7 +1278,7 @@ func Test_codec_MergeResponse(t *testing.T) {
 					},
 				},
 			},
-			false,
+			"",
 		},
 		{
 			"loki labels",
@@ -1295,15 +1305,14 @@ func Test_codec_MergeResponse(t *testing.T) {
 				Version:    1,
 				Data:       []string{"foo", "bar", "buzz", "blip", "blop"},
 			},
-			false,
+			"",
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			got, err := DefaultCodec.MergeResponse(tt.responses...)
-			if (err != nil) != tt.wantErr {
-				t.Errorf("codec.MergeResponse() error = %v, wantErr %v", err, tt.wantErr)
-				return
+			if tt.errorMessage != "" {
+				require.ErrorContains(t, err, tt.errorMessage)
 			}
 			require.Equal(t, tt.want, got)
 		})
diff --git a/pkg/querier/queryrange/views_test.go b/pkg/querier/queryrange/views_test.go
index a3d41ad91a9a..eacf157074f1 100644
--- a/pkg/querier/queryrange/views_test.go
+++ b/pkg/querier/queryrange/views_test.go
@@ -217,7 +217,6 @@ func TestMergedViewMaterialize(t *testing.T) {
 	}
 	expected := []string{`{baz="woof", i="1"}`, `{baz="woof", i="3"}`, `{foo="bar", i="2"}`}
 	require.ElementsMatch(t, series, expected)
-
 }
 
 func TestMergedViewJSON(t *testing.T) {