webmail-client/backend/scripts/reindex_search.py

"""Rebuild the search index for every message, from Postgres alone.

Both backends index the same :class:`SearchDoc` (SPEC §5.2), and every field of
it is already stored in the database — subject/addresses/flags on ``messages``,
the plain-text body on ``body_parts``, filenames on ``attachments``. So the
index is fully derivable and never needs its own backup: if the Meilisearch
volume is lost, or you switch SEARCH_BACKEND, run this.

    python scripts/reindex_search.py                 # everything
    python scripts/reindex_search.py --user <uuid>   # one user
    python scripts/reindex_search.py --purge         # drop stale docs first

In Docker:  ./deploy/webmail.sh reindex [--purge]

``--purge`` clears each user's existing documents before rebuilding, which is
what you want after deleting mail outside the app or when the index has drifted;
without it the rebuild is a pure upsert (safe to re-run, safe to interrupt).
"""

from __future__ import annotations

import argparse
import asyncio
import uuid
from collections import defaultdict
from datetime import datetime

from sqlalchemy import select

from app.db import get_sessionmaker
from app.models import Attachment, BodyPart, Mailbox, Message
from app.search.base import SearchBackend, SearchDoc
from app.search.factory import get_search_backend

# Rows pulled per round-trip, and how many documents are pushed to the backend
# concurrently. Deliberately modest: this runs on a NAS alongside live sync.
BATCH = 500
CONCURRENCY = 8


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


async def _index_batch(backend: SearchBackend, docs: list[SearchDoc]) -> None:
    sem = asyncio.Semaphore(CONCURRENCY)

    async def one(doc: SearchDoc) -> None:
        async with sem:
            await backend.index(doc)

    await asyncio.gather(*(one(d) for d in docs))


async def main() -> None:
    parser = argparse.ArgumentParser(description="Rebuild the search index from Postgres")
    parser.add_argument("--user", help="only this user_id (UUID)")
    parser.add_argument(
        "--purge",
        action="store_true",
        help="delete each user's existing documents before rebuilding",
    )
    args = parser.parse_args()

    only_user = uuid.UUID(args.user) if args.user else None
    backend = get_search_backend()
    sm = get_sessionmaker()

    if args.purge:
        async with sm() as session:
            users = (
                [only_user]
                if only_user
                else list(await session.scalars(select(Message.user_id).distinct()))
            )
        for user_id in users:
            await backend.delete_user(user_id)
        print(f"purged existing documents for {len(users)} user(s)")

    total = 0
    after: uuid.UUID | None = None
    while True:
        async with sm() as session:
            stmt = (
                select(Message, Mailbox.special_use)
                .join(Mailbox, Mailbox.id == Message.mailbox_id)
                .order_by(Message.id)
                .limit(BATCH)
            )
            if after is not None:
                stmt = stmt.where(Message.id > after)
            if only_user is not None:
                stmt = stmt.where(Message.user_id == only_user)
            rows = list(await session.execute(stmt))
            if not rows:
                break

            ids = [row[0].id for row in rows]
            after = ids[-1]

            # Two batched lookups instead of 2N per-message queries.
            bodies: dict[uuid.UUID, str] = {}
            for message_id, content in await session.execute(
                select(BodyPart.message_id, BodyPart.content).where(
                    BodyPart.message_id.in_(ids), BodyPart.kind == "text"
                )
            ):
                bodies.setdefault(message_id, content)

            names: dict[uuid.UUID, list[str]] = defaultdict(list)
            for message_id, filename in await session.execute(
                select(Attachment.message_id, Attachment.filename).where(
                    Attachment.message_id.in_(ids),
                    Attachment.content_id.is_(None),  # real attachments, not inline images
                )
            ):
                names[message_id].append(filename)

        docs = [
            SearchDoc(
                message_id=msg.id,
                thread_id=msg.thread_id or msg.id,
                user_id=msg.user_id,
                account_id=msg.account_id,
                mailbox_special_use=special_use,
                from_=msg.from_addr,
                to=msg.to_addr,
                cc=msg.cc_addr,
                subject=msg.subject,
                body_text=bodies.get(msg.id, ""),
                attachment_names=names.get(msg.id, []),
                has_attachments=msg.has_attachments,
                is_unread="\\Seen" not in msg.flags,
                is_starred="\\Flagged" in msg.flags,
                date=_epoch(msg.internal_date),
                size_bytes=msg.size_bytes,
            )
            for msg, special_use in rows
        ]

        await _index_batch(backend, docs)
        total += len(docs)
        print(f"indexed {total} messages", flush=True)

    print(f"done: {total} messages reindexed")


if __name__ == "__main__":
    asyncio.run(main())