Bugfix
This commit is contained in:
+11
-13
@@ -38,12 +38,11 @@ async def get_upload_token(file_name: str, file_type: str, file_size: int, respo
|
||||
|
||||
user_ip = request.client.host
|
||||
file_uuid = ''.join(random.choices(string.ascii_letters + string.digits, k=6))
|
||||
async with s3_worker as worker:
|
||||
try:
|
||||
post_data = await worker.generate_upload_post(file_name, content_type=file_type)
|
||||
except Exception as exception:
|
||||
response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
|
||||
return BaseResponse(result="Error generating S3 access token. Error: " + str(exception), error=True)
|
||||
try:
|
||||
post_data = await s3_worker.generate_upload_post(file_uuid, content_type=file_type)
|
||||
except Exception as exception:
|
||||
response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
|
||||
return BaseResponse(result="Error generating S3 access token. Error: " + str(exception), error=True)
|
||||
|
||||
redis_worker.create_record(user_ip, file_name, file_uuid, file_type, datetime.now().isoformat(), file_size)
|
||||
print(post_data)
|
||||
@@ -64,18 +63,17 @@ async def get_file_by_uuid(file_uuid:str, response: Response, request: Request)
|
||||
response.status_code = status.HTTP_400_BAD_REQUEST
|
||||
return BaseResponse(result={"data": None, "comment": "File with this UUID not found"}, error=True)
|
||||
|
||||
async with s3_worker as worker:
|
||||
try:
|
||||
download_url = await worker.generate_download_url(file_uuid, redis_data["file_name"])
|
||||
except Exception as exception:
|
||||
response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
|
||||
return BaseResponse(result="Error generating S3 access token. Error: " + str(exception), error=True)
|
||||
try:
|
||||
download_url = await s3_worker.generate_download_url(file_uuid, redis_data["file_name"])
|
||||
except Exception as exception:
|
||||
response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
|
||||
return BaseResponse(result="Error generating S3 access token. Error: " + str(exception), error=True)
|
||||
|
||||
return BaseResponse(result={"data": {"url": download_url, "file_name": redis_data["file_name"], "file_size": redis_data["file_size"]}, "comment": "Ok"})
|
||||
|
||||
|
||||
async def main():
|
||||
config = uvicorn.Config("main:app", port=5000, host="0.0.0.0", log_level="debug")
|
||||
config = uvicorn.Config("main:app", port=int(os.getenv('BACKEND_PORT')), host="0.0.0.0", log_level="debug")
|
||||
server = uvicorn.Server(config)
|
||||
await server.serve()
|
||||
|
||||
|
||||
+43
-64
@@ -1,6 +1,7 @@
|
||||
import os
|
||||
|
||||
import asyncio
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from aiobotocore.session import get_session
|
||||
@@ -10,84 +11,62 @@ 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_config = {
|
||||
"aws_access_key_id": os.getenv('S3_ACCESS_KEY_ID'),
|
||||
"aws_secret_access_key": os.getenv('S3_SECRET_ACCESS_KEY'),
|
||||
"endpoint_url": os.getenv('S3_ENDPOINT_URL')
|
||||
}
|
||||
|
||||
self._s3_session = get_session()
|
||||
self._s3_client = None
|
||||
|
||||
self.bucket = os.getenv('BUCKET_NAME')
|
||||
self.max_file_size = os.getenv('MAX_FILES_SIZE')
|
||||
|
||||
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,
|
||||
)
|
||||
@asynccontextmanager
|
||||
async def get_client(self):
|
||||
async with self._s3_session.create_client("s3", **self._s3_config) as client:
|
||||
yield client
|
||||
|
||||
async def generate_upload_post(self, key: str, content_type: str, expires_in: int = 300) -> dict:
|
||||
return await self._client.generate_presigned_post(
|
||||
Bucket=self.bucket,
|
||||
Key=key,
|
||||
Fields={
|
||||
"Content-Type": content_type,
|
||||
"acl": "private",
|
||||
},
|
||||
Conditions=[
|
||||
["content-length-range", 0, self.max_file_size],
|
||||
{"acl": "private"},
|
||||
],
|
||||
ExpiresIn=expires_in,
|
||||
)
|
||||
async with self.get_client() as client:
|
||||
return await client.generate_presigned_post(
|
||||
Bucket=self.bucket,
|
||||
Key=key,
|
||||
Fields={
|
||||
"Content-Type": content_type,
|
||||
"acl": "private",
|
||||
},
|
||||
Conditions=[
|
||||
["content-length-range", 0, self.max_file_size],
|
||||
{"acl": "private"},
|
||||
],
|
||||
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 with self.get_client() as client:
|
||||
return await client.generate_presigned_url("get_object",
|
||||
Params={
|
||||
"Bucket": self.bucket,
|
||||
"Key": key,
|
||||
"ResponseContentDisposition": (
|
||||
f'attachment; filename="{filename}"'
|
||||
),
|
||||
},
|
||||
ExpiresIn=expires_in,
|
||||
)
|
||||
|
||||
|
||||
s3_worker = S3Worker()
|
||||
async def test_run():
|
||||
async with s3_worker 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)
|
||||
worker = S3Worker()
|
||||
|
||||
url = await worker.generate_download_url("test.txt", "hello.txt")
|
||||
print(url)
|
||||
url = await worker.generate_upload_post("some.jpg", content_type="image/jpeg")
|
||||
print(url)
|
||||
url = await worker.generate_download_url("some.jpg", "image.jpg")
|
||||
print(url)
|
||||
url = await worker.generate_download_url("test.txt", "hello.txt")
|
||||
print(url)
|
||||
url = await worker.generate_upload_post("some.jpg", content_type="image/jpeg")
|
||||
print(url)
|
||||
url = await worker.generate_download_url("some.jpg", "image.jpg")
|
||||
print(url)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user