refactor query decompression

Signed-off-by: Ahmed Hassan <[email protected]>

Author: afhassan

pkg/querier/tripperware/query.go

@@ -454,35 +454,9 @@ func BodyBytes(res *http.Response, logger log.Logger) ([]byte, error) {
 		}
 	}
 
-	// if the response is gzipped, lets unzip it here
-	if strings.EqualFold(res.Header.Get("Content-Encoding"), "gzip") {
-		gReader, err := gzip.NewReader(buf)
-		if err != nil {
-			return nil, err
-		}
-		defer runutil.CloseWithLogOnErr(logger, gReader, "close gzip reader")
-
-		return io.ReadAll(gReader)
-	}
-
-	// if the response is snappy compressed, decode it here
-	if strings.EqualFold(res.Header.Get("Content-Encoding"), "snappy") {
-		sReader := snappy.NewReader(buf)
-		return io.ReadAll(sReader)
-	}
-
-	// if the response is zstd compressed, decode it here
-	if strings.EqualFold(res.Header.Get("Content-Encoding"), "zstd") {
-		zReader, err := zstd.NewReader(buf)
-		if err != nil {
-			return nil, err
-		}
-		defer runutil.CloseWithLogOnErr(logger, io.NopCloser(zReader), "close zstd decoder")
-
-		return io.ReadAll(zReader)
-	}
-
-	return buf.Bytes(), nil
+	// Handle decoding response if it was compressed
+	encoding := res.Header.Get("Content-Encoding")
+	return decode(buf, encoding, logger)
 }
 
 func BodyBytesFromHTTPGRPCResponse(res *httpgrpc.HTTPResponse, logger log.Logger) ([]byte, error) {
@@ -491,9 +465,16 @@ func BodyBytesFromHTTPGRPCResponse(res *httpgrpc.HTTPResponse, logger log.Logger
 		headers[h.Key] = h.Values
 	}
 
+	// Handle decoding response if it was compressed
+	encoding := headers.Get("Content-Encoding")
+	buf := bytes.NewBuffer(res.Body)
+	return decode(buf, encoding, logger)
+}
+
+func decode(buf *bytes.Buffer, encoding string, logger log.Logger) ([]byte, error) {
 	// if the response is gzipped, lets unzip it here
-	if strings.EqualFold(headers.Get("Content-Encoding"), "gzip") {
-		gReader, err := gzip.NewReader(bytes.NewBuffer(res.Body))
+	if strings.EqualFold(encoding, "gzip") {
+		gReader, err := gzip.NewReader(buf)
 		if err != nil {
 			return nil, err
 		}
@@ -503,14 +484,14 @@ func BodyBytesFromHTTPGRPCResponse(res *httpgrpc.HTTPResponse, logger log.Logger
 	}
 
 	// if the response is snappy compressed, decode it here
-	if strings.EqualFold(headers.Get("Content-Encoding"), "snappy") {
-		sReader := snappy.NewReader(bytes.NewBuffer(res.Body))
+	if strings.EqualFold(encoding, "snappy") {
+		sReader := snappy.NewReader(buf)
 		return io.ReadAll(sReader)
 	}
 
 	// if the response is zstd compressed, decode it here
-	if strings.EqualFold(headers.Get("Content-Encoding"), "zstd") {
-		zReader, err := zstd.NewReader(bytes.NewBuffer(res.Body))
+	if strings.EqualFold(encoding, "zstd") {
+		zReader, err := zstd.NewReader(buf)
 		if err != nil {
 			return nil, err
 		}
@@ -519,7 +500,7 @@ func BodyBytesFromHTTPGRPCResponse(res *httpgrpc.HTTPResponse, logger log.Logger
 		return io.ReadAll(zReader)
 	}
 
-	return res.Body, nil
+	return buf.Bytes(), nil
 }
 
 // UnmarshalJSON implements json.Unmarshaler.