"""Local-filesystem object store with sha256 content-addressed dedup. This is the M1 storage backend for media blobs (image / voice / video / file downloads referenced by :class:`app.db.models.Media`). The interface is the seam a future MinIO/S3 backend will swap in behind: * :meth:`ObjectStore.put_bytes` — write bytes under an opaque key. * :meth:`ObjectStore.put_stream` — write an async-iterable stream under a key. * :meth:`ObjectStore.put_by_hash` — content-addressed put with sha256 dedup: computes the hash, and only writes when no file with that hash exists. The path layout is ``base_dir//`` (a simple fan-out directory so no single directory holds millions of entries). * :meth:`ObjectStore.exists` / :meth:`get_path` — predicate / pure-path helpers. All disk I/O is dispatched via :func:`asyncio.to_thread`, so the methods are genuinely async and safe to call from a FastAPI handler / scraper loop without blocking the event loop on a slow disk. The default base directory is :attr:`Settings.media_dir`. """ from __future__ import annotations import asyncio import hashlib from collections.abc import AsyncIterator from pathlib import Path from app.config import get_settings __all__ = ["ObjectStore"] class ObjectStore: """Async local-filesystem object store with sha256-based dedup.""" def __init__(self, base_dir: Path | None = None) -> None: """Create a store rooted at ``base_dir``. Defaults to :attr:`Settings.media_dir` when omitted so callers can ``ObjectStore()`` and land in the configured media location. """ self.base_dir: Path = Path(base_dir) if base_dir is not None else get_settings().media_dir # ------------------------------------------------------------------ # Path helpers (pure, no I/O) # ------------------------------------------------------------------ def get_path(self, key: str) -> Path: """Return ``base_dir / `` (a pure-path computation; no I/O). ``key`` may contain ``/`` separators (e.g. ``"foo/bar.txt"``); they become real subdirectories on disk. The key is treated as opaque — no hash is computed. """ return self.base_dir / key def _hash_path(self, sha256_hex: str, ext: str) -> Path: """Return ``base_dir / / ``.""" if ext and not ext.startswith("."): ext = "." + ext return self.base_dir / sha256_hex[:2] / f"{sha256_hex}{ext}" # ------------------------------------------------------------------ # Writes # ------------------------------------------------------------------ async def put_bytes(self, key: str, data: bytes) -> Path: """Write ``data`` under ``key`` (overwrites if present). Returns the resolved path. Always writes; for content-addressed dedup use :meth:`put_by_hash`. """ def _write() -> Path: path = self.get_path(key) path.parent.mkdir(parents=True, exist_ok=True) path.write_bytes(data) return path return await asyncio.to_thread(_write) async def put_stream( self, key: str, stream: AsyncIterator[bytes] ) -> Path: """Concatenate an async-byte-stream under ``key`` (overwrites). Chunks are consumed in the event-loop thread; the file open + each write is dispatched to a worker thread to avoid blocking on disk. """ path = self.get_path(key) def _mkdir() -> None: path.parent.mkdir(parents=True, exist_ok=True) await asyncio.to_thread(_mkdir) def _write_chunk(buf: list[bytes], f) -> None: for chunk in buf: f.write(chunk) buf.clear() # Stream chunks in the event loop but push each disk write to a # thread so we never block on slow storage mid-iteration. def _open_and_prep(): return open(path, "wb") f = await asyncio.to_thread(_open_and_prep) try: buffer: list[bytes] = [] async for chunk in stream: buffer.append(chunk) # Flush once we have a meaningful amount buffered. if sum(len(c) for c in buffer) >= 65536: await asyncio.to_thread(_write_chunk, buffer, f) if buffer: await asyncio.to_thread(_write_chunk, buffer, f) finally: await asyncio.to_thread(f.close) return path async def put_by_hash( self, data: bytes, *, ext: str = "" ) -> tuple[str, Path]: """Content-addressed put with sha256 dedup. Computes ``sha256(data)``, derives ``base_dir/<2 chars>/``, and writes the blob **only if no file with that path already exists** (deduplication across messages that reference the same blob via md5). Returns ``(sha256_hex, path)``. On a hit the existing file is left untouched (no rewrite), so dedup is observable. """ sha256_hex = hashlib.sha256(data).hexdigest() path = self._hash_path(sha256_hex, ext) def _write_if_missing() -> Path: if path.exists(): return path path.parent.mkdir(parents=True, exist_ok=True) path.write_bytes(data) return path await asyncio.to_thread(_write_if_missing) return sha256_hex, path # ------------------------------------------------------------------ # Reads / predicates # ------------------------------------------------------------------ async def exists(self, key: str) -> bool: """True iff ``key`` exists on disk under ``base_dir``.""" return await asyncio.to_thread(self.get_path(key).exists)