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

"""Indexer orchestration (SPEC §6.4).

Turns a :class:`RemoteMessage` into persisted rows: raw blob (deduped), parsed
MIME, sanitized HTML, ``messages``/``body_parts``/``attachments`` rows, JWZ
threading, and a ``known_recipients`` upsert. Returns a :class:`SearchDoc` for
the caller to index *after commit* (FR-30 inline indexing) and to drive the
``message.new`` event. Idempotent on ``(mailbox_id, uid)``.
"""

from __future__ import annotations

import uuid
from dataclasses import dataclass
from datetime import UTC, datetime

from sqlalchemy import select
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.ext.asyncio import AsyncSession

from app.concurrency import loop_semaphore
from app.config import get_settings
from app.indexer import blobstore, rawstore
from app.indexer.mime import ParsedEmail, parse_message, parse_message_file
from app.indexer.sanitize import sanitize_html
from app.models import Account, Attachment, BodyPart, KnownRecipient, Message
from app.providers.base import RemoteMessage
from app.search.base import SearchDoc
from app.threading_.normalize import synthesize_message_id
from app.threading_.threader import thread_message

_NO_RECIPIENT_HARVEST = {"junk", "sent", "drafts", "trash"}


@dataclass
class PersistResult:
    message: Message
    search_doc: SearchDoc
    has_remote_images: bool
    gmail_labels: tuple[str, ...] = ()


def _epoch(value: datetime) -> int:
    return int(value.timestamp())


def _build_search_doc(message: Message, parsed: ParsedEmail, special_use: str | None) -> SearchDoc:
    return SearchDoc(
        message_id=message.id,
        thread_id=message.thread_id or message.id,
        user_id=message.user_id,
        account_id=message.account_id,
        mailbox_special_use=special_use,
        from_=parsed.from_addr,
        to=parsed.to_addr,
        cc=parsed.cc_addr,
        subject=parsed.subject,
        body_text=parsed.text_body,
        attachment_names=[a.filename for a in parsed.attachments if a.content_id is None],
        has_attachments=parsed.has_attachments,
        is_unread="\\Seen" not in message.flags,
        is_starred="\\Flagged" in message.flags,
        date=_epoch(message.internal_date),
        size_bytes=message.size_bytes,
    )


async def _harvest_recipient(session: AsyncSession, message: Message, parsed: ParsedEmail) -> None:
    """Upsert the From address into known_recipients (weight 1) (SPEC §6.4)."""
    if not parsed.from_addr:
        return
    sender = parsed.from_addr[0]
    email = (sender.get("email") or "").lower()
    if not email:
        return
    stmt = pg_insert(KnownRecipient).values(
        user_id=message.user_id,
        email=email,
        name=sender.get("name", ""),
        times_used=1,
        last_used_at=message.internal_date,
    )
    stmt = stmt.on_conflict_do_update(
        index_elements=[KnownRecipient.user_id, KnownRecipient.email],
        set_={
            "times_used": KnownRecipient.times_used + 1,
            "last_used_at": stmt.excluded.last_used_at,
            "name": stmt.excluded.name,
        },
    )
    await session.execute(stmt)


async def _resolve_message_id(
    session: AsyncSession,
    *,
    user_id: uuid.UUID,
    account_id: uuid.UUID,
    parsed_id: str | None,
    raw: bytes,
) -> str:
    """Resolve the Message-ID to store, applying SPEC §7.1 duplicate handling.

    Missing id -> synthesize. A *duplicate* id already seen in a **different
    account** is a separate delivery (e.g. a bulk sender that reuses one
    Message-ID for every recipient) -> synthesize a non-authoritative id so the
    two do not merge. A same-account duplicate (e.g. Gmail INBOX + All Mail
    holding one message) keeps the real id so the copies stay in one thread.
    """
    if parsed_id is None:
        return synthesize_message_id(raw)
    other_account = await session.scalar(
        select(Message.account_id)
        .where(Message.user_id == user_id, Message.message_id_hdr == parsed_id)
        .limit(1)
    )
    if other_account is not None and other_account != account_id:
        return synthesize_message_id(raw)
    return parsed_id


async def _keeps_raw(session: AsyncSession, account_id: uuid.UUID) -> bool:
    """Whether the raw message is kept for good: not for a hosted mailbox, whose
    mail server holds it (it is only cached, App. D §D.6 item 2)."""
    kind = await session.scalar(select(Account.kind).where(Account.id == account_id))
    return kind != "native"


