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

"""Attachments, kept once: inside their message (SPEC App. D §D.6 item 1).

An attachment is not stored apart from its message any more — it was, which
kept most of a mailbox's bytes twice (once base64 in the raw message, once
decoded on its own). It is taken out of the raw message when wanted:

* :func:`load` — its bytes (downloads, zip, inline images, forwarding, thumbnails);
* :func:`file` — a file holding it, for streaming with HTTP Range (videos):
  taken out into ``{user}/cache/`` and kept a day after its last use;
* :func:`dedupe_existing` — the one-time pass removing the separate copies
  stored before, each only after checking it can be taken out of its message.

Blocking (parsing, disk): call these in a thread — after
:func:`app.indexer.rawstore.ensure` for the message, which fetches a hosted
mailbox's message from its mail server when it is not cached.
"""

from __future__ import annotations

import contextlib
import hashlib
import logging
import time
import uuid
from collections import defaultdict
from pathlib import Path

from app.indexer import blobstore, rawstore
from app.indexer.mime import load_message, part_payload

logger = logging.getLogger(__name__)

CACHE_KEEP_S = 24 * 3600  # a taken-out file is kept this long after its last use
CACHE_MAX_BYTES = 1024**3  # per user; the least recently used go first beyond it
# part_index of a row whose content is not in its message (kept as its own file).
NOT_IN_MESSAGE = -1


def _legacy(user_id: uuid.UUID, sha: str) -> Path | None:
    """The separate copy stored before attachments were kept once, if still there."""
    path = blobstore.blob_path(user_id, sha)
    return path if path.exists() else None


def load(user_id: uuid.UUID, message_sha: str, sha: str, part_index: int | None) -> bytes:
    """An attachment's bytes. Raises FileNotFoundError when they cannot be had."""
    legacy = _legacy(user_id, sha)
    if legacy is not None:
        return legacy.read_bytes()
    raw = rawstore.locate(user_id, message_sha)
    if part_index == NOT_IN_MESSAGE or raw is None:
        raise FileNotFoundError(sha)
    payload = part_payload(load_message(path=raw), part_index, sha)
    if payload is None:
        raise FileNotFoundError(sha)
    return payload


def materialize(user_id: uuid.UUID, message_sha: str, sha: str, part_index: int | None) -> bool:
    """Make sure the attachment exists as a blob of its own — for a compose
    attachment (forward, a draft reopened), which is read by its hash. False
    when its bytes cannot be had."""
    if blobstore.exists(user_id, sha):
        return True
    try:
        blobstore.store(user_id, load(user_id, message_sha, sha, part_index))
    except FileNotFoundError:
        return False
    return True


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


def prune_cache(cache: Path) -> None:
    """Drop streaming copies past their keep time, and the oldest beyond the cap."""
    files = [(p, p.stat()) for p in cache.iterdir() if p.is_file()]
    now = time.time()
    total = 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 > CACHE_KEEP_S or total > CACHE_MAX_BYTES:
            path.unlink(missing_ok=True)


def has_file(user_id: uuid.UUID, sha: str) -> bool:
    """Whether :func:`file` can answer without reading the message."""
    return _legacy(user_id, sha) is not None or (_cache_dir(user_id) / sha).exists()


def file(user_id: uuid.UUID, message_sha: str, sha: str, part_index: int | None) -> Path:
    """A file holding the attachment, to stream from. Raises FileNotFoundError."""
    legacy = _legacy(user_id, sha)
    if legacy is not None:
        return legacy
    cache = _cache_dir(user_id)
    path = cache / sha
    if path.exists():
        path.touch()  # used again: kept longer
        return path
    data = load(user_id, message_sha, sha, part_index)
    cache.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(f".{uuid.uuid4().hex}.tmp")
    tmp.write_bytes(data)
    tmp.replace(path)
    with contextlib.suppress(OSError):  # a file removed meanwhile by another request
        prune_cache(cache)
    return path


def drop_cached(user_id: uuid.UUID) -> None:
    """Forget every taken-out file of a user (account removed)."""
    cache = _cache_dir(user_id)
    if cache.exists():
        for path in cache.iterdir():
            path.unlink(missing_ok=True)


# ---- removing the copies stored before ------------------------------------------------

# A file this recent may belong to work in progress (a sent message waiting
# for its Sent copy is referenced only by its job): left for the next pass.
_MIN_AGE_S = 3600


