Merge pull request #432 from discordianfish/prometheus-request-headers

Support setting headers on requests to Prometheus

cmd/adapter/adapter.go

@@ -25,6 +25,7 @@ import (
 	"net/http"
 	"net/url"
 	"os"
+	"strings"
 	"time"
 
 	basecmd "github.com/kubernetes-sigs/custom-metrics-apiserver/pkg/cmd"
@@ -71,6 +72,8 @@ type PrometheusAdapter struct {
 	PrometheusClientTLSKeyFile string
 	// PrometheusTokenFile points to the file that contains the bearer token when connecting with Prometheus
 	PrometheusTokenFile string
+	// PrometheusHeaders is a k=v list of headers to set on requests to PrometheusURL
+	PrometheusHeaders []string
 	// AdapterConfigFile points to the file containing the metrics discovery configuration.
 	AdapterConfigFile string
 	// MetricsRelistInterval is the interval at which to relist the set of available metrics
@@ -112,8 +115,7 @@ func (cmd *PrometheusAdapter) makePromClient() (prom.Client, error) {
 		}
 		httpClient.Transport = transport.NewBearerAuthRoundTripper(string(data), httpClient.Transport)
 	}
-
-	genericPromClient := prom.NewGenericAPIClient(httpClient, baseURL)
+	genericPromClient := prom.NewGenericAPIClient(httpClient, baseURL, parseHeaderArgs(cmd.PrometheusHeaders))
 	instrumentedGenericPromClient := mprom.InstrumentGenericAPIClient(genericPromClient, baseURL.String())
 	return prom.NewClientForAPI(instrumentedGenericPromClient), nil
 }
@@ -133,6 +135,8 @@ func (cmd *PrometheusAdapter) addFlags() {
 		"Optional client TLS key file to use when connecting with Prometheus, auto-renewal is not supported")
 	cmd.Flags().StringVar(&cmd.PrometheusTokenFile, "prometheus-token-file", cmd.PrometheusTokenFile,
 		"Optional file containing the bearer token to use when connecting with Prometheus")
+	cmd.Flags().StringArrayVar(&cmd.PrometheusHeaders, "prometheus-header", cmd.PrometheusHeaders,
+		"Optional header to set on requests to prometheus-url. Can be repeated")
 	cmd.Flags().StringVar(&cmd.AdapterConfigFile, "config", cmd.AdapterConfigFile,
 		"Configuration file containing details of how to transform between Prometheus metrics "+
 			"and custom metrics API resources")
@@ -405,3 +409,16 @@ func makePrometheusCAClient(caFilePath string, tlsCertFilePath string, tlsKeyFil
 		},
 	}, nil
 }
+
+func parseHeaderArgs(args []string) http.Header {
+	headers := make(http.Header, len(args))
+	for _, h := range args {
+		parts := strings.SplitN(h, "=", 2)
+		value := ""
+		if len(parts) > 1 {
+			value = parts[1]
+		}
+		headers.Add(parts[0], value)
+	}
+	return headers
+}

cmd/adapter/adapter_test.go

@@ -20,6 +20,7 @@ import (
 	"net/http"
 	"os"
 	"path/filepath"
+	"reflect"
 	"testing"
 )
 
@@ -146,3 +147,41 @@ func TestMakePrometheusCAClient(t *testing.T) {
 		}
 	}
 }
+
+func TestParseHeaderArgs(t *testing.T) {
+
+	tests := []struct {
+		args    []string
+		headers http.Header
+	}{
+		{
+			headers: http.Header{},
+		},
+		{
+			args: []string{"foo=bar"},
+			headers: http.Header{
+				"Foo": []string{"bar"},
+			},
+		},
+		{
+			args: []string{"foo"},
+			headers: http.Header{
+				"Foo": []string{""},
+			},
+		},
+		{
+			args: []string{"foo=bar", "foo=baz", "bux=baz=23"},
+			headers: http.Header{
+				"Foo": []string{"bar", "baz"},
+				"Bux": []string{"baz=23"},
+			},
+		},
+	}
+
+	for _, test := range tests {
+		got := parseHeaderArgs(test.args)
+		if !reflect.DeepEqual(got, test.headers) {
+			t.Errorf("Expected %#v but got %#v", test.headers, got)
+		}
+	}
+}

pkg/client/api.go

@@ -47,6 +47,7 @@ type GenericAPIClient interface {
 type httpAPIClient struct {
 	client  *http.Client
 	baseURL *url.URL
+	headers http.Header
 }
 
 func (c *httpAPIClient) Do(ctx context.Context, verb, endpoint string, query url.Values) (APIResponse, error) {
@@ -58,6 +59,7 @@ func (c *httpAPIClient) Do(ctx context.Context, verb, endpoint string, query url
 		return APIResponse{}, fmt.Errorf("error constructing HTTP request to Prometheus: %v", err)
 	}
 	req.WithContext(ctx)
+	req.Header = c.headers
 
 	resp, err := c.client.Do(req)
 	defer func() {
@@ -113,10 +115,11 @@ func (c *httpAPIClient) Do(ctx context.Context, verb, endpoint string, query url
 }
 
 // NewGenericAPIClient builds a new generic Prometheus API client for the given base URL and HTTP Client.
-func NewGenericAPIClient(client *http.Client, baseURL *url.URL) GenericAPIClient {
+func NewGenericAPIClient(client *http.Client, baseURL *url.URL, headers http.Header) GenericAPIClient {
 	return &httpAPIClient{
 		client:  client,
 		baseURL: baseURL,
+		headers: headers,
 	}
 }
 
@@ -139,8 +142,8 @@ func NewClientForAPI(client GenericAPIClient) Client {
 }
 
 // NewClient creates a Client for the given HTTP client and base URL (the location of the Prometheus server).
-func NewClient(client *http.Client, baseURL *url.URL) Client {
-	genericClient := NewGenericAPIClient(client, baseURL)
+func NewClient(client *http.Client, baseURL *url.URL, headers http.Header) Client {
+	genericClient := NewGenericAPIClient(client, baseURL, headers)
 	return NewClientForAPI(genericClient)
 }