From c8bbb27e46489b61d68151b88b2fc21ba848fb2c Mon Sep 17 00:00:00 2001
From: Nejc Habjan <nejc.habjan@siemens.com>
Date: Tue, 22 Nov 2022 16:32:25 +0100
Subject: [PATCH] fix: respect retry-after header with 429 responses

---
 src/poetry/publishing/uploader.py |  1 +
 src/poetry/utils/authenticator.py | 12 +++++++++++-
 src/poetry/utils/constants.py     |  4 +++-
 tests/utils/test_authenticator.py | 28 ++++++++++++++++++++++++++++
 4 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/src/poetry/publishing/uploader.py b/src/poetry/publishing/uploader.py
index 7e6cd6d9b6a..256fd235259 100644
--- a/src/poetry/publishing/uploader.py
+++ b/src/poetry/publishing/uploader.py
@@ -69,6 +69,7 @@ def adapter(self) -> adapters.HTTPAdapter:
             connect=5,
             total=10,
             allowed_methods=["GET"],
+            respect_retry_after_header=True,
             status_forcelist=STATUS_FORCELIST,
         )
 
diff --git a/src/poetry/utils/authenticator.py b/src/poetry/utils/authenticator.py
index 2b1011e4d56..0fb238fb56c 100644
--- a/src/poetry/utils/authenticator.py
+++ b/src/poetry/utils/authenticator.py
@@ -24,6 +24,7 @@
 from poetry.config.config import Config
 from poetry.exceptions import PoetryException
 from poetry.utils.constants import REQUESTS_TIMEOUT
+from poetry.utils.constants import RETRY_AFTER_HEADER
 from poetry.utils.constants import STATUS_FORCELIST
 from poetry.utils.password_manager import HTTPAuthCredential
 from poetry.utils.password_manager import PasswordManager
@@ -251,6 +252,7 @@ def request(
         send_kwargs.update(settings)
 
         attempt = 0
+        resp = None
 
         while True:
             is_last_attempt = attempt >= 5
@@ -267,7 +269,7 @@ def request(
 
             if not is_last_attempt:
                 attempt += 1
-                delay = 0.5 * attempt
+                delay = self._get_backoff(resp, attempt)
                 logger.debug("Retrying HTTP request in %s seconds.", delay)
                 time.sleep(delay)
                 continue
@@ -275,6 +277,14 @@ def request(
         # this should never really be hit under any sane circumstance
         raise PoetryException("Failed HTTP {} request", method.upper())
 
+    def _get_backoff(self, response: requests.Response | None, attempt: int) -> float:
+        if response is not None:
+            retry_after = response.headers.get(RETRY_AFTER_HEADER, "")
+            if retry_after:
+                return float(retry_after)
+
+        return 0.5 * attempt
+
     def get(self, url: str, **kwargs: Any) -> requests.Response:
         return self.request("get", url, **kwargs)
 
diff --git a/src/poetry/utils/constants.py b/src/poetry/utils/constants.py
index b755fb68e5e..56bec540ae2 100644
--- a/src/poetry/utils/constants.py
+++ b/src/poetry/utils/constants.py
@@ -4,5 +4,7 @@
 # Timeout for HTTP requests using the requests library.
 REQUESTS_TIMEOUT = 15
 
+RETRY_AFTER_HEADER = "retry-after"
+
 # Server response codes to retry requests on.
-STATUS_FORCELIST = [500, 501, 502, 503, 504]
+STATUS_FORCELIST = [429, 500, 501, 502, 503, 504]
diff --git a/tests/utils/test_authenticator.py b/tests/utils/test_authenticator.py
index 545b73f53bf..05c5490ac11 100644
--- a/tests/utils/test_authenticator.py
+++ b/tests/utils/test_authenticator.py
@@ -242,6 +242,33 @@ def callback(*_: Any, **___: Any) -> None:
     assert sleep.call_count == 5
 
 
+def test_authenticator_request_respects_retry_header(
+    mocker: MockerFixture,
+    config: Config,
+    http: type[httpretty.httpretty],
+):
+    sleep = mocker.patch("time.sleep")
+    sdist_uri = f"https://foo.bar/files/{uuid.uuid4()!s}/foo-0.1.0.tar.gz"
+    content = str(uuid.uuid4())
+    seen = []
+
+    def callback(
+        request: requests.Request, uri: str, response_headers: dict
+    ) -> list[int | dict | str]:
+        if not seen.count(uri):
+            seen.append(uri)
+            return [429, {"Retry-After": "42"}, "Retry later"]
+
+        return [200, response_headers, content]
+
+    http.register_uri(httpretty.GET, sdist_uri, body=callback)
+    authenticator = Authenticator(config, NullIO())
+
+    response = authenticator.request("get", sdist_uri)
+    assert sleep.call_args[0] == (42.0,)
+    assert response.text == content
+
+
 @pytest.mark.parametrize(
     ["status", "attempts"],
     [
@@ -249,6 +276,7 @@ def callback(*_: Any, **___: Any) -> None:
         (401, 0),
         (403, 0),
         (404, 0),
+        (429, 5),
         (500, 5),
         (501, 5),
         (502, 5),