Repository: piskvorky/smart_open
Pull request #785 (2023-09-22 07:05:37)

Retry finalizing multipart s3 upload
#### Title

Please **pick a concise, informative and complete title** for your PR.
The title is important because it will appear in [our change log](https://github.com/RaRe-Technologies/smart_open/blob/master/CHANGELOG.md).

#### Motivation

Please explain the motivation behind this PR in the description.

If you're fixing a bug, link to the issue number like so:

- Fixes #628

If you're adding a new feature, then consider opening a ticket and discussing it with the maintainers before you actually do the hard work.

#### Tests

If you're fixing a bug, consider [test-driven development](https://en.wikipedia.org/wiki/Test-driven_development):

1. Create a unit test that demonstrates the bug. The test should **fail**.
2. Implement your bug fix.
3. The test you created should now **pass**.

If you're implementing a new feature, include unit tests for it.

Make sure all existing unit tests pass.
You can run them locally using:

    pytest smart_open

If there are any failures, please fix them before creating the PR (or mark it as WIP, see below).

#### Work in progress

If you're still working on your PR, include "WIP" in the title.
We'll skip reviewing it for the time being.
Once you're ready to review, remove the "WIP" from the title, and ping one of the maintainers (e.g. mpenkov).

#### Checklist

Before you create the PR, please make sure you have:

- [ ] Picked a concise, informative and complete title
- [ ] Clearly explained the motivation behind the PR
- [ ] Linked to any existing issues that your PR will be solving
- [ ] Included tests for any new functionality
- [ ] Checked that all unit tests pass

#### Workflow

Please avoid rebasing and force-pushing to the branch of the PR once a review is in progress.
Rebasing can make your commits look a bit cleaner, but it also makes life more difficult from the reviewer, because they are no longer able to distinguish between code that has already been reviewed, and unreviewed code.



Patch:
diff --git a/smart_open/s3.py b/smart_open/s3.py
index 60f3ef0..7b6dfe9 100644
--- a/smart_open/s3.py
+++ b/smart_open/s3.py
@@ -28,6 +28,11 @@ import smart_open.utils
 
 from smart_open import constants
 
+from typing import (
+    Callable,
+    List,
+)
+
 logger = logging.getLogger(__name__)
 
 DEFAULT_MIN_PART_SIZE = 50 * 1024**2
@@ -47,13 +52,52 @@ URI_EXAMPLES = (
     's3://my_key:my_secret@my_server:my_port@my_bucket/my_key',
 )
 
-_UPLOAD_ATTEMPTS = 6
-_SLEEP_SECONDS = 10
-
 # Returned by AWS when we try to seek beyond EOF.
 _OUT_OF_RANGE = 'InvalidRange'
 
 
+class Retry:
+    def __init__(self):
+        self.attempts: int = 6
+        self.sleep_seconds: int = 10
+        self.exceptions: List[Exception] = [botocore.exceptions.EndpointConnectionError]
+        self.client_error_codes: List[str] = ['NoSuchUpload']
+
+    def _do(self, fn: Callable):
+        for attempt in range(self.attempts):
+            try:
+                return fn()
+            except tuple(self.exceptions) as err:
+                logger.critical(
+                    'Caught non-fatal %s, retrying %d more times',
+                    err,
+                    self.attempts - attempt - 1,
+                )
+                logger.exception(err)
+                time.sleep(self.sleep_seconds)
+            except botocore.exceptions.ClientError as err:
+                error_code = err.response['Error'].get('Code')
+                if error_code not in self.client_error_codes:
+                    raise
+                logger.critical(
+                    'Caught non-fatal ClientError (%s), retrying %d more times',
+                    error_code,
+                    self.attempts - attempt - 1,
+                )
+                logger.exception(err)
+                time.sleep(self.sleep_seconds)
+        else:
+            logger.critical('encountered too many non-fatal errors, giving up')
+            raise IOError('%s failed after %d attempts', fn.func, self.attempts)
+
+
+#
+# The retry mechanism for this submodule.  Client code may modify it, e.g. by
+# updating RETRY.sleep_seconds and friends.
+#
+RETRY = Retry()
+
+
 class _ClientWrapper:
     """Wraps a client to inject the appropriate keyword args into each method call.
 
@@ -803,7 +847,7 @@ multipart upload may fail")
                 Bucket=bucket,
                 Key=key,
             )
-            self._upload_id = _retry_if_failed(partial)['UploadId']
+            self._upload_id = RETRY._do(partial)['UploadId']
         except botocore.client.ClientError as error:
             raise ValueError(
                 'the bucket %r does not exist, or is forbidden for access (%r)' % (
@@ -843,7 +887,7 @@ multipart upload may fail")
                 UploadId=self._upload_id,
                 MultipartUpload={'Parts': self._parts},
             )
-            _retry_if_failed(partial)
+            RETRY._do(partial)
             logger.debug('%s: completed multipart upload', self)
         elif self._upload_id:
             #
@@ -954,7 +998,7 @@ multipart upload may fail")
         # of a temporary connection problem, so this part needs to be
         # especially robust.
         #
-        upload = _retry_if_failed(
+        upload = RETRY._do(
             functools.partial(
                 self._client.upload_part,
                 Bucket=self._bucket,
@@ -1119,28 +1163,6 @@ class SinglepartWriter(io.BufferedIOBase):
         return "smart_open.s3.SinglepartWriter(bucket=%r, key=%r)" % (self._bucket, self._key)
 
 
-def _retry_if_failed(
-        partial,
-        attempts=_UPLOAD_ATTEMPTS,
-        sleep_seconds=_SLEEP_SECONDS,
-        exceptions=None):
-    if exceptions is None:
-        exceptions = (botocore.exceptions.EndpointConnectionError, )
-    for attempt in range(attempts):
-        try:
-            return partial()
-        except exceptions:
-            logger.critical(
-                'Unable to connect to the endpoint. Check your network connection. '
-                'Sleeping and retrying %d more times '
-                'before giving up.' % (attempts - attempt - 1)
-            )
-            time.sleep(sleep_seconds)
-    else:
-        logger.critical('Unable to connect to the endpoint. Giving up.')
-        raise IOError('Unable to connect to the endpoint after %d attempts' % attempts)
-
-
 def _accept_all(key):
     return True
 


Test patch:
diff --git a/smart_open/tests/test_s3.py b/smart_open/tests/test_s3.py
index b2ab87b..d6e8c76 100644
--- a/smart_open/tests/test_s3.py
+++ b/smart_open/tests/test_s3.py
@@ -21,6 +21,7 @@ import sys
 import boto3
 import botocore.client
 import botocore.endpoint
+import botocore.exceptions
 import pytest
 
 # See https://github.com/piskvorky/smart_open/issues/800
@@ -952,19 +953,32 @@ def populate_bucket(num_keys=10):
 
 
 class RetryIfFailedTest(unittest.TestCase):
+    def setUp(self):
+        self.retry = smart_open.s3.Retry()
+        self.retry.attempts = 3
+        self.retry.sleep_seconds = 0
+
     def test_success(self):
         partial = mock.Mock(return_value=1)
-        result = smart_open.s3._retry_if_failed(partial, attempts=3, sleep_seconds=0)
+        result = self.retry._do(partial)
         self.assertEqual(result, 1)
         self.assertEqual(partial.call_count, 1)
 
-    def test_failure(self):
+    def test_failure_exception(self):
         partial = mock.Mock(side_effect=ValueError)
-        exceptions = (ValueError, )
-
+        self.retry.exceptions = {ValueError: 'Let us retry ValueError'}
         with self.assertRaises(IOError):
-            smart_open.s3._retry_if_failed(partial, attempts=3, sleep_seconds=0, exceptions=exceptions)
+            self.retry._do(partial)
+        self.assertEqual(partial.call_count, 3)
 
+    def test_failure_client_error(self):
+        partial = mock.Mock(
+            side_effect=botocore.exceptions.ClientError(
+                {'Error': {'Code': 'NoSuchUpload'}}, 'NoSuchUpload'
+            )
+        )
+        with self.assertRaises(IOError):
+            self.retry._do(partial)
         self.assertEqual(partial.call_count, 3)