From 77e898322142adfff29d039e22da44647b6ed3d9 Mon Sep 17 00:00:00 2001
From: Fabian Reinartz <fab.reinartz@gmail.com>
Date: Thu, 2 Jul 2015 19:12:04 +0200
Subject: [PATCH] promql: add MarshalJSON method for SamplePair

---
 promql/engine.go        | 17 +++++++++
 web/api/legacy/query.go | 84 ++++++++++++++++++++++++++++++++++++++---
 2 files changed, 95 insertions(+), 6 deletions(-)

diff --git a/promql/engine.go b/promql/engine.go
index 3ed8d7ded..9f81b52a8 100644
--- a/promql/engine.go
+++ b/promql/engine.go
@@ -14,6 +14,7 @@
 package promql
 
 import (
+	"encoding/json"
 	"fmt"
 	"math"
 	"runtime"
@@ -48,6 +49,22 @@ type Scalar struct {
 	Timestamp clientmodel.Timestamp   `json:"timestamp"`
 }
 
+// MarshalJSON implements json.Marshaler.
+func (s *Sample) MarshalJSON() ([]byte, error) {
+	v := struct {
+		Metric clientmodel.COWMetric `json:"metric"`
+		Value  metric.SamplePair     `json:"value"`
+	}{
+		Metric: s.Metric,
+		Value: metric.SamplePair{
+			Timestamp: s.Timestamp,
+			Value:     s.Value,
+		},
+	}
+
+	return json.Marshal(&v)
+}
+
 func (s *Scalar) String() string {
 	return fmt.Sprintf("scalar: %v @[%v]", s.Value, s.Timestamp)
 }
diff --git a/web/api/legacy/query.go b/web/api/legacy/query.go
index e1db462f9..de7279005 100644
--- a/web/api/legacy/query.go
+++ b/web/api/legacy/query.go
@@ -17,7 +17,9 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"net/http"
+	"net/url"
 	"sort"
 	"strconv"
 	"time"
@@ -26,7 +28,7 @@ import (
 
 	clientmodel "github.com/prometheus/client_golang/model"
 
-	"github.com/prometheus/prometheus/util/httputil"
+	"github.com/prometheus/prometheus/promql"
 )
 
 // Enables cross-site script calls.
@@ -39,7 +41,7 @@ func setAccessControlHeaders(w http.ResponseWriter) {
 
 func httpJSONError(w http.ResponseWriter, err error, code int) {
 	w.WriteHeader(code)
-	httputil.ErrorJSON(w, err)
+	errorJSON(w, err)
 }
 
 func parseTimestampOrNow(t string, now clientmodel.Timestamp) (clientmodel.Timestamp, error) {
@@ -67,7 +69,7 @@ func (api *API) Query(w http.ResponseWriter, r *http.Request) {
 	setAccessControlHeaders(w)
 	w.Header().Set("Content-Type", "application/json")
 
-	params := httputil.GetQueryParams(r)
+	params := getQueryParams(r)
 	expr := params.Get("expr")
 
 	timestamp, err := parseTimestampOrNow(params.Get("timestamp"), api.Now())
@@ -88,7 +90,35 @@ func (api *API) Query(w http.ResponseWriter, r *http.Request) {
 	}
 	log.Debugf("Instant query: %s\nQuery stats:\n%s\n", expr, query.Stats())
 
-	httputil.RespondJSON(w, res.Value)
+	if vec, ok := res.Value.(promql.Vector); ok {
+		respondJSON(w, plainVec(vec))
+		return
+	}
+
+	respondJSON(w, res.Value)
+}
+
+// plainVec is an indirection that hides the original MarshalJSON methods
+// which does not fit the response format for the legacy API.
+type plainVec promql.Vector
+
+func (pv plainVec) MarshalJSON() ([]byte, error) {
+	type plainSmpl promql.Sample
+
+	v := make([]*plainSmpl, len(pv))
+	for i, sv := range pv {
+		v[i] = (*plainSmpl)(sv)
+	}
+
+	return json.Marshal(&v)
+}
+
+func (pv plainVec) Type() promql.ExprType {
+	return promql.ExprVector
+}
+
+func (pv plainVec) String() string {
+	return ""
 }
 
 // QueryRange handles the /api/query_range endpoint.
@@ -96,7 +126,7 @@ func (api *API) QueryRange(w http.ResponseWriter, r *http.Request) {
 	setAccessControlHeaders(w)
 	w.Header().Set("Content-Type", "application/json")
 
-	params := httputil.GetQueryParams(r)
+	params := getQueryParams(r)
 	expr := params.Get("expr")
 
 	duration, err := parseDuration(params.Get("range"))
@@ -148,7 +178,7 @@ func (api *API) QueryRange(w http.ResponseWriter, r *http.Request) {
 	}
 
 	log.Debugf("Range query: %s\nQuery stats:\n%s\n", expr, query.Stats())
-	httputil.RespondJSON(w, matrix)
+	respondJSON(w, matrix)
 }
 
 // Metrics handles the /api/metrics endpoint.
@@ -166,3 +196,45 @@ func (api *API) Metrics(w http.ResponseWriter, r *http.Request) {
 	}
 	w.Write(resultBytes)
 }
+
+// GetQueryParams calls r.ParseForm and returns r.Form.
+func getQueryParams(r *http.Request) url.Values {
+	r.ParseForm()
+	return r.Form
+}
+
+var jsonFormatVersion = 1
+
+// ErrorJSON writes the given error JSON-formatted to w.
+func errorJSON(w io.Writer, err error) error {
+	data := struct {
+		Type    string `json:"type"`
+		Value   string `json:"value"`
+		Version int    `json:"version"`
+	}{
+		Type:    "error",
+		Value:   err.Error(),
+		Version: jsonFormatVersion,
+	}
+	enc := json.NewEncoder(w)
+	return enc.Encode(data)
+}
+
+// RespondJSON converts the given data value to JSON and writes it to w.
+func respondJSON(w io.Writer, val promql.Value) error {
+	data := struct {
+		Type    string      `json:"type"`
+		Value   interface{} `json:"value"`
+		Version int         `json:"version"`
+	}{
+		Type:    val.Type().String(),
+		Value:   val,
+		Version: jsonFormatVersion,
+	}
+	// TODO(fabxc): Adding MarshalJSON to promql.Values might be a good idea.
+	if sc, ok := val.(*promql.Scalar); ok {
+		data.Value = sc.Value
+	}
+	enc := json.NewEncoder(w)
+	return enc.Encode(data)
+}