Skip to content

Commit

Permalink
Fix multipart S3 uploads
Browse files Browse the repository at this point in the history
While reworking the retries the nested-loop was removed and this now
causes only the first part to be uploaded.

This adds a test to ensure that it keeps working and also fixes the
issue.
  • Loading branch information
packi committed Mar 20, 2024
1 parent 3bcc58f commit 641fe51
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 1 deletion.
1 change: 0 additions & 1 deletion rohmu/object_storage/s3.py
Original file line number Diff line number Diff line change
Expand Up @@ -447,7 +447,6 @@ def multipart_upload_file_object(
if progress_fn:
# TODO: change this to incremental progress. Size parameter is currently unused.
progress_fn(bytes_sent, size) # type: ignore[arg-type]
break

self.stats.operation(StorageOperation.multipart_complete)
try:
Expand Down
23 changes: 23 additions & 0 deletions test/object_storage/test_s3.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,29 @@ def test_store_file_from_disk(infra: S3Infra) -> None:
)


def test_store_file_object_large(infra: S3Infra) -> None:
test_data = b"test-data" * 2
chunk_size = len(test_data) // 2
file_object = BytesIO(test_data)

infra.transfer.multipart_chunk_size = chunk_size # simulate smaller chunk size to force multiple chunks

metadata = {"Content-Length": len(test_data), "some-date": datetime(2022, 11, 15, 18, 30, 58, 486644)}
infra.transfer.store_file_object(key="test_key2", fd=file_object, metadata=metadata, multipart=True)

notifier = infra.notifier
s3_client = infra.s3_client

s3_client.create_multipart_upload.assert_called()
assert s3_client.upload_part.call_count == 2
s3_client.complete_multipart_upload.assert_called()
notifier.object_created.assert_called_once_with(
key="test_key2",
size=len(test_data),
metadata={"Content-Length": "18", "some-date": "2022-11-15 18:30:58.486644"},
)


@pytest.mark.parametrize("multipart", [False, None, True])
def test_store_file_object(infra: S3Infra, multipart: Optional[bool]) -> None:
test_data = b"test-data"
Expand Down

0 comments on commit 641fe51

Please sign in to comment.