from __future__ import annotations import base64 import hashlib from dataclasses import dataclass from pathlib import Path import anyio from sqlalchemy.ext.asyncio import AsyncSession from app.core.config import get_settings from app.core.exceptions import UploadConflictError, UploadNotFoundError, ValidationError from app.core.ids import new_id from app.models.entities import UploadSession from app.repositories.file_repository import FileRepository from app.repositories.upload_repository import UploadRepository from app.services.file_service import CreateFileMetadataInput, FileService from app.services.job_service import JobService from app.services.local_storage_service import LocalStorageService from app.services.storage_service import StorageService @dataclass class CreateUploadInput: upload_length: int upload_metadata: str | None = None @dataclass class FinalizeUploadInput: directory_id: str filename: str class UploadService: def __init__(self, session: AsyncSession) -> None: self.session = session self.repository = UploadRepository(session) self.file_service = FileService(session) self.file_repository = FileRepository(session) self.local_storage_service = LocalStorageService(session) self.storage_service = StorageService(session) self.job_service = JobService(session) self.settings = get_settings() async def ensure_temp_dir(self) -> Path: temp_dir = Path(self.settings.upload_temp_dir) await anyio.to_thread.run_sync(lambda: temp_dir.mkdir(parents=True, exist_ok=True)) return temp_dir async def create_upload(self, payload: CreateUploadInput) -> UploadSession: if payload.upload_length <= 0: raise ValidationError("Upload-Length must be greater than zero.") temp_dir = await self.ensure_temp_dir() upload_id = new_id("upl") temp_path = temp_dir / f"{upload_id}.bin" await anyio.Path(temp_path).write_bytes(b"") upload_session = UploadSession( id=upload_id, total_size_bytes=payload.upload_length, received_size_bytes=0, status="created", temp_path=str(temp_path), tus_upload_url=f"/files/{upload_id}", upload_metadata_json=payload.upload_metadata, ) created = await self.repository.create(upload_session) await self.session.commit() return created async def get_upload(self, upload_id: str) -> UploadSession: upload_session = await self.repository.get_by_id(upload_id) if upload_session is None: raise UploadNotFoundError() return upload_session async def append_upload_bytes(self, upload_id: str, offset: int, chunk: bytes) -> UploadSession: upload_session = await self.get_upload(upload_id) if upload_session.status in {"finalized", "finalizing", "expired"}: raise UploadConflictError("Upload session is not writable.") if offset != upload_session.received_size_bytes: raise UploadConflictError("Upload offset does not match the current stored offset.") if upload_session.received_size_bytes + len(chunk) > upload_session.total_size_bytes: raise UploadConflictError("Upload would exceed declared length.") temp_path = anyio.Path(upload_session.temp_path) async with await anyio.open_file(temp_path, "ab") as file_handle: await file_handle.write(chunk) upload_session.received_size_bytes += len(chunk) upload_session.status = ( "uploaded" if upload_session.received_size_bytes == upload_session.total_size_bytes else "uploading" ) saved = await self.repository.save(upload_session) await self.session.commit() return saved async def finalize_upload(self, upload_id: str, payload: FinalizeUploadInput): upload_session = await self.get_upload(upload_id) if upload_session.received_size_bytes != upload_session.total_size_bytes: raise UploadConflictError("Upload is not complete.") if upload_session.status == "finalized": raise UploadConflictError("Upload has already been finalized.") upload_session.status = "finalizing" await self.repository.save(upload_session) await self.session.flush() content_hash = await self._compute_sha256(upload_session.temp_path) mime_type = self._guess_mime_type(payload.filename) file_entry = await self.file_service.create_file_metadata( CreateFileMetadataInput( directory_id=payload.directory_id, name=payload.filename, mime_type=mime_type, size_bytes=upload_session.total_size_bytes, content_hash=content_hash, ) ) blob = await self.file_repository.get_primary_blob_for_file(file_entry.id) if blob is None: raise UploadConflictError("Primary blob was not created for finalized upload.") await self.local_storage_service.persist_blob_from_temp(blob, upload_session.temp_path) await self.job_service.enqueue_file_reconcile(file_entry.id) if mime_type is not None: await self.job_service.enqueue_preview_generation(file_entry.id) upload_session.directory_id = payload.directory_id upload_session.filename = payload.filename upload_session.status = "finalized" await self.repository.save(upload_session) await self.session.commit() await anyio.Path(upload_session.temp_path).unlink(missing_ok=True) return upload_session, file_entry async def _compute_sha256(self, temp_path: str) -> str: def _hash_file() -> str: digest = hashlib.sha256() with open(temp_path, "rb") as file_handle: while True: chunk = file_handle.read(1024 * 1024) if not chunk: break digest.update(chunk) return digest.hexdigest() return await anyio.to_thread.run_sync(_hash_file) @staticmethod def parse_upload_metadata(raw_metadata: str | None) -> dict[str, str]: if not raw_metadata: return {} parsed: dict[str, str] = {} for item in raw_metadata.split(","): item = item.strip() if not item: continue parts = item.split(" ", 1) key = parts[0] value = "" if len(parts) == 2 and parts[1]: value = base64.b64decode(parts[1]).decode("utf-8") parsed[key] = value return parsed @staticmethod def _guess_mime_type(filename: str) -> str | None: lowered = filename.lower() if lowered.endswith(".mp4"): return "video/mp4" if lowered.endswith(".jpg") or lowered.endswith(".jpeg"): return "image/jpeg" if lowered.endswith(".png"): return "image/png" if lowered.endswith(".pdf"): return "application/pdf" return None