SG-20154 Add Retries on 503 Errors when uploading to S3 in Shotgun python-api (shotgunsoftware#263)

NorberMV · web-flow · commit e640baed8d2a · 2022-05-31T10:31:33.000-05:00
* Add the 503 retries in the internal function who maker the S3 uploads.
* Mock the S3 response with a expected HTTPError exception, for test 503 responses are retried when uploading to S3.
* Adding side_effect parameter to the mock object.
* Add comment in the _setup_mock() _make_upload_request definition.
* Rename test and add expected HTTPError exception error message to the assertion.
diff --git a/shotgun_api3/shotgun.py b/shotgun_api3/shotgun.py
@@ -3588,6 +3588,18 @@ def _http_request(self, verb, path, body, headers):
 
         return (http_status, resp_headers, resp_body)
 
+    def _make_upload_request(self, request, opener):
+        """
+        Open the given request object, return the
+        response, raises URLError on protocol errors.
+        """
+        try:
+            result = opener.open(request)
+
+        except urllib.error.HTTPError:
+            raise
+        return result
+
     def _parse_http_status(self, status):
         """
         Parse the status returned from the http request.
@@ -4049,21 +4061,40 @@ def _upload_data_to_storage(self, data, content_type, size, storage_url):
         :returns: upload url.
         :rtype: str
         """
-        try:
-            opener = self._build_opener(urllib.request.HTTPHandler)
+        opener = self._build_opener(urllib.request.HTTPHandler)
+
+        request = urllib.request.Request(storage_url, data=data)
+        request.add_header("Content-Type", content_type)
+        request.add_header("Content-Length", size)
+        request.get_method = lambda: "PUT"
+
+        attempt = 1
+        max_attempts = 4  # Three retries on failure
+        backoff = 0.75  # Seconds to wait before retry, times the attempt number
+
+        while attempt <= max_attempts:
+            try:
+                result = self._make_upload_request(request, opener)
+
+                LOG.debug("Completed request to %s" % request.get_method())
+
+            except urllib.error.HTTPError as e:
+                if e.code == 500:
+                    raise ShotgunError("Server encountered an internal error.\n%s\n%s\n\n" % (storage_url, e))
+                elif attempt != max_attempts and e.code == 503:
+                    LOG.debug("Got a 503 response. Waiting and retrying...")
+                    time.sleep(float(attempt) * backoff)
+                    attempt += 1
+                    continue
+                else:
+                    if e.code == 503:
+                        raise ShotgunError("Got a 503 response when uploading to %s: %s" % (storage_url, e))
+                    raise ShotgunError("Unanticipated error occurred uploading to %s: %s" % (storage_url, e))
 
-            request = urllib.request.Request(storage_url, data=data)
-            request.add_header("Content-Type", content_type)
-            request.add_header("Content-Length", size)
-            request.get_method = lambda: "PUT"
-            result = opener.open(request)
-            etag = result.info()["Etag"]
-        except urllib.error.HTTPError as e:
-            if e.code == 500:
-                raise ShotgunError("Server encountered an internal error.\n%s\n%s\n\n" % (storage_url, e))
             else:
-                raise ShotgunError("Unanticipated error occurred uploading to %s: %s" % (storage_url, e))
+                break
 
+        etag = result.info()["Etag"]
         LOG.debug("Part upload completed successfully.")
         return etag
 
diff --git a/tests/base.py b/tests/base.py
@@ -9,6 +9,7 @@
 from shotgun_api3.shotgun import json
 from shotgun_api3.shotgun import ServerCapabilities
 from shotgun_api3.lib import six
+from shotgun_api3.lib.six.moves import urllib
 
 if six.PY2:
     from shotgun_api3.lib.six.moves.configparser import SafeConfigParser as ConfigParser
@@ -128,7 +129,17 @@ def _setup_mock(self):
         # eaiser than mocking the http connection + response
         self.sg._http_request = mock.Mock(spec=api.Shotgun._http_request,
                                           return_value=((200, "OK"), {}, None))
-
+        # Replace the function used to make the final call to the S3 server, and simulate
+        # the exception HTTPError raised with 503 status errors
+        self.sg._make_upload_request = mock.Mock(spec=api.Shotgun._make_upload_request,
+                                                 side_effect = urllib.error.HTTPError(
+                                                     "url",
+                                                     503,
+                                                     "The server is currently down or to busy to reply."
+                                                     "Please try again later.",
+                                                     {},
+                                                     None
+                                                 ))
         # also replace the function that is called to get the http connection
         # to avoid calling the server. OK to return a mock as we will not use
         # it
diff --git a/tests/test_client.py b/tests/test_client.py
@@ -37,7 +37,6 @@
 from shotgun_api3.shotgun import ServerCapabilities, SG_TIMEZONE
 from . import base
 
-
 if six.PY3:
     from base64 import encodebytes as base64encode
 else:
@@ -196,7 +195,6 @@ def test_split_url(self):
         sg = api.Shotgun("https://ci.shotgunstudio.com",
                          "foo", "bar", connect=False)
 
-
         base_url = "https://ci.shotgunstudio.com"
         expected_server = "ci.shotgunstudio.com"
         expected_auth = None
@@ -439,6 +437,29 @@ def test_call_rpc(self):
         self._mock_http(d, status=(502, "bad gateway"))
         self.assertRaises(api.ProtocolError, self.sg._call_rpc, "list", a)
 
+    def test_upload_s3(self):
+        """
+        Test 503 response is retried when uploading to S3.
+        """
+        this_dir, _ = os.path.split(__file__)
+        storage_url = "http://foo.com/"
+        path = os.path.abspath(os.path.expanduser(
+            os.path.join(this_dir, "sg_logo.jpg")))
+        max_attempts = 4  # Max retries to S3 server attempts
+        # Expected HTTPError exception error message
+        expected = "The server is currently down or to busy to reply." \
+                   "Please try again later."
+
+        # Test the Internal function that is used to upload each
+        # data part in the context of multi-part uploads to S3, we
+        # simulate the HTTPError exception raised with 503 status errors
+        with self.assertRaises(api.ShotgunError, msg=expected):
+            self.sg._upload_file_to_storage(path, storage_url)
+        # Test the max retries attempt
+        self.assertTrue(
+            max_attempts == self.sg._make_upload_request.call_count,
+            "Call is repeated up to 3 times")
+
     def test_transform_data(self):
         """Outbound data is transformed"""
         timestamp = time.time()
