decorate _select_object_content with a retry (#1780)

Author: kukushking

awswrangler/_utils.py

@@ -8,7 +8,8 @@
 import random
 import time
 from concurrent.futures import FIRST_COMPLETED, Future, wait
-from typing import Any, Callable, Dict, Generator, List, Optional, Sequence, Tuple, Union, cast
+from functools import partial, wraps
+from typing import Any, Callable, Dict, Generator, List, Optional, Sequence, Tuple, Type, Union, cast
 
 import boto3
 import botocore.config
@@ -333,9 +334,42 @@ def check_duplicated_columns(df: pd.DataFrame) -> Any:
         )
 
 
+def retry(
+    ex: Type[Exception],
+    ex_code: Optional[str] = None,
+    base: float = 1.0,
+    max_num_tries: int = 3,
+) -> Callable[..., Any]:
+    """
+    Decorate function with decorrelated Jitter retries.
+
+    Parameters
+    ----------
+    ex : Exception
+        Exception to retry on
+    ex_code : Optional[str]
+        Response error code
+    base : float
+        Base delay
+    max_num_tries : int
+        Maximum number of retries
+
+    Returns
+    -------
+    Callable[..., Any]
+        Function
+    """
+
+    def wrapper(f: Callable[..., Any]) -> Any:
+        return wraps(f)(partial(try_it, f, ex, ex_code=ex_code, base=base, max_num_tries=max_num_tries))
+
+    return wrapper
+
+
 def try_it(
     f: Callable[..., Any],
     ex: Any,
+    *args: Any,
     ex_code: Optional[str] = None,
     base: float = 1.0,
     max_num_tries: int = 3,
@@ -348,7 +382,7 @@ def try_it(
     delay: float = base
     for i in range(max_num_tries):
         try:
-            return f(**kwargs)
+            return f(*args, **kwargs)
         except ex as exception:
             if ex_code is not None and hasattr(exception, "response"):
                 if exception.response["Error"]["Code"] != ex_code:

awswrangler/exceptions.py

@@ -115,3 +115,7 @@ class FailedQualityCheck(Exception):
 
 class AlreadyExists(Exception):
     """AlreadyExists."""
+
+
+class S3SelectRequestIncomplete(Exception):
+    """S3SelectRequestIncomplete."""

awswrangler/s3/_select.py

@@ -35,6 +35,9 @@ def _gen_scan_range(obj_size: int, scan_range_chunk_size: Optional[int] = None)
 
 
 @engine.dispatch_on_engine
+@_utils.retry(
+    ex=exceptions.S3SelectRequestIncomplete,
+)
 def _select_object_content(
     boto3_session: Optional[boto3.Session],
     args: Dict[str, Any],
@@ -64,7 +67,9 @@ def _select_object_content(
             request_complete = True
     # If the End Event is not received, the results may be incomplete
     if not request_complete:
-        raise Exception(f"S3 Select request for path {args['Key']} is incomplete as End Event was not received")
+        raise exceptions.S3SelectRequestIncomplete(
+            f"S3 Select request for path {args['Key']} is incomplete as End Event was not received"
+        )
 
     return _utils.list_to_arrow_table(mapping=payload_records)