Skip to content

Commit 78ce2cb

Browse files
fix(streaming): _S3SeekableIO.__next__/__iter__ don't advance self._position, corrupting subsequent seek()/read() (#8389)
* fix(streaming): _S3SeekableIO.__next__/__iter__ don't advance self._position, corrupting subsequent seek()/read() __next__ and __iter__ delegated straight to raw_stream.__next__()/raw_stream.__iter__(), bypassing self._position bookkeeping entirely -- unlike read()/readline()/readlines(), which all correctly advance self._position by the number of bytes actually consumed. seek() uses self._position as ground truth to compute the S3 GetObject Range header on the next raw_stream access. Consuming any data via the iterator protocol (e.g. "for line in s3_object: ...", a fully supported documented usage) left self._position stuck at its pre-iteration value. A subsequent seek()/read() would then reopen the S3 stream with a Range header computed from that stale position -- silently returning the wrong slice of the object (skipped or duplicated bytes), with no exception raised. Fix: route __next__ through the already position-tracked readline(), and have __iter__ return self, matching the standard Python file-iterator protocol. * fix(streaming): preserve iteration chunk behavior --------- Co-authored-by: Leandro <lcdama@amazon.pt>
1 parent e0566b0 commit 78ce2cb

2 files changed

Lines changed: 60 additions & 2 deletions

File tree

‎aws_lambda_powertools/utilities/streaming/_s3_seekable_io.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -169,10 +169,12 @@ def closed(self) -> bool:
169169
return self._closed
170170

171171
def __next__(self):
172-
return self.raw_stream.__next__()
172+
chunk = next(self.raw_stream)
173+
self._position += len(chunk)
174+
return chunk
173175

174176
def __iter__(self):
175-
return self.raw_stream.__iter__()
177+
return self
176178

177179
def __enter__(self):
178180
return self

‎tests/functional/streaming/_boto3/test_s3_seekable_io.py‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,62 @@ def test_next(s3_seekable_obj, s3_client_stub):
163163
next(s3_seekable_obj)
164164

165165

166+
def test_next_advances_position_without_changing_chunk_size(s3_seekable_obj, s3_client_stub):
167+
payload = b"a" * 1024 + b"b" * 481
168+
streaming_body = PowertoolsStreamingBody(raw_stream=io.BytesIO(payload), content_length=len(payload))
169+
170+
s3_client_stub.add_response(
171+
"get_object",
172+
{"Body": streaming_body},
173+
{"Bucket": s3_seekable_obj.bucket, "Key": s3_seekable_obj.key, "Range": "bytes=0-"},
174+
)
175+
176+
assert next(s3_seekable_obj) == payload[:1024]
177+
assert s3_seekable_obj.tell() == 1024
178+
179+
assert next(s3_seekable_obj) == payload[1024:]
180+
assert s3_seekable_obj.tell() == len(payload)
181+
182+
with pytest.raises(StopIteration):
183+
next(s3_seekable_obj)
184+
185+
186+
def test_iter_returns_self(s3_seekable_obj):
187+
assert iter(s3_seekable_obj) is s3_seekable_obj
188+
189+
190+
def test_seek_after_partial_iteration_reads_from_correct_position(s3_seekable_obj, s3_client_stub):
191+
payload = bytes(range(256)) * 8
192+
streaming_body = PowertoolsStreamingBody(raw_stream=io.BytesIO(payload), content_length=len(payload))
193+
194+
s3_client_stub.add_response(
195+
"get_object",
196+
{"Body": streaming_body},
197+
{"Bucket": s3_seekable_obj.bucket, "Key": s3_seekable_obj.key, "Range": "bytes=0-"},
198+
)
199+
200+
assert next(s3_seekable_obj) == payload[:1024]
201+
assert s3_seekable_obj.seek(5, io.SEEK_CUR) == 1029
202+
203+
remaining_payload = payload[1029:]
204+
resumed_streaming_body = PowertoolsStreamingBody(
205+
raw_stream=io.BytesIO(remaining_payload),
206+
content_length=len(remaining_payload),
207+
)
208+
s3_client_stub.add_response(
209+
"get_object",
210+
{"Body": resumed_streaming_body},
211+
{
212+
"Bucket": s3_seekable_obj.bucket,
213+
"Key": s3_seekable_obj.key,
214+
"Range": "bytes=1029-",
215+
},
216+
)
217+
218+
assert s3_seekable_obj.read(7) == payload[1029:1036]
219+
assert s3_seekable_obj.tell() == 1036
220+
221+
166222
def test_context_manager(s3_seekable_obj, s3_client_stub):
167223
payload = b"test"
168224
streaming_body = PowertoolsStreamingBody(raw_stream=io.BytesIO(payload), content_length=len(payload))

0 commit comments

Comments
 (0)