From f79b61e2a141d6b1c104818dbc01eae76db500bb Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@gmail.com>
Date: Wed, 11 Dec 2019 00:42:44 +0200
Subject: [PATCH] app/vmselect/promql: return matrix instead of vector on
 subqueries to `/api/v1/query` like Prometheus does

---
 app/vmselect/prometheus/prometheus.go | 68 +++++++++++++++++++++------
 app/vmselect/promql/parser.go         | 16 +++++++
 2 files changed, 69 insertions(+), 15 deletions(-)

diff --git a/app/vmselect/prometheus/prometheus.go b/app/vmselect/prometheus/prometheus.go
index 035be9389..fcb1483c2 100644
--- a/app/vmselect/prometheus/prometheus.go
+++ b/app/vmselect/prometheus/prometheus.go
@@ -560,21 +560,13 @@ func QueryHandler(at *auth.Token, w http.ResponseWriter, r *http.Request) error
 		start = ct - queryOffset
 	}
 	if childQuery, windowStr, offsetStr := promql.IsMetricSelectorWithRollup(query); childQuery != "" {
-		var window int64
-		if len(windowStr) > 0 {
-			var err error
-			window, err = promql.PositiveDurationValue(windowStr, step)
-			if err != nil {
-				return err
-			}
+		window, err := parsePositiveDuration(windowStr, step)
+		if err != nil {
+			return fmt.Errorf("cannot parse window: %s", err)
 		}
-		var offset int64
-		if len(offsetStr) > 0 {
-			var err error
-			offset, err = promql.DurationValue(offsetStr, step)
-			if err != nil {
-				return err
-			}
+		offset, err := parseDuration(offsetStr, step)
+		if err != nil {
+			return fmt.Errorf("cannot parse offset: %s", err)
 		}
 		start -= offset
 		end := start
@@ -585,6 +577,31 @@ func QueryHandler(at *auth.Token, w http.ResponseWriter, r *http.Request) error
 		queryDuration.UpdateDuration(startTime)
 		return nil
 	}
+	if childQuery, windowStr, stepStr, offsetStr := promql.IsRollup(query); childQuery != "" {
+		newStep, err := parsePositiveDuration(stepStr, step)
+		if err != nil {
+			return fmt.Errorf("cannot parse step: %s", err)
+		}
+		if newStep > 0 {
+			step = newStep
+		}
+		window, err := parsePositiveDuration(windowStr, step)
+		if err != nil {
+			return fmt.Errorf("cannot parse window: %s", err)
+		}
+		offset, err := parseDuration(offsetStr, step)
+		if err != nil {
+			return fmt.Errorf("cannot parse offset: %s", err)
+		}
+		start -= offset
+		end := start
+		start = end - window
+		if err := queryRangeHandler(at, w, childQuery, start, end, step, r, ct); err != nil {
+			return err
+		}
+		queryDuration.UpdateDuration(startTime)
+		return nil
+	}
 
 	ec := promql.EvalConfig{
 		AuthToken:     at,
@@ -609,6 +626,20 @@ func QueryHandler(at *auth.Token, w http.ResponseWriter, r *http.Request) error
 
 var queryDuration = metrics.NewSummary(`vm_request_duration_seconds{path="/api/v1/query"}`)
 
+func parseDuration(s string, step int64) (int64, error) {
+	if len(s) == 0 {
+		return 0, nil
+	}
+	return promql.DurationValue(s, step)
+}
+
+func parsePositiveDuration(s string, step int64) (int64, error) {
+	if len(s) == 0 {
+		return 0, nil
+	}
+	return promql.PositiveDurationValue(s, step)
+}
+
 // QueryRangeHandler processes /api/v1/query_range request.
 //
 // See https://prometheus.io/docs/prometheus/latest/querying/api/#range-queries
@@ -632,6 +663,14 @@ func QueryRangeHandler(at *auth.Token, w http.ResponseWriter, r *http.Request) e
 	if err != nil {
 		return err
 	}
+	if err := queryRangeHandler(at, w, query, start, end, step, r, ct); err != nil {
+		return err
+	}
+	queryRangeDuration.UpdateDuration(startTime)
+	return nil
+}
+
+func queryRangeHandler(at *auth.Token, w http.ResponseWriter, query string, start, end, step int64, r *http.Request, ct int64) error {
 	deadline := getDeadline(r)
 	mayCache := !getBool(r, "nocache")
 	lookbackDelta, err := getMaxLookback(r)
@@ -679,7 +718,6 @@ func QueryRangeHandler(at *auth.Token, w http.ResponseWriter, r *http.Request) e
 
 	w.Header().Set("Content-Type", "application/json")
 	WriteQueryRangeResponse(w, result)
-	queryRangeDuration.UpdateDuration(startTime)
 	return nil
 }
 
diff --git a/app/vmselect/promql/parser.go b/app/vmselect/promql/parser.go
index 528893b0d..caaa48f14 100644
--- a/app/vmselect/promql/parser.go
+++ b/app/vmselect/promql/parser.go
@@ -1290,6 +1290,22 @@ func (p *parser) parseIdentExpr() (expr, error) {
 	}
 }
 
+// IsRollup verifies whether s is a rollup with non-empty window.
+//
+// It returns the wrapped query with the corresponding window, step and offset.
+func IsRollup(s string) (childQuery string, window, step, offset string) {
+	expr, err := parsePromQLWithCache(s)
+	if err != nil {
+		return
+	}
+	re, ok := expr.(*rollupExpr)
+	if !ok || len(re.Window) == 0 {
+		return
+	}
+	wrappedQuery := re.Expr.AppendString(nil)
+	return string(wrappedQuery), re.Window, re.Step, re.Offset
+}
+
 // IsMetricSelectorWithRollup verifies whether s contains PromQL metric selector
 // wrapped into rollup.
 //