Add humble retry logic for package downloads

poetry/installation/authenticator.py

@@ -1,5 +1,12 @@
+import time
+
 from typing import TYPE_CHECKING
 
+import requests
+import requests.auth
+import requests.exceptions
+
+from poetry.exceptions import PoetryException
 from poetry.utils._compat import urlparse
 from poetry.utils.password_manager import PasswordManager
 
@@ -10,10 +17,6 @@
     from typing import Tuple
 
     from clikit.api.io import IO
-    from requests import Request  # noqa
-    from requests import Response  # noqa
-    from requests import Session  # noqa
-
     from poetry.config.config import Config
 
 
@@ -26,24 +29,22 @@ def __init__(self, config, io):  # type: (Config, IO) -> None
         self._password_manager = PasswordManager(self._config)
 
     @property
-    def session(self):  # type: () -> Session
-        from requests import Session  # noqa
-
+    def session(self):  # type: () -> requests.Session
         if self._session is None:
-            self._session = Session()
+            self._session = requests.Session()
 
         return self._session
 
-    def request(self, method, url, **kwargs):  # type: (str, str, Any) -> Response
-        from requests import Request  # noqa
-        from requests.auth import HTTPBasicAuth
-
-        request = Request(method, url)
+    def request(
+        self, method, url, **kwargs
+    ):  # type: (str, str, Any) -> requests.Response
+        request = requests.Request(method, url)
+        io = kwargs.get("io", self._io)
 
         username, password = self._get_credentials_for_url(url)
 
         if username is not None and password is not None:
-            request = HTTPBasicAuth(username, password)(request)
+            request = requests.auth.HTTPBasicAuth(username, password)(request)
 
         session = self.session
         prepared_request = session.prepare_request(request)
@@ -63,11 +64,31 @@ def request(self, method, url, **kwargs):  # type: (str, str, Any) -> Response
             "allow_redirects": kwargs.get("allow_redirects", True),
         }
         send_kwargs.update(settings)
-        resp = session.send(prepared_request, **send_kwargs)
 
-        resp.raise_for_status()
+        attempt = 0
+
+        while True:
+            is_last_attempt = attempt >= 5
+            try:
+                resp = session.send(prepared_request, **send_kwargs)
+                if resp.status_code not in [502, 503, 504] or is_last_attempt:
+                    resp.raise_for_status()
+                    return resp
+            except (requests.exceptions.ConnectionError, ConnectionError) as e:
+                if is_last_attempt:
+                    raise e
+
+            if not is_last_attempt:
+                attempt += 1
+                delay = 0.5 * attempt
+                io.write_line(
+                    "<debug>Retrying HTTP request in {} seconds.</debug>".format(delay)
+                )
+                time.sleep(delay)
+                continue
 
-        return resp
+        # this should never really be hit under any sane circumstance
+        raise PoetryException("Failed HTTP {} request", method.upper())
 
     def _get_credentials_for_url(
         self, url

poetry/installation/executor.py

@@ -603,7 +603,9 @@ def _download_link(self, operation, link):
         return archive
 
     def _download_archive(self, operation, link):  # type: (Operation, Link) -> Path
-        response = self._authenticator.request("get", link.url, stream=True)
+        response = self._authenticator.request(
+            "get", link.url, stream=True, io=self._sections.get(id(operation))
+        )
         wheel_size = response.headers.get("content-length")
         operation_message = self.get_operation_message(operation)
         message = "  <fg=blue;options=bold>•</> {message}: <info>Downloading...</>".format(

tests/installation/test_authenticator.py

@@ -1,6 +1,9 @@
 import re
+import uuid
 
+import httpretty
 import pytest
+import requests
 
 from poetry.installation.authenticator import Authenticator
 from poetry.io.null_io import NullIO
@@ -113,3 +116,66 @@ def test_authenticator_uses_empty_strings_as_default_username(
     request = http.last_request()
 
     assert "Basic OmJhcg==" == request.headers["Authorization"]
+
+
+@httpretty.activate(allow_net_connect=False)
+def test_authenticator_request_retries_on_exception(mocker, config):
+    sleep = mocker.patch("time.sleep")
+    sdist_uri = "https://foo.bar/files/{}/foo-0.1.0.tar.gz".format(str(uuid.uuid4()))
+    content = str(uuid.uuid4())
+    seen = list()
+
+    def callback(request, uri, response_headers):
+        if seen.count(uri) < 2:
+            seen.append(uri)
+            raise requests.exceptions.ConnectionError("Disconnected")
+        return [200, response_headers, content]
+
+    httpretty.register_uri(httpretty.GET, sdist_uri, body=callback)
+
+    authenticator = Authenticator(config, NullIO())
+    response = authenticator.request("get", sdist_uri)
+    assert response.text == content
+    assert sleep.call_count == 2
+
+
+@httpretty.activate(allow_net_connect=False)
+def test_authenticator_request_raises_exception_when_attempts_exhausted(mocker, config):
+    sleep = mocker.patch("time.sleep")
+    sdist_uri = "https://foo.bar/files/{}/foo-0.1.0.tar.gz".format(str(uuid.uuid4()))
+
+    def callback(*_, **__):
+        raise requests.exceptions.ConnectionError(str(uuid.uuid4()))
+
+    httpretty.register_uri(httpretty.GET, sdist_uri, body=callback)
+    authenticator = Authenticator(config, NullIO())
+
+    with pytest.raises(requests.exceptions.ConnectionError):
+        authenticator.request("get", sdist_uri)
+
+    assert sleep.call_count == 5
+
+
+@httpretty.activate(allow_net_connect=False)
+@pytest.mark.parametrize(
+    "status, attempts",
+    [(400, 0), (401, 0), (403, 0), (404, 0), (500, 0), (502, 5), (503, 5), (504, 5)],
+)
+def test_authenticator_request_retries_on_status_code(mocker, config, status, attempts):
+    sleep = mocker.patch("time.sleep")
+    sdist_uri = "https://foo.bar/files/{}/foo-0.1.0.tar.gz".format(str(uuid.uuid4()))
+    content = str(uuid.uuid4())
+
+    def callback(request, uri, response_headers):
+        return [status, response_headers, content]
+
+    httpretty.register_uri(httpretty.GET, sdist_uri, body=callback)
+    authenticator = Authenticator(config, NullIO())
+
+    with pytest.raises(requests.exceptions.HTTPError) as excinfo:
+        authenticator.request("get", sdist_uri)
+
+    assert excinfo.value.response.status_code == status
+    assert excinfo.value.response.text == content
+
+    assert sleep.call_count == attempts