Skip to content

Commit 6cd421e

Browse files
committed
Add multigraphql package for batching several GraphQL queries in one request
Some of the upcoming gh features depend on querying information about an arbitrary number of repositories (determined at runtime). GraphQL (in theory) allows us to perform all those lookups in a single query, rather than over N individual queries, but we don't yet have a great mechanism for combining several GraphQL queries into one and parsing the combined result. This implements a potential approach: 1. The `Parse(str)` function returns a Query; 2. A PreparedQuery combines that with scalar values ("variables") for the query; 3. `Merge(queries...)` returns the final GraphQL query string + combined variables for the request; 4. `Decode(body, destinations)` segments the JSON response into corresponding destination interfaces.
1 parent bce5d21 commit 6cd421e

File tree

6 files changed

+540
-0
lines changed

6 files changed

+540
-0
lines changed

pkg/multigraphql/decode.go

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
package multigraphql
2+
3+
import (
4+
"bytes"
5+
"encoding/json"
6+
"fmt"
7+
"io"
8+
"strconv"
9+
"strings"
10+
)
11+
12+
type graphqlResponse struct {
13+
Data map[string]*json.RawMessage
14+
Errors []struct {
15+
Message string
16+
}
17+
}
18+
19+
// Decode parses the GraphQL JSON response
20+
func Decode(r io.Reader, destinations []interface{}) error {
21+
resp := graphqlResponse{}
22+
if err := json.NewDecoder(r).Decode(&resp); err != nil {
23+
return err
24+
}
25+
26+
if len(resp.Errors) > 0 {
27+
messages := []string{}
28+
for _, e := range resp.Errors {
29+
messages = append(messages, e.Message)
30+
}
31+
return fmt.Errorf("GraphQL error: %s", strings.Join(messages, "; "))
32+
}
33+
34+
for alias, value := range resp.Data {
35+
if !strings.HasPrefix(alias, "multi_") {
36+
continue
37+
}
38+
i, _ := strconv.Atoi(strings.TrimPrefix(alias, "multi_"))
39+
dec := json.NewDecoder(bytes.NewReader([]byte(*value)))
40+
if err := dec.Decode(destinations[i]); err != nil {
41+
return err
42+
}
43+
}
44+
45+
return nil
46+
}

pkg/multigraphql/decode_test.go

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
package multigraphql
2+
3+
import (
4+
"bytes"
5+
"testing"
6+
)
7+
8+
func TestDecode(t *testing.T) {
9+
buf := bytes.NewBufferString(`
10+
{ "extensions": [],
11+
"data": {
12+
"multi_000": { "world": true },
13+
"multi_001": { "machines": "are learning" }
14+
} }
15+
`)
16+
17+
hello := struct {
18+
World bool
19+
}{}
20+
ai := struct {
21+
Machines string
22+
}{}
23+
24+
err := Decode(buf, []interface{}{&hello, &ai})
25+
if err != nil {
26+
t.Fatalf("got error: %v", err)
27+
}
28+
29+
if !hello.World {
30+
t.Errorf("expected World to be true")
31+
}
32+
if ai.Machines != "are learning" {
33+
t.Errorf("expected machines to be learning, got %q", ai.Machines)
34+
}
35+
}
36+
37+
func TestDecode_errors(t *testing.T) {
38+
buf := bytes.NewBufferString(`
39+
{ "extensions": [],
40+
"errors": [
41+
{ "message": "boom" },
42+
{ "message": "shutting down" }
43+
] }
44+
`)
45+
46+
hello := struct {
47+
World bool
48+
}{}
49+
50+
err := Decode(buf, []interface{}{&hello})
51+
if err == nil || err.Error() != "GraphQL error: boom; shutting down" {
52+
t.Fatalf("got error: %v", err)
53+
}
54+
}

