From 5999d115dd0b159193a98d4924e73628404d84e8 Mon Sep 17 00:00:00 2001
From: Richard87 <richard.hagen@gmail.com>
Date: Mon, 14 Oct 2024 16:11:38 +0200
Subject: [PATCH] feat: support any query

---
 Dockerfile   |  1 +
 config.go    |  8 +++++---
 main.go      | 36 +++++++++++++++++++++++++++++-------
 queries.yaml |  4 ++++
 4 files changed, 39 insertions(+), 10 deletions(-)
 create mode 100644 queries.yaml

diff --git a/Dockerfile b/Dockerfile
index 4a8656a..40fa9b0 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -21,6 +21,7 @@ RUN go build -ldflags "-s -w" -a -installsuffix cgo -o /radix-prometheus-proxy
 FROM gcr.io/distroless/static
 
 COPY --from=builder /radix-prometheus-proxy /radix-prometheus-proxy
+COPY queries.yaml /queries.yaml
 
 EXPOSE 8000
 USER 1000
diff --git a/config.go b/config.go
index f9878be..36b33e0 100644
--- a/config.go
+++ b/config.go
@@ -8,9 +8,10 @@ import (
 )
 
 type Config struct {
-	LogLevel  string `envconfig:"LOG_LEVEL" default:"info"`
-	LogPretty bool   `envconfig:"LOG_PRETTY" default:"false"`
-	Port      int    `envconfig:"PORT" default:"8000"`
+	LogLevel    string `envconfig:"LOG_LEVEL" default:"info"`
+	LogPretty   bool   `envconfig:"LOG_PRETTY" default:"false"`
+	Port        int    `envconfig:"PORT" default:"8000"`
+	QueriesFile string `envconfig:"QUERIES" default:"queries.yaml"`
 
 	Prometheus url.URL `envconfig:"PROMETHEUS" required:"true"`
 }
@@ -29,6 +30,7 @@ func MustParseConfig() Config {
 	log.Info().Str("Log level", c.LogLevel).Send()
 	log.Info().Bool("Log pretty", c.LogPretty).Send()
 	log.Info().Stringer("Prometheus", &c.Prometheus).Send()
+	log.Info().Str("Queries file", c.QueriesFile).Send()
 
 	return c
 }
diff --git a/main.go b/main.go
index 08ac284..bc199ca 100644
--- a/main.go
+++ b/main.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"net/http"
+	"os"
 	"os/signal"
 	"time"
 
@@ -12,32 +13,53 @@ import (
 	"github.com/prometheus/common/model"
 	"github.com/rs/zerolog/log"
 	"golang.org/x/sys/unix"
+	"gopkg.in/yaml.v2"
 )
 
-var query = `min_over_time(probe_success{instance="https://api.dev.radix.equinor.com/health/"}[5m])`
-
 func main() {
 	ctx, cancel := signal.NotifyContext(context.Background(), unix.SIGTERM, unix.SIGINT)
 	defer cancel()
 
 	config := MustParseConfig()
-	promController := NewPrometheusController(config)
+	queries := mustParseQueryFile(config.QueriesFile)
+	promController := NewPrometheusController(config.Prometheus.String(), queries)
+
 	router := NewRouter(promController)
 
-	log.Ctx(ctx).Info().Msgf("Starting server on http://localhost:%d/query", config.Port)
+	log.Ctx(ctx).Info().Msgf("Starting server on http://localhost:%d/query/{query}", config.Port)
 	err := Serve(ctx, config.Port, router)
 	log.Err(err).Msg("Terminated")
 }
 
-func NewPrometheusController(config Config) RouteMapper {
-	apiClient, err := prometheusApi.NewClient(prometheusApi.Config{Address: config.Prometheus.String()})
+func mustParseQueryFile(queriesFile string) map[string]string {
+	queries := map[string]string{}
+	content, err := os.ReadFile(queriesFile)
+	if err != nil {
+		log.Fatal().Err(err).Msg("unable to read queries file")
+	}
+
+	if err = yaml.Unmarshal(content, &queries); err != nil {
+		log.Fatal().Err(err).Msg("unable to parse queries file")
+	}
+	return queries
+}
+
+func NewPrometheusController(prometheusUrl string, queries map[string]string) RouteMapper {
+	apiClient, err := prometheusApi.NewClient(prometheusApi.Config{Address: prometheusUrl})
 	if err != nil {
 		log.Fatal().Err(err).Msg("failed to create the Prometheus API client")
 	}
 	api := prometheusV1.NewAPI(apiClient)
 
 	return func(mux *http.ServeMux) {
-		mux.HandleFunc("/query", func(w http.ResponseWriter, r *http.Request) {
+		mux.HandleFunc("/query/{query}", func(w http.ResponseWriter, r *http.Request) {
+			query, ok := queries[r.PathValue("query")]
+			if !ok {
+				w.WriteHeader(http.StatusNotFound)
+				_, _ = w.Write([]byte("Query not found"))
+				return
+			}
+
 			logger := log.Ctx(r.Context())
 			//
 			end := time.Now()
diff --git a/queries.yaml b/queries.yaml
new file mode 100644
index 0000000..525b6df
--- /dev/null
+++ b/queries.yaml
@@ -0,0 +1,4 @@
+dev: min_over_time(probe_success{instance="https://api.dev.radix.equinor.com/health/"}[5m])
+playground: min_over_time(probe_success{instance="https://api.playground.radix.equinor.com/health/"}[5m])
+platform: min_over_time(probe_success{instance="https://api.radix.equinor.com/health/"}[5m])
+c2: min_over_time(probe_success{instance="https://api.c2.radix.equinor.com/health/"}[5m])