Queries in context

relay/relay.go

@@ -1,6 +1,7 @@
 package relay
 
 import (
+	"context"
 	"encoding/base64"
 	"encoding/json"
 	"errors"
@@ -47,18 +48,36 @@ type Handler struct {
 	Schema *graphql.Schema
 }
 
+type Params struct {
+	Query         string                 `json:"query"`
+	OperationName string                 `json:"operationName"`
+	Variables     map[string]interface{} `json:"variables"`
+}
+
 func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	var params struct {
-		Query         string                 `json:"query"`
-		OperationName string                 `json:"operationName"`
-		Variables     map[string]interface{} `json:"variables"`
-	}
+	var params Params
 	if err := json.NewDecoder(r.Body).Decode(&params); err != nil {
 		http.Error(w, err.Error(), http.StatusBadRequest)
 		return
 	}
 
-	response := h.Schema.Exec(r.Context(), params.Query, params.OperationName, params.Variables)
+	// set parameters to context
+	ctx := WithParams(r.Context(), params)
+
+	response := h.Schema.Exec(ctx, params.Query, params.OperationName, params.Variables)
+
+	// check additional queries from context
+	if queries, ok := GetQueries(ctx); ok {
+		for topic, params := range queries {
+			addResponse := h.Schema.Exec(r.Context(), params.Query, params.OperationName, params.Variables)
+			// TODO: don't ignore error
+			if response.Extensions == nil {
+				response.Extensions = make(map[string]interface{})
+			}
+			response.Extensions[topic] = addResponse.Data
+		}
+	}
+
 	responseJSON, err := json.Marshal(response)
 	if err != nil {
 		http.Error(w, err.Error(), http.StatusInternalServerError)
@@ -68,3 +87,28 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	w.Header().Set("Content-Type", "application/json")
 	w.Write(responseJSON)
 }
+
+type ctxKey int
+
+const (
+	paramsKey ctxKey = iota
+	queriesKey
+)
+
+func WithParams(ctx context.Context, params Params) context.Context {
+	return context.WithValue(ctx, paramsKey, params)
+}
+
+func GetParams(ctx context.Context) (Params, bool) {
+	v, ok := ctx.Value(paramsKey).(Params)
+	return v, ok
+}
+
+func WithQueries(ctx context.Context, queries map[string]Params) context.Context {
+	return context.WithValue(ctx, queriesKey, queries)
+}
+
+func GetQueries(ctx context.Context) (map[string]Params, bool) {
+	v, ok := ctx.Value(queriesKey).(map[string]Params)
+	return v, ok
+}

relay/relay_test.go

@@ -10,14 +10,19 @@ import (
 	"github.com/lygo/graphql-go/relay"
 )
 
-var starwarsSchema = graphql.MustParseSchema(starwars.Schema, &starwars.Resolver{})
-
 func TestServeHTTP(t *testing.T) {
+	starwarsSchema := graphql.MustParseSchema(starwars.Schema, &starwars.Resolver{})
 	w := httptest.NewRecorder()
 	r := httptest.NewRequest("POST", "/some/path/here", strings.NewReader(`{"query":"{ hero { name } }", "operationName":"", "variables": null}`))
 	h := relay.Handler{Schema: starwarsSchema}
 
-	h.ServeHTTP(w, r)
+	ctx := relay.WithQueries(r.Context(), map[string]relay.Params{
+		"addition": {
+			Query: `{ hero(episode: EMPIRE) { name } }`,
+		},
+	})
+
+	h.ServeHTTP(w, r.WithContext(ctx))
 
 	if w.Code != 200 {
 		t.Fatalf("Expected status code 200, got %d.", w.Code)
@@ -28,7 +33,7 @@ func TestServeHTTP(t *testing.T) {
 		t.Fatalf("Invalid content-type. Expected [application/json], but instead got [%s]", contentType)
 	}
 
-	expectedResponse := `{"data":{"hero":{"name":"R2-D2"}}}`
+	expectedResponse := `{"data":{"hero":{"name":"R2-D2"}},"extensions":{"addition":{"hero":{"name":"Luke Skywalker"}}}}`
 	actualResponse := w.Body.String()
 	if expectedResponse != actualResponse {
 		t.Fatalf("Invalid response. Expected [%s], but instead got [%s]", expectedResponse, actualResponse)