✨Streaming utils for zipping and reading/wiring to S3 (#7186)

Co-authored-by: Andrei Neagu <[email protected]>

Author: GitHK

packages/aws-library/requirements/_base.txt

@@ -244,6 +244,8 @@ protobuf==5.29.3
     #   opentelemetry-proto
 psutil==6.1.1
     # via -r requirements/../../../packages/service-library/requirements/_base.in
+pycryptodome==3.21.0
+    # via stream-zip
 pydantic==2.10.6
     # via
     #   -c requirements/../../../packages/common-library/requirements/../../../requirements/constraints.txt
@@ -368,6 +370,8 @@ six==1.17.0
     # via python-dateutil
 sniffio==1.3.1
     # via anyio
+stream-zip==0.0.83
+    # via -r requirements/../../../packages/service-library/requirements/_base.in
 tenacity==9.0.0
     # via -r requirements/../../../packages/service-library/requirements/_base.in
 toolz==1.0.0

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

@@ -14,9 +14,12 @@
 from botocore import exceptions as botocore_exc
 from botocore.client import Config
 from models_library.basic_types import SHA256Str
+from models_library.bytes_iters import BytesIter, DataSize
 from models_library.storage_schemas import ETag, S3BucketName, UploadedPart
 from pydantic import AnyUrl, ByteSize, TypeAdapter
+from servicelib.bytes_iters import DEFAULT_READ_CHUNK_SIZE, BytesStreamer
 from servicelib.logging_utils import log_catch, log_context
+from servicelib.s3_utils import FileLikeReader
 from servicelib.utils import limited_gather
 from settings_library.s3 import S3Settings
 from types_aiobotocore_s3 import S3Client
@@ -470,6 +473,57 @@ async def copy_objects_recursively(
             limit=_MAX_CONCURRENT_COPY,
         )
 
+    async def get_bytes_streamer_from_object(
+        self,
+        bucket_name: S3BucketName,
+        object_key: S3ObjectKey,
+        *,
+        chunk_size: int = DEFAULT_READ_CHUNK_SIZE,
+    ) -> BytesStreamer:
+        """stream read an object from S3 chunk by chunk"""
+
+        # NOTE `download_fileobj` cannot be used to implement this because
+        # it will buffer the entire file in memory instead of reading it
+        # chunk by chunk
+
+        # below is a quick call
+        head_response = await self._client.head_object(
+            Bucket=bucket_name, Key=object_key
+        )
+        data_size = DataSize(head_response["ContentLength"])
+
+        async def _() -> BytesIter:
+            # Download the file in chunks
+            position = 0
+            while position < data_size:
+                # Calculate the range for this chunk
+                end = min(position + chunk_size - 1, data_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
+
+        return BytesStreamer(data_size, _)
+
+    @s3_exception_handler(_logger)
+    async def upload_object_from_file_like(
+        self,
+        bucket_name: S3BucketName,
+        object_key: S3ObjectKey,
+        file_like_reader: FileLikeReader,
+    ) -> None:
+        """streams write an object in S3 from an AsyncIterable[bytes]"""
+        await self._client.upload_fileobj(file_like_reader, bucket_name, object_key)  # type: ignore[arg-type]
+
     @staticmethod
     def is_multipart(file_size: ByteSize) -> bool:
         return file_size >= MULTIPART_UPLOADS_MIN_TOTAL_SIZE