Add Header method to Pusher for custom header (#1218)

Signed-off-by: songjiayang <[email protected]>

Author: songjiayang

prometheus/push/push.go

@@ -77,6 +77,7 @@ type Pusher struct {
 	registerer prometheus.Registerer
 
 	client             HTTPDoer
+	header             http.Header
 	useBasicAuth       bool
 	username, password string
 
@@ -201,6 +202,13 @@ func (p *Pusher) Client(c HTTPDoer) *Pusher {
 	return p
 }
 
+// Header sets a custom HTTP header for the Pusher's client. For convenience, this method
+// returns a pointer to the Pusher itself.
+func (p *Pusher) Header(header http.Header) *Pusher {
+	p.header = header
+	return p
+}
+
 // BasicAuth configures the Pusher to use HTTP Basic Authentication with the
 // provided username and password. For convenience, this method returns a
 // pointer to the Pusher itself.
@@ -236,6 +244,9 @@ func (p *Pusher) Delete() error {
 	if err != nil {
 		return err
 	}
+	if p.header != nil {
+		req.Header = p.header
+	}
 	if p.useBasicAuth {
 		req.SetBasicAuth(p.username, p.password)
 	}
@@ -286,6 +297,9 @@ func (p *Pusher) push(ctx context.Context, method string) error {
 	if err != nil {
 		return err
 	}
+	if p.header != nil {
+		req.Header = p.header
+	}
 	if p.useBasicAuth {
 		req.SetBasicAuth(p.username, p.password)
 	}

prometheus/push/push_test.go

@@ -31,13 +31,15 @@ func TestPush(t *testing.T) {
 		lastMethod string
 		lastBody   []byte
 		lastPath   string
+		lastHeader http.Header
 	)
 
 	// Fake a Pushgateway that responds with 202 to DELETE and with 200 in
 	// all other cases.
 	pgwOK := httptest.NewServer(
 		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			lastMethod = r.Method
+			lastHeader = r.Header
 			var err error
 			lastBody, err = io.ReadAll(r.Body)
 			if err != nil {
@@ -281,4 +283,27 @@ func TestPush(t *testing.T) {
 	if lastPath != "/metrics/job/testjob/a/x/b/y" && lastPath != "/metrics/job/testjob/b/y/a/x" {
 		t.Error("unexpected path:", lastPath)
 	}
+
+	// Push some Collectors with custom header, all good.
+	header := make(http.Header)
+	header.Set("Authorization", "Bearer Token")
+	if err := New(pgwOK.URL, "testjob").
+		Collector(metric1).
+		Collector(metric2).
+		Header(header).
+		Push(); err != nil {
+		t.Fatal(err)
+	}
+	if lastMethod != http.MethodPut {
+		t.Errorf("got method %q for Add, want %q", lastMethod, http.MethodPut)
+	}
+	if !bytes.Equal(lastBody, wantBody) {
+		t.Errorf("got body %v, want %v", lastBody, wantBody)
+	}
+	if lastPath != "/metrics/job/testjob" {
+		t.Error("unexpected path:", lastPath)
+	}
+	if lastHeader == nil || lastHeader.Get("Authorization") == "" {
+		t.Error("empty Authorization header")
+	}
 }