"""Atomic download cache with hash verification for Mosaic.""" import os import hashlib import shutil import logging from pathlib import Path from typing import Optional logger = logging.getLogger(__name__) class DownloadCache: """Manages atomic downloads with hash-sharded storage and integrity verification.""" def __init__(self, base_dir: str = "output/raw"): self.base_dir = Path(base_dir) self.base_dir.mkdir(parents=True, exist_ok=True) def get_shard_path(self, content_hash: str, collection: str, ext: str = "") -> Path: """Get hash-sharded storage path: raw/{collection}/{ab}/{cd}/{hash}.{ext}""" if len(content_hash) < 4: raise ValueError(f"Invalid hash: {content_hash}") ab = content_hash[:2] cd = content_hash[2:4] filename = content_hash + (f".{ext}" if ext else "") return self.base_dir / collection / ab / cd / filename def store_atomic(self, tmp_path: str, content_hash: str, collection: str, ext: str = "", expected_size: Optional[int] = None) -> str: """Atomically move a temp file to its cache location. Args: tmp_path: Path to the temporary download file content_hash: SHA256 hash of the content collection: Source collection name ext: File extension expected_size: Expected file size (from Content-Length) Returns: Final path of the cached file Raises: ValueError: If hash or size verification fails """ tmp = Path(tmp_path) if not tmp.exists(): raise FileNotFoundError(f"Temp file not found: {tmp_path}") # Verify size if expected if expected_size is not None: actual_size = tmp.stat().st_size if actual_size != expected_size: raise ValueError( f"Size mismatch: expected {expected_size}, got {actual_size}" ) # Verify hash actual_hash = self.compute_hash(tmp_path) if actual_hash != content_hash: raise ValueError( f"Hash mismatch: expected {content_hash}, got {actual_hash}" ) # Atomic move to final location final_path = self.get_shard_path(content_hash, collection, ext) final_path.parent.mkdir(parents=True, exist_ok=True) shutil.move(str(tmp), str(final_path)) logger.debug("Cached: %s → %s", tmp_path, final_path) return str(final_path) def is_cached(self, content_hash: str, collection: str, ext: str = "") -> bool: """Check if a file exists in cache with valid hash.""" path = self.get_shard_path(content_hash, collection, ext) if not path.exists(): return False actual_hash = self.compute_hash(str(path)) if actual_hash != content_hash: logger.warning("Cache corruption detected: %s (expected %s, got %s)", path, content_hash, actual_hash) return False return True def get_cached_path(self, content_hash: str, collection: str, ext: str = "") -> Optional[str]: """Get path to cached file, or None if not cached/corrupted.""" if self.is_cached(content_hash, collection, ext): return str(self.get_shard_path(content_hash, collection, ext)) return None def verify_file(self, file_path: str, expected_hash: str) -> bool: """Verify a file's hash matches expected.""" actual = self.compute_hash(file_path) return actual == expected_hash def verify_all(self, collection: str) -> tuple[int, int, list[str]]: """Verify all files in a collection. Returns: (verified_count, mismatch_count, list_of_mismatched_paths) """ collection_dir = self.base_dir / collection if not collection_dir.exists(): return 0, 0, [] verified = 0 mismatched = 0 bad_paths = [] for root, _, files in os.walk(collection_dir): for fname in files: if fname.startswith('.') or fname.endswith('.tmp'): continue fpath = os.path.join(root, fname) # Extract expected hash from filename expected_hash = Path(fname).stem actual_hash = self.compute_hash(fpath) if actual_hash == expected_hash: verified += 1 else: mismatched += 1 bad_paths.append(fpath) return verified, mismatched, bad_paths @staticmethod def compute_hash(file_path: str) -> str: """Compute SHA256 hash of a file.""" sha256 = hashlib.sha256() with open(file_path, 'rb') as f: for chunk in iter(lambda: f.read(8192), b''): sha256.update(chunk) return sha256.hexdigest() @staticmethod def hash_content(content: bytes) -> str: """Compute SHA256 hash of bytes content.""" return hashlib.sha256(content).hexdigest() def check_disk_space(self, required_mb: int = 100) -> bool: """Check if enough disk space is available.""" stat = shutil.disk_usage(str(self.base_dir)) free_mb = stat.free // (1024 * 1024) if free_mb < required_mb: logger.warning("Low disk space: %d MB free, %d MB required", free_mb, required_mb) return False return True