101 lines
3.1 KiB
Python
101 lines
3.1 KiB
Python
import hashlib
|
|
import json
|
|
import os
|
|
import uuid
|
|
from datetime import date
|
|
from pathlib import Path
|
|
|
|
import aiofiles
|
|
from PIL import Image
|
|
from fastapi import UploadFile, HTTPException
|
|
from starlette.responses import FileResponse
|
|
|
|
from core.utils import is_valid_uuid
|
|
from core.service.mongo_db import use_collection
|
|
from root import ROOT_DIR
|
|
from upload.models.upload import FileModel
|
|
|
|
COLLECTION = 'files'
|
|
ThisModel = FileModel
|
|
|
|
|
|
def create_dir(path: str) -> None:
|
|
"""Create a directory if it does not exist."""
|
|
if not os.path.exists(path):
|
|
os.makedirs(path)
|
|
|
|
|
|
def _fix_str_uuid(data: dict) -> dict:
|
|
for key, value in data.items():
|
|
if isinstance(value, str) and is_valid_uuid(value):
|
|
data[key] = uuid.UUID(value)
|
|
return data
|
|
|
|
|
|
def storage_path(dest: str = '', prefix='storage/') -> Path:
|
|
return Path(__file__).parent.parent.parent.joinpath(f'{prefix}{dest}')
|
|
|
|
|
|
class UploadRepository:
|
|
@staticmethod
|
|
async def upload(file: UploadFile, data: str, module):
|
|
if data:
|
|
data = json.JSONDecoder().decode(data)
|
|
|
|
if isinstance(data, dict):
|
|
data = _fix_str_uuid(data)
|
|
|
|
content = await file.read() # async read
|
|
file_size = len(content)
|
|
ext = file.filename.split('.')[-1]
|
|
file_name = '{0}.{1}'.format(hashlib.md5(content).hexdigest(), ext)
|
|
|
|
today = date.today()
|
|
dir = Path(__file__).parent.parent.parent.joinpath(
|
|
'storage/' + today.strftime("%Y%m"))
|
|
|
|
create_dir(dir.__str__())
|
|
|
|
path = dir.joinpath(file_name)
|
|
|
|
async with aiofiles.open(path, 'wb') as out_file:
|
|
await out_file.write(content) # async write
|
|
|
|
if file.content_type.startswith('image'):
|
|
UploadRepository._generate_thumbnail(dir, file_name)
|
|
UploadRepository._generate_thumbnail(dir, file_name, (512, 512))
|
|
UploadRepository._generate_thumbnail(dir, file_name, (48, 48))
|
|
|
|
model = ThisModel(
|
|
file_name=file.filename,
|
|
file_type=file.content_type,
|
|
file_size=file_size / 1024,
|
|
module=module,
|
|
payload=data,
|
|
storage=path.relative_to(ROOT_DIR).as_posix().__str__(),
|
|
).init_id().cu_at()
|
|
|
|
result = await use_collection(COLLECTION).insert_one(model.mongo_dict())
|
|
assert result.acknowledged
|
|
return model.mongo_dict()
|
|
|
|
@staticmethod
|
|
def _generate_thumbnail(dir: Path, file_name, size=(256, 256)):
|
|
try:
|
|
im = Image.open(dir.joinpath(file_name))
|
|
im.thumbnail(size)
|
|
thumb_dir = dir.joinpath('thumb').__str__()
|
|
create_dir(thumb_dir)
|
|
thumb_name = '{0}_{1}'.format(str(size[0]), file_name)
|
|
im.save('{0}/{1}'.format(thumb_dir, thumb_name, "JPEG"))
|
|
except IOError:
|
|
print("cannot create thumbnail for")
|
|
|
|
@staticmethod
|
|
async def get_file(_id: uuid.UUID):
|
|
result = await use_collection(COLLECTION).find_one({"_id": _id})
|
|
if not result:
|
|
raise HTTPException(status_code=404, detail='File Not Found')
|
|
file = storage_path(result.get('storage'), '')
|
|
return FileResponse(file)
|