Support setHeaders in Query & Mutation

graphql.go

@@ -33,19 +33,19 @@ func NewClient(url string, httpClient *http.Client) *Client {
 // Query executes a single GraphQL query request,
 // with a query derived from q, populating the response into it.
 // q should be a pointer to struct that corresponds to the GraphQL schema.
-func (c *Client) Query(ctx context.Context, q interface{}, variables map[string]interface{}) error {
-	return c.do(ctx, queryOperation, q, variables)
+func (c *Client) Query(ctx context.Context, q interface{}, variables map[string]interface{}, setHeaders func(*http.Request)) error {
+	return c.do(ctx, queryOperation, q, variables, setHeaders)
 }
 
 // Mutate executes a single GraphQL mutation request,
 // with a mutation derived from m, populating the response into it.
 // m should be a pointer to struct that corresponds to the GraphQL schema.
-func (c *Client) Mutate(ctx context.Context, m interface{}, variables map[string]interface{}) error {
-	return c.do(ctx, mutationOperation, m, variables)
+func (c *Client) Mutate(ctx context.Context, m interface{}, variables map[string]interface{}, setHeaders func(*http.Request)) error {
+	return c.do(ctx, mutationOperation, m, variables, setHeaders)
 }
 
 // do executes a single GraphQL operation.
-func (c *Client) do(ctx context.Context, op operationType, v interface{}, variables map[string]interface{}) error {
+func (c *Client) do(ctx context.Context, op operationType, v interface{}, variables map[string]interface{}, setHeaders func(*http.Request)) error {
 	var query string
 	switch op {
 	case queryOperation:
@@ -65,7 +65,16 @@ func (c *Client) do(ctx context.Context, op operationType, v interface{}, variab
 	if err != nil {
 		return err
 	}
-	resp, err := ctxhttp.Post(ctx, c.httpClient, c.url, "application/json", &buf)
+	req, err := http.NewRequest(http.MethodPost, c.url, &buf)
+	if err != nil {
+		return err
+	}
+	req.Header.Set("Content-Type", "application/json")
+	if setHeaders != nil {
+		setHeaders(req)
+	}
+
+	resp, err := ctxhttp.Do(ctx, c.httpClient, req)
 	if err != nil {
 		return err
 	}

graphql_test.go

@@ -2,6 +2,7 @@ package graphql_test
 
 import (
 	"context"
+	"fmt"
 	"io"
 	"io/ioutil"
 	"net/http"
@@ -49,7 +50,7 @@ func TestClient_Query_partialDataWithErrorResponse(t *testing.T) {
 			ID graphql.ID
 		} `graphql:"node2: node(id: \"NotExist\")"`
 	}
-	err := client.Query(context.Background(), &q, nil)
+	err := client.Query(context.Background(), &q, nil, nil)
 	if err == nil {
 		t.Fatal("got error: nil, want: non-nil")
 	}
@@ -89,7 +90,7 @@ func TestClient_Query_noDataWithErrorResponse(t *testing.T) {
 			Name graphql.String
 		}
 	}
-	err := client.Query(context.Background(), &q, nil)
+	err := client.Query(context.Background(), &q, nil, nil)
 	if err == nil {
 		t.Fatal("got error: nil, want: non-nil")
 	}
@@ -113,7 +114,7 @@ func TestClient_Query_errorStatusCode(t *testing.T) {
 			Name graphql.String
 		}
 	}
-	err := client.Query(context.Background(), &q, nil)
+	err := client.Query(context.Background(), &q, nil, nil)
 	if err == nil {
 		t.Fatal("got error: nil, want: non-nil")
 	}
@@ -144,7 +145,7 @@ func TestClient_Query_emptyVariables(t *testing.T) {
 			Name string
 		}
 	}
-	err := client.Query(context.Background(), &q, map[string]interface{}{})
+	err := client.Query(context.Background(), &q, map[string]interface{}{}, nil)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -153,6 +154,77 @@ func TestClient_Query_emptyVariables(t *testing.T) {
 	}
 }
 
+// Test that nil setHeaders has no impact
+func TestClient_Query_setHeadersNil(t *testing.T) {
+	mux := http.NewServeMux()
+	mux.HandleFunc("/graphql", func(w http.ResponseWriter, req *http.Request) {
+		body := mustRead(req.Body)
+		if got, want := body, `{"query":"{user{name}}"}`+"\n"; got != want {
+			t.Errorf("got body: %v, want %v", got, want)
+		}
+		w.Header().Set("Content-Type", "application/json")
+		name := req.Header.Get("TestUser")
+		if name == "" {
+			name = "Gopher"
+		}
+		outBody := fmt.Sprintf("{\"data\": {\"user\": {\"name\": \"%s\"}}}\n", name)
+		mustWrite(w, outBody)
+	})
+
+	client := graphql.NewClient("/graphql", &http.Client{Transport: localRoundTripper{handler: mux}})
+
+	var q struct {
+		User struct {
+			Name string
+		}
+	}
+
+	err := client.Query(context.Background(), &q, map[string]interface{}{}, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got, want := q.User.Name, "Gopher"; got != want {
+		t.Errorf("got q.User.Name: %q, want: %q", got, want)
+	}
+}
+
+// Test that setHeaders does send headers
+func TestClient_Query_setHeaders(t *testing.T) {
+	mux := http.NewServeMux()
+	mux.HandleFunc("/graphql", func(w http.ResponseWriter, req *http.Request) {
+		body := mustRead(req.Body)
+		if got, want := body, `{"query":"{user{name}}"}`+"\n"; got != want {
+			t.Errorf("got body: %v, want %v", got, want)
+		}
+		w.Header().Set("Content-Type", "application/json")
+		name := req.Header.Get("TestUser")
+		if name == "" {
+			name = "Gopher"
+		}
+		outBody := fmt.Sprintf("{\"data\": {\"user\": {\"name\": \"%s\"}}}\n", name)
+		mustWrite(w, outBody)
+	})
+
+	client := graphql.NewClient("/graphql", &http.Client{Transport: localRoundTripper{handler: mux}})
+
+	var q struct {
+		User struct {
+			Name string
+		}
+	}
+	setHeaders := func(req *http.Request) {
+		req.Header.Set("TestUser", "testUser")
+	}
+
+	err := client.Query(context.Background(), &q, map[string]interface{}{}, setHeaders)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got, want := q.User.Name, "testUser"; got != want {
+		t.Errorf("got q.User.Name: %q, want: %q", got, want)
+	}
+}
+
 // localRoundTripper is an http.RoundTripper that executes HTTP transactions
 // by using handler directly, instead of going over an HTTP connection.
 type localRoundTripper struct {