This commit is contained in:
IgorVolochay
2026-01-06 17:27:59 +03:00
parent c11fb3d9b7
commit 6920ff381c
3 changed files with 59 additions and 80 deletions
+11 -13
View File
@@ -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
View File
@@ -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__":