async def _parse_bounded(user_id: uuid.UUID, sha: str, raw: bytes) -> ParsedEmail:
    """Parse under the PARSE_CONCURRENCY cap; stream large messages from disk (§7).

    Peak RAM stays ~PARSE_CONCURRENCY x STREAM_PARSE_MB: only a few parses run at
    once, and a message above the threshold is parsed from its (already written)
    blob file instead of being held as an in-memory tree beside its raw bytes.
    """
    settings = get_settings()
    threshold = settings.stream_parse_mb * 1024 * 1024
    async with loop_semaphore("mime_parse", "parse", settings.parse_concurrency):
        path = rawstore.locate(user_id, sha) if len(raw) > threshold else None
        if path is not None:
            return parse_message_file(path)
        return parse_message(raw)


async def persist_message(
    session: AsyncSession,
    *,
    user_id: uuid.UUID,
    account_id: uuid.UUID,
    mailbox_id: uuid.UUID,
    mailbox_special_use: str | None,
    remote: RemoteMessage,
    dedupe_account_scope: bool = False,
) -> PersistResult | None:
    """Persist one message. Returns None if it already exists (idempotent).

    ``dedupe_account_scope`` skips a message already stored under a *different*
    mailbox of the same account (same Message-ID) — used for Gmail's All Mail pass so
    a message already synced from INBOX/Sent isn't duplicated (SPEC: Gmail sync)."""
    existing = await session.scalar(
        select(Message.id).where(Message.mailbox_id == mailbox_id, Message.uid == remote.uid)
    )
    if existing is not None:
        return None

    if await _keeps_raw(session, account_id):
        sha = blobstore.store(user_id, remote.raw)
    else:
        sha = rawstore.put_cached(user_id, remote.raw)
    parsed = await _parse_bounded(user_id, sha, remote.raw)

    message_id_hdr = await _resolve_message_id(
        session, user_id=user_id, account_id=account_id, parsed_id=parsed.message_id, raw=remote.raw
    )
    if dedupe_account_scope and message_id_hdr is not None:
        dup = await session.scalar(
            select(Message.id)
            .where(
                Message.account_id == account_id,
                Message.message_id_hdr == message_id_hdr,
                Message.mailbox_id != mailbox_id,
            )
            .limit(1)
        )
        if dup is not None:
            return None
    message = Message(
        id=uuid.uuid4(),
        user_id=user_id,
        account_id=account_id,
        mailbox_id=mailbox_id,
        uid=remote.uid,
        message_id_hdr=message_id_hdr,
        # Keep the real delivered Message-ID (may differ from the stored, possibly
        # synthesized, message_id_hdr) so cross-account copies collapse in threads.
        delivery_id=parsed.message_id,
        in_reply_to=parsed.in_reply_to,
        references_hdr=parsed.references,
        subject=parsed.subject,
        from_addr=parsed.from_addr,
        to_addr=parsed.to_addr,
        cc_addr=parsed.cc_addr,
        bcc_addr=parsed.bcc_addr,
        date_hdr=parsed.date,
        internal_date=remote.internal_date or datetime.now(UTC),
        flags=sorted(remote.flags),
        size_bytes=len(remote.raw),
        has_attachments=parsed.has_attachments,
        snippet=parsed.snippet,
        blob_sha256=sha,
    )
    session.add(message)
    await session.flush()

    # Body parts: plain text (for indexing) + sanitized HTML (for display).
    has_remote_images = False
    if parsed.text_body:
        session.add(BodyPart(message_id=message.id, kind="text", content=parsed.text_body))
    if parsed.html_body is not None:
        result = sanitize_html(parsed.html_body, message.id)
        has_remote_images = result.has_remote_images
        session.add(
            BodyPart(
                message_id=message.id,
                kind="html_sanitized",
                content=result.html,
                sanitizer_v=result.sanitizer_v,
            )
        )

    # Attachments are not stored apart: they stay inside the raw blob, and are
    # taken out of it when wanted (app/indexer/attachments.py).
    for att in parsed.attachments:
        session.add(
            Attachment(
                message_id=message.id,
                user_id=user_id,
                filename=att.filename,
                mime_type=att.mime_type,
                size_bytes=len(att.payload),
                content_id=att.content_id,
                blob_sha256=blobstore.sha256_hex(att.payload),
                part_index=att.part_index,
            )
        )

    message.has_remote_images = has_remote_images
    # Drafts are not conversation messages: keep them out of threads so they don't
    # pollute the reading pane or thread list. They surface only via the drafts UI.
    if mailbox_special_use != "drafts":
        await thread_message(session, message)

    if mailbox_special_use not in _NO_RECIPIENT_HARVEST:
        await _harvest_recipient(session, message, parsed)

    await session.flush()
    doc = _build_search_doc(message, parsed, mailbox_special_use)
    return PersistResult(
        message=message,
        search_doc=doc,
        has_remote_images=has_remote_images,
        gmail_labels=remote.gmail_labels,
    )