Support extra custom query parameters in requests to Prometheus backend

pkg/prometheus/config/config.go

@@ -19,10 +19,11 @@ type Configuration struct {
 	TokenFilePath            string `mapstructure:"token_file_path"`
 	TokenOverrideFromContext bool   `mapstructure:"token_override_from_context"`
 
-	MetricNamespace   string `mapstructure:"metric_namespace"`
-	LatencyUnit       string `mapstructure:"latency_unit"`
-	NormalizeCalls    bool   `mapstructure:"normalize_calls"`
-	NormalizeDuration bool   `mapstructure:"normalize_duration"`
+	MetricNamespace      string            `mapstructure:"metric_namespace"`
+	LatencyUnit          string            `mapstructure:"latency_unit"`
+	NormalizeCalls       bool              `mapstructure:"normalize_calls"`
+	NormalizeDuration    bool              `mapstructure:"normalize_duration"`
+	AdditionalParameters map[string]string `mapstructure:"additional_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
+		params 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.params {
+		query.Set(k, v)
+	}
+	u.RawQuery = query.Encode()
+
+	return u
+}
+
+func createPromClient(logger *zap.Logger, cfg config.Configuration) (api.Client, error) {
 	roundTripper, err := getHTTPRoundTripper(&cfg, logger)
 	if err != nil {
 		return nil, err
 	}
+
 	client, err := api.NewClient(api.Config{
 		Address:      cfg.ServerURL,
 		RoundTripper: roundTripper,
@@ -78,10 +95,27 @@ 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,
+		params: cfg.AdditionalParameters,
+	}
+
+	return customClient, nil
+}
+
+// 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))
+
 	const operationLabel = "span_name"
 
+	promClient, err := createPromClient(logger, 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"),
 

plugin/metricstore/prometheus/metricstore/reader_test.go

@@ -863,6 +863,31 @@ func TestInvalidCertFile(t *testing.T) {
 	assert.Nil(t, reader)
 }
 
+func TestCreatePromClientWithAdditionalParameters(t *testing.T) {
+	expParams := map[string]string{
+		"param1": "value1",
+		"param2": "value2",
+	}
+
+	logger := zap.NewNop()
+	cfg := config.Configuration{
+		ServerURL:            "http://localhost:1234",
+		AdditionalParameters: expParams,
+	}
+
+	customClient, err := createPromClient(logger, cfg)
+	require.NoError(t, err)
+
+	u := customClient.URL("", nil)
+
+	q := u.Query()
+
+	for k, v := range expParams {
+		recV := q.Get(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 {

plugin/metricstore/prometheus/options.go

@@ -51,10 +51,11 @@ func DefaultConfig() config.Configuration {
 		ServerURL:      defaultServerURL,
 		ConnectTimeout: defaultConnectTimeout,
 
-		MetricNamespace:   defaultMetricNamespace,
-		LatencyUnit:       defaultLatencyUnit,
-		NormalizeCalls:    defaultNormalizeCalls,
-		NormalizeDuration: defaultNormalizeCalls,
+		MetricNamespace:      defaultMetricNamespace,
+		LatencyUnit:          defaultLatencyUnit,
+		NormalizeCalls:       defaultNormalizeCalls,
+		NormalizeDuration:    defaultNormalizeCalls,
+		AdditionalParameters: nil,
 	}
 }