diff --git a/main.py b/main.py index 43f4cd46..5b3e5c49 100644 --- a/main.py +++ b/main.py @@ -16,6 +16,7 @@ from auth.oauth import oauth_login, oauth_authorize from base.redis import redis from base.resolvers import resolvers from resolvers.auth import confirm_email_handler +from resolvers.upload import upload_handler from services.main import storages_init from services.stat.viewed import ViewedStorage from services.zine.gittask import GitTask @@ -69,7 +70,8 @@ routes = [ # Route("/messages", endpoint=sse_messages), Route("/oauth/{provider}", endpoint=oauth_login), Route("/oauth-authorize", endpoint=oauth_authorize), - Route("/confirm/{token}", endpoint=confirm_email_handler) + Route("/confirm/{token}", endpoint=confirm_email_handler), + Route("/upload", endpoint=upload_handler, methods=['POST']) ] app = Starlette( diff --git a/requirements.txt b/requirements.txt index a1a718ed..06974f05 100644 --- a/requirements.txt +++ b/requirements.txt @@ -32,3 +32,6 @@ graphql-ws nltk~=3.8.1 pymystem3~=0.2.0 transformers~=4.28.1 +boto3~=1.28.2 +botocore~=1.31.2 +python-multipart~=0.0.6 diff --git a/resolvers/upload.py b/resolvers/upload.py new file mode 100644 index 00000000..441824cc --- /dev/null +++ b/resolvers/upload.py @@ -0,0 +1,63 @@ +import os +import shutil +import tempfile +import uuid +import boto3 +from botocore.exceptions import BotoCoreError, ClientError +from starlette.responses import JSONResponse + +STORJ_ACCESS_KEY = os.environ.get('STORJ_ACCESS_KEY') +STORJ_SECRET_KEY = os.environ.get('STORJ_SECRET_KEY') +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') + + +async def upload_handler(request): + form = await request.form() + file = form.get('file') + + if file is None: + return JSONResponse({'error': 'No file uploaded'}, status_code=400) + + file_name, file_extension = os.path.splitext(file.filename) + + key = str(uuid.uuid4()) + file_extension + + # Create an S3 client with Storj configuration + s3 = boto3.client('s3', + aws_access_key_id=STORJ_ACCESS_KEY, + aws_secret_access_key=STORJ_SECRET_KEY, + endpoint_url=STORJ_END_POINT) + + try: + # Save the uploaded file to a temporary file + with tempfile.NamedTemporaryFile() as tmp_file: + shutil.copyfileobj(file.file, tmp_file) + + file_stats = os.stat(tmp_file.name) + file_size = file_stats.st_size + + # 25MB + if file_size > 26214400: + return JSONResponse({'error': 'File is too large'}, status_code=400) + + s3.upload_file( + Filename=tmp_file.name, + Bucket=STORJ_BUCKET_NAME, + Key=key, + ExtraArgs={ + "ContentType": file.content_type + } + ) + + url = 'http://' + CDN_DOMAIN + '/' + key + + return JSONResponse({'url': url, 'originalFilename': file.filename, 'size': file_size}) + + except (BotoCoreError, ClientError) as e: + print(e) + return JSONResponse({'error': 'Failed to upload file'}, status_code=500) + + +