Merge branch 'develop' into fix/dynamodb-number-magnitude

Author: leandrodamascena

aws_lambda_powertools/utilities/streaming/_s3_seekable_io.py

@@ -5,8 +5,10 @@
 from typing import IO, TYPE_CHECKING, Any, TypeVar
 
 import boto3
+from botocore.exceptions import ClientError
 
 from aws_lambda_powertools.shared import user_agent
+from aws_lambda_powertools.utilities.streaming.compat import PowertoolsStreamingBody
 from aws_lambda_powertools.utilities.streaming.constants import MESSAGE_STREAM_NOT_WRITABLE
 
 if TYPE_CHECKING:
@@ -15,8 +17,6 @@
 
     from mypy_boto3_s3.client import S3Client
 
-    from aws_lambda_powertools.utilities.streaming.compat import PowertoolsStreamingBody
-
     _CData = TypeVar("_CData")
 
 logger = logging.getLogger(__name__)
@@ -98,10 +98,21 @@ def raw_stream(self) -> PowertoolsStreamingBody:
         """
         Returns the boto3 StreamingBody, starting the stream from the sought position.
         """
+        if self._closed:
+            raise ValueError("I/O operation on closed file.")
+
         if self._raw_stream is None:
             range_header = f"bytes={self._position}-"
             logger.debug(f"Starting new stream at {range_header}")
-            self._raw_stream = self.s3_client.get_object(Range=range_header, **self._sdk_options).get("Body")
+            try:
+                self._raw_stream = self.s3_client.get_object(Range=range_header, **self._sdk_options).get("Body")
+            except ClientError as exc:
+                # S3 rejects a range that starts at or past the end of the object, which includes any range
+                # on an empty object. A file returns no data at that position instead of raising.
+                if exc.response.get("Error", {}).get("Code") != "InvalidRange":
+                    raise
+                logger.debug(f"Position {self._position} is at or past the end of the object")
+                self._raw_stream = PowertoolsStreamingBody(raw_stream=io.BytesIO(b""), content_length=0)
             self._closed = False
 
         return self._raw_stream
@@ -183,7 +194,9 @@ def __exit__(self, *kwargs):
         self.close()
 
     def close(self) -> None:
-        self.raw_stream.close()
+        # Only close a stream that is already open, rather than opening a new one just to close it
+        if self._raw_stream is not None:
+            self._raw_stream.close()
         self._closed = True
 
     def fileno(self) -> int:

tests/functional/streaming/_boto3/test_s3_seekable_io.py

@@ -5,7 +5,9 @@
 import boto3
 import pytest
 from botocore import stub
+from botocore.exceptions import ClientError
 
+from aws_lambda_powertools.utilities.streaming import S3Object
 from aws_lambda_powertools.utilities.streaming._s3_seekable_io import _S3SeekableIO
 from aws_lambda_powertools.utilities.streaming.compat import PowertoolsStreamingBody
 
@@ -134,17 +136,93 @@ def test_readlines(s3_seekable_obj, s3_client_stub):
     assert s3_seekable_obj.tell() == len(payload)
 
 
-def test_closed(s3_seekable_obj, s3_client_stub):
-    payload = b"test"
-    streaming_body = PowertoolsStreamingBody(raw_stream=io.BytesIO(payload), content_length=len(payload))
+def test_read_at_end_of_object_returns_empty_bytes(s3_seekable_obj, s3_client_stub):
+    s3_client_stub.add_response("head_object", {"ContentLength": 4})
+    # S3 rejects a range that starts at the end of the object
+    s3_client_stub.add_client_error(
+        "get_object",
+        service_error_code="InvalidRange",
+        http_status_code=416,
+        expected_params={"Bucket": s3_seekable_obj.bucket, "Key": s3_seekable_obj.key, "Range": "bytes=4-"},
+    )
 
-    s3_client_stub.add_response(
+    s3_seekable_obj.seek(0, io.SEEK_END)
+
+    assert s3_seekable_obj.read() == b""
+    assert s3_seekable_obj.tell() == 4
+
+
+def test_read_empty_object_returns_empty_bytes(s3_seekable_obj, s3_client_stub):
+    # S3 rejects any range on an empty object
+    s3_client_stub.add_client_error(
         "get_object",
-        {"Body": streaming_body},
-        {"Bucket": s3_seekable_obj.bucket, "Key": s3_seekable_obj.key, "Range": "bytes=0-"},
+        service_error_code="InvalidRange",
+        http_status_code=416,
+        expected_params={"Bucket": s3_seekable_obj.bucket, "Key": s3_seekable_obj.key, "Range": "bytes=0-"},
     )
 
+    assert s3_seekable_obj.read() == b""
+    assert list(s3_seekable_obj) == []
+    assert s3_seekable_obj.tell() == 0
+
+
+def test_raw_stream_raises_other_client_errors(s3_seekable_obj, s3_client_stub):
+    s3_client_stub.add_client_error("get_object", service_error_code="NoSuchKey", http_status_code=404)
+
+    with pytest.raises(ClientError, match="NoSuchKey"):
+        s3_seekable_obj.read()
+
+
+def test_closed(s3_seekable_obj, s3_client_stub):
     s3_seekable_obj.close()
+
+    assert s3_seekable_obj.closed is True
+    # Closing an object that was never read must not open a stream just to close it
+    s3_client_stub.assert_no_pending_responses()
+
+
+@pytest.mark.parametrize("stream_class", [_S3SeekableIO, S3Object])
+@pytest.mark.parametrize("read_method", ["read", "readline", "readlines", "__next__"])
+@pytest.mark.parametrize("initial_state", ["unread", "partially_read", "seeked", "empty"])
+def test_reads_after_close_do_not_reopen_stream(s3_client, s3_client_stub, stream_class, read_method, initial_state):
+    stream = stream_class(bucket="bucket", key="key", boto3_client=s3_client)
+    expected_params = {"Bucket": "bucket", "Key": "key", "Range": "bytes=0-"}
+
+    if initial_state == "empty":
+        s3_client_stub.add_client_error(
+            "get_object",
+            service_error_code="InvalidRange",
+            http_status_code=416,
+            expected_params=expected_params,
+        )
+        assert stream.read() == b""
+    elif initial_state != "unread":
+        payload = b"hello\nworld"
+        body = PowertoolsStreamingBody(raw_stream=io.BytesIO(payload), content_length=len(payload))
+        s3_client_stub.add_response("get_object", {"Body": body}, expected_params)
+        assert stream.read(1) == b"h"
+        if initial_state == "seeked":
+            stream.seek(3)
+
+    position = stream.tell()
+    stream.close()
+    stream.close()
+
+    with pytest.raises(ValueError, match="I/O operation on closed file"):
+        getattr(stream, read_method)()
+
+    assert stream.closed is True
+    assert stream.tell() == position
+    # The stub has no queued responses, so any attempt to reopen the stream would fail the test.
+    s3_client_stub.assert_no_pending_responses()
+
+
+def test_context_manager_at_end_of_object(s3_seekable_obj, s3_client_stub):
+    s3_client_stub.add_response("head_object", {"ContentLength": 4})
+
+    with s3_seekable_obj as f:
+        f.seek(0, io.SEEK_END)
+
     assert s3_seekable_obj.closed is True