pkg/multigraphql/merge.go

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
package multigraphql
2+
3+
import (
4+
"bytes"
5+
"fmt"
6+
"io"
7+
"regexp"
8+
"sort"
9+
"strings"
10+
)
11+
12+
// PreparedQuery represents a Query with associated variable values
13+
type PreparedQuery struct {
14+
variableValues map[string]interface{}
15+
Query
16+
}
17+
18+
var identifier = regexp.MustCompile(`\$[a-zA-Z]\w*`)
19+
20+
// Merge combines multiple queries into one while avoiding variable collisions
21+
func Merge(queries ...PreparedQuery) (string, map[string]interface{}) {
22+
out := &bytes.Buffer{}
23+
queryStrings := []string{}
24+
allVariables := map[string]string{}
25+
allValues := map[string]interface{}{}
26+
seenFragments := map[string]struct{}{"": {}}
27+
28+
for i, q := range queries {
29+
renames := mergeVariables(allVariables, q.variables, func(k string) string {
30+
return fmt.Sprintf("%s_%03d", k, i)
31+
})
32+
33+
for key, value := range q.variableValues {
34+
if newKey, exists := renames[key]; exists {
35+
key = newKey
36+
}
37+
allValues[key] = value
38+
}
39+
40+
if _, seen := seenFragments[q.fragments]; !seen {
41+
fmt.Fprintln(out, q.fragments)
42+
seenFragments[q.fragments] = struct{}{}
43+
}
44+
45+
finalQuery := renameVariables(q.query, renames)
46+
queryStrings = append(queryStrings, fmt.Sprintf("multi_%03d: %s", i, finalQuery))
47+
}
48+
49+
fmt.Fprint(out, "query")
50+
writeVariables(out, allVariables)
51+
fmt.Fprintf(out, " {\n\t%s\n}", strings.Join(queryStrings, "\n\t"))
52+
53+
return out.String(), allValues
54+
}
55+
56+
func mergeVariables(dest, src map[string]string, keyGen func(string) string) map[string]string {
57+
renames := map[string]string{}
58+
for key, value := range src {
59+
if _, exists := dest[key]; exists {
60+
newKey := keyGen(key)
61+
renames[key] = newKey
62+
key = newKey
63+
}
64+
dest[key] = value
65+
}
66+
return renames
67+
}
68+
69+
func renameVariables(q string, dictionary map[string]string) string {
70+
return identifier.ReplaceAllStringFunc(q, func(v string) string {
71+
if newName, exists := dictionary[v[1:]]; exists {
72+
return "$" + newName
73+
}
74+
return v
75+
})
76+
}
77+
78+
func writeVariables(out io.Writer, variables map[string]string) {
79+
if len(variables) == 0 {
80+
return
81+
}
82+
83+
vars := []string{}
84+
for key, value := range variables {
85+
vars = append(vars, fmt.Sprintf("$%s: %s", key, value))
86+
}
87+
sort.Sort(sort.StringSlice(vars))
88+
89+
fmt.Fprint(out, "(\n")
90+
for _, v := range vars {
91+
fmt.Fprintf(out, "\t%s\n", v)
92+
}
93+
fmt.Fprint(out, ")")
94+
}

pkg/multigraphql/merge_test.go

