webmail-client/backend/app/indexer/rawstore.py

"""Where a message's raw source is (SPEC App. D §D.6 item 2).

An external account's messages are kept whole in the blob store, as always.
A hosted mailbox's are not: the mail server beside the app holds them, so the
webmail keeps only what it reads them for (headers, the sanitized body, the
snippet, the search index) and a **cache** of raw messages — the newest ones
on arrival, the ones opened since — bounded per user (``RAW_CACHE_MB``,
least recently used out first). Anything else is fetched from the mail server
when needed: downloading an attachment, "show original", forwarding.

Layout: ``${BLOB_ROOT}/{user_id}/rawcache/{sha256}``, the same content hash
``messages.blob_sha256`` names.

* :func:`locate` — the file, if it is here (blob store or cache); cheap, sync;
* :func:`ensure` — the file, fetched from the mail server if need be (async):
  call it before handing the work to a thread that reads the file.
"""

from __future__ import annotations

import contextlib
import logging
import time
import uuid
from pathlib import Path

from app.config import get_settings
from app.indexer import blobstore

logger = logging.getLogger(__name__)

CACHE_DIR = "rawcache"
KEEP_S = 30 * 24 * 3600  # unused this long: dropped whatever the room
PRUNE_EVERY = 50
_writes: dict[uuid.UUID, int] = {}


def cache_dir(user_id: uuid.UUID) -> Path:
    return blobstore.blob_path(user_id, "00").parent.parent / CACHE_DIR


def cache_path(user_id: uuid.UUID, sha: str) -> Path:
    return cache_dir(user_id) / sha


def locate(user_id: uuid.UUID, sha: str) -> Path | None:
    """The message's raw file, if it is kept or cached here."""
    kept = blobstore.blob_path(user_id, sha)
    if kept.exists():
        return kept
    cached = cache_path(user_id, sha)
    if cached.exists():
        with contextlib.suppress(OSError):
            cached.touch()  # used again: kept longer
        return cached
    return None


def load(user_id: uuid.UUID, sha: str) -> bytes:
    """The raw bytes if they are here. Raises FileNotFoundError."""
    path = locate(user_id, sha)
    if path is None:
        raise FileNotFoundError(sha)
    return path.read_bytes()


def put_cached(user_id: uuid.UUID, data: bytes, sha: str | None = None) -> str:
    """Cache a raw message (``sha`` names it; its content hash by default)."""
    sha = sha or blobstore.sha256_hex(data)
    if blobstore.exists(user_id, sha):
        return sha  # kept for good already (another account has the same message)
    path = cache_path(user_id, sha)
    if not path.exists():
        path.parent.mkdir(parents=True, exist_ok=True)
        tmp = path.with_suffix(f".{uuid.uuid4().hex}.tmp")
        tmp.write_bytes(data)
        tmp.replace(path)
    else:
        path.touch()
    # Pruning lists the whole cache: done every PRUNE_EVERY writes, not each
    # (a backfill writes thousands); meanwhile it overshoots by that much at most.
    writes = _writes.get(user_id, 0) + 1
    _writes[user_id] = writes % PRUNE_EVERY
    if writes >= PRUNE_EVERY:
        with contextlib.suppress(OSError):  # a file removed meanwhile by another request
            prune(user_id)
    return sha


def prune(user_id: uuid.UUID, *, max_bytes: int | None = None) -> int:
    """Drop cached messages unused for KEEP_S, and the least recently used
    beyond the cap. Returns bytes freed."""
    cache = cache_dir(user_id)
    if not cache.is_dir():
        return 0
    cap = max_bytes if max_bytes is not None else get_settings().raw_cache_mb * 1024 * 1024
    files = [(p, p.stat()) for p in cache.iterdir() if p.is_file() and not p.name.endswith(".tmp")]
    now = time.time()
    total = freed = 0
    for path, st in sorted(files, key=lambda f: f[1].st_mtime, reverse=True):
        total += st.st_size
        if now - st.st_mtime > KEEP_S or total > cap:
            path.unlink(missing_ok=True)
            freed += st.st_size
    return freed


def drop(user_id: uuid.UUID, sha: str) -> None:
    cache_path(user_id, sha).unlink(missing_ok=True)


async def ensure(message: object) -> Path:
    """The raw file of ``message`` (a :class:`~app.models.Message`), fetched
    from the hosted mail server into the cache when it is not here. Raises
    FileNotFoundError when it cannot be had (an external account's message
    missing from the blob store; a native one gone from the server)."""
    from app.models import Message

    assert isinstance(message, Message)
    found = locate(message.user_id, message.blob_sha256)
    if found is not None:
        return found
    raw = await _fetch_native(message)
    if raw is None:
        raise FileNotFoundError(message.blob_sha256)
    got = blobstore.sha256_hex(raw)
    if got != message.blob_sha256:
        # The server's copy differs from what was indexed (it should not):
        # cached under the indexed name, so the message can still be read.
        logger.warning("raw of message %s changed on the mail server", message.id)
    put_cached(message.user_id, raw, message.blob_sha256)
    return cache_path(message.user_id, message.blob_sha256)


async def ensure_bytes(message: object) -> bytes:
    path = await ensure(message)
    import asyncio

    return await asyncio.to_thread(path.read_bytes)


async def _fetch_native(message: object) -> bytes | None:
    from sqlalchemy import select

    from app.db import session_scope
    from app.models import Account, Mailbox, Message

    assert isinstance(message, Message)
    async with session_scope() as db:
        row = (
            await db.execute(
                select(Account.kind, Account.email, Mailbox.path)
                .join(Mailbox, Mailbox.account_id == Account.id)
                .where(Account.id == message.account_id, Mailbox.id == message.mailbox_id)
            )
        ).first()
    if row is None or row.kind != "native":
        return None
    from app.hosted.native import native_provider_config
    from app.providers.base import ProviderError
    from app.providers.imap_smtp import ImapSmtpProvider

    provider = ImapSmtpProvider(native_provider_config(row.email))
    try:
        return await provider.fetch_raw(row.path, message.uid)
    except ProviderError as exc:
        logger.warning("cannot fetch message %s from the mail server: %s", message.id, exc)
        return None
    finally:
        await provider.close()