yunke/yunke-server/app/storage/object_store.py

"""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/<first 2 hash chars>/<hash><ext>`` (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 / <key>`` (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 / <first 2 hash chars> / <hash><ext>``."""
        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>/<hash><ext>``,
        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)