webmail-client/backend/app/search/pg.py

"""Postgres full-text SearchBackend (SPEC §5.2, §8; M0-M1).

Indexing populates ``messages.body_indexed`` (a ``tsvector``) from the doc's
searchable fields. Querying translates the parsed §8.2 groups (AND of ORs) into
SQL over messages joined to mailboxes/accounts, always filtered by ``user_id``
(invariant I-7). Relevance ranking / typo tolerance is Meilisearch's job (M2);
here results are date-ordered with keyset pagination.
"""

from __future__ import annotations

import base64
import datetime as dt
import html
import re
from uuid import UUID

from sqlalchemy import Text, any_, cast, exists, func, literal, not_, or_, select, tuple_, update
from sqlalchemy.sql.elements import ColumnElement

from app.db import get_sessionmaker
from app.models import Attachment, Mailbox, Message
from app.search.base import SearchBackend, SearchDoc, SearchHit, SearchQuery

_TS_CONFIG = "simple"  # language-neutral; avoids wrong-language stemming (M0)
_RELATIVE_MULT = {"d": 1, "w": 7, "m": 30, "y": 365}


def _doc_text_blob(doc: SearchDoc) -> str:
    parts = [doc.subject, doc.body_text]
    for field in (doc.from_, doc.to, doc.cc):
        for addr in field:
            parts.append(addr.get("name", ""))
            parts.append(addr.get("email", ""))
    parts.extend(doc.attachment_names)
    return " ".join(p for p in parts if p)


def _parse_date(value: str) -> dt.datetime | None:
    try:
        return dt.datetime.strptime(value, "%Y-%m-%d").replace(tzinfo=dt.UTC)
    except ValueError:
        return None


def _parse_relative(value: str) -> dt.timedelta | None:
    m = re.fullmatch(r"(\d+)([dwmy])", value.strip().lower())
    if not m:
        return None
    n = int(m.group(1)) * _RELATIVE_MULT[m.group(2)]
    return dt.timedelta(days=n)


def _parse_size(value: str) -> int | None:
    m = re.fullmatch(r"(\d+)([km])", value.strip().lower())
    if not m:
        return None
    return int(m.group(1)) * (1024 if m.group(2) == "k" else 1024 * 1024)


def _flag_condition(flag: str) -> ColumnElement[bool]:
    # flag = ANY(messages.flags) — the array membership operator, not func any().
    return literal(flag) == any_(Message.flags)


class PgSearchBackend(SearchBackend):
    async def index(self, doc: SearchDoc) -> None:
        blob = _doc_text_blob(doc)
        async with get_sessionmaker()() as session:
            await session.execute(
                update(Message)
                .where(Message.id == doc.message_id, Message.user_id == doc.user_id)
                .values(body_indexed=func.to_tsvector(_TS_CONFIG, blob))
            )
            await session.commit()

    async def delete(self, message_id: UUID) -> None:
        async with get_sessionmaker()() as session:
            await session.execute(
                update(Message).where(Message.id == message_id).values(body_indexed=None)
            )
            await session.commit()

    async def delete_user(self, user_id: UUID) -> None:
        # User deletion cascades message rows (FK ON DELETE CASCADE); clearing the
        # tsvector here keeps the backend contract satisfied for partial re-index.
        async with get_sessionmaker()() as session:
            await session.execute(
                update(Message).where(Message.user_id == user_id).values(body_indexed=None)
            )
            await session.commit()

    # --- querying ------------------------------------------------------------

    def _clause_condition(self, clause: dict[str, object]) -> ColumnElement[bool] | None:
        value = str(clause.get("value", ""))
        if clause.get("type") == "text":
            return Message.body_indexed.op("@@")(func.websearch_to_tsquery(_TS_CONFIG, value))
        key = clause.get("key")
        cond = self._filter_condition(str(key), value)
        return cond

    def _filter_condition(self, key: str, value: str) -> ColumnElement[bool] | None:
        like = f"%{value}%"
        if key == "from":
            return cast(Message.from_addr, Text).ilike(like)
        if key == "to":
            return cast(Message.to_addr, Text).ilike(like)
        if key == "cc":
            return cast(Message.cc_addr, Text).ilike(like)
        if key == "subject":
            return Message.subject.ilike(like)
        if key == "in":
            return or_(Mailbox.special_use == value.lower(), Mailbox.path.ilike(like))
        if key == "is":
            v = value.lower()
            if v == "unread":
                return not_(_flag_condition("\\Seen"))
            if v == "read":
                return _flag_condition("\\Seen")
            if v == "starred":
                return _flag_condition("\\Flagged")
            if v == "answered":
                return _flag_condition("\\Answered")
            return None
        if key == "has":
            return Message.has_attachments.is_(True) if value.lower() == "attachment" else None
        if key == "filename":
            return exists().where(
                Attachment.message_id == Message.id, Attachment.filename.ilike(like)
            )
        if key == "before":
            d = _parse_date(value)
            return Message.internal_date < d if d else None
        if key == "after":
            d = _parse_date(value)
            return Message.internal_date >= d if d else None
        if key == "newer_than":
            delta = _parse_relative(value)
            return Message.internal_date >= func.now() - delta if delta else None
        if key == "older_than":
            delta = _parse_relative(value)
            return Message.internal_date < func.now() - delta if delta else None
        if key == "larger":
            n = _parse_size(value)
            return Message.size_bytes > n if n is not None else None
        if key == "smaller":
            n = _parse_size(value)
            return Message.size_bytes < n if n is not None else None
        return None

    def _build_conditions(self, groups: list[list[dict[str, object]]]) -> list[ColumnElement[bool]]:
        conditions: list[ColumnElement[bool]] = []
        for group in groups:
            ors: list[ColumnElement[bool]] = []
            for clause in group:
                cond = self._clause_condition(clause)
                if cond is None:
                    continue
                if clause.get("negate"):
                    cond = not_(cond)
                ors.append(cond)
            if ors:
                conditions.append(or_(*ors))
        return conditions

    async def query(
        self, q: SearchQuery, *, limit: int, cursor: str | None
    ) -> tuple[list[SearchHit], str | None]:
        groups = q.filters.get("groups", []) if isinstance(q.filters, dict) else []
        assert isinstance(groups, list)
        conditions = self._build_conditions(groups)

        rank = (
            func.ts_rank_cd(Message.body_indexed, func.websearch_to_tsquery(_TS_CONFIG, q.text))
            if q.text
            else literal(0.0)
        )
        stmt = (
            select(
                Message.id,
                Message.thread_id,
                Message.subject,
                Message.snippet,
                Message.internal_date,
                rank.label("score"),
            )
            .select_from(Message)
            .join(Mailbox, Mailbox.id == Message.mailbox_id)
            .where(Message.user_id == q.user_id, Message.thread_id.is_not(None))
        )
        for cond in conditions:
            stmt = stmt.where(cond)

        decoded = _decode_cursor(cursor)
        if decoded is not None:
            cur_date, cur_id = decoded
            stmt = stmt.where(
                tuple_(Message.internal_date, Message.id)
                < tuple_(literal(cur_date), literal(cur_id))
            )
        stmt = stmt.order_by(Message.internal_date.desc(), Message.id.desc()).limit(limit + 1)

        terms = _positive_text_terms(groups)
        async with get_sessionmaker()() as session:
            rows = (await session.execute(stmt)).all()

        has_more = len(rows) > limit
        rows = rows[:limit]
        hits = [
            SearchHit(
                message_id=row.id,
                thread_id=row.thread_id,
                score=float(row.score or 0.0),
                highlights=_highlights(row.subject, row.snippet, terms),
            )
            for row in rows
        ]
        next_cursor = (
            _encode_cursor(rows[-1].internal_date, rows[-1].id) if has_more and rows else None
        )
        return hits, next_cursor


def _positive_text_terms(groups: list[list[dict[str, object]]]) -> list[str]:
    terms: list[str] = []
    for group in groups:
        for clause in group:
            if clause.get("type") == "text" and not clause.get("negate"):
                terms.append(str(clause.get("value", "")))
    return terms


def _highlights(subject: str, snippet: str, terms: list[str]) -> dict[str, str]:
    out: dict[str, str] = {}
    if not terms:
        return out
    marked_subject = _mark(subject, terms)
    if "<mark>" in marked_subject:
        out["subject"] = marked_subject
    marked_body = _mark(snippet, terms)
    if "<mark>" in marked_body:
        out["body"] = marked_body
    return out


def _mark(value: str, terms: list[str]) -> str:
    """HTML-escape ``value`` and wrap term matches in <mark> (SPEC §5.2 highlights)."""
    escaped = html.escape(value)
    for term in terms:
        if not term:
            continue
        pattern = re.compile(re.escape(html.escape(term)), re.IGNORECASE)
        escaped = pattern.sub(lambda m: f"<mark>{m.group(0)}</mark>", escaped)
    return escaped


def _encode_cursor(date: dt.datetime, message_id: UUID) -> str:
    raw = f"{date.isoformat()}|{message_id}"
    return base64.urlsafe_b64encode(raw.encode()).decode()


def _decode_cursor(cursor: str | None) -> tuple[dt.datetime, UUID] | None:
    if not cursor:
        return None
    try:
        raw = base64.urlsafe_b64decode(cursor.encode()).decode()
        date_str, id_str = raw.split("|", 1)
        return dt.datetime.fromisoformat(date_str), UUID(id_str)
    except (ValueError, TypeError):
        return None