From 2771bcc5da1c3bde389e1905c360d830875567e1 Mon Sep 17 00:00:00 2001
From: songjiayang <songjiayang1@gmail.com>
Date: Tue, 7 Feb 2023 23:16:37 +0800
Subject: [PATCH] Add `Header` method to Pusher for custom header (#1218)

Signed-off-by: songjiayang <songjiayang1@gmail.com>
---
 prometheus/push/push.go      | 14 ++++++++++++++
 prometheus/push/push_test.go | 25 +++++++++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/prometheus/push/push.go b/prometheus/push/push.go
index 29f6cd3..d3a3cef 100644
--- a/prometheus/push/push.go
+++ b/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)
 	}
diff --git a/prometheus/push/push_test.go b/prometheus/push/push_test.go
index cc061cf..f907408 100644
--- a/prometheus/push/push_test.go
+++ b/prometheus/push/push_test.go
@@ -31,6 +31,7 @@ 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
@@ -38,6 +39,7 @@ func TestPush(t *testing.T) {
 	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")
+	}
 }