Modify respond method to return proper label values

Author: Blagoj Atanasovski

pkg/api/label_values.go

@@ -6,6 +6,7 @@ import (
 	"github.com/NYTimes/gziphandler"
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/common/route"
+	"github.com/timescale/timescale-prometheus/pkg/promql"
 	"github.com/timescale/timescale-prometheus/pkg/query"
 	"math"
 	"net/http"
@@ -26,13 +27,16 @@ func LabelValues(queriable *query.Queryable) http.Handler {
 			respondError(w, http.StatusInternalServerError, err, "internal")
 			return
 		}
-		names, warnings, err := querier.LabelValues(name)
+		var values labelsValue
+		values, warnings, err := querier.LabelValues(name)
 		if err != nil {
 			respondError(w, http.StatusInternalServerError, err, "internal")
 			return
 		}
 
-		respond(w, names, warnings)
+		respondLabels(w, &promql.Result{
+			Value: values,
+		}, warnings)
 	})
 
 	return gziphandler.GzipHandler(hf)

pkg/api/labels.go

@@ -2,12 +2,27 @@ package api
 
 import (
 	"context"
+	"encoding/json"
 	"github.com/NYTimes/gziphandler"
+	"github.com/prometheus/prometheus/promql/parser"
+	"github.com/prometheus/prometheus/storage"
+	"github.com/timescale/timescale-prometheus/pkg/promql"
 	"github.com/timescale/timescale-prometheus/pkg/query"
 	"math"
 	"net/http"
+	"strings"
 )
 
+type labelsValue []string
+
+func (l labelsValue) Type() parser.ValueType {
+	return parser.ValueTypeNone
+}
+
+func (l labelsValue) String() string {
+	return strings.Join(l, "\n")
+}
+
 func Labels(queriable *query.Queryable) http.Handler {
 	hf := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Access-Control-Allow-Origin", "*")
@@ -17,14 +32,25 @@ func Labels(queriable *query.Queryable) http.Handler {
 			respondError(w, http.StatusInternalServerError, err, "internal")
 			return
 		}
+		var names labelsValue
 		names, warnings, err := querier.LabelNames()
 		if err != nil {
 			respondError(w, http.StatusInternalServerError, err, "internal")
 			return
 		}
-
-		respond(w, names, warnings)
+		respondLabels(w, &promql.Result{
+			Value: names,
+		}, warnings)
 	})
 
 	return gziphandler.GzipHandler(hf)
 }
+
+func respondLabels(w http.ResponseWriter, res *promql.Result, warnings storage.Warnings) {
+	setHeaders(w, res, warnings)
+	resp := &response{
+		Status: "success",
+		Data:   res.Value,
+	}
+	_ = json.NewEncoder(w).Encode(resp)
+}

pkg/api/query.go

@@ -14,7 +14,6 @@ import (
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/promql/parser"
 	"github.com/prometheus/prometheus/storage"
-	"github.com/prometheus/prometheus/util/stats"
 	"github.com/timescale/timescale-prometheus/pkg/log"
 	"github.com/timescale/timescale-prometheus/pkg/promql"
 	"github.com/timescale/timescale-prometheus/pkg/query"
@@ -74,19 +73,26 @@ func Query(queryEngine *promql.Engine, queryable *query.Queryable) http.Handler
 			return
 		}
 
-		respond(w, res, res.Warnings)
+		respondQuery(w, res, res.Warnings)
 	})
 
 	return gziphandler.GzipHandler(hf)
 }
 
-func respond(w http.ResponseWriter, res *promql.Result, warnings storage.Warnings) {
+func setHeaders(w http.ResponseWriter, res *promql.Result, warnings storage.Warnings) {
 	w.Header().Set("Content-Type", "application/json")
 	if warnings != nil && len(warnings) > 0 {
 		w.Header().Set("Cache-Control", "no-store")
 	}
-	w.WriteHeader(http.StatusOK)
+	if res != nil && res.Value != nil {
+		w.WriteHeader(http.StatusOK)
+	} else {
+		w.WriteHeader(http.StatusNoContent)
+	}
+}
 
+func respondQuery(w http.ResponseWriter, res *promql.Result, warnings storage.Warnings) {
+	setHeaders(w, res, warnings)
 	switch resVal := res.Value.(type) {
 	case promql.Vector:
 		warnings := make([]string, 0, len(res.Warnings))
@@ -113,7 +119,6 @@ func respond(w http.ResponseWriter, res *promql.Result, warnings storage.Warning
 		}
 		_ = json.NewEncoder(w).Encode(resp)
 	}
-
 }
 
 func respondError(w http.ResponseWriter, status int, err error, errType string) {

pkg/api/query_range.go

@@ -2,7 +2,6 @@ package api
 
 import (
 	"context"
-	"github.com/prometheus/prometheus/util/stats"
 	"net/http"
 
 	"github.com/NYTimes/gziphandler"
@@ -101,15 +100,7 @@ func QueryRange(queryEngine *promql.Engine, queriable *query.Queryable) http.Han
 			return
 		}
 
-		var qs *stats.QueryStats
-		if r.FormValue("stats") != "" {
-			qs = stats.NewQueryStats(qry.Stats())
-		}
-		respond(w, &queryData{
-			Result:     res.Value,
-			ResultType: res.Value.Type(),
-			Stats:      qs,
-		}, res.Warnings)
+		respondQuery(w, res, res.Warnings)
 	})
 
 	return gziphandler.GzipHandler(hf)

pkg/pgmodel/end_to_end_tests/main_test.go

@@ -87,7 +87,7 @@ func TestMain(m *testing.M) {
 
 func withDB(t testing.TB, DBName string, f func(db *pgxpool.Pool, t testing.TB)) {
 	testhelpers.WithDB(t, DBName, testhelpers.NoSuperuser, func(db *pgxpool.Pool, t testing.TB, connectURL string) {
-		performMigrate(t, DBName, connectURL)
+		performMigrate(t, connectURL)
 
 		//need to get a new pool after the Migrate to catch any GUC changes made during Migrate
 		db, err := pgxpool.Connect(context.Background(), connectURL)
@@ -101,7 +101,7 @@ func withDB(t testing.TB, DBName string, f func(db *pgxpool.Pool, t testing.TB))
 	})
 }
 
-func performMigrate(t testing.TB, DBName string, connectURL string) {
+func performMigrate(t testing.TB, connectURL string) {
 	dbStd, err := sql.Open("pgx", connectURL)
 	defer func() {
 		err := dbStd.Close()

pkg/pgmodel/end_to_end_tests/migrate_test.go

@@ -41,7 +41,7 @@ func TestMigrateTwice(t *testing.T) {
 		t.Skip("skipping integration test")
 	}
 	testhelpers.WithDB(t, *testDatabase, testhelpers.NoSuperuser, func(db *pgxpool.Pool, t testing.TB, connectURL string) {
-		performMigrate(t, *testDatabase, connectURL)
-		performMigrate(t, *testDatabase, connectURL)
+		performMigrate(t, connectURL)
+		performMigrate(t, connectURL)
 	})
 }