Add error translation from S3ResponseError to Botocore errors

Co-authored-by: Michael Graeb <graebm@amazon.com>

s3transfer/crt.py

@@ -29,6 +29,7 @@
     S3Client,
     S3RequestTlsMode,
     S3RequestType,
+    S3ResponseError,
     get_recommended_throughput_target_gbps,
 )
 from botocore import UNSIGNED
@@ -196,6 +197,9 @@ def __init__(self, crt_s3_client, crt_request_serializer, osutil=None):
         self._s3_args_creator = S3ClientArgsCreator(
             crt_request_serializer, self._osutil
         )
+        self._crt_exception_translator = (
+            crt_request_serializer.translate_crt_exception
+        )
         self._future_coordinators = []
         self._semaphore = threading.Semaphore(128)  # not configurable
         # A counter to create unique id's for each transfer submitted.
@@ -299,7 +303,10 @@ def _release_semaphore(self, **kwargs):
 
     def _submit_transfer(self, request_type, call_args):
         on_done_after_calls = [self._release_semaphore]
-        coordinator = CRTTransferCoordinator(transfer_id=self._id_counter)
+        coordinator = CRTTransferCoordinator(
+            transfer_id=self._id_counter,
+            exception_translator=self._crt_exception_translator,
+        )
         components = {
             'meta': CRTTransferMeta(self._id_counter, call_args),
             'coordinator': coordinator,
@@ -410,6 +417,9 @@ def serialize_http_request(self, transfer_type, future):
         """
         raise NotImplementedError('serialize_http_request()')
 
+    def translate_crt_exception(self, exception):
+        raise NotImplementedError('translate_crt_exception()')
+
 
 class BotocoreCRTRequestSerializer(BaseCRTRequestSerializer):
     def __init__(self, session, client_kwargs=None):
@@ -533,6 +543,40 @@ def serialize_http_request(self, transfer_type, future):
         crt_request = self._convert_to_crt_http_request(botocore_http_request)
         return crt_request
 
+    def translate_crt_exception(self, exception):
+        if isinstance(exception, S3ResponseError):
+            return self._translate_crt_s3_response_error(exception)
+        else:
+            return None
+
+    def _translate_crt_s3_response_error(self, s3_response_error):
+        status_code = s3_response_error.status_code
+        if status_code < 301:
+            # Botocore's exception parsing only
+            # runs on status codes >= 301
+            return None
+
+        headers = {k: v for k, v in s3_response_error.headers}
+        operation_name = s3_response_error.operation_name
+        if operation_name is not None:
+            service_model = self._client.meta.service_model
+            shape = service_model.operation_model(operation_name).output_shape
+        else:
+            shape = None
+
+        response_dict = {
+            'headers': botocore.awsrequest.HeadersDict(headers),
+            'status_code': status_code,
+            'body': s3_response_error.body,
+        }
+        parsed_response = self._client._response_parser.parse(
+            response_dict, shape=shape
+        )
+
+        error_code = parsed_response.get("Error", {}).get("Code")
+        error_class = self._client.exceptions.from_code(error_code)
+        return error_class(parsed_response, operation_name=operation_name)
+
 
 class FakeRawResponse(BytesIO):
     def stream(self, amt=1024, decode_content=None):
@@ -565,8 +609,11 @@ def _get_credentials(self):
 class CRTTransferCoordinator:
     """A helper class for managing CRTTransferFuture"""
 
-    def __init__(self, transfer_id=None, s3_request=None):
+    def __init__(
+        self, transfer_id=None, s3_request=None, exception_translator=None
+    ):
         self.transfer_id = transfer_id
+        self._exception_translator = exception_translator
         self._s3_request = s3_request
         self._lock = threading.Lock()
         self._exception = None
@@ -599,11 +646,27 @@ def result(self, timeout=None):
             self._crt_future.result(timeout)
         except KeyboardInterrupt:
             self.cancel()
+            self._crt_future.result(timeout)
             raise
+        except Exception as e:
+            self.handle_exception(e)
         finally:
             if self._s3_request:
                 self._s3_request = None
-            self._crt_future.result(timeout)
+
+    def handle_exception(self, exc):
+        translated_exc = None
+        if self._exception_translator:
+            try:
+                translated_exc = self._exception_translator(exc)
+            except Exception:
+                # Bail out if we hit an issue translating
+                # and raise the original error.
+                pass
+        if translated_exc is not None:
+            raise translated_exc from exc
+        else:
+            raise exc
 
     def done(self):
         if self._crt_future is None:

tests/integration/test_crt.py

@@ -511,3 +511,20 @@ def test_download_cancel(self):
         possible_matches = glob.glob('%s*' % download_path)
         self.assertEqual(possible_matches, [])
         self._assert_subscribers_called()
+
+    def test_exception_translation(self):
+        # Test that CRT's S3ResponseError translates to botocore error
+        transfer = self._create_s3_transfer()
+        download_path = os.path.join(
+            self.files.rootdir, 'obviously-no-such-key.txt'
+        )
+        with self.assertRaises(self.client.exceptions.NoSuchKey) as cm:
+            future = transfer.download(
+                self.bucket_name,
+                'obviously-no-such-key.txt',
+                download_path,
+                subscribers=[self.record_subscriber],
+            )
+            future.result()
+
+        self.assertEqual(cm.exception.response['Error']['Code'], 'NoSuchKey')

tests/unit/test_crt.py

@@ -14,7 +14,7 @@
 
 import pytest
 from botocore.credentials import Credentials, ReadOnlyCredentials
-from botocore.exceptions import NoCredentialsError
+from botocore.exceptions import ClientError, NoCredentialsError
 from botocore.session import Session
 
 from s3transfer.exceptions import TransferNotDoneError
@@ -164,6 +164,60 @@ def test_delete_request(self):
         self.assertEqual(self.expected_host, crt_request.headers.get("host"))
         self.assertIsNone(crt_request.headers.get("Authorization"))
 
+    def _create_crt_response_error(
+        self, status_code, body, operation_name=None
+    ):
+        return awscrt.s3.S3ResponseError(
+            code=14343,
+            name='AWS_ERROR_S3_INVALID_RESPONSE_STATUS',
+            message='Invalid response status from request',
+            status_code=status_code,
+            headers=[
+                ('x-amz-request-id', 'QSJHJJZR2EDYD4GQ'),
+                (
+                    'x-amz-id-2',
+                    'xDbgdKdvYZTjgpOTzm7yNP2JPrOQl+eaQvUkFdOjdJoWkIC643fgHxdsHpUKvVAfjKf5F6otEYA=',
+                ),
+                ('Content-Type', 'application/xml'),
+                ('Transfer-Encoding', 'chunked'),
+                ('Date', 'Fri, 10 Nov 2023 23:22:47 GMT'),
+                ('Server', 'AmazonS3'),
+            ],
+            body=body,
+            operation_name=operation_name,
+        )
+
+    def test_translate_get_object_404(self):
+        body = (
+            b'<?xml version="1.0" encoding="UTF-8"?>\n<Error>'
+            b'<Code>NoSuchKey</Code>'
+            b'<Message>The specified key does not exist.</Message>'
+            b'<Key>obviously-no-such-key.txt</Key>'
+            b'<RequestId>SBJ7ZQY03N1WDW9T</RequestId>'
+            b'<HostId>SomeHostId</HostId></Error>'
+        )
+        crt_exc = self._create_crt_response_error(404, body, 'GetObject')
+        boto_err = self.request_serializer.translate_crt_exception(crt_exc)
+        self.assertIsInstance(
+            boto_err, self.session.create_client('s3').exceptions.NoSuchKey
+        )
+
+    def test_translate_head_object_404(self):
+        # There's no body in a HEAD response, so we can't map it to a modeled S3 exception.
+        # But it should still map to a botocore ClientError
+        body = None
+        crt_exc = self._create_crt_response_error(
+            404, body, operation_name='HeadObject'
+        )
+        boto_err = self.request_serializer.translate_crt_exception(crt_exc)
+        self.assertIsInstance(boto_err, ClientError)
+
+    def test_translate_unknown_operation_404(self):
+        body = None
+        crt_exc = self._create_crt_response_error(404, body)
+        boto_err = self.request_serializer.translate_crt_exception(crt_exc)
+        self.assertIsInstance(boto_err, ClientError)
+
 
 @requires_crt_pytest
 class TestBotocoreCRTCredentialsWrapper: