feat(gensupport): pass in headers via context

internal/gensupport/send.go

@@ -15,6 +15,7 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/googleapis/gax-go/v2"
+	"github.com/googleapis/gax-go/v2/callctx"
 )
 
 // Use this error type to return an error which allows introspection of both
@@ -43,6 +44,12 @@ func (e wrappedCallErr) Is(target error) bool {
 // req.WithContext, then calls any functions returned by the hooks in
 // reverse order.
 func SendRequest(ctx context.Context, client *http.Client, req *http.Request) (*http.Response, error) {
+	// Add headers set in context metadata.
+	headers := callctx.HeadersFromContext(ctx)
+	for k, v := range headers {
+		req.Header[k] = append(req.Header[k], v...)
+	}
+
 	// Disallow Accept-Encoding because it interferes with the automatic gzip handling
 	// done by the default http.Transport. See https://github.com/google/google-api-go-client/issues/219.
 	if _, ok := req.Header["Accept-Encoding"]; ok {
@@ -77,6 +84,12 @@ func send(ctx context.Context, client *http.Client, req *http.Request) (*http.Re
 // req.WithContext, then calls any functions returned by the hooks in
 // reverse order.
 func SendRequestWithRetry(ctx context.Context, client *http.Client, req *http.Request, retry *RetryConfig) (*http.Response, error) {
+	// Add headers set in context metadata.
+	headers := callctx.HeadersFromContext(ctx)
+	for k, v := range headers {
+		req.Header[k] = append(req.Header[k], v...)
+	}
+
 	// Disallow Accept-Encoding because it interferes with the automatic gzip handling
 	// done by the default http.Transport. See https://github.com/google/google-api-go-client/issues/219.
 	if _, ok := req.Header["Accept-Encoding"]; ok {

internal/gensupport/send_test.go

@@ -7,6 +7,9 @@ package gensupport
 import (
 	"context"
 	"errors"
+	"fmt"
+	"github.com/google/go-cmp/cmp"
+	"github.com/googleapis/gax-go/v2/callctx"
 	"net/http"
 	"testing"
 )
@@ -31,6 +34,39 @@ func TestSendRequestWithRetry(t *testing.T) {
 	}
 }
 
+type headerRoundTripper struct {
+	wantHeader http.Header
+}
+
+func (rt *headerRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
+	// Ignore x-goog headers sent by SendRequestWithRetry
+	delete(r.Header, "X-Goog-Api-Client")
+	delete(r.Header, "X-Goog-Gcs-Idempotency-Token")
+	if diff := cmp.Diff(r.Header, rt.wantHeader); diff != "" {
+		return nil, fmt.Errorf("headers don't match: %v", diff)
+	}
+	return &http.Response{StatusCode: 200}, nil
+}
+
+// Ensure that headers set via the context are passed through to the request as expected.
+func TestSendRequestHeader(t *testing.T) {
+	ctx := context.Background()
+	ctx = callctx.SetHeaders(ctx, "foo", "100", "bar", "200")
+	client := http.Client{
+		Transport: &headerRoundTripper{
+			wantHeader: map[string][]string{"foo": {"100"}, "bar": {"200"}},
+		},
+	}
+	req, _ := http.NewRequest("GET", "url", nil)
+	if _, err := SendRequest(ctx, &client, req); err != nil {
+		t.Errorf("SendRequest: %v", err)
+	}
+	req2, _ := http.NewRequest("GET", "url", nil)
+	if _, err := SendRequestWithRetry(ctx, &client, req2, nil); err != nil {
+		t.Errorf("SendRequest: %v", err)
+	}
+}
+
 type brokenRoundTripper struct{}
 
 func (t *brokenRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {