improve: configurable transfermanager and handle max retries error

Author: FelixAbrahamsson

fast_s3/fetcher.py

@@ -3,6 +3,7 @@
 from typing import List, Union
 
 from botocore.exceptions import ClientError
+from s3transfer.exceptions import RetriesExceededError
 
 from .file import File, Status
 from .transfer_manager import transfer_manager
@@ -20,6 +21,7 @@ def __init__(
         ordered=True,
         buffer_size=1024,
         n_workers=32,
+        **transfer_manager_kwargs,
     ):
         self.paths = paths
         self.ordered = ordered
@@ -30,6 +32,7 @@ def __init__(
             aws_secret_access_key=aws_secret_access_key,
             region_name=region_name,
             n_workers=n_workers,
+            **transfer_manager_kwargs,
         )
         self.bucket_name = bucket_name
         self.files: List[File] = []
@@ -60,8 +63,8 @@ def process_index(self, index):
         try:
             file.future.result()
             return file.with_status(Status.done)
-        except ClientError:
-            return file.with_status(Status.error)
+        except (ClientError, RetriesExceededError) as e:
+            return file.with_status(Status.error, exception=e)
 
     def queue_download_(self):
         if self.current_path_index < len(self):

fast_s3/file.py

@@ -1,7 +1,7 @@
 import io
 from enum import Enum
 from pathlib import Path
-from typing import Union
+from typing import Optional, Union
 
 from pydantic import BaseModel
 from s3transfer.futures import TransferFuture
@@ -18,11 +18,12 @@ class File(BaseModel):
     future: TransferFuture
     path: Union[str, Path]
     status: Status = Status.pending
+    exception: Optional[Exception] = None
 
     class Config:
         arbitrary_types_allowed = True
 
-    def with_status(self, status: Status):
+    def with_status(self, status: Status, exception: Optional[Exception] = None):
         attributes = self.dict()
-        attributes.update(status=status)
+        attributes.update(status=status, exception=exception)
         return File(**attributes)

fast_s3/transfer_manager.py

@@ -10,6 +10,7 @@ def transfer_manager(
     aws_secret_access_key: str,
     region_name: str,
     n_workers=32,
+    **kwargs,
 ) -> TransferManager:
     session = boto3.Session(
         aws_access_key_id=aws_access_key_id,
@@ -22,6 +23,6 @@ def transfer_manager(
         config=botocore.config.Config(max_pool_connections=n_workers),
     )
     transfer_config = s3transfer.TransferConfig(
-        use_threads=True, max_concurrency=n_workers
+        use_threads=True, max_concurrency=n_workers, **kwargs
     )
     return s3transfer.create_transfer_manager(client, transfer_config)

fast_s3/uploader.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import List, Tuple, Union
+from typing import List, Union
 
 from s3transfer.futures import TransferFuture
 
@@ -15,13 +15,15 @@ def __init__(
         region_name: str,
         bucket_name: str,
         n_workers=32,
+        **transfer_manager_kwargs,
     ):
         self.transfer_manager = transfer_manager(
             endpoint_url=endpoint_url,
             aws_access_key_id=aws_access_key_id,
             aws_secret_access_key=aws_secret_access_key,
             region_name=region_name,
             n_workers=n_workers,
+            **transfer_manager_kwargs,
         )
         self.bucket_name = bucket_name
         self.futures: List[TransferFuture] = []