88 lines
2.6 KiB
Python
88 lines
2.6 KiB
Python
import os
|
|
|
|
import asyncio
|
|
|
|
from dotenv import load_dotenv
|
|
from aiobotocore.session import get_session
|
|
|
|
|
|
class S3Worker:
|
|
def __init__(self):
|
|
load_dotenv()
|
|
|
|
self._S3_ACCESS_KEY_ID = os.getenv('S3_ACCESS_KEY_ID')
|
|
self._S3_SECRET_ACCESS_KEY = os.getenv('S3_SECRET_ACCESS_KEY')
|
|
self._S3_ENDPOINT_URL = os.getenv('S3_ENDPOINT_URL')
|
|
|
|
self._s3_session = get_session()
|
|
self._s3_client = None
|
|
|
|
self.bucket = os.getenv('BUCKET_NAME')
|
|
|
|
async def __aenter__(self):
|
|
self._client = await self._s3_session.create_client(
|
|
"s3",
|
|
region_name="us-east-1",
|
|
aws_access_key_id=self._S3_ACCESS_KEY_ID,
|
|
aws_secret_access_key=self._S3_SECRET_ACCESS_KEY,
|
|
endpoint_url=self._S3_ENDPOINT_URL,
|
|
).__aenter__()
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc, tb):
|
|
await self._client.__aexit__(exc_type, exc, tb)
|
|
|
|
|
|
async def upload_file(self, key: str, data: bytes):
|
|
await self._client.put_object(
|
|
Bucket=self.bucket,
|
|
Key=key,
|
|
Body=data,
|
|
)
|
|
|
|
async def generate_upload_url(self, key: str, content_type: str, expires_in: int = 300) -> str:
|
|
return await self._client.generate_presigned_url("put_object",
|
|
Params={
|
|
"Bucket": self.bucket,
|
|
"Key": key,
|
|
"ContentType": content_type,
|
|
},
|
|
ExpiresIn=expires_in,
|
|
)
|
|
|
|
async def download_file(self, key: str):
|
|
return await self._client.get_object(
|
|
Bucket=self.bucket,
|
|
Key=key,
|
|
)
|
|
|
|
async def generate_download_url(self, key: str, filename: str, expires_in: int = 300) -> str:
|
|
return await self._client.generate_presigned_url("get_object",
|
|
Params={
|
|
"Bucket": self.bucket,
|
|
"Key": key,
|
|
"ResponseContentDisposition": (
|
|
f'attachment; filename="{filename}"'
|
|
),
|
|
},
|
|
ExpiresIn=expires_in,
|
|
)
|
|
|
|
|
|
async def test_run():
|
|
async with S3Worker() as worker:
|
|
await worker.upload_file("test.txt", b"hello")
|
|
file = await worker.download_file("test.txt")
|
|
file_text = await file["Body"].read()
|
|
print(file_text)
|
|
|
|
url = await worker.generate_download_url("test.txt", "hello.txt")
|
|
print(url)
|
|
url = await worker.generate_upload_url("some.jpg", content_type="image/jpeg")
|
|
print(url)
|
|
url = await worker.generate_download_url("some.jpg", "image.jpg")
|
|
print(url)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(test_run()) |