api command: add GraphQL support for --paginate

pkg/cmd/api/api.go

@@ -76,7 +76,11 @@ on the format of the value:
 Raw request body may be passed from the outside via a file specified by '--input'.
 Pass "-" to read from standard input. In this mode, parameters specified via
 '--field' flags are serialized into URL query parameters.
-`,
+
+In '--paginate' mode, all pages of results will sequentially be requested until
+there are no more pages of results. For GraphQL requests, this requires that the
+original query accepts an '$endCursor: String' variable and that it fetches the
+'pageInfo{ hasNextPage, endCursor }' set of fields from a collection.`,
 		Example: heredoc.Doc(`
 			$ gh api repos/:owner/:repo/releases
 
@@ -89,6 +93,20 @@ Pass "-" to read from standard input. In this mode, parameters specified via
 					}
 				}
 			'
+			
+			$ gh api graphql --paginate -f query='
+				query($endCursor: String) {
+					viewer {
+						repositories(first: 100, after: $endCursor) {
+							nodes { nameWithOwner }
+							pageInfo {
+								hasNextPage
+								endCursor
+							}
+						}
+					}
+				}
+			'
 		`),
 		Args: cobra.ExactArgs(1),
 		RunE: func(c *cobra.Command, args []string) error {
@@ -162,7 +180,7 @@ func apiRun(opts *ApiOptions) error {
 			return err
 		}
 
-		err = processResponse(resp, opts)
+		endCursor, err := processResponse(resp, opts)
 		if err != nil {
 			return err
 		}
@@ -170,7 +188,15 @@ func apiRun(opts *ApiOptions) error {
 		if !opts.Paginate {
 			break
 		}
-		requestPath, hasNextPage = findNextPage(resp)
+
+		if opts.RequestPath == "graphql" {
+			hasNextPage = endCursor != ""
+			if hasNextPage {
+				params["endCursor"] = endCursor
+			}
+		} else {
+			requestPath, hasNextPage = findNextPage(resp)
+		}
 
 		if hasNextPage && opts.ShowResponseHeaders {
 			fmt.Fprint(opts.IO.Out, "\n")
@@ -180,21 +206,7 @@ func apiRun(opts *ApiOptions) error {
 	return nil
 }
 
-var linkRE = regexp.MustCompile(`<([^>]+)>;\s*rel="([^"]+)"`)
-
-func findNextPage(resp *http.Response) (string, bool) {
-	for _, m := range linkRE.FindAllStringSubmatch(resp.Header.Get("Link"), -1) {
-		if len(m) < 2 {
-			continue
-		}
-		if m[2] == "next" {
-			return m[1], true
-		}
-	}
-	return "", false
-}
-
-func processResponse(resp *http.Response, opts *ApiOptions) error {
+func processResponse(resp *http.Response, opts *ApiOptions) (endCursor string, err error) {
 	if opts.ShowResponseHeaders {
 		fmt.Fprintln(opts.IO.Out, resp.Proto, resp.Status)
 		printHeaders(opts.IO.Out, resp.Header, opts.IO.ColorEnabled())
@@ -202,43 +214,55 @@ func processResponse(resp *http.Response, opts *ApiOptions) error {
 	}
 
 	if resp.StatusCode == 204 {
-		return nil
+		return
 	}
 	var responseBody io.Reader = resp.Body
 	defer resp.Body.Close()
 
 	isJSON, _ := regexp.MatchString(`[/+]json(;|$)`, resp.Header.Get("Content-Type"))
 
-	var err error
 	var serverError string
 	if isJSON && (opts.RequestPath == "graphql" || resp.StatusCode >= 400) {
 		responseBody, serverError, err = parseErrorResponse(responseBody, resp.StatusCode)
 		if err != nil {
-			return err
+			return
 		}
 	}
 
+	var bodyCopy *bytes.Buffer
+	isGraphQLPaginate := isJSON && resp.StatusCode == 200 && opts.Paginate && opts.RequestPath == "graphql"
+	if isGraphQLPaginate {
+		bodyCopy = &bytes.Buffer{}
+		responseBody = io.TeeReader(responseBody, bodyCopy)
+	}
+
 	if isJSON && opts.IO.ColorEnabled() {
 		err = jsoncolor.Write(opts.IO.Out, responseBody, "  ")
 		if err != nil {
-			return err
+			return
 		}
 	} else {
 		_, err = io.Copy(opts.IO.Out, responseBody)
 		if err != nil {
-			return err
+			return
 		}
 	}
 
 	if serverError != "" {
 		fmt.Fprintf(opts.IO.ErrOut, "gh: %s\n", serverError)
-		return cmdutil.SilentError
+		err = cmdutil.SilentError
+		return
 	} else if resp.StatusCode > 299 {
 		fmt.Fprintf(opts.IO.ErrOut, "gh: HTTP %d\n", resp.StatusCode)
-		return cmdutil.SilentError
+		err = cmdutil.SilentError
+		return
 	}
 
-	return nil
+	if isGraphQLPaginate {
+		endCursor = findEndCursor(bodyCopy)
+	}
+
+	return
 }
 
 var placeholderRE = regexp.MustCompile(`\:(owner|repo)\b`)

pkg/cmd/api/api_test.go

@@ -2,6 +2,7 @@ package api
 
 import (
 	"bytes"
+	"encoding/json"
 	"fmt"
 	"io/ioutil"
 	"net/http"
@@ -13,6 +14,7 @@ import (
 	"github.com/cli/cli/pkg/iostreams"
 	"github.com/google/shlex"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func Test_NewCmdApi(t *testing.T) {
@@ -293,7 +295,7 @@ func Test_apiRun(t *testing.T) {
 	}
 }
 
-func Test_apiRun_pagination(t *testing.T) {
+func Test_apiRun_paginationREST(t *testing.T) {
 	io, _, stdout, stderr := iostreams.Test()
 
 	requestCount := 0
@@ -346,6 +348,83 @@ func Test_apiRun_pagination(t *testing.T) {
 	assert.Equal(t, "https://api.github.com/repositories/1227/issues?page=3", responses[2].Request.URL.String())
 }
 
+func Test_apiRun_paginationGraphQL(t *testing.T) {
+	io, _, stdout, stderr := iostreams.Test()
+
+	requestCount := 0
+	responses := []*http.Response{
+		{
+			StatusCode: 200,
+			Header:     http.Header{"Content-Type": []string{`application/json`}},
+			Body: ioutil.NopCloser(bytes.NewBufferString(`{
+				"data": {
+					"nodes": ["page one"],
+					"pageInfo": {
+						"endCursor": "PAGE1_END",
+						"hasNextPage": true
+					}
+				}
+			}`)),
+		},
+		{
+			StatusCode: 200,
+			Header:     http.Header{"Content-Type": []string{`application/json`}},
+			Body: ioutil.NopCloser(bytes.NewBufferString(`{
+				"data": {
+					"nodes": ["page two"],
+					"pageInfo": {
+						"endCursor": "PAGE2_END",
+						"hasNextPage": false
+					}
+				}
+			}`)),
+		},
+	}
+
+	options := ApiOptions{
+		IO: io,
+		HttpClient: func() (*http.Client, error) {
+			var tr roundTripper = func(req *http.Request) (*http.Response, error) {
+				resp := responses[requestCount]
+				resp.Request = req
+				requestCount++
+				return resp, nil
+			}
+			return &http.Client{Transport: tr}, nil
+		},
+
+		RequestMethod: "POST",
+		RequestPath:   "graphql",
+		Paginate:      true,
+	}
+
+	err := apiRun(&options)
+	require.NoError(t, err)
+
+	assert.Contains(t, stdout.String(), `"page one"`)
+	assert.Contains(t, stdout.String(), `"page two"`)
+	assert.Equal(t, "", stderr.String(), "stderr")
+
+	var requestData struct {
+		Variables map[string]interface{}
+	}
+
+	bb, err := ioutil.ReadAll(responses[0].Request.Body)
+	require.NoError(t, err)
+	err = json.Unmarshal(bb, &requestData)
+	require.NoError(t, err)
+	_, hasCursor := requestData.Variables["endCursor"].(string)
+	assert.Equal(t, false, hasCursor)
+
+	bb, err = ioutil.ReadAll(responses[1].Request.Body)
+	require.NoError(t, err)
+	err = json.Unmarshal(bb, &requestData)
+	require.NoError(t, err)
+	endCursor, hasCursor := requestData.Variables["endCursor"].(string)
+	assert.Equal(t, true, hasCursor)
+	assert.Equal(t, "PAGE1_END", endCursor)
+}
+
 func Test_apiRun_inputFile(t *testing.T) {
 	tests := []struct {
 		name          string

pkg/cmd/api/pagination.go

@@ -0,0 +1,87 @@
+package api
+
+import (
+	"encoding/json"
+	"io"
+	"net/http"
+	"regexp"
+)
+
+var linkRE = regexp.MustCompile(`<([^>]+)>;\s*rel="([^"]+)"`)
+
+func findNextPage(resp *http.Response) (string, bool) {
+	for _, m := range linkRE.FindAllStringSubmatch(resp.Header.Get("Link"), -1) {
+		if len(m) >= 2 && m[2] == "next" {
+			return m[1], true
+		}
+	}
+	return "", false
+}
+
+func findEndCursor(r io.Reader) string {
+	dec := json.NewDecoder(r)
+
+	var idx int
+	var stack []json.Delim
+	var lastKey string
+	var contextKey string
+
+	var endCursor string
+	var hasNextPage bool
+	var foundEndCursor bool
+	var foundNextPage bool
+
+loop:
+	for {
+		t, err := dec.Token()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return ""
+		}
+
+		switch tt := t.(type) {
+		case json.Delim:
+			switch tt {
+			case '{', '[':
+				stack = append(stack, tt)
+				contextKey = lastKey
+				idx = 0
+			case '}', ']':
+				stack = stack[:len(stack)-1]
+				contextKey = ""
+				idx = 0
+			}
+		default:
+			isKey := len(stack) > 0 && stack[len(stack)-1] == '{' && idx%2 == 0
+			idx++
+
+			switch tt := t.(type) {
+			case string:
+				if isKey {
+					lastKey = tt
+				} else if contextKey == "pageInfo" && lastKey == "endCursor" {
+					endCursor = tt
+					foundEndCursor = true
+					if foundNextPage {
+						break loop
+					}
+				}
+			case bool:
+				if contextKey == "pageInfo" && lastKey == "hasNextPage" {
+					hasNextPage = tt
+					foundNextPage = true
+					if foundEndCursor {
+						break loop
+					}
+				}
+			}
+		}
+	}
+
+	if hasNextPage {
+		return endCursor
+	}
+	return ""
+}

pkg/cmd/api/pagination_test.go

@@ -0,0 +1,118 @@
+package api
+
+import (
+	"bytes"
+	"io"
+	"net/http"
+	"testing"
+)
+
+func Test_findNextPage(t *testing.T) {
+	tests := []struct {
+		name  string
+		resp  *http.Response
+		want  string
+		want1 bool
+	}{
+		{
+			name:  "no Link header",
+			resp:  &http.Response{},
+			want:  "",
+			want1: false,
+		},
+		{
+			name: "no next page in Link",
+			resp: &http.Response{
+				Header: http.Header{
+					"Link": []string{`<https://api.github.com/issues?page=3>; rel="last"`},
+				},
+			},
+			want:  "",
+			want1: false,
+		},
+		{
+			name: "has next page",
+			resp: &http.Response{
+				Header: http.Header{
+					"Link": []string{`<https://api.github.com/issues?page=2>; rel="next", <https://api.github.com/issues?page=3>; rel="last"`},
+				},
+			},
+			want:  "https://api.github.com/issues?page=2",
+			want1: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, got1 := findNextPage(tt.resp)
+			if got != tt.want {
+				t.Errorf("findNextPage() got = %v, want %v", got, tt.want)
+			}
+			if got1 != tt.want1 {
+				t.Errorf("findNextPage() got1 = %v, want %v", got1, tt.want1)
+			}
+		})
+	}
+}
+
+func Test_findEndCursor(t *testing.T) {
+	tests := []struct {
+		name string
+		json io.Reader
+		want string
+	}{
+		{
+			name: "blank",
+			json: bytes.NewBufferString(`{}`),
+			want: "",
+		},
+		{
+			name: "unrelated fields",
+			json: bytes.NewBufferString(`{
+				"hasNextPage": true,
+				"endCursor": "THE_END"
+			}`),
+			want: "",
+		},
+		{
+			name: "has next page",
+			json: bytes.NewBufferString(`{
+				"pageInfo": {
+					"hasNextPage": true,
+					"endCursor": "THE_END"
+				}
+			}`),
+			want: "THE_END",
+		},
+		{
+			name: "more pageInfo blocks",
+			json: bytes.NewBufferString(`{
+				"pageInfo": {
+					"hasNextPage": true,
+					"endCursor": "THE_END"
+				},
+				"pageInfo": {
+					"hasNextPage": true,
+					"endCursor": "NOT_THIS"
+				}
+			}`),
+			want: "THE_END",
+		},
+		{
+			name: "no next page",
+			json: bytes.NewBufferString(`{
+				"pageInfo": {
+					"hasNextPage": false,
+					"endCursor": "THE_END"
+				}
+			}`),
+			want: "",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := findEndCursor(tt.json); got != tt.want {
+				t.Errorf("findEndCursor() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}