diff --git a/aws_lambda_powertools/utilities/streaming/_s3_seekable_io.py b/aws_lambda_powertools/utilities/streaming/_s3_seekable_io.py index 244e81dfd56..c3773e80d47 100644 --- a/aws_lambda_powertools/utilities/streaming/_s3_seekable_io.py +++ b/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: diff --git a/tests/functional/streaming/_boto3/test_s3_seekable_io.py b/tests/functional/streaming/_boto3/test_s3_seekable_io.py index 9ad81361d99..abf1e593bcc 100644 --- a/tests/functional/streaming/_boto3/test_s3_seekable_io.py +++ b/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