inspirations

adapters/s3_adapter.py

@@ -1,5 +1,5 @@
 from contextlib import asynccontextmanager
-from typing import Optional, BinaryIO
+from typing import Optional, BinaryIO, AsyncGenerator
 import aioboto3
 from botocore.exceptions import ClientError
 import os
@@ -56,11 +56,25 @@ class S3Adapter:
             print(f"Error downloading from S3: {e}")
             return None
 
-    async def stream_file(self, object_name: str, chunk_size: int = 65536):
+    async def get_file_size(self, object_name: str) -> Optional[int]:
+        """Returns the size of the file in bytes."""
+        try:
+            async with self._get_client() as client:
+                response = await client.head_object(Bucket=self.bucket_name, Key=object_name)
+                return response['ContentLength']
+        except ClientError as e:
+            print(f"Error getting file size from S3: {e}")
+            return None
+
+    async def stream_file(self, object_name: str, range_header: Optional[str] = None, chunk_size: int = 65536) -> AsyncGenerator[bytes, None]:
         """Streams a file from S3 yielding chunks. Memory-efficient for large files."""
         try:
             async with self._get_client() as client:
-                response = await client.get_object(Bucket=self.bucket_name, Key=object_name)
+                args = {'Bucket': self.bucket_name, 'Key': object_name}
+                if range_header:
+                    args['Range'] = range_header
+
+                response = await client.get_object(**args)
                 # aioboto3 Body is an aiohttp StreamReader wrapper
                 body = response['Body']