def _part_indexes(raw_path: Path) -> dict[str, int]:
    """Content hash → MIME part position, for every leaf part of a message."""
    out: dict[str, int] = {}
    for index, part in enumerate(load_message(path=raw_path).walk()):
        if part.is_multipart():
            continue
        payload = part.get_payload(decode=True)
        if isinstance(payload, bytes):
            out.setdefault(hashlib.sha256(payload).hexdigest(), index)
    return out


async def dedupe_existing(batch: int = 200, *, user_id: uuid.UUID | None = None) -> dict[str, int]:
    """Record where each older attachment sits in its message, and delete its
    separate copy once it is known to be there. Idempotent; resumable."""
    import asyncio

    from sqlalchemy import select, update

    from app.db import session_scope
    from app.models import Attachment, ComposeAttachment, Message

    stats = {"rows": 0, "found": 0, "not_found": 0, "files_removed": 0, "bytes_freed": 0}
    while True:
        async with session_scope() as db:
            rows = (
                await db.execute(
                    select(
                        Attachment.id,
                        Attachment.user_id,
                        Attachment.blob_sha256,
                        Message.blob_sha256,
                    )
                    .join(Message, Message.id == Attachment.message_id)
                    .where(
                        Attachment.part_index.is_(None),
                        *([Attachment.user_id == user_id] if user_id is not None else []),
                    )
                    .order_by(Attachment.message_id)
                    .limit(batch)
                )
            ).all()
        if not rows:
            return stats
        by_message: dict[tuple[uuid.UUID, str], list[tuple[uuid.UUID, str]]] = defaultdict(list)
        for att_id, owner, sha, msg_sha in rows:
            by_message[(owner, msg_sha)].append((att_id, sha))

        located: dict[uuid.UUID, int] = {}
        for (owner, msg_sha), atts in by_message.items():
            raw = rawstore.locate(owner, msg_sha)
            try:
                indexes = await asyncio.to_thread(_part_indexes, raw) if raw is not None else {}
            except Exception:  # noqa: BLE001 - unreadable message: keep the copies
                logger.warning("attachments: cannot parse %s", raw, exc_info=True)
                indexes = {}
            for att_id, sha in atts:
                located[att_id] = indexes.get(sha, NOT_IN_MESSAGE)

        async with session_scope() as db:
            for att_id, index in located.items():
                await db.execute(
                    update(Attachment).where(Attachment.id == att_id).values(part_index=index)
                )
        stats["rows"] += len(located)
        stats["found"] += sum(1 for i in located.values() if i != NOT_IN_MESSAGE)
        stats["not_found"] += sum(1 for i in located.values() if i == NOT_IN_MESSAGE)

        # A separate copy goes only when every attachment row with that content
        # is found in its message and nothing else uses the file.
        candidates = {
            (owner, sha)
            for att_id, owner, sha, _ in rows
            if located.get(att_id, NOT_IN_MESSAGE) != NOT_IN_MESSAGE
        }
        async with session_scope() as db:
            for owner, sha in candidates:
                path = blobstore.blob_path(owner, sha)
                if not path.exists() or time.time() - path.stat().st_mtime < _MIN_AGE_S:
                    continue
                in_use = (
                    await db.scalar(
                        select(Message.id)
                        .where(Message.user_id == owner, Message.blob_sha256 == sha)
                        .limit(1)
                    )
                    or await db.scalar(
                        select(ComposeAttachment.id)
                        .where(
                            ComposeAttachment.user_id == owner,
                            ComposeAttachment.blob_sha256 == sha,
                        )
                        .limit(1)
                    )
                    or await db.scalar(
                        select(Attachment.id)
                        .where(
                            Attachment.user_id == owner,
                            Attachment.blob_sha256 == sha,
                            (Attachment.part_index.is_(None))
                            | (Attachment.part_index == NOT_IN_MESSAGE),
                        )
                        .limit(1)
                    )
                )
                if in_use:
                    continue
                size = path.stat().st_size
                path.unlink(missing_ok=True)
                stats["files_removed"] += 1
                stats["bytes_freed"] += size


async def dedupe_existing_job(ctx: dict[str, object]) -> dict[str, int]:
    """arq job (bulk queue): :func:`dedupe_existing`, run once after an update."""
    stats = await dedupe_existing()
    logger.info("attachments kept once: %s", stats)
    return stats