prom: removes limitation of one query per tsdb call

pkg/tsdb/prometheus/prometheus.go

@@ -83,41 +83,48 @@ func (e *PrometheusExecutor) getClient(dsInfo *models.DataSource) (apiv1.API, er
 }
 
 func (e *PrometheusExecutor) Query(ctx context.Context, dsInfo *models.DataSource, tsdbQuery *tsdb.TsdbQuery) (*tsdb.Response, error) {
-	result := &tsdb.Response{}
+	result := &tsdb.Response{
+		Results: map[string]*tsdb.QueryResult{},
+	}
 
 	client, err := e.getClient(dsInfo)
 	if err != nil {
 		return nil, err
 	}
 
-	query, err := parseQuery(dsInfo, tsdbQuery.Queries, tsdbQuery)
+	querys, err := parseQuery(dsInfo, tsdbQuery.Queries, tsdbQuery)
 	if err != nil {
 		return nil, err
 	}
 
-	timeRange := apiv1.Range{
-		Start: query.Start,
-		End:   query.End,
-		Step:  query.Step,
-	}
+	for _, query := range querys {
+		timeRange := apiv1.Range{
+			Start: query.Start,
+			End:   query.End,
+			Step:  query.Step,
+		}
 
-	span, ctx := opentracing.StartSpanFromContext(ctx, "alerting.prometheus")
-	span.SetTag("expr", query.Expr)
-	span.SetTag("start_unixnano", int64(query.Start.UnixNano()))
-	span.SetTag("stop_unixnano", int64(query.End.UnixNano()))
-	defer span.Finish()
+		plog.Debug("Sending query", "start", timeRange.Start, "end", timeRange.End, "step", timeRange.Step, "query", query.Expr)
 
-	value, err := client.QueryRange(ctx, query.Expr, timeRange)
+		span, ctx := opentracing.StartSpanFromContext(ctx, "alerting.prometheus")
+		span.SetTag("expr", query.Expr)
+		span.SetTag("start_unixnano", int64(query.Start.UnixNano()))
+		span.SetTag("stop_unixnano", int64(query.End.UnixNano()))
+		defer span.Finish()
 
-	if err != nil {
-		return nil, err
-	}
+		value, err := client.QueryRange(ctx, query.Expr, timeRange)
 
-	queryResult, err := parseResponse(value, query)
-	if err != nil {
-		return nil, err
+		if err != nil {
+			return nil, err
+		}
+
+		queryResult, err := parseResponse(value, query)
+		if err != nil {
+			return nil, err
+		}
+		result.Results[query.RefId] = queryResult
 	}
-	result.Results = queryResult
+
 	return result, nil
 }
 
@@ -140,51 +147,54 @@ func formatLegend(metric model.Metric, query *PrometheusQuery) string {
 	return string(result)
 }
 
-func parseQuery(dsInfo *models.DataSource, queries []*tsdb.Query, queryContext *tsdb.TsdbQuery) (*PrometheusQuery, error) {
-	queryModel := queries[0]
+func parseQuery(dsInfo *models.DataSource, queries []*tsdb.Query, queryContext *tsdb.TsdbQuery) ([]*PrometheusQuery, error) {
+	qs := []*PrometheusQuery{}
+	for _, queryModel := range queries {
+		expr, err := queryModel.Model.Get("expr").String()
+		if err != nil {
+			return nil, err
+		}
 
-	expr, err := queryModel.Model.Get("expr").String()
-	if err != nil {
-		return nil, err
-	}
+		format := queryModel.Model.Get("legendFormat").MustString("")
 
-	format := queryModel.Model.Get("legendFormat").MustString("")
+		start, err := queryContext.TimeRange.ParseFrom()
+		if err != nil {
+			return nil, err
+		}
 
-	start, err := queryContext.TimeRange.ParseFrom()
-	if err != nil {
-		return nil, err
-	}
+		end, err := queryContext.TimeRange.ParseTo()
+		if err != nil {
+			return nil, err
+		}
 
-	end, err := queryContext.TimeRange.ParseTo()
-	if err != nil {
-		return nil, err
-	}
+		dsInterval, err := tsdb.GetIntervalFrom(dsInfo, queryModel.Model, time.Second*15)
+		if err != nil {
+			return nil, err
+		}
 
-	dsInterval, err := tsdb.GetIntervalFrom(dsInfo, queryModel.Model, time.Second*15)
-	if err != nil {
-		return nil, err
-	}
+		intervalFactor := queryModel.Model.Get("intervalFactor").MustInt64(1)
+		interval := intervalCalculator.Calculate(queryContext.TimeRange, dsInterval)
+		step := time.Duration(int64(interval.Value) * intervalFactor)
 
-	intervalFactor := queryModel.Model.Get("intervalFactor").MustInt64(1)
-	interval := intervalCalculator.Calculate(queryContext.TimeRange, dsInterval)
-	step := time.Duration(int64(interval.Value) * intervalFactor)
+		qs = append(qs, &PrometheusQuery{
+			Expr:         expr,
+			Step:         step,
+			LegendFormat: format,
+			Start:        start,
+			End:          end,
+			RefId:        queryModel.RefId,
+		})
+	}
 
-	return &PrometheusQuery{
-		Expr:         expr,
-		Step:         step,
-		LegendFormat: format,
-		Start:        start,
-		End:          end,
-	}, nil
+	return qs, nil
 }
 
-func parseResponse(value model.Value, query *PrometheusQuery) (map[string]*tsdb.QueryResult, error) {
-	queryResults := make(map[string]*tsdb.QueryResult)
+func parseResponse(value model.Value, query *PrometheusQuery) (*tsdb.QueryResult, error) {
 	queryRes := tsdb.NewQueryResult()
 
 	data, ok := value.(model.Matrix)
 	if !ok {
-		return queryResults, fmt.Errorf("Unsupported result format: %s", value.Type().String())
+		return queryRes, fmt.Errorf("Unsupported result format: %s", value.Type().String())
 	}
 
 	for _, v := range data {
@@ -204,6 +214,5 @@ func parseResponse(value model.Value, query *PrometheusQuery) (map[string]*tsdb.
 		queryRes.Series = append(queryRes.Series, &series)
 	}
 
-	queryResults["A"] = queryRes
-	return queryResults, nil
+	return queryRes, nil
 }

pkg/tsdb/prometheus/prometheus_test.go

@@ -60,9 +60,10 @@ func TestPrometheus(t *testing.T) {
 			Convey("with 48h time range", func() {
 				queryContext.TimeRange = tsdb.NewTimeRange("12h", "now")
 
-				model, err := parseQuery(dsInfo, queryModels, queryContext)
-
+				models, err := parseQuery(dsInfo, queryModels, queryContext)
 				So(err, ShouldBeNil)
+
+				model := models[0]
 				So(model.Step, ShouldEqual, time.Second*30)
 			})
 		})
@@ -83,18 +84,22 @@ func TestPrometheus(t *testing.T) {
 			Convey("with 48h time range", func() {
 				queryContext.TimeRange = tsdb.NewTimeRange("48h", "now")
 
-				model, err := parseQuery(dsInfo, queryModels, queryContext)
+				models, err := parseQuery(dsInfo, queryModels, queryContext)
 
 				So(err, ShouldBeNil)
+
+				model := models[0]
 				So(model.Step, ShouldEqual, time.Minute*2)
 			})
 
 			Convey("with 1h time range", func() {
 				queryContext.TimeRange = tsdb.NewTimeRange("1h", "now")
 
-				model, err := parseQuery(dsInfo, queryModels, queryContext)
+				models, err := parseQuery(dsInfo, queryModels, queryContext)
 
 				So(err, ShouldBeNil)
+
+				model := models[0]
 				So(model.Step, ShouldEqual, time.Second*15)
 			})
 		})
@@ -116,9 +121,11 @@ func TestPrometheus(t *testing.T) {
 				Convey("with 48h time range", func() {
 					queryContext.TimeRange = tsdb.NewTimeRange("48h", "now")
 
-					model, err := parseQuery(dsInfo, queryModels, queryContext)
+					models, err := parseQuery(dsInfo, queryModels, queryContext)
 
 					So(err, ShouldBeNil)
+
+					model := models[0]
 					So(model.Step, ShouldEqual, time.Minute*20)
 				})
 			})
@@ -139,9 +146,11 @@ func TestPrometheus(t *testing.T) {
 				Convey("with 48h time range", func() {
 					queryContext.TimeRange = tsdb.NewTimeRange("48h", "now")
 
-					model, err := parseQuery(dsInfo, queryModels, queryContext)
+					models, err := parseQuery(dsInfo, queryModels, queryContext)
 
 					So(err, ShouldBeNil)
+
+					model := models[0]
 					So(model.Step, ShouldEqual, time.Minute*2)
 				})
 			})

pkg/tsdb/prometheus/types.go

@@ -8,4 +8,5 @@ type PrometheusQuery struct {
 	LegendFormat string
 	Start        time.Time
 	End          time.Time
+	RefId        string
 }