Skip to content

Commit

Permalink
Fix multipart S3 uploads
Browse files Browse the repository at this point in the history
Not sure how this worked before but something broke multipart uploads.
This adds a test to ensure that it keeps working and also tries to fix
the issue.
  • Loading branch information
packi committed Mar 20, 2024
1 parent 3bcc58f commit 6eedae1
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 6eedae1

Please sign in to comment.