From 0a515c3bd622e03f905315a8d47acfd85e8819fc Mon Sep 17 00:00:00 2001
From: Thomas Jackson <jacksontj.89@gmail.com>
Date: Fri, 8 Sep 2023 21:35:36 -0700
Subject: [PATCH] Add error wrapping for the servergroup and APIClient

Fixes #301
---
 pkg/promclient/error_wrap.go   | 85 ++++++++++++++++++++++++++++++++++
 pkg/proxystorage/proxy.go      |  1 +
 pkg/servergroup/config.go      |  1 +
 pkg/servergroup/servergroup.go |  8 ++--
 4 files changed, 92 insertions(+), 3 deletions(-)
 create mode 100644 pkg/promclient/error_wrap.go

diff --git a/pkg/promclient/error_wrap.go b/pkg/promclient/error_wrap.go
new file mode 100644
index 000000000..c048516bc
--- /dev/null
+++ b/pkg/promclient/error_wrap.go
@@ -0,0 +1,85 @@
+package promclient
+
+import (
+	"context"
+	"time"
+
+	"github.com/pkg/errors"
+	v1 "github.com/prometheus/client_golang/api/prometheus/v1"
+	"github.com/prometheus/common/model"
+	"github.com/prometheus/prometheus/model/labels"
+)
+
+type ErrorWrap struct {
+	A   API
+	Msg string
+}
+
+func (e *ErrorWrap) LabelNames(ctx context.Context, matchers []string, startTime time.Time, endTime time.Time) (v []string, w v1.Warnings, err error) {
+	defer func() {
+		if err != nil {
+			err = errors.Wrap(err, e.Msg)
+		}
+	}()
+	return e.A.LabelNames(ctx, matchers, startTime, endTime)
+}
+
+// LabelValues performs a query for the values of the given label.
+func (e *ErrorWrap) LabelValues(ctx context.Context, label string, matchers []string, startTime time.Time, endTime time.Time) (v model.LabelValues, w v1.Warnings, err error) {
+	defer func() {
+		if err != nil {
+			err = errors.Wrap(err, e.Msg)
+		}
+	}()
+	return e.A.LabelValues(ctx, label, matchers, startTime, endTime)
+}
+
+// Query performs a query for the given time.
+func (e *ErrorWrap) Query(ctx context.Context, query string, ts time.Time) (v model.Value, w v1.Warnings, err error) {
+	defer func() {
+		if err != nil {
+			err = errors.Wrap(err, e.Msg)
+		}
+	}()
+	return e.A.Query(ctx, query, ts)
+}
+
+// QueryRange performs a query for the given range.
+func (e *ErrorWrap) QueryRange(ctx context.Context, query string, r v1.Range) (v model.Value, w v1.Warnings, err error) {
+	defer func() {
+		if err != nil {
+			err = errors.Wrap(err, e.Msg)
+		}
+	}()
+	return e.A.QueryRange(ctx, query, r)
+}
+
+// Series finds series by label matchers.
+func (e *ErrorWrap) Series(ctx context.Context, matches []string, startTime time.Time, endTime time.Time) (v []model.LabelSet, w v1.Warnings, err error) {
+	defer func() {
+		if err != nil {
+			err = errors.Wrap(err, e.Msg)
+		}
+	}()
+	return e.A.Series(ctx, matches, startTime, endTime)
+}
+
+// GetValue loads the raw data for a given set of matchers in the time range
+func (e *ErrorWrap) GetValue(ctx context.Context, start, end time.Time, matchers []*labels.Matcher) (v model.Value, w v1.Warnings, err error) {
+	defer func() {
+		if err != nil {
+			err = errors.Wrap(err, e.Msg)
+		}
+	}()
+	return e.A.GetValue(ctx, start, end, matchers)
+}
+
+// Metadata returns metadata about metrics currently scraped by the metric name.
+func (e *ErrorWrap) Metadata(ctx context.Context, metric, limit string) (v map[string][]v1.Metadata, err error) {
+	defer func() {
+		if err != nil {
+			err = errors.Wrap(err, e.Msg)
+		}
+	}()
+	return e.A.Metadata(ctx, metric, limit)
+}
diff --git a/pkg/proxystorage/proxy.go b/pkg/proxystorage/proxy.go
index 6ffad44a9..f2453ddbb 100644
--- a/pkg/proxystorage/proxy.go
+++ b/pkg/proxystorage/proxy.go
@@ -105,6 +105,7 @@ func (p *ProxyStorage) ApplyConfig(c *proxyconfig.Config) error {
 		cfg: c,
 	}
 	for i, sgCfg := range c.ServerGroups {
+		sgCfg.Ordinal = i
 		tmp, err := servergroup.NewServerGroup()
 		if err != nil {
 			failed = true
diff --git a/pkg/servergroup/config.go b/pkg/servergroup/config.go
index d6faedaa8..4f5bd0d02 100644
--- a/pkg/servergroup/config.go
+++ b/pkg/servergroup/config.go
@@ -35,6 +35,7 @@ const (
 // Config is the configuration for a ServerGroup that promxy will talk to.
 // This is where the vast majority of options exist.
 type Config struct {
+	Ordinal int `yaml:"-"`
 	// RemoteRead directs promxy to load RAW data (meaning matrix selectors such as `foo[1h]`)
 	// through the RemoteRead API on prom.
 	// Pros:
diff --git a/pkg/servergroup/servergroup.go b/pkg/servergroup/servergroup.go
index f36a583e0..5150c46f0 100644
--- a/pkg/servergroup/servergroup.go
+++ b/pkg/servergroup/servergroup.go
@@ -270,7 +270,8 @@ func (s *ServerGroup) loadTargetGroupMap(targetGroupMap map[string][]*targetgrou
 					}
 				}
 
-				apiClients = append(apiClients, apiClient)
+				// Add wrap for the specific target, and add to the list
+				apiClients = append(apiClients, &promclient.ErrorWrap{apiClient, "error in target=" + u.String()})
 			}
 		}
 	}
@@ -286,8 +287,9 @@ func (s *ServerGroup) loadTargetGroupMap(targetGroupMap map[string][]*targetgrou
 	}
 
 	newState := &ServerGroupState{
-		Targets:   targets,
-		apiClient: apiClient,
+		Targets: targets,
+		// Add error wrap for this specific servergroup
+		apiClient: &promclient.ErrorWrap{apiClient, fmt.Sprintf("error in servergroup ord=%d", s.Cfg.Ordinal)},
 		ctx:       ctx,
 		ctxCancel: ctxCancel,
 	}