diff --git a/github_backup/github_backup.py b/github_backup/github_backup.py index 64c08c3..fdf2e22 100644 --- a/github_backup/github_backup.py +++ b/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) diff --git a/tests/test_app_auth.py b/tests/test_app_auth.py index a206bc6..d10c3e6 100644 --- a/tests/test_app_auth.py +++ b/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"]