Cast volume response on receive.

pkg/querier/queryrange/volume.go

@@ -2,6 +2,7 @@ package queryrange
 
 import (
 	"context"
+	"fmt"
 	"sort"
 	"time"
 
@@ -48,18 +49,23 @@ func NewVolumeMiddleware() queryrangebase.Middleware {
 				}
 			})
 
-			type f func(context.Context) (time.Time, definitions.Response, error)
+			type f func(context.Context) (time.Time, *VolumeResponse, error)
 			var jobs []f
 
 			for bucket, req := range reqs {
 				b, r := bucket, req
-				jobs = append(jobs, f(func(ctx context.Context) (time.Time, definitions.Response, error) {
+				jobs = append(jobs, f(func(ctx context.Context) (time.Time, *VolumeResponse, error) {
 					resp, err := next.Do(ctx, r)
 					if err != nil {
 						return b, nil, err
 					}
 
-					return b, resp, nil
+					volumeResponse, ok := resp.(*VolumeResponse)
+					if !ok {
+						return b, nil, fmt.Errorf("unexpected response type: want (*VolumeResponse), got (%T)", resp)
+					}
+
+					return b, volumeResponse, nil
 				}))
 			}
 
@@ -80,7 +86,7 @@ func NewVolumeMiddleware() queryrangebase.Middleware {
 					}
 
 					collector <- &bucketedVolumeResponse{
-						bucket, resp.(*VolumeResponse),
+						bucket, resp,
 					}
 					return err
 				})

pkg/util/marshal/marshal_test.go

@@ -714,6 +714,34 @@ func Test_WriteSeriesResponseJSON(t *testing.T) {
 			require.NoError(t, err)
 
 			require.JSONEqf(t, tc.expected, b.String(), "Series Test %d failed", i)
+
+		})
+	}
+}
+
+func Test_WriteVolumeResponseJSON(t *testing.T) {
+	for _, tc := range []struct {
+		name     string
+		input    *logproto.VolumeResponse
+		expected string
+	}{
+		{
+			name:     "empty",
+			input:    &logproto.VolumeResponse{Volumes: []logproto.Volume{}},
+			expected: `{"volumes":[]}`,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			var b bytes.Buffer
+			err := WriteVolumeResponseJSON(tc.input, &b)
+			require.NoError(t, err)
+
+			require.JSONEq(t, tc.expected, b.String())
+
+			var resp logproto.VolumeResponse
+			err = json.Unmarshal([]byte(tc.expected), &resp)
+			require.NoError(t, err)
+			require.Equal(t, tc.input, &resp)
 		})
 	}
 }