added S3 streaming and integration test

Author: Andrei Neagu

packages/aws-library/src/aws_library/s3/__init__.py

@@ -3,7 +3,11 @@
     SimcoreS3API,
     UploadedBytesTransferredCallback,
 )
-from ._constants import PRESIGNED_LINK_MAX_SIZE, S3_MAX_FILE_SIZE
+from ._constants import (
+    MIN_MULTIPART_UPLOAD_CHUNK_SIZE,
+    PRESIGNED_LINK_MAX_SIZE,
+    S3_MAX_FILE_SIZE,
+)
 from ._errors import (
     S3AccessError,
     S3BucketInvalidError,
@@ -23,6 +27,7 @@
 
 __all__: tuple[str, ...] = (
     "CopiedBytesTransferredCallback",
+    "MIN_MULTIPART_UPLOAD_CHUNK_SIZE",
     "MultiPartUploadLinks",
     "PRESIGNED_LINK_MAX_SIZE",
     "S3_MAX_FILE_SIZE",

packages/aws-library/src/aws_library/s3/_client.py

@@ -3,7 +3,7 @@
 import functools
 import logging
 import urllib.parse
-from collections.abc import AsyncGenerator, Sequence
+from collections.abc import AsyncGenerator, Callable, Sequence
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Final, Protocol, cast
@@ -18,6 +18,7 @@
 from pydantic import AnyUrl, ByteSize, TypeAdapter
 from servicelib.logging_utils import log_catch, log_context
 from servicelib.utils import limited_gather
+from servicelib.zip_stream import DEFAULT_CHUNK_SIZE, FileStream
 from settings_library.s3 import S3Settings
 from types_aiobotocore_s3 import S3Client
 from types_aiobotocore_s3.literals import BucketLocationConstraintType
@@ -57,6 +58,14 @@ def __call__(self, total_bytes_copied: int, *, file_name: str) -> None:
         ...
 
 
+class AsyncFileProtocol(Protocol):
+    async def read(self, chunk_size: int) -> bytes:
+        ...
+
+    async def write(self, data: bytes) -> None:
+        ...
+
+
 @dataclass(frozen=True)
 class SimcoreS3API:  # pylint: disable=too-many-public-methods
     _client: S3Client
@@ -470,6 +479,79 @@ async def copy_objects_recursively(
             limit=_MAX_CONCURRENT_COPY,
         )
 
+    async def get_object_file_stream(
+        self,
+        bucket_name: S3BucketName,
+        object_key: S3ObjectKey,
+        *,
+        chunk_size: int = DEFAULT_CHUNK_SIZE,
+    ) -> FileStream:
+        response = await self._client.head_object(Bucket=bucket_name, Key=object_key)
+        file_size = response["ContentLength"]
+
+        # Download the file in chunks
+        position = 0
+        while position < file_size:
+            # Calculate the range for this chunk
+            end = min(position + chunk_size - 1, file_size - 1)
+            range_header = f"bytes={position}-{end}"
+
+            # Download the chunk
+            response = await self._client.get_object(
+                Bucket=bucket_name, Key=object_key, Range=range_header
+            )
+
+            chunk = await response["Body"].read()
+
+            # Yield the chunk for processing
+            yield chunk
+
+            position += chunk_size
+
+    @s3_exception_handler(_logger)
+    async def upload_object_from_file_stream(
+        self,
+        bucket_name: S3BucketName,
+        object_key: S3ObjectKey,
+        file_stream: Callable[[], FileStream],
+    ) -> None:
+        # Create a multipart upload
+        multipart_response = await self._client.create_multipart_upload(
+            Bucket=bucket_name, Key=object_key
+        )
+        upload_id = multipart_response["UploadId"]
+
+        try:
+            parts = []
+            part_number = 1
+
+            async for chunk in file_stream():
+                print(f"partsizze={len(chunk)}")
+
+                part_response = await self._client.upload_part(
+                    Bucket=bucket_name,
+                    Key=object_key,
+                    PartNumber=part_number,
+                    UploadId=upload_id,
+                    Body=chunk,
+                )
+                parts.append({"ETag": part_response["ETag"], "PartNumber": part_number})
+                part_number += 1
+
+            # Complete the multipart upload
+            await self._client.complete_multipart_upload(
+                Bucket=bucket_name,
+                Key=object_key,
+                UploadId=upload_id,
+                MultipartUpload={"Parts": parts},
+            )
+        except Exception:
+            # Abort the multipart upload if something goes wrong
+            await self._client.abort_multipart_upload(
+                Bucket=bucket_name, Key=object_key, UploadId=upload_id
+            )
+            raise
+
     @staticmethod
     def is_multipart(file_size: ByteSize) -> bool:
         return file_size >= MULTIPART_UPLOADS_MIN_TOTAL_SIZE

packages/aws-library/src/aws_library/s3/_constants.py

@@ -6,6 +6,9 @@
 MULTIPART_UPLOADS_MIN_TOTAL_SIZE: Final[ByteSize] = TypeAdapter(
     ByteSize
 ).validate_python("100MiB")
+MIN_MULTIPART_UPLOAD_CHUNK_SIZE: Final[int] = TypeAdapter(ByteSize).validate_python(
+    "5MiB"
+)
 MULTIPART_COPY_THRESHOLD: Final[ByteSize] = TypeAdapter(ByteSize).validate_python(
     "100MiB"
 )

packages/aws-library/src/aws_library/s3/_errors.py

@@ -27,3 +27,10 @@ class S3UploadNotFoundError(S3AccessError):
 
 class S3DestinationNotEmptyError(S3AccessError):
     msg_template: str = "The destination {dst_prefix} is not empty"
+
+
+class S3MultipartUploadMinChunkSizeError(S3RuntimeError):
+    msg_template: str = (
+        "chunk_size='{current_chunk_size}' must be grater of equal "
+        "to min_chunk_size='{min_chunk_size}'"
+    )