Resolve the app token again when retrying after a rate-limit wait

github_backup/github_backup.py

@@ -901,23 +901,31 @@ def _extract_legal_url(response_body_bytes):
 
         next_url = None
 
+        def build_request(auth):
+            return _construct_request(
+                per_page=per_page if paginated else None,
+                query_args=query_args,
+                template=next_url or template,
+                auth=auth,
+                as_app=args.as_app,
+                fine=args.token_fine is not None,
+            )
+
         while True:
             # FIRST: Fetch response
 
             for attempt in range(args.max_retries + 1):
-                # Resolved per request so long runs pick up refreshed app tokens
+                # Resolved per request (and again before each retry) so long
+                # runs pick up refreshed app tokens
                 auth = get_auth(args, encode=not args.as_app)
-                request = _construct_request(
-                    per_page=per_page if paginated else None,
-                    query_args=query_args,
-                    template=next_url or template,
-                    auth=auth,
-                    as_app=args.as_app,
-                    fine=args.token_fine is not None,
-                )
                 try:
                     http_response = make_request_with_retry(
-                        request, auth, args.max_retries
+                        build_request(auth),
+                        auth,
+                        args.max_retries,
+                        rebuild_request=lambda: build_request(
+                            get_auth(args, encode=not args.as_app)
+                        ),
                     )
                 except HTTPError as exc:
                     if exc.code == 451:
@@ -1028,17 +1036,29 @@ def retrieve_graphql_data(args, query, variables=None, log_context=None):
     ).encode("utf-8")
     endpoint = get_github_graphql_url(args)
 
-    for attempt in range(args.max_retries + 1):
+    def build_request(auth):
         request = Request(endpoint, data=payload, method="POST")
         request.add_header("Accept", "application/json")
         request.add_header("Content-Type", "application/json")
         request.add_header("Authorization", "bearer " + auth)
+        return request
+
+    for attempt in range(args.max_retries + 1):
+        if attempt:
+            # Resolved again after a retry wait so long runs pick up refreshed
+            # app tokens
+            auth = get_graphql_auth(args)
         log_url = endpoint
         if log_context:
             log_url = "{0} ({1})".format(log_url, log_context)
         logger.info("Requesting {0}".format(log_url))
 
-        http_response = make_request_with_retry(request, auth, args.max_retries)
+        http_response = make_request_with_retry(
+            build_request(auth),
+            auth,
+            args.max_retries,
+            rebuild_request=lambda: build_request(get_graphql_auth(args)),
+        )
 
         status = http_response.getcode()
         if status != 200:
@@ -1097,8 +1117,14 @@ def retrieve_graphql_data(args, query, variables=None, log_context=None):
     )  # pragma: no cover
 
 
-def make_request_with_retry(request, auth, max_retries=5):
-    """Make HTTP request with automatic retry for transient errors."""
+def make_request_with_retry(request, auth, max_retries=5, rebuild_request=None):
+    """Make HTTP request with automatic retry for transient errors.
+
+    If given, rebuild_request is called to build the request again before each
+    retry. A rate-limit wait can last longer than a GitHub App installation
+    token, so callers that resolve credentials per request use it to resend
+    with the token that is current after the wait.
+    """
 
     def is_retryable_status(status_code, headers):
         # Server errors are always retryable
@@ -1110,6 +1136,8 @@ def is_retryable_status(status_code, headers):
         return False
 
     for attempt in range(max_retries + 1):
+        if attempt and rebuild_request is not None:
+            request = rebuild_request()
         try:
             return urlopen(request, context=https_ctx)
 

tests/test_app_auth.py

@@ -337,3 +337,89 @@ def mock_urlopen(request, *args, **kwargs):
 
     assert results == [{"id": 1}, {"id": 2}]
     assert page_auth == [b"token ghs_1", b"token ghs_2"]
+
+
+def rate_limited_403():
+    reset = int(time.time()) + 900
+    return HTTPError(
+        url="https://api.github.com/test",
+        code=403,
+        msg="rate limit exceeded",
+        hdrs={"x-ratelimit-remaining": "0", "x-ratelimit-reset": str(reset)},
+        fp=None,
+    )
+
+
+@pytest.fixture
+def token_rotates_during_wait():
+    """The installation token is replaced while a retry is waiting, as happens
+    when the rate limit resets after the token has expired."""
+    current = {"token": "ghs_1"}
+
+    def sleep(seconds):
+        current["token"] = "ghs_2"
+
+    with patch(
+        "github_backup.github_backup.get_app_installation_token",
+        side_effect=lambda args: current["token"],
+    ), patch("github_backup.github_backup.time.sleep", side_effect=sleep):
+        yield
+
+
+def test_rate_limit_retry_picks_up_refreshed_token(app_args, token_rotates_during_wait):
+    """Retrying after a rate-limit wait sends the token that is current after the
+    wait, not the one the request was first built with."""
+    sent_auth = []
+
+    def mock_urlopen(request, *args, **kwargs):
+        sent_auth.append(request.get_header("Authorization"))
+        if len(sent_auth) == 1:
+            raise rate_limited_403()
+        return MockHTTPResponse([{"id": 1}])
+
+    with patch("github_backup.github_backup.urlopen", side_effect=mock_urlopen):
+        results = github_backup.retrieve_data(
+            app_args, "https://api.github.com/orgs/testorg/repos"
+        )
+
+    assert results == [{"id": 1}]
+    assert sent_auth == [b"token ghs_1", b"token ghs_2"]
+
+
+def test_graphql_rate_limit_retry_picks_up_refreshed_token(
+    app_args, token_rotates_during_wait
+):
+    sent_auth = []
+
+    def mock_urlopen(request, *args, **kwargs):
+        sent_auth.append(request.get_header("Authorization"))
+        if len(sent_auth) == 1:
+            raise rate_limited_403()
+        return MockHTTPResponse({"data": {"ok": True}})
+
+    with patch("github_backup.github_backup.urlopen", side_effect=mock_urlopen):
+        data = github_backup.retrieve_graphql_data(app_args, "query { ok }")
+
+    assert data == {"ok": True}
+    assert sent_auth == ["bearer ghs_1", "bearer ghs_2"]
+
+
+def test_graphql_rate_limited_error_retry_picks_up_refreshed_token(
+    app_args, token_rotates_during_wait
+):
+    """GraphQL can also report the rate limit as a RATE_LIMITED error in a 200."""
+    sent_auth = []
+
+    def mock_urlopen(request, *args, **kwargs):
+        sent_auth.append(request.get_header("Authorization"))
+        if len(sent_auth) == 1:
+            return MockHTTPResponse(
+                {"errors": [{"type": "RATE_LIMITED", "message": "rate limited"}]}
+            )
+        return MockHTTPResponse({"data": {"ok": True}})
+
+    with patch("github_backup.github_backup.urlopen", side_effect=mock_urlopen):
+        data = github_backup.retrieve_graphql_data(app_args, "query { ok }")
+
+    assert data == {"ok": True}
+    assert sent_auth == ["bearer ghs_1", "bearer ghs_2"]