Support extra custom query parameters in requests to Prometheus backend

pkg/prometheus/config/config.go

@@ -23,6 +23,10 @@ type Configuration struct {
 	LatencyUnit       string `mapstructure:"latency_unit"`
 	NormalizeCalls    bool   `mapstructure:"normalize_calls"`
 	NormalizeDuration bool   `mapstructure:"normalize_duration"`
+	// ExtraQueryParams is used to provide extra parameters to be appended
+	// to the URL of queries going out to the metrics backend.
+	// is being called for metric read/write
+	ExtraQueryParams map[string]string `mapstructure:"extra_query_parameters"`
 }
 
 func (c *Configuration) Validate() error {

plugin/metricstore/prometheus/metricstore/reader.go

@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"net"
 	"net/http"
+	"net/url"
 	"os"
 	"path/filepath"
 	"strings"
@@ -60,16 +61,32 @@ type (
 		metricDesc        string
 		buildPromQuery    func(p promQueryParams) string
 	}
+
+	promClient struct {
+		api.Client
+		extraParams map[string]string
+	}
 )
 
-// NewMetricsReader returns a new MetricsReader.
-func NewMetricsReader(cfg config.Configuration, logger *zap.Logger, tracer trace.TracerProvider) (*MetricsReader, error) {
-	logger.Info("Creating metrics reader", zap.Any("configuration", cfg))
+// URL decorator enables adding additional query parameters to the request sent to prometheus backend
+func (p promClient) URL(ep string, args map[string]string) *url.URL {
+	u := p.Client.URL(ep, args)
+
+	query := u.Query()
+	for k, v := range p.extraParams {
+		query.Add(k, v)
+	}
+	u.RawQuery = query.Encode()
+
+	return u
+}
 
-	roundTripper, err := getHTTPRoundTripper(&cfg, logger)
+func createPromClient(cfg config.Configuration) (api.Client, error) {
+	roundTripper, err := getHTTPRoundTripper(&cfg)
 	if err != nil {
 		return nil, err
 	}
+
 	client, err := api.NewClient(api.Config{
 		Address:      cfg.ServerURL,
 		RoundTripper: roundTripper,
@@ -78,10 +95,25 @@ func NewMetricsReader(cfg config.Configuration, logger *zap.Logger, tracer trace
 		return nil, fmt.Errorf("failed to initialize prometheus client: %w", err)
 	}
 
+	customClient := promClient{
+		Client:      client,
+		extraParams: cfg.ExtraQueryParams,
+	}
+
+	return customClient, nil
+}
+
+// NewMetricsReader returns a new MetricsReader.
+func NewMetricsReader(cfg config.Configuration, logger *zap.Logger, tracer trace.TracerProvider) (*MetricsReader, error) {
 	const operationLabel = "span_name"
 
+	promClient, err := createPromClient(cfg)
+	if err != nil {
+		return nil, err
+	}
+
 	mr := &MetricsReader{
-		client: promapi.NewAPI(client),
+		client: promapi.NewAPI(promClient),
 		logger: logger,
 		tracer: tracer.Tracer("prom-metrics-reader"),
 
@@ -314,7 +346,7 @@ func logErrorToSpan(span trace.Span, err error) {
 	span.SetStatus(codes.Error, err.Error())
 }
 
-func getHTTPRoundTripper(c *config.Configuration, _ *zap.Logger) (rt http.RoundTripper, err error) {
+func getHTTPRoundTripper(c *config.Configuration) (rt http.RoundTripper, err error) {
 	ctlsConfig, err := c.TLS.LoadTLSConfig(context.Background())
 	if err != nil {
 		return nil, err

plugin/metricstore/prometheus/metricstore/reader_test.go

@@ -11,6 +11,7 @@ import (
 	"net/http/httptest"
 	"net/url"
 	"os"
+	"sort"
 	"strings"
 	"sync/atomic"
 	"testing"
@@ -737,13 +738,12 @@ func TestGetRoundTripperTLSConfig(t *testing.T) {
 		{"tls disabled", false},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
-			logger := zap.NewNop()
 			config := &config.Configuration{
 				ConnectTimeout:           9 * time.Millisecond,
 				TLS:                      configtls.ClientConfig{},
 				TokenOverrideFromContext: true,
 			}
-			rt, err := getHTTPRoundTripper(config, logger)
+			rt, err := getHTTPRoundTripper(config)
 			require.NoError(t, err)
 
 			server := newFakePromServer(t)
@@ -781,7 +781,7 @@ func TestGetRoundTripperTokenFile(t *testing.T) {
 		ConnectTimeout:           time.Second,
 		TokenFilePath:            file.Name(),
 		TokenOverrideFromContext: false,
-	}, nil)
+	})
 	require.NoError(t, err)
 
 	server := newFakePromServer(t)
@@ -815,7 +815,7 @@ func TestGetRoundTripperTokenFromContext(t *testing.T) {
 		ConnectTimeout:           time.Second,
 		TokenFilePath:            file.Name(),
 		TokenOverrideFromContext: true,
-	}, nil)
+	})
 	require.NoError(t, err)
 
 	server := newFakePromServer(t)
@@ -842,7 +842,7 @@ func TestGetRoundTripperTokenError(t *testing.T) {
 
 	_, err := getHTTPRoundTripper(&config.Configuration{
 		TokenFilePath: tokenFilePath,
-	}, nil)
+	})
 	assert.ErrorContains(t, err, "failed to get token from file")
 }
 
@@ -863,6 +863,36 @@ func TestInvalidCertFile(t *testing.T) {
 	assert.Nil(t, reader)
 }
 
+func TestCreatePromClientWithExtraQueryParameters(t *testing.T) {
+	extraParams := map[string]string{
+		"param1": "value1",
+		"param2": "value2",
+	}
+
+	cfg := config.Configuration{
+		ServerURL:        "http://localhost:1234?param1=value0",
+		ExtraQueryParams: extraParams,
+	}
+
+	expParams := map[string][]string{
+		"param1": {"value0", "value1"},
+		"param2": {"value2"},
+	}
+
+	customClient, err := createPromClient(cfg)
+	require.NoError(t, err)
+
+	u := customClient.URL("", nil)
+
+	q := u.Query()
+
+	for k, v := range expParams {
+		sort.Strings(q[k])
+		recV := q[k]
+		require.Equal(t, v, recV)
+	}
+}
+
 func startMockPrometheusServer(t *testing.T, wantPromQlQuery string, wantWarnings []string) *httptest.Server {
 	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if len(wantWarnings) > 0 {