From b24af037c304badfadffc1707bba7c4c6a321d23 Mon Sep 17 00:00:00 2001
From: Trevor Whitney <trevorjwhitney@gmail.com>
Date: Mon, 7 Oct 2024 14:03:26 -0600
Subject: [PATCH] feat: add range support to agg metric volume queries

---
 pkg/querier/queryrange/roundtrip.go | 58 ++++++++++++++++++-----------
 1 file changed, 36 insertions(+), 22 deletions(-)

diff --git a/pkg/querier/queryrange/roundtrip.go b/pkg/querier/queryrange/roundtrip.go
index 0ead0aaf4c67..50948d3ef7b1 100644
--- a/pkg/querier/queryrange/roundtrip.go
+++ b/pkg/querier/queryrange/roundtrip.go
@@ -1207,15 +1207,33 @@ func aggMetricsVolumeHandler(
 		return nil, httpgrpc.Errorf(http.StatusBadRequest, "%s", err.Error())
 	}
 
-	lokiReq := &LokiInstantRequest{
-		Query:     expr.String(),
-		Limit:     1000,
-		Direction: logproto.BACKWARD,
-		TimeTs:    r.GetEnd().UTC(),
-		Path:      "/loki/api/v1/query",
-		Plan: &plan.QueryPlan{
-			AST: expr,
-		},
+	var lokiReq base.Request
+	if r.GetStep() <= 0 {
+		lokiReq = &LokiInstantRequest{
+			Query:     expr.String(),
+			Limit:     1000,
+			Direction: logproto.BACKWARD,
+			TimeTs:    r.GetEnd().UTC(),
+			Path:      "/loki/api/v1/query",
+			Plan: &plan.QueryPlan{
+				AST: expr,
+			},
+			CachingOptions: r.GetCachingOptions(),
+		}
+	} else {
+		lokiReq = &LokiRequest{
+			Query:     expr.String(),
+			Limit:     1000,
+			Step:      r.GetStep(),
+			StartTs:   r.GetStart().UTC(),
+			EndTs:     r.GetEnd().UTC(),
+			Direction: logproto.BACKWARD,
+			Path:      "/loki/api/v1/query_range",
+			Plan: &plan.QueryPlan{
+				AST: expr,
+			},
+			CachingOptions: r.GetCachingOptions(),
+		}
 	}
 
 	resp, err := logHandler.Do(ctx, lokiReq)
@@ -1232,7 +1250,7 @@ func aggMetricsVolumeHandler(
 	sortableResult := make([]sortableSampleStream, 0, len(result))
 	resultType := loghttp.ResultTypeVector
 
-  // sort the response to match the index volume respsone
+	// sort the response to match the index volume respsone
 	for _, stream := range result {
 		if resultType == loghttp.ResultTypeVector && len(stream.Samples) > 1 {
 			resultType = loghttp.ResultTypeMatrix
@@ -1262,19 +1280,15 @@ func aggMetricsVolumeHandler(
 		})
 	}
 
-	
-
-	
-
 	return &LokiPromResponse{
-		Response:   &base.PrometheusResponse{
-		Status:  loghttp.QueryStatusSuccess,
-		Data:    base.PrometheusData{
-		ResultType: resultType,
-		Result:     respStreams,
-	},
-		Headers: re.Response.Headers,
-	},
+		Response: &base.PrometheusResponse{
+			Status: loghttp.QueryStatusSuccess,
+			Data: base.PrometheusData{
+				ResultType: resultType,
+				Result:     respStreams,
+			},
+			Headers: re.Response.Headers,
+		},
 		Statistics: re.Statistics,
 	}, nil
 }