"""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. """ 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 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 = blobstore.blob_path(user_id, message_sha) if part_index == NOT_IN_MESSAGE or not raw.exists(): 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 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 = blobstore.blob_path(owner, msg_sha) try: indexes = await asyncio.to_thread(_part_indexes, raw) if raw.exists() 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