99 lines
3.7 KiB
Python
99 lines
3.7 KiB
Python
from contextlib import asynccontextmanager
|
|
from typing import Optional, BinaryIO
|
|
import aioboto3
|
|
from botocore.exceptions import ClientError
|
|
import os
|
|
|
|
class S3Adapter:
|
|
def __init__(self,
|
|
endpoint_url: str,
|
|
aws_access_key_id: str,
|
|
aws_secret_access_key: str,
|
|
bucket_name: str):
|
|
self.endpoint_url = endpoint_url
|
|
self.aws_access_key_id = aws_access_key_id
|
|
self.aws_secret_access_key = aws_secret_access_key
|
|
self.bucket_name = bucket_name
|
|
self.session = aioboto3.Session()
|
|
|
|
@asynccontextmanager
|
|
async def _get_client(self):
|
|
async with self.session.client( # type: ignore[reportGeneralTypeIssues]
|
|
"s3",
|
|
endpoint_url=self.endpoint_url,
|
|
aws_access_key_id=self.aws_access_key_id,
|
|
aws_secret_access_key=self.aws_secret_access_key,
|
|
) as client:
|
|
yield client
|
|
|
|
async def upload_file(self, object_name: str, data: bytes, content_type: Optional[str] = None):
|
|
"""Uploads bytes data to S3."""
|
|
try:
|
|
extra_args = {}
|
|
if content_type:
|
|
extra_args["ContentType"] = content_type
|
|
|
|
async with self._get_client() as client:
|
|
await client.put_object(
|
|
Bucket=self.bucket_name,
|
|
Key=object_name,
|
|
Body=data,
|
|
**extra_args
|
|
)
|
|
return True
|
|
except ClientError as e:
|
|
# logging.error(e)
|
|
print(f"Error uploading to S3: {e}")
|
|
return False
|
|
|
|
async def get_file(self, object_name: str) -> Optional[bytes]:
|
|
"""Downloads a file from S3 and returns bytes."""
|
|
try:
|
|
async with self._get_client() as client:
|
|
response = await client.get_object(Bucket=self.bucket_name, Key=object_name)
|
|
return await response['Body'].read()
|
|
except ClientError as e:
|
|
print(f"Error downloading from S3: {e}")
|
|
return None
|
|
|
|
async def stream_file(self, object_name: str, chunk_size: int = 65536):
|
|
"""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)
|
|
# aioboto3 Body is an aiohttp StreamReader wrapper
|
|
body = response['Body']
|
|
|
|
while True:
|
|
chunk = await body.read(chunk_size)
|
|
if not chunk:
|
|
break
|
|
yield chunk
|
|
except ClientError as e:
|
|
print(f"Error streaming from S3: {e}")
|
|
return
|
|
|
|
async def delete_file(self, object_name: str):
|
|
"""Deletes a file from S3."""
|
|
try:
|
|
async with self._get_client() as client:
|
|
await client.delete_object(Bucket=self.bucket_name, Key=object_name)
|
|
return True
|
|
except ClientError as e:
|
|
print(f"Error deleting from S3: {e}")
|
|
return False
|
|
|
|
async def get_presigned_url(self, object_name: str, expiration: int = 3600) -> Optional[str]:
|
|
"""Generate a presigned URL to share an S3 object."""
|
|
try:
|
|
async with self._get_client() as client:
|
|
response = await client.generate_presigned_url(
|
|
'get_object',
|
|
Params={'Bucket': self.bucket_name, 'Key': object_name},
|
|
ExpiresIn=expiration
|
|
)
|
|
return response
|
|
except ClientError as e:
|
|
print(f"Error generating presigned URL: {e}")
|
|
return None
|