Add TransferConfig support for CRTTransferManager

Author: hssyoo

.changes/next-release/enhancement-awscrt-49742.json

@@ -0,0 +1,5 @@
+{
+  "type": "enhancement",
+  "category": "``awscrt``",
+  "description": "``CRTTransferManager`` now supports the following ``TransferConfig`` options - ``multipart_threshold``, ``multipart_chunksize``"
+}

s3transfer/crt.py

@@ -193,7 +193,9 @@ class CRTTransferManager:
 
     _UNSUPPORTED_BUCKET_PATTERNS = TransferManager._UNSUPPORTED_BUCKET_PATTERNS
 
-    def __init__(self, crt_s3_client, crt_request_serializer, osutil=None):
+    def __init__(
+        self, crt_s3_client, crt_request_serializer, osutil=None, config=None
+    ):
         """A transfer manager interface for Amazon S3 on CRT s3 client.
 
         :type crt_s3_client: awscrt.s3.S3Client
@@ -207,12 +209,18 @@ def __init__(self, crt_s3_client, crt_request_serializer, osutil=None):
         :type osutil: s3transfer.utils.OSUtils
         :param osutil: OSUtils object to use for os-related behavior when
             using with transfer manager.
+
+        :type config: s3transfer.manager.TransferConfig
+        :param config: The transfer configuration to be used when
+            making CRT S3 client requests.
         """
         if osutil is None:
             self._osutil = OSUtils()
         self._crt_s3_client = crt_s3_client
         self._s3_args_creator = S3ClientArgsCreator(
-            crt_request_serializer, self._osutil
+            crt_request_serializer,
+            self._osutil,
+            config,
         )
         self._crt_exception_translator = (
             crt_request_serializer.translate_crt_exception
@@ -732,9 +740,22 @@ def set_s3_request(self, s3_request):
 
 
 class S3ClientArgsCreator:
-    def __init__(self, crt_request_serializer, os_utils):
+    def __init__(self, crt_request_serializer, os_utils, config=None):
         self._request_serializer = crt_request_serializer
         self._os_utils = os_utils
+        self._config = config
+
+    def _get_crt_transfer_config_options(self):
+        part_size = self._config.multipart_chunksize
+        if (
+            self._config.get_deep_attr('multipart_chunksize')
+            is self._config.UNSET_DEFAULT
+        ):
+            # Let CRT dynamically calculate part size.
+            part_size = None
+        return {
+            'part_size': part_size,
+        }
 
     def get_make_request_args(
         self, request_type, call_args, coordinator, future, on_done_after_calls
@@ -823,6 +844,11 @@ def _get_make_request_args_put_object(
         )
         make_request_args['send_filepath'] = send_filepath
         make_request_args['checksum_config'] = checksum_config
+        if self._config is not None:
+            make_request_args.update(self._get_crt_transfer_config_options())
+            make_request_args['multipart_upload_threshold'] = (
+                self._config.multipart_threshold
+            )
         return make_request_args
 
     def _get_make_request_args_get_object(
@@ -859,6 +885,8 @@ def _get_make_request_args_get_object(
         make_request_args['recv_filepath'] = recv_filepath
         make_request_args['on_body'] = on_body
         make_request_args['checksum_config'] = checksum_config
+        if self._config is not None:
+            make_request_args.update(self._get_crt_transfer_config_options())
         return make_request_args
 
     def _default_get_make_request_args(

s3transfer/manager.py

@@ -50,6 +50,8 @@
 
 
 class TransferConfig:
+    UNSET_DEFAULT = object()
+
     def __init__(
         self,
         multipart_threshold=8 * MB,
@@ -152,12 +154,19 @@ def __init__(
 
     def _validate_attrs_are_nonzero(self):
         for attr, attr_val in self.__dict__.items():
-            if attr_val is not None and attr_val <= 0:
+            if (
+                attr_val is not None
+                and attr_val is not self.UNSET_DEFAULT
+                and attr_val <= 0
+            ):
                 raise ValueError(
                     f'Provided parameter {attr} of value {attr_val} must '
                     'be greater than 0.'
                 )
 
+    def get_deep_attr(self, item):
+        return object.__getattribute__(self, item)
+
 
 class TransferManager:
     ALLOWED_DOWNLOAD_ARGS = ALLOWED_DOWNLOAD_ARGS

tests/functional/test_crt.py

@@ -18,6 +18,8 @@
 
 from botocore.session import Session
 
+from s3transfer.constants import MB
+from s3transfer.manager import TransferConfig
 from s3transfer.subscribers import BaseSubscriber
 from tests import (
     HAS_CRT,
@@ -769,3 +771,59 @@ def test_crt_s3_client_error_handling(self):
         )
         with self.assertRaises(awscrt.exceptions.AwsCrtError):
             future.result()
+
+    def test_transfer_config_used_in_upload_request(self):
+        config = TransferConfig(
+            multipart_threshold=4 * MB,
+            multipart_chunksize=2 * MB,
+        )
+        transfer_manager = s3transfer.crt.CRTTransferManager(
+            crt_s3_client=self.s3_crt_client,
+            crt_request_serializer=self.request_serializer,
+            config=config,
+        )
+        future = transfer_manager.upload(
+            self.filename, self.bucket, self.key, {}, []
+        )
+        future.result()
+
+        callargs_kwargs = self.s3_crt_client.make_request.call_args[1]
+        assert callargs_kwargs['multipart_upload_threshold'] == 4 * MB
+        assert callargs_kwargs['part_size'] == 2 * MB
+
+    def test_transfer_config_used_in_download_request(self):
+        config = TransferConfig(
+            multipart_threshold=4 * MB,
+            multipart_chunksize=2 * MB,
+        )
+        transfer_manager = s3transfer.crt.CRTTransferManager(
+            crt_s3_client=self.s3_crt_client,
+            crt_request_serializer=self.request_serializer,
+            config=config,
+        )
+        future = transfer_manager.download(
+            self.bucket, self.key, self.filename, {}, []
+        )
+        future.result()
+
+        callargs_kwargs = self.s3_crt_client.make_request.call_args[1]
+        assert callargs_kwargs['part_size'] == 2 * MB
+        # Config option only used for PUT requests.
+        assert 'multipart_upload_threshold' not in callargs_kwargs
+
+    def test_unset_part_size_defaults_to_none_in_upload_request(self):
+        config = TransferConfig(
+            multipart_chunksize=TransferConfig.UNSET_DEFAULT,
+        )
+        transfer_manager = s3transfer.crt.CRTTransferManager(
+            crt_s3_client=self.s3_crt_client,
+            crt_request_serializer=self.request_serializer,
+            config=config,
+        )
+        future = transfer_manager.upload(
+            self.filename, self.bucket, self.key, {}, []
+        )
+        future.result()
+
+        callargs_kwargs = self.s3_crt_client.make_request.call_args[1]
+        assert callargs_kwargs['part_size'] is None