From 3e73be8d2e34ec09f1b75516b23e576628f2a376 Mon Sep 17 00:00:00 2001
From: bergquist <carl.bergquist@gmail.com>
Date: Wed, 21 Sep 2016 11:31:56 +0200
Subject: [PATCH] feat(prometheus): improve error handling

---
 pkg/tsdb/prometheus/prometheus.go      | 46 +++++++++++++++++++-------
 pkg/tsdb/prometheus/prometheus_test.go |  2 +-
 2 files changed, 35 insertions(+), 13 deletions(-)

diff --git a/pkg/tsdb/prometheus/prometheus.go b/pkg/tsdb/prometheus/prometheus.go
index d835d5a822c..b32a99220b2 100644
--- a/pkg/tsdb/prometheus/prometheus.go
+++ b/pkg/tsdb/prometheus/prometheus.go
@@ -49,14 +49,17 @@ func (e *PrometheusExecutor) Execute(queries tsdb.QuerySlice, queryContext *tsdb
 
 	client, err := e.getClient()
 	if err != nil {
-		result.Error = err
-		return result
+		return resultWithError(result, err)
 	}
 
 	from, _ := queryContext.TimeRange.FromTime()
 	to, _ := queryContext.TimeRange.ToTime()
 
-	query := parseQuery(queries)
+	query, err := parseQuery(queries)
+
+	if err != nil {
+		return resultWithError(result, err)
+	}
 
 	timeRange := prometheus.Range{
 		Start: from,
@@ -67,15 +70,14 @@ func (e *PrometheusExecutor) Execute(queries tsdb.QuerySlice, queryContext *tsdb
 	value, err := client.QueryRange(context.Background(), query.Expr, timeRange)
 
 	if err != nil {
-		result.Error = err
-		return result
+		return resultWithError(result, err)
 	}
 
 	result.QueryResults = parseResponse(value, query)
 	return result
 }
 
-func formatLegend(metric pmodel.Metric, query PrometheusQuery) string {
+func formatLegend(metric pmodel.Metric, query *PrometheusQuery) string {
 	r, _ := regexp.Compile(`\{\{\s*(.+?)\s*\}\}`)
 
 	result := r.ReplaceAllFunc([]byte(query.LegendFormat), func(in []byte) []byte {
@@ -90,17 +92,32 @@ func formatLegend(metric pmodel.Metric, query PrometheusQuery) string {
 	return string(result)
 }
 
-func parseQuery(queries tsdb.QuerySlice) PrometheusQuery {
+func parseQuery(queries tsdb.QuerySlice) (*PrometheusQuery, error) {
 	queryModel := queries[0]
 
-	return PrometheusQuery{
-		Expr:         queryModel.Model.Get("expr").MustString(),
-		Step:         time.Second * time.Duration(queryModel.Model.Get("step").MustInt64(1)),
-		LegendFormat: queryModel.Model.Get("legendFormat").MustString(),
+	expr, err := queryModel.Model.Get("expr").String()
+	if err != nil {
+		return nil, err
 	}
+
+	step, err := queryModel.Model.Get("step").Int64()
+	if err != nil {
+		return nil, err
+	}
+
+	format, err := queryModel.Model.Get("legendFormat").String()
+	if err != nil {
+		return nil, err
+	}
+
+	return &PrometheusQuery{
+		Expr:         expr,
+		Step:         time.Second * time.Duration(step),
+		LegendFormat: format,
+	}, nil
 }
 
-func parseResponse(value pmodel.Value, query PrometheusQuery) map[string]*tsdb.QueryResult {
+func parseResponse(value pmodel.Value, query *PrometheusQuery) map[string]*tsdb.QueryResult {
 	queryResults := make(map[string]*tsdb.QueryResult)
 	queryRes := &tsdb.QueryResult{}
 
@@ -123,3 +140,8 @@ func parseResponse(value pmodel.Value, query PrometheusQuery) map[string]*tsdb.Q
 	queryResults["A"] = queryRes
 	return queryResults
 }
+
+func resultWithError(result *tsdb.BatchResult, err error) *tsdb.BatchResult {
+	result.Error = err
+	return result
+}
diff --git a/pkg/tsdb/prometheus/prometheus_test.go b/pkg/tsdb/prometheus/prometheus_test.go
index 6edb6260a09..f7489ae9afc 100644
--- a/pkg/tsdb/prometheus/prometheus_test.go
+++ b/pkg/tsdb/prometheus/prometheus_test.go
@@ -16,7 +16,7 @@ func TestPrometheus(t *testing.T) {
 				p.LabelName("device"): p.LabelValue("mobile"),
 			}
 
-			query := PrometheusQuery{
+			query := &PrometheusQuery{
 				LegendFormat: "legend {{app}} {{device}} {{broken}}",
 			}