api: output a single JSON array in REST pagination mode (cli#7190)

When fetching N pages, avoid printing N separate JSON arrays to the output stream. Instead, massage the output so that the N pages of data are merged into a single JSON array. This is achieved by omitting the final `]` for the first page, and omitting the initial `[` for all subsequent pages.
kousikmitra · Jun 9, 2023 · 63a4319 · 63a4319
1 parent ad4a489
commit 63a4319
Show file tree

Hide file tree

Showing 3 changed files with 129 additions and 6 deletions.
diff --git a/pkg/cmd/api/api.go b/pkg/cmd/api/api.go
@@ -321,17 +321,24 @@ func apiRun(opts *ApiOptions) error {
 		return err
 	}
 
+	isFirstPage := true
 	hasNextPage := true
 	for hasNextPage {
 		resp, err := httpRequest(httpClient, host, method, requestPath, requestBody, requestHeaders)
 		if err != nil {
 			return err
 		}
 
-		endCursor, err := processResponse(resp, opts, bodyWriter, headersWriter, tmpl)
+		if !isGraphQL {
+			requestPath, hasNextPage = findNextPage(resp)
+			requestBody = nil // prevent repeating GET parameters
+		}
+
+		endCursor, err := processResponse(resp, opts, bodyWriter, headersWriter, tmpl, isFirstPage, !hasNextPage)
 		if err != nil {
 			return err
 		}
+		isFirstPage = false
 
 		if !opts.Paginate {
 			break
@@ -342,9 +349,6 @@ func apiRun(opts *ApiOptions) error {
 			if hasNextPage {
 				params["endCursor"] = endCursor
 			}
-		} else {
-			requestPath, hasNextPage = findNextPage(resp)
-			requestBody = nil // prevent repeating GET parameters
 		}
 
 		if hasNextPage && opts.ShowResponseHeaders {
@@ -355,7 +359,7 @@ func apiRun(opts *ApiOptions) error {
 	return tmpl.Flush()
 }
 
-func processResponse(resp *http.Response, opts *ApiOptions, bodyWriter, headersWriter io.Writer, template *template.Template) (endCursor string, err error) {
+func processResponse(resp *http.Response, opts *ApiOptions, bodyWriter, headersWriter io.Writer, template *template.Template, isFirstPage, isLastPage bool) (endCursor string, err error) {
 	if opts.ShowResponseHeaders {
 		fmt.Fprintln(headersWriter, resp.Proto, resp.Status)
 		printHeaders(headersWriter, resp.Header, opts.IO.ColorEnabled())
@@ -403,6 +407,13 @@ func processResponse(resp *http.Response, opts *ApiOptions, bodyWriter, headersW
 	} else if isJSON && opts.IO.ColorEnabled() {
 		err = jsoncolor.Write(bodyWriter, responseBody, "  ")
 	} else {
+		if isJSON && opts.Paginate && !isGraphQLPaginate && !opts.ShowResponseHeaders {
+			responseBody = &paginatedArrayReader{
+				Reader:      responseBody,
+				isFirstPage: isFirstPage,
+				isLastPage:  isLastPage,
+			}
+		}
 		_, err = io.Copy(bodyWriter, responseBody)
 	}
 	if err != nil {

diff --git a/pkg/cmd/api/api_test.go b/pkg/cmd/api/api_test.go
@@ -668,6 +668,78 @@ func Test_apiRun_paginationREST(t *testing.T) {
 	assert.Equal(t, "https://api.github.com/repositories/1227/issues?page=3", responses[2].Request.URL.String())
 }
 
+func Test_apiRun_arrayPaginationREST(t *testing.T) {
+	ios, _, stdout, stderr := iostreams.Test()
+	ios.SetStdoutTTY(false)
+
+	requestCount := 0
+	responses := []*http.Response{
+		{
+			StatusCode: 200,
+			Body:       io.NopCloser(bytes.NewBufferString(`[{"item":1},{"item":2}]`)),
+			Header: http.Header{
+				"Content-Type": []string{"application/json"},
+				"Link":         []string{`<https://api.github.com/repositories/1227/issues?page=2>; rel="next", <https://api.github.com/repositories/1227/issues?page=4>; rel="last"`},
+			},
+		},
+		{
+			StatusCode: 200,
+			Body:       io.NopCloser(bytes.NewBufferString(`[{"item":3},{"item":4}]`)),
+			Header: http.Header{
+				"Content-Type": []string{"application/json"},
+				"Link":         []string{`<https://api.github.com/repositories/1227/issues?page=3>; rel="next", <https://api.github.com/repositories/1227/issues?page=4>; rel="last"`},
+			},
+		},
+		{
+			StatusCode: 200,
+			Body:       io.NopCloser(bytes.NewBufferString(`[{"item":5}]`)),
+			Header: http.Header{
+				"Content-Type": []string{"application/json"},
+				"Link":         []string{`<https://api.github.com/repositories/1227/issues?page=4>; rel="next", <https://api.github.com/repositories/1227/issues?page=4>; rel="last"`},
+			},
+		},
+		{
+			StatusCode: 200,
+			Body:       io.NopCloser(bytes.NewBufferString(`[]`)),
+			Header: http.Header{
+				"Content-Type": []string{"application/json"},
+			},
+		},
+	}
+
+	options := ApiOptions{
+		IO: ios,
+		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
+		},
+		Config: func() (config.Config, error) {
+			return config.NewBlankConfig(), nil
+		},
+
+		RequestMethod:       "GET",
+		RequestMethodPassed: true,
+		RequestPath:         "issues",
+		Paginate:            true,
+		RawFields:           []string{"per_page=50", "page=1"},
+	}
+
+	err := apiRun(&options)
+	assert.NoError(t, err)
+
+	assert.Equal(t, `[{"item":1},{"item":2},{"item":3},{"item":4},{"item":5} ]`, stdout.String(), "stdout")
+	assert.Equal(t, "", stderr.String(), "stderr")
+
+	assert.Equal(t, "https://api.github.com/issues?page=1&per_page=50", responses[0].Request.URL.String())
+	assert.Equal(t, "https://api.github.com/repositories/1227/issues?page=2", responses[1].Request.URL.String())
+	assert.Equal(t, "https://api.github.com/repositories/1227/issues?page=3", responses[2].Request.URL.String())
+}
+
 func Test_apiRun_paginationGraphQL(t *testing.T) {
 	ios, _, stdout, stderr := iostreams.Test()
 
@@ -1236,7 +1308,7 @@ func Test_processResponse_template(t *testing.T) {
 	tmpl := template.New(ios.Out, ios.TerminalWidth(), ios.ColorEnabled())
 	err := tmpl.Parse(opts.Template)
 	require.NoError(t, err)
-	_, err = processResponse(&resp, &opts, ios.Out, io.Discard, tmpl)
+	_, err = processResponse(&resp, &opts, ios.Out, io.Discard, tmpl, true, true)
 	require.NoError(t, err)
 	err = tmpl.Flush()
 	require.NoError(t, err)

diff --git a/pkg/cmd/api/pagination.go b/pkg/cmd/api/pagination.go
@@ -106,3 +106,43 @@ func addPerPage(p string, perPage int, params map[string]interface{}) string {
 
 	return fmt.Sprintf("%s%sper_page=%d", p, sep, perPage)
 }
+
+// paginatedArrayReader wraps a Reader to omit the opening and/or the closing square bracket of a
+// JSON array in order to apply pagination context between multiple API requests.
+type paginatedArrayReader struct {
+	io.Reader
+	isFirstPage bool
+	isLastPage  bool
+
+	isSubsequentRead bool
+	cachedByte       byte
+}
+
+func (r *paginatedArrayReader) Read(p []byte) (int, error) {
+	var n int
+	var err error
+	if r.cachedByte != 0 && len(p) > 0 {
+		p[0] = r.cachedByte
+		n, err = r.Reader.Read(p[1:])
+		n += 1
+		r.cachedByte = 0
+	} else {
+		n, err = r.Reader.Read(p)
+	}
+	if !r.isSubsequentRead && !r.isFirstPage && n > 0 && p[0] == '[' {
+		if n > 1 && p[1] == ']' {
+			// empty array case
+			p[0] = ' '
+		} else {
+			// avoid starting a new array and continue with a comma instead
+			p[0] = ','
+		}
+	}
+	if !r.isLastPage && n > 0 && p[n-1] == ']' {
+		// avoid closing off an array in case we determine we are at EOF
+		r.cachedByte = p[n-1]
+		n -= 1
+	}
+	r.isSubsequentRead = true
+	return n, err
+}