Merge pull request #2317 from mdaniel/api-rate-limit

Honor the rate limit headers in "hub api"

Author: mislav

commands/api.go

@@ -9,6 +9,7 @@ import (
 	"regexp"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/github/hub/github"
 	"github.com/github/hub/ui"
@@ -73,7 +74,10 @@ var cmdApi = &Command{
 		resource as indicated in the "Link" response header. For GraphQL queries,
 		this utilizes 'pageInfo' that must be present in the query; see EXAMPLES.
 
-		Note that multiple JSON documents will be output as a result.
+		Note that multiple JSON documents will be output as a result. If the API
+		rate limit has been reached, the final document that is output will be the
+		HTTP 403 notice, and the process will exit with a non-zero status. One way
+		this can be avoided is by enabling '--obey-ratelimit'.
 
 	--color[=<WHEN>]
 		Enable colored output even if stdout is not a terminal. <WHEN> can be one
@@ -86,6 +90,11 @@ var cmdApi = &Command{
 		requests as well. Just make sure to not use '--cache' for any GraphQL
 		mutations.
 
+	--obey-ratelimit
+		After exceeding the API rate limit, pause the process until the reset time
+		of the current rate limit window and retry the request. Note that this may
+		cause the process to hang for a long time (maximum of 1 hour).
+
 	<ENDPOINT>
 		The GitHub API endpoint to send the HTTP request to (default: "/").
 		
@@ -136,7 +145,7 @@ func init() {
 	CmdRunner.Use(cmdApi)
 }
 
-func apiCommand(cmd *Command, args *Args) {
+func apiCommand(_ *Command, args *Args) {
 	path := ""
 	if !args.IsParamsEmpty() {
 		path = args.GetParam(0)
@@ -235,15 +244,20 @@ func apiCommand(cmd *Command, args *Args) {
 	parseJSON := args.Flag.Bool("--flat")
 	includeHeaders := args.Flag.Bool("--include")
 	paginate := args.Flag.Bool("--paginate")
+	rateLimitWait := args.Flag.Bool("--obey-ratelimit")
 
 	args.NoForward()
 
-	requestLoop := true
-	for requestLoop {
+	for {
 		response, err := gh.GenericAPIRequest(method, path, body, headers, cacheTTL)
 		utils.Check(err)
-		success := response.StatusCode < 300
 
+		if rateLimitWait && response.StatusCode == 403 && response.RateLimitRemaining() == 0 {
+			pauseUntil(response.RateLimitReset())
+			continue
+		}
+
+		success := response.StatusCode < 300
 		jsonType := true
 		if !success {
 			jsonType, _ = regexp.MatchString(`[/+]json(?:;|$)`, response.Header.Get("Content-Type"))
@@ -273,7 +287,6 @@ func apiCommand(cmd *Command, args *Args) {
 			os.Exit(22)
 		}
 
-		requestLoop = false
 		if paginate {
 			if isGraphQL && hasNextPage && endCursor != "" {
 				if v, ok := params["variables"]; ok {
@@ -283,15 +296,31 @@ func apiCommand(cmd *Command, args *Args) {
 					variables := map[string]interface{}{"endCursor": endCursor}
 					params["variables"] = variables
 				}
-				requestLoop = true
+				goto next
 			} else if nextLink := response.Link("next"); nextLink != "" {
 				path = nextLink
-				requestLoop = true
+				goto next
 			}
 		}
-		if requestLoop && !parseJSON {
+
+		break
+	next:
+		if !parseJSON {
 			fmt.Fprintf(out, "\n")
 		}
+
+		if rateLimitWait && response.RateLimitRemaining() == 0 {
+			pauseUntil(response.RateLimitReset())
+		}
+	}
+}
+
+func pauseUntil(timestamp int) {
+	rollover := time.Unix(int64(timestamp)+1, 0)
+	duration := time.Until(rollover)
+	if duration > 0 {
+		ui.Errorf("API rate limit exceeded; pausing until %v ...\n", rollover)
+		time.Sleep(duration)
 	}
 }
 

features/api.feature

@@ -421,3 +421,73 @@ Feature: hub api
     Given I am "octocat" on github.com with OAuth token "TOKEN2"
     When I run `hub api -t count --cache 5`
     Then it should pass with ".count	2"
+
+  Scenario: Honor rate limit with pagination
+    Given the GitHub API server:
+      """
+      get('/hello') {
+        page = (params[:page] || 1).to_i
+        if page < 2
+          response.headers['X-Ratelimit-Remaining'] = '0'
+          response.headers['X-Ratelimit-Reset'] = Time.now.utc.to_i.to_s
+          response.headers['Link'] = %(</hello?page=#{page+1}>; rel="next")
+        end
+        json [{}]
+      }
+      """
+    When I successfully run `hub api --obey-ratelimit --paginate hello`
+    Then the stderr should contain "API rate limit exceeded; pausing until "
+
+  Scenario: Succumb to rate limit with pagination
+    Given the GitHub API server:
+      """
+      get('/hello') {
+        page = (params[:page] || 1).to_i
+        response.headers['X-Ratelimit-Remaining'] = '0'
+        response.headers['X-Ratelimit-Reset'] = Time.now.utc.to_i.to_s
+        if page == 2
+          status 403
+          json :message => "API rate limit exceeded"
+        else
+          response.headers['Link'] = %(</hello?page=#{page+1}>; rel="next")
+          json [{page:page}]
+        end
+      }
+      """
+    When I run `hub api --paginate -t hello`
+    Then the exit status should be 22
+    And the stderr should not contain "API rate limit exceeded"
+    And the stdout should contain exactly:
+      """
+      .[0].page	1
+      .message	API rate limit exceeded\n
+      """
+
+  Scenario: Honor rate limit for 403s
+    Given the GitHub API server:
+      """
+      count = 0
+      get('/hello') {
+        count += 1
+        if count == 1
+          response.headers['X-Ratelimit-Remaining'] = '0'
+          response.headers['X-Ratelimit-Reset'] = Time.now.utc.to_i.to_s
+          halt 403
+        end
+        json [{}]
+      }
+      """
+    When I successfully run `hub api --obey-ratelimit hello`
+    Then the stderr should contain "API rate limit exceeded; pausing until "
+
+  Scenario: 403 unrelated to rate limit
+    Given the GitHub API server:
+      """
+      get('/hello') {
+        response.headers['X-Ratelimit-Remaining'] = '1'
+        status 403
+      }
+      """
+    When I run `hub api --obey-ratelimit hello`
+    Then the exit status should be 22
+    Then the stderr should not contain "API rate limit exceeded"

github/http.go

@@ -32,6 +32,11 @@ const checksType = "application/vnd.github.antiope-preview+json;charset=utf-8"
 const draftsType = "application/vnd.github.shadow-cat-preview+json;charset=utf-8"
 const cacheVersion = 2
 
+const (
+	rateLimitRemainingHeader = "X-Ratelimit-Remaining"
+	rateLimitResetHeader     = "X-Ratelimit-Reset"
+)
+
 var inspectHeaders = []string{
 	"Authorization",
 	"X-GitHub-OTP",
@@ -517,3 +522,21 @@ func (res *simpleResponse) Link(name string) string {
 	}
 	return ""
 }
+
+func (res *simpleResponse) RateLimitRemaining() int {
+	if v := res.Header.Get(rateLimitRemainingHeader); len(v) > 0 {
+		if num, err := strconv.Atoi(v); err == nil {
+			return num
+		}
+	}
+	return -1
+}
+
+func (res *simpleResponse) RateLimitReset() int {
+	if v := res.Header.Get(rateLimitResetHeader); len(v) > 0 {
+		if ts, err := strconv.Atoi(v); err == nil {
+			return ts
+		}
+	}
+	return -1
+}