import os import uuid import hashlib import struct from pathlib import Path from typing import Tuple, Dict, Any, Optional from app.models.MigrationModel import MediaItemStatusEnum project_root = Path(__file__).resolve().parents[4] MIGRATION_BASE_DIR = str(project_root / "uploads" / "migrations") os.makedirs(MIGRATION_BASE_DIR, exist_ok=True) class StorageManager: """ Handles streaming disk storage, deterministic path resolution (/media/{sha256[0:2]}/{sha256}.ext), and atomic POSIX file writes (.tmp_xyz -> fsync -> rename) to guarantee zero corrupt files. """ @staticmethod def get_job_dir(job_id: str) -> str: job_dir = os.path.join(MIGRATION_BASE_DIR, job_id) os.makedirs(os.path.join(job_dir, "dataset"), exist_ok=True) os.makedirs(os.path.join(job_dir, "archives"), exist_ok=True) os.makedirs(os.path.join(job_dir, "media"), exist_ok=True) os.makedirs(os.path.join(job_dir, "logs"), exist_ok=True) return job_dir @classmethod def save_job_config(cls, job_id: str, config: Dict[str, Any]): import json job_dir = cls.get_job_dir(job_id) config_path = os.path.join(job_dir, "job_config.json") with open(config_path, "w", encoding="utf-8") as f: json.dump(config, f) @classmethod def get_job_config(cls, job_id: str) -> Dict[str, Any]: import json job_dir = cls.get_job_dir(job_id) config_path = os.path.join(job_dir, "job_config.json") if os.path.exists(config_path): try: with open(config_path, "r", encoding="utf-8") as f: return json.load(f) except Exception: pass return {} @classmethod async def save_upload_stream_async(cls, upload_file, dest_path: str, chunk_size: int = 256 * 1024) -> Tuple[int, str]: """ Asynchronously streams HTTP upload file to disk using threadpool offloading and 256KB buffer. Prevents blocking asyncio event loop during multi-hundred MB uploads. """ import asyncio def _write_sync(): sha256_hash = hashlib.sha256() total_size = 0 os.makedirs(os.path.dirname(dest_path), exist_ok=True) temp_path = f"{dest_path}.tmp_{uuid.uuid4().hex}" with open(temp_path, "wb") as f: while True: chunk = upload_file.file.read(chunk_size) if not chunk: break f.write(chunk) sha256_hash.update(chunk) total_size += len(chunk) f.flush() os.fsync(f.fileno()) os.rename(temp_path, dest_path) return total_size, sha256_hash.hexdigest() return await asyncio.to_thread(_write_sync) @classmethod def save_upload_stream_sync(cls, file_obj, dest_path: str, chunk_size: int = 65536) -> Tuple[int, str]: """ Synchronous 64KB streaming file saver. """ sha256_hash = hashlib.sha256() total_size = 0 os.makedirs(os.path.dirname(dest_path), exist_ok=True) temp_path = f"{dest_path}.tmp_{uuid.uuid4().hex}" with open(temp_path, "wb") as f: while chunk := file_obj.read(chunk_size): f.write(chunk) sha256_hash.update(chunk) total_size += len(chunk) f.flush() os.fsync(f.fileno()) os.rename(temp_path, dest_path) return total_size, sha256_hash.hexdigest() @classmethod def write_media_file_atomically(cls, job_id: str, file_bytes: bytes, original_filename: str, sha256_hash: str) -> Tuple[str, str, int, int]: """ Writes image file atomically: write to .tmp_xyz -> fsync -> rename to /media/{sha256[0:2]}/{sha256}.ext Returns: (storage_path, cdn_url, width, height) """ ext = os.path.splitext(original_filename)[1].lower() or ".jpg" if not ext.startswith("."): ext = f".{ext}" sub_dir = sha256_hash[:2] media_dir = os.path.join(cls.get_job_dir(job_id), "media", sub_dir) os.makedirs(media_dir, exist_ok=True) filename = f"{sha256_hash}{ext}" final_disk_path = os.path.join(media_dir, filename) temp_disk_path = os.path.join(media_dir, f".tmp_{uuid.uuid4().hex}") # Fast binary PNG width/height header inspection width, height = 0, 0 if ext == ".png" and len(file_bytes) >= 24: if file_bytes[:8] == b'\x89PNG\r\n\x1a\n' and file_bytes[12:16] == b'IHDR': try: width, height = struct.unpack(">II", file_bytes[16:24]) except Exception: pass if width == 0 or height == 0: try: from PIL import Image as PIL_Image import io with PIL_Image.open(io.BytesIO(file_bytes)) as img: width, height = img.size except Exception: width, height = 0, 0 # Write to temporary file -> fsync -> rename if not os.path.exists(final_disk_path): with open(temp_disk_path, "wb") as f: f.write(file_bytes) f.flush() os.fsync(f.fileno()) os.rename(temp_disk_path, final_disk_path) storage_path = os.path.join("uploads", "migrations", job_id, "media", sub_dir, filename) cdn_url = f"/uploads/migrations/{job_id}/media/{sub_dir}/{filename}" return storage_path, cdn_url, width, height @classmethod def exists(cls, storage_path: str) -> bool: """ Verifies whether a physical storage file exists on disk. """ if not storage_path: return False if storage_path.startswith("/"): return os.path.exists(storage_path) disk_path = str(project_root / storage_path) return os.path.exists(disk_path)