Lines changed: 156 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,156 @@
1+
package multigraphql
2+
3+
import (
4+
"reflect"
5+
"testing"
6+
)
7+
8+
func TestMerge(t *testing.T) {
9+
type args struct {
10+
queries []PreparedQuery
11+
}
12+
tests := []struct {
13+
name string
14+
args args
15+
wantQuery string
16+
wantValues map[string]interface{}
17+
}{
18+
{
19+
name: "A single query",
20+
args: args{
21+
queries: []PreparedQuery{
22+
PreparedQuery{
23+
variableValues: map[string]interface{}{
24+
"owner": "monalisa",
25+
"repo": "hello-world",
26+
},
27+
Query: Query{
28+
query: `repository(owner: $owner, name: $repo) { id }`,
29+
variables: map[string]string{"owner": "String!", "repo": "String"},
30+
},
31+
},
32+
},
33+
},
34+
wantQuery: `query(
35+
$owner: String!
36+
$repo: String
37+
) {
38+
multi_000: repository(owner: $owner, name: $repo) { id }
39+
}`,
40+
wantValues: map[string]interface{}{
41+
"owner": "monalisa",
42+
"repo": "hello-world",
43+
},
44+
},
45+
{
46+
name: "Multiple queries",
47+
args: args{
48+
queries: []PreparedQuery{
49+
PreparedQuery{
50+
variableValues: map[string]interface{}{
51+
"owner": "monalisa",
52+
"repo": "hello-world",
53+
},
54+
Query: Query{
55+
query: `repository(owner: $owner, name: $repo) { id }`,
56+
variables: map[string]string{"owner": "String!", "repo": "String"},
57+
},
58+
},
59+
PreparedQuery{
60+
variableValues: map[string]interface{}{
61+
"owner": "hubot",
62+
"repo": "chatops",
63+
"user": "octocat",
64+
},
65+
Query: Query{
66+
query: `repository(owner: $owner, name: $repo, assignee: $user) { id }`,
67+
variables: map[string]string{"owner": "String", "repo": "String!", "user": "String"},
68+
},
69+
},
70+
PreparedQuery{
71+
variableValues: map[string]interface{}{
72+
"owner": "github",
73+
"user": "ghost",
74+
},
75+
Query: Query{
76+
query: `repository(owner: $owner, assignee: $user) { id }`,
77+
variables: map[string]string{"owner": "String", "user": "String!"},
78+
},
79+
},
80+
},
81+
},
82+
wantQuery: `query(
83+
$owner: String!
84+
$owner_001: String
85+
$owner_002: String
86+
$repo: String
87+
$repo_001: String!
88+
$user: String
89+
$user_002: String!
90+
) {
91+
multi_000: repository(owner: $owner, name: $repo) { id }
92+
multi_001: repository(owner: $owner_001, name: $repo_001, assignee: $user) { id }
93+
multi_002: repository(owner: $owner_002, assignee: $user_002) { id }
94+
}`,
95+
wantValues: map[string]interface{}{
96+
"owner": "monalisa",
97+
"repo": "hello-world",
98+
"owner_001": "hubot",
99+
"repo_001": "chatops",
100+
"user": "octocat",
101+
"owner_002": "github",
102+
"user_002": "ghost",
103+
},
104+
},
105+
{
106+
name: "Queries with fragments",
107+
args: args{
108+
queries: []PreparedQuery{
109+
PreparedQuery{
110+
variableValues: map[string]interface{}{},
111+
Query: Query{
112+
query: `a { ...b }`,
113+
fragments: `fragment b on B { boo }`,
114+
variables: map[string]string{},
115+
},
116+
},
117+
PreparedQuery{
118+
variableValues: map[string]interface{}{},
119+
Query: Query{
120+
query: `c { ...b }`,
121+
fragments: `fragment b on B { boo }`,
122+
variables: map[string]string{},
123+
},
124+
},
125+
PreparedQuery{
126+
variableValues: map[string]interface{}{},
127+
Query: Query{
128+
query: `d { ...e }`,
129+
fragments: `fragment e on E { eeek }`,
130+
variables: map[string]string{},
131+
},
132+
},
133+
},
134+
},
135+
wantQuery: `fragment b on B { boo }
136+
fragment e on E { eeek }
137+
query {
138+
multi_000: a { ...b }
139+
multi_001: c { ...b }
140+
multi_002: d { ...e }
141+
}`,
142+
wantValues: map[string]interface{}{},
143+
},
144+
}
145+
for _, tt := range tests {
146+
t.Run(tt.name, func(t *testing.T) {
147+
got, got1 := Merge(tt.args.queries...)
148+
if got != tt.wantQuery {
149+
t.Errorf("Merge() got = %#v, want %#v", got, tt.wantQuery)
150+
}
151+
if !reflect.DeepEqual(got1, tt.wantValues) {
152+
t.Errorf("Merge() got1 = %v, want %v", got1, tt.wantValues)
153+
}
154+
})
155+
}
156+
}

0 commit comments

Comments
 (0)