From ef2c902b32e823b683fba7b24f4bf5805f5a19e0 Mon Sep 17 00:00:00 2001
From: Stepan Vladovskiy <stepan.vladovskiy@gmail.com>
Date: Mon, 6 May 2024 04:50:19 -0300
Subject: [PATCH] feat: aioboto3 for instead of boto3

---
 main.py        | 33 +++++++++++----------------------
 pyproject.toml |  3 +--
 2 files changed, 12 insertions(+), 24 deletions(-)

diff --git a/main.py b/main.py
index dfba7c8..69fe0b8 100644
--- a/main.py
+++ b/main.py
@@ -1,8 +1,7 @@
 import os
 import tempfile
 import uuid
-import boto3
-from botocore.exceptions import BotoCoreError, ClientError
+import aioboto3  # Ensure this is only imported once
 from starlette.applications import Starlette
 from starlette.responses import JSONResponse
 from starlette.routing import Route
@@ -16,7 +15,6 @@ STORJ_END_POINT = os.environ.get('STORJ_END_POINT')
 STORJ_BUCKET_NAME = os.environ.get('STORJ_BUCKET_NAME')
 CDN_DOMAIN = os.environ.get('CDN_DOMAIN')
 
-
 @check_auth
 async def upload_handler(request: Request):
     form = await request.form()
@@ -28,36 +26,28 @@ async def upload_handler(request: Request):
     file_name, file_extension = os.path.splitext(file.filename)
     key = str(uuid.uuid4()) + file_extension
 
-    s3 = boto3.client('s3',
-                      aws_access_key_id=STORJ_ACCESS_KEY,
-                      aws_secret_access_key=STORJ_SECRET_KEY,
-                      endpoint_url=STORJ_END_POINT)
-
-    try:
+    async with aioboto3.client('s3',
+                               aws_access_key_id=STORJ_ACCESS_KEY,
+                               aws_secret_access_key=STORJ_SECRET_KEY,
+                               endpoint_url=STORJ_END_POINT) as s3:
         with tempfile.NamedTemporaryFile() as tmp_file:
             while True:
-                chunk = await file.read(8192)  # 8192 bytes by default.
+                chunk = await file.read(8192)
                 if not chunk:
                     break
                 tmp_file.write(chunk)
+            tmp_file.flush()
 
-            s3.upload_file(
+            await s3.upload_file(
                 Filename=tmp_file.name,
                 Bucket=STORJ_BUCKET_NAME,
                 Key=key,
-                ExtraArgs={
-                    "ContentType": file.content_type
-                }
+                ExtraArgs={"ContentType": file.content_type}
             )
 
-        url = 'http://' + CDN_DOMAIN + '/' + key
-
+        url = f'http://{CDN_DOMAIN}/{key}'
         return JSONResponse({'url': url, 'originalFilename': file.filename})
 
-    except (BotoCoreError, ClientError) as e:
-        print(e)
-        return JSONResponse({'error': 'Failed to upload file'}, status_code=500)
-
 async def home(request: Request):
     return JSONResponse({'message': 'Hello World!'})
 
@@ -70,5 +60,4 @@ app = Starlette(debug=True, routes=routes)
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host='0.0.0.0', port=8080)
-    
\ No newline at end of file
+    uvicorn.run(app, host='0.0.0.0', port=8080)
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index ed67076..b81f107 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -11,8 +11,7 @@ python = "^3.12"
 aiohttp = "^3.9.1"
 uvicorn = "^0.24.0.post1"
 starlette = "^0.33.0"
-boto3 = "^1.33.6"
-botocore = "^1.33.6"
+aioboto3 = "^9.0.0"
 
 [tool.poetry.dev-dependencies]
 black = "^23.10.1"