Making test_head_object_with_retry public again

Author: igorborgest

awswrangler/pandas.py

@@ -111,7 +111,7 @@ def _read_csv_iterator(self, bucket_name, key_path, max_result_size=200_000_000,
         if pd_additional_kwargs['compression'] is not None:
             raise InvalidParameters("max_result_size currently does not support compressed files")
 
-        metadata = S3._head_object_with_retry(client_s3=self._client_s3, bucket=bucket_name, key=key_path)
+        metadata = S3.head_object_with_retry(client_s3=self._client_s3, bucket=bucket_name, key=key_path)
         total_size = metadata["ContentLength"]
         logger.debug(f"total_size: {total_size}")
         if total_size <= 0:

awswrangler/s3.py

@@ -323,7 +323,17 @@ def list_objects(self, path: str) -> List[str]:
                     stop=tenacity.stop_after_attempt(max_attempt_number=10),
                     reraise=True,
                     after=tenacity.after_log(logger, INFO))
-    def _head_object_with_retry(client_s3: client, bucket: str, key: str) -> Dict[str, Any]:
+    def head_object_with_retry(client_s3: client, bucket: str, key: str) -> Dict[str, Any]:
+        """
+        Executes the Boto3 head_object() function with an extra layer of random exponential back-off.
+
+        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/s3.html#S3.Client.head_object
+
+        :param client_s3: Boto3 S3 client
+        :param bucket: Bucket name.
+        :param key: Key
+        :return: Boto3 head_object() regular response dictionary
+        """
         return client_s3.head_object(Bucket=bucket, Key=key)
 
     @staticmethod
@@ -336,7 +346,7 @@ def _get_objects_head_remote(send_pipe, session_primitives: "SessionPrimitives",
         logger.debug(f"len(objects_paths): {len(objects_paths)}")
         for object_path in objects_paths:
             bucket, key = object_path.replace("s3://", "").split("/", 1)
-            res = S3._head_object_with_retry(client_s3=client_s3, bucket=bucket, key=key)
+            res = S3.head_object_with_retry(client_s3=client_s3, bucket=bucket, key=key)
             size = res["ContentLength"]
             objects_sizes[object_path] = size
         logger.debug(f"len(objects_sizes): {len(objects_sizes)}")

testing/build-image.sh

@@ -1,5 +1,5 @@
 #!/usr/bin/env bash
-set -e
+set -ex
 
 cp ../requirements.txt .
 cp ../requirements-dev.txt .

testing/deploy-cloudformation.sh

@@ -1,5 +1,5 @@
 #!/usr/bin/env bash
-set -e
+set -ex
 
 aws cloudformation deploy \
   --template-file template.yaml \

testing/start-jupyter.sh

@@ -1,5 +1,5 @@
 #!/usr/bin/env bash
-set -e
+set -ex
 
 cd ..
 pip install -e .

testing/test_awswrangler/test_s3.py

@@ -231,3 +231,14 @@ def test_copy_listed_objects(session, bucket, database, mode, procs_io_bound):
 def test_wait_object_exists(bucket):
     with pytest.raises(S3WaitObjectTimeout):
         wr.s3.wait_object_exists(path=f"s3://{bucket}/test_wait_object_exists.txt", timeout=5.0)
+
+
+def test_head_object_with_retry(bucket):
+    key = "test_head_object_with_retry"
+    boto3.resource("s3").Object(bucket, key).put(Body=str("Hello!"))
+    res = wr.s3.head_object_with_retry(
+        client_s3=boto3.client("s3"),
+        bucket=bucket,
+        key=key
+    )
+    assert res["ResponseMetadata"]["HTTPHeaders"]["content-length"] == "6"