Merge pull request #571 from pyrra-dev/api-prometheus

Add API for querying Prometheus "directly"

Author: metalmatze

go.mod

@@ -17,6 +17,7 @@ require (
 	github.com/prometheus/common v0.39.0
 	github.com/prometheus/prometheus v0.41.0
 	github.com/stretchr/testify v1.8.1
+	golang.org/x/exp v0.0.0-20221212164502-fae10dda9338
 	golang.org/x/net v0.4.0
 	google.golang.org/protobuf v1.28.1
 	k8s.io/api v0.26.0
@@ -68,7 +69,6 @@ require (
 	go.uber.org/goleak v1.2.0 // indirect
 	go.uber.org/multierr v1.6.0 // indirect
 	go.uber.org/zap v1.24.0 // indirect
-	golang.org/x/exp v0.0.0-20221212164502-fae10dda9338 // indirect
 	golang.org/x/oauth2 v0.3.0 // indirect
 	golang.org/x/sys v0.3.0 // indirect
 	golang.org/x/term v0.3.0 // indirect

go.sum

@@ -26,7 +26,7 @@ import (
 	"github.com/go-kit/log/level"
 	connectprometheus "github.com/polarsignals/connect-go-prometheus"
 	"github.com/prometheus/client_golang/api"
-	prometheusv1 "github.com/prometheus/client_golang/api/prometheus/v1"
+	prometheusapiv1 "github.com/prometheus/client_golang/api/prometheus/v1"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/collectors"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
@@ -40,6 +40,7 @@ import (
 
 	objectivesv1alpha1 "github.com/pyrra-dev/pyrra/proto/objectives/v1alpha1"
 	"github.com/pyrra-dev/pyrra/proto/objectives/v1alpha1/objectivesv1alpha1connect"
+	"github.com/pyrra-dev/pyrra/proto/prometheus/v1/prometheusv1connect"
 	"github.com/pyrra-dev/pyrra/slo"
 )
 
@@ -185,7 +186,7 @@ func cmdAPI(logger log.Logger, reg *prometheus.Registry, promClient api.Client,
 	defer cache.Close()
 	promAPI := &promCache{
 		api: &promLogger{
-			api:    prometheusv1.NewAPI(promClient),
+			api:    prometheusapiv1.NewAPI(promClient),
 			logger: logger,
 		},
 		cache: cache,
@@ -204,22 +205,32 @@ func cmdAPI(logger log.Logger, reg *prometheus.Registry, promClient api.Client,
 
 	r.Route(routePrefix, func(r chi.Router) {
 		objectiveService := &objectiveServer{
-			logger:  logger,
+			logger:  log.WithPrefix(logger, "service", "objective"),
 			promAPI: promAPI,
 			client: objectivesv1alpha1connect.NewObjectiveBackendServiceClient(
 				http.DefaultClient,
 				apiURL.String(),
 				connect.WithInterceptors(prometheusInterceptor),
 			),
 		}
+
 		objectivePath, objectiveHandler := objectivesv1alpha1connect.NewObjectiveServiceHandler(
 			objectiveService,
 			connect.WithInterceptors(prometheusInterceptor),
 		)
+
+		prometheusService := &prometheusServer{
+			logger:  log.WithPrefix(logger, "service", "prometheus"),
+			promAPI: promAPI,
+		}
+		prometheusPath, prometheusHandler := prometheusv1connect.NewPrometheusServiceHandler(prometheusService)
+
 		if routePrefix != "/" {
 			r.Mount(objectivePath, http.StripPrefix(routePrefix, objectiveHandler))
+			r.Mount(prometheusPath, http.StripPrefix(routePrefix, prometheusHandler))
 		} else {
 			r.Mount(objectivePath, objectiveHandler)
+			r.Mount(prometheusPath, prometheusHandler)
 		}
 
 		r.Handle("/metrics", promhttp.HandlerFor(reg, promhttp.HandlerOpts{}))
@@ -333,17 +344,17 @@ func (c *thanosClient) Do(ctx context.Context, r *http.Request) (*http.Response,
 
 type prometheusAPI interface {
 	// Query performs a query for the given time.
-	Query(ctx context.Context, query string, ts time.Time, opts ...prometheusv1.Option) (model.Value, prometheusv1.Warnings, error)
+	Query(ctx context.Context, query string, ts time.Time, opts ...prometheusapiv1.Option) (model.Value, prometheusapiv1.Warnings, error)
 	// QueryRange performs a query for the given range.
-	QueryRange(ctx context.Context, query string, r prometheusv1.Range, opts ...prometheusv1.Option) (model.Value, prometheusv1.Warnings, error)
+	QueryRange(ctx context.Context, query string, r prometheusapiv1.Range, opts ...prometheusapiv1.Option) (model.Value, prometheusapiv1.Warnings, error)
 }
 
 type promLogger struct {
 	api    prometheusAPI
 	logger log.Logger
 }
 
-func (l *promLogger) Query(ctx context.Context, query string, ts time.Time, opts ...prometheusv1.Option) (model.Value, prometheusv1.Warnings, error) {
+func (l *promLogger) Query(ctx context.Context, query string, ts time.Time, opts ...prometheusapiv1.Option) (model.Value, prometheusapiv1.Warnings, error) {
 	level.Debug(l.logger).Log(
 		"msg", "running instant query",
 		"query", query,
@@ -352,7 +363,7 @@ func (l *promLogger) Query(ctx context.Context, query string, ts time.Time, opts
 	return l.api.Query(ctx, query, ts, opts...)
 }
 
-func (l *promLogger) QueryRange(ctx context.Context, query string, r prometheusv1.Range, opts ...prometheusv1.Option) (model.Value, prometheusv1.Warnings, error) {
+func (l *promLogger) QueryRange(ctx context.Context, query string, r prometheusapiv1.Range, opts ...prometheusapiv1.Option) (model.Value, prometheusapiv1.Warnings, error) {
 	level.Debug(l.logger).Log(
 		"msg", "running range query",
 		"query", query,
@@ -383,7 +394,7 @@ func contextGetPromCache(ctx context.Context) time.Duration {
 	return 0
 }
 
-func (p *promCache) Query(ctx context.Context, query string, ts time.Time) (model.Value, prometheusv1.Warnings, error) {
+func (p *promCache) Query(ctx context.Context, query string, ts time.Time) (model.Value, prometheusapiv1.Warnings, error) {
 	if value, exists := p.cache.Get(query); exists {
 		return value.(model.Value), nil, nil
 	}
@@ -410,7 +421,7 @@ func (p *promCache) Query(ctx context.Context, query string, ts time.Time) (mode
 	return value, warnings, nil
 }
 
-func (p *promCache) QueryRange(ctx context.Context, query string, r prometheusv1.Range) (model.Value, prometheusv1.Warnings, error) {
+func (p *promCache) QueryRange(ctx context.Context, query string, r prometheusapiv1.Range) (model.Value, prometheusapiv1.Warnings, error) {
 	// Get the full time range of this query from start to end.
 	// We round by 10s to adjust for small imperfections to increase cache hits.
 	timeRange := r.End.Sub(r.Start).Round(10 * time.Second)
@@ -477,6 +488,61 @@ func (s *objectiveServer) List(ctx context.Context, req *connect.Request[objecti
 		return nil, err
 	}
 
+	groupingMatchers := map[string]*labels.Matcher{}
+	if req.Msg.Grouping != "" {
+		ms, err := parser.ParseMetricSelector(req.Msg.Grouping)
+		if err != nil {
+			return nil, connect.NewError(connect.CodeInvalidArgument, err)
+		}
+		for _, m := range ms {
+			groupingMatchers[m.Name] = m
+		}
+	}
+
+	for _, o := range resp.Msg.Objectives {
+		oi := objectivesv1alpha1.ToInternal(o)
+
+		// If specific grouping was selected we need to merge the label matchers for the queries.
+		if len(groupingMatchers) > 0 {
+			if oi.Indicator.Ratio != nil {
+				for _, m := range oi.Indicator.Ratio.Errors.LabelMatchers {
+					if rm, replace := groupingMatchers[m.Name]; replace {
+						m.Type = rm.Type
+						m.Value = rm.Value
+					}
+				}
+				for _, m := range oi.Indicator.Ratio.Total.LabelMatchers {
+					if rm, replace := groupingMatchers[m.Name]; replace {
+						m.Type = rm.Type
+						m.Value = rm.Value
+					}
+				}
+			}
+			if oi.Indicator.Latency != nil {
+				for _, m := range oi.Indicator.Latency.Success.LabelMatchers {
+					if rm, replace := groupingMatchers[m.Name]; replace {
+						m.Type = rm.Type
+						m.Value = rm.Value
+					}
+				}
+				for _, m := range oi.Indicator.Latency.Total.LabelMatchers {
+					if rm, replace := groupingMatchers[m.Name]; replace {
+						m.Type = rm.Type
+						m.Value = rm.Value
+					}
+				}
+			}
+		}
+
+		o.Queries = &objectivesv1alpha1.Queries{
+			CountTotal:       oi.QueryTotal(oi.Window),
+			CountErrors:      oi.QueryErrors(oi.Window),
+			GraphErrorBudget: oi.QueryErrorBudget(),
+			GraphRequests:    oi.RequestRange(5 * time.Minute),
+			GraphErrors:      oi.ErrorsRange(5 * time.Minute),
+		}
+	}
+
 	return connect.NewResponse(&objectivesv1alpha1.ListResponse{
 		Objectives: resp.Msg.Objectives,
 	}), nil
@@ -647,7 +713,7 @@ func (s *objectiveServer) GraphErrorBudget(ctx context.Context, req *connect.Req
 	step := end.Sub(start) / 1000
 
 	query := objective.QueryErrorBudget()
-	value, _, err := s.promAPI.QueryRange(contextSetPromCache(ctx, 15*time.Second), query, prometheusv1.Range{
+	value, _, err := s.promAPI.QueryRange(contextSetPromCache(ctx, 15*time.Second), query, prometheusapiv1.Range{
 		Start: start,
 		End:   end,
 		Step:  step,
@@ -1022,7 +1088,7 @@ func (s *objectiveServer) GraphRate(ctx context.Context, req *connect.Request[ob
 
 	query := objective.RequestRange(timeRange)
 
-	value, _, err := s.promAPI.QueryRange(contextSetPromCache(ctx, cacheDuration), query, prometheusv1.Range{
+	value, _, err := s.promAPI.QueryRange(contextSetPromCache(ctx, cacheDuration), query, prometheusapiv1.Range{
 		Start: start,
 		End:   end,
 		Step:  step,
@@ -1117,7 +1183,7 @@ func (s *objectiveServer) GraphErrors(ctx context.Context, req *connect.Request[
 	cacheDuration := rangeCache(start, end)
 
 	query := objective.ErrorsRange(timeRange)
-	value, _, err := s.promAPI.QueryRange(contextSetPromCache(ctx, cacheDuration), query, prometheusv1.Range{
+	value, _, err := s.promAPI.QueryRange(contextSetPromCache(ctx, cacheDuration), query, prometheusapiv1.Range{
 		Start: start,
 		End:   end,
 		Step:  step,
@@ -1235,7 +1301,7 @@ func (s *objectiveServer) GraphDuration(ctx context.Context, req *connect.Reques
 	for _, percentile := range objectivePercentiles {
 		if objective.Target >= percentile {
 			query := objective.DurationRange(timeRange, percentile)
-			value, _, err := s.promAPI.QueryRange(contextSetPromCache(ctx, cacheDuration), query, prometheusv1.Range{
+			value, _, err := s.promAPI.QueryRange(contextSetPromCache(ctx, cacheDuration), query, prometheusapiv1.Range{
 				Start: start,
 				End:   end,
 				Step:  step,

main.go

@@ -0,0 +1,126 @@
+package main
+
+import (
+	"context"
+	"time"
+
+	"github.com/bufbuild/connect-go"
+	"github.com/go-kit/log"
+	prometheusapiv1 "github.com/prometheus/client_golang/api/prometheus/v1"
+	"github.com/prometheus/common/model"
+
+	v1 "github.com/pyrra-dev/pyrra/proto/prometheus/v1"
+)
+
+type prometheusServer struct {
+	logger  log.Logger
+	promAPI *promCache
+}
+
+func (ps *prometheusServer) Query(ctx context.Context, req *connect.Request[v1.QueryRequest]) (*connect.Response[v1.QueryResponse], error) {
+	value, warnings, err := ps.promAPI.Query(ctx, req.Msg.Query, time.Unix(req.Msg.Time, 0))
+	if err != nil {
+		return nil, err
+	}
+
+	switch v := value.(type) {
+	case *model.String:
+		return connect.NewResponse(&v1.QueryResponse{
+			Warnings: warnings,
+			Options: &v1.QueryResponse_String_{
+				String_: &v1.String{
+					Time:  v.Timestamp.Unix(),
+					Value: v.Value,
+				},
+			},
+		}), err
+	case *model.Scalar:
+		return connect.NewResponse(&v1.QueryResponse{
+			Warnings: warnings,
+			Options: &v1.QueryResponse_Scalar{
+				Scalar: &v1.SamplePair{
+					Time:  v.Timestamp.Unix(),
+					Value: float64(v.Value),
+				},
+			},
+		}), nil
+	case model.Vector:
+		vector := convertVector(v)
+		return connect.NewResponse(&v1.QueryResponse{
+			Warnings: warnings,
+			Options: &v1.QueryResponse_Vector{
+				Vector: vector,
+			},
+		}), nil
+	}
+
+	return connect.NewResponse(&v1.QueryResponse{
+		Warnings: warnings,
+		Options:  nil,
+	}), nil
+}
+
+func (ps *prometheusServer) QueryRange(ctx context.Context, req *connect.Request[v1.QueryRangeRequest]) (*connect.Response[v1.QueryRangeResponse], error) {
+	value, warnings, err := ps.promAPI.QueryRange(ctx, req.Msg.GetQuery(), prometheusapiv1.Range{
+		Start: time.Unix(req.Msg.GetStart(), 0),
+		End:   time.Unix(req.Msg.GetEnd(), 0),
+		Step:  time.Duration(req.Msg.GetStep()) * time.Second,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	switch v := value.(type) {
+	case model.Matrix:
+		return connect.NewResponse(&v1.QueryRangeResponse{
+			Warnings: warnings,
+			Options: &v1.QueryRangeResponse_Matrix{
+				Matrix: convertMatrix(v),
+			},
+		}), nil
+	}
+
+	return connect.NewResponse(&v1.QueryRangeResponse{
+		Warnings: warnings,
+	}), nil
+}
+
+func convertVector(in model.Vector) *v1.Vector {
+	samples := make([]*v1.Sample, 0, len(in))
+	for _, si := range in {
+		samples = append(samples,
+			&v1.Sample{
+				Metric: convertLabelSet(si.Metric),
+				Time:   si.Timestamp.Unix(),
+				Value:  float64(si.Value),
+			},
+		)
+	}
+	return &v1.Vector{Samples: samples}
+}
+
+func convertLabelSet(metric model.Metric) map[string]string {
+	out := make(map[string]string, len(metric))
+	for name, value := range metric {
+		out[string(name)] = string(value)
+	}
+	return out
+}
+
+func convertMatrix(in model.Matrix) *v1.Matrix {
+	samples := make([]*v1.SampleStream, in.Len())
+	for i, sampleStream := range in {
+		values := make([]*v1.SamplePair, len(sampleStream.Values))
+		for j, pair := range sampleStream.Values {
+			values[j] = &v1.SamplePair{
+				Time:  pair.Timestamp.Unix(),
+				Value: float64(pair.Value),
+			}
+		}
+		samples[i] = &v1.SampleStream{
+			Metric: convertLabelSet(sampleStream.Metric),
+			Values: values,
+		}
+	}
+	return &v1.Matrix{Samples: samples}
+}