yunke/yunke-server/app/sync/orchestrator.py

"""The sync brain — coordinates scrapers, repos, watermarks, and the SyncRun log.

Each public ``sync_*`` function follows the same shape:

    read watermark -> scrape -> upsert -> (advance watermark) -> log SyncRun

Two flavors of "incremental" live here, per Task 2.7:

* **Time-window incremental** (messages only). The message endpoint pages
  backward over a ``(start, end]`` timestamp window, so a real incremental query
  is possible. ``sync_wechat_messages`` reads ``last_timestamp`` for
  ``(message, wechat, "<wechat_id>:<friend_id>")``, queries
  ``[last_timestamp, now]`` via :func:`scrape_wechat_messages`, upserts each
  record, **best-effort downloads media** for media-bearing messages, and
  advances the watermark to ``max(timestamp)``. A re-run with no new messages
  therefore re-queries the empty window and fetches nothing — the watermark
  advanced past the previously-seen rows.
* **Full re-scrape + idempotent upsert** (every other entity). The list
  endpoints (friends / accounts / groups / voice_calls / departments) have no
  ``since`` param, so "incremental" here = re-scrape everything and rely on the
  repos' idempotent Core upsert (same PK twice -> one refreshed row). The
  watermark records ``last_synced_at`` only — no cursor advance, because there
  is no cursor to advance.

Failure policy: a scraper that raises is caught, recorded as a ``SyncRun`` row
with ``status='failed'`` and the error string, and a failed :class:`SyncResult`
is returned — the orchestrator does NOT re-raise, so a single broken entity
does not abort a multi-entity batch. The caller decides whether to surface /
retry. Media-download failures are best-effort: a media error is logged and the
sync continues (the message is still upserted; the watermark still advances).

Transactions: the orchestrator never commits. Every write is flushed into the
caller's ``AsyncSession``; the caller owns the commit / rollback. This keeps a
multi-entity batch atomic at the caller's discretion and lets a test roll a
failed advance back by simply not committing.
"""

from __future__ import annotations

import dataclasses
import datetime as dt
import logging
from collections.abc import AsyncIterator, Callable
from typing import Any

from sqlalchemy.ext.asyncio import AsyncSession

from app.db.models import SyncRun
from app.scraper.media import download_message_media
from app.scraper.wechat import (
    scrape_departments,
    scrape_wechat_accounts,
    scrape_wechat_friends,
    scrape_wechat_groups,
    scrape_wechat_messages,
    scrape_wechat_voice_calls,
)
from app.session.yunke_client import YunkeClient
from app.storage.object_store import ObjectStore
from app.storage.repos import (
    upsert_chat_group,
    upsert_customer,
    upsert_department,
    upsert_messages,
    upsert_voice_call,
    upsert_wechat_account,
)
from app.sync.watermarks import get_last_timestamp, set_watermark

logger = logging.getLogger(__name__)

__all__ = [
    "SyncResult",
    "sync_all_wechat",
    "sync_departments",
    "sync_wechat_accounts",
    "sync_wechat_friends",
    "sync_wechat_groups",
    "sync_wechat_messages",
    "sync_wechat_voice_calls",
]

# ---------------------------------------------------------------------------
# Stable SyncRun / SyncState keys (entity_type / channel labels)
# ---------------------------------------------------------------------------

# entity_type labels — kept as module constants so the SyncRun / SyncState keys
# are stable strings shared between the orchestrator and any dashboard / ops
# consumer (e.g. Task 6.4 monitoring).
ENTITY_DEPARTMENT = "department"
ENTITY_WECHAT_ACCOUNT = "wechat_account"
ENTITY_WECHAT_FRIEND = "wechat_friend"
ENTITY_WECHAT_GROUP = "wechat_group"
ENTITY_WECHAT_VOICE_CALL = "wechat_voice_call"
ENTITY_MESSAGE = "message"

CHANNEL_WECHAT = "wechat"


# ---------------------------------------------------------------------------
# SyncResult — small per-run outcome returned to the caller
# ---------------------------------------------------------------------------


@dataclasses.dataclass(frozen=True)
class SyncResult:
    """Outcome of one sync run (one entity / one conversation).

    Mirrors the essential fields of the persisted ``SyncRun`` row so callers can
    react to a run without re-querying the DB. ``status`` is ``"success"`` or
    ``"failed"``; on failure ``error`` carries the exception text.
    """

    entity: str
    status: str
    fetched: int
    upserted: int
    channel: str | None = None
    scope_id: str | None = None
    error: str | None = None


# ---------------------------------------------------------------------------
# Time helpers
# ---------------------------------------------------------------------------


def _now() -> dt.datetime:
    """Aware UTC now (never naive). Used for SyncRun timestamps."""
    return dt.datetime.now(dt.UTC)


def _now_ms() -> int:
    """Current time as ms-epoch (the message scraper's timestamp unit)."""
    return int(_now().timestamp() * 1000)


# ---------------------------------------------------------------------------
# SyncRun bookkeeping
# ---------------------------------------------------------------------------


async def _start_run(
    session: AsyncSession,
    entity_type: str,
    *,
    channel: str | None = None,
    scope_id: str | None = None,
) -> SyncRun:
    """Open a ``running`` SyncRun row and flush it (assigning run_id)."""
    run = SyncRun(
        entity_type=entity_type,
        channel=channel,
        scope_id=scope_id,
        status="running",
        started_at=_now(),
    )
    session.add(run)
    await session.flush()
    return run


async def _finish_run(
    session: AsyncSession,
    run: SyncRun,
    *,
    status: str,
    fetched: int,
    upserted: int,
    error: str | None = None,
) -> None:
    """Stamp a SyncRun row with final status / counts / finished_at and flush."""
    run.status = status
    run.records_fetched = fetched
    run.records_upserted = upserted
    run.error = error
    run.finished_at = _now()
    await session.flush()


# ---------------------------------------------------------------------------
# Messages — time-window incremental (the workhorse)
# ---------------------------------------------------------------------------


def _message_scope(wechat_id: str, friend_id: str) -> str:
    """Scope-id key for a 1:1 (or room) WeChat conversation."""
    return f"{wechat_id}:{friend_id}"


async def sync_wechat_messages(
    session: AsyncSession,
    client: YunkeClient,
    store: ObjectStore,
    wechat_id: str,
    friend_id: str,
    *,
    types: str = "",
) -> SyncResult:
    """Incremental sync of one WeChat conversation's messages.

    Reads ``last_timestamp`` for ``(message, wechat, "<wechat_id>:<friend_id>")``
    (default ``0`` on first sync), queries ``[last_timestamp, now]`` via the
    scraper, upserts every yielded message, best-effort downloads media for
    each, and advances the watermark to ``max(timestamp)``.

    Returns a :class:`SyncResult`. On a scraper failure, records a
    ``status='failed'`` SyncRun with the error and returns a failed result
    (does NOT re-raise).
    """
    scope_id = _message_scope(wechat_id, friend_id)
    last_ts = await get_last_timestamp(
        session,
        ENTITY_MESSAGE,
        channel=CHANNEL_WECHAT,
        scope_id=scope_id,
        default=0,
    )
    end = _now_ms()

    run = await _start_run(
        session,
        ENTITY_MESSAGE,
        channel=CHANNEL_WECHAT,
        scope_id=scope_id,
    )

    # --- Scrape (may raise) ---
    try:
        msgs: list[dict[str, Any]] = [
            msg
            async for msg in scrape_wechat_messages(
                client, wechat_id, friend_id, last_ts, end, types=types
            )
        ]
    except Exception as exc:
        logger.exception(
            "scrape_wechat_messages failed for %s", scope_id
        )
        error = repr(exc)
        await _finish_run(
            session, run, status="failed", fetched=0, upserted=0, error=error
        )
        return SyncResult(
            entity=ENTITY_MESSAGE,
            status="failed",
            fetched=0,
            upserted=0,
            channel=CHANNEL_WECHAT,
            scope_id=scope_id,
            error=error,
        )

    # --- Upsert (idempotent; same msg_id twice -> one refreshed row) ---
    upserted = await upsert_messages(session, msgs) if msgs else 0

    # --- Best-effort media download (never aborts the sync) ---
    for msg in msgs:
        try:
            await download_message_media(client, msg, store, session)
        except Exception:
            logger.warning(
                "media download failed for msg %s; sync continues",
                msg.get("msg_id"),
                exc_info=True,
            )

    # --- Advance watermark to max(timestamp) seen in this batch ---
    timestamps = [m.get("timestamp") for m in msgs if m.get("timestamp") is not None]
    if timestamps:
        await set_watermark(
            session,
            ENTITY_MESSAGE,
            channel=CHANNEL_WECHAT,
            scope_id=scope_id,
            last_timestamp=max(timestamps),
        )

    await _finish_run(
        session, run, status="success", fetched=len(msgs), upserted=upserted
    )
    return SyncResult(
        entity=ENTITY_MESSAGE,
        status="success",
        fetched=len(msgs),
        upserted=upserted,
        channel=CHANNEL_WECHAT,
        scope_id=scope_id,
    )


# ---------------------------------------------------------------------------
# Full re-scrape entities — idempotent upsert, no since-param
# ---------------------------------------------------------------------------


# Type aliases for readability (a scrape factory takes the client captured at
# call time and returns an async generator of records; an upsert takes the
# session + one record).
_ScrapeFactory = Callable[[], AsyncIterator[dict[str, Any]]]
_UpsertOne = Callable[[AsyncSession, dict[str, Any]], Any]


async def _sync_full_rescrape(
    session: AsyncSession,
    *,
    entity_type: str,
    channel: str | None,
    scope_id: str | None,
    scrape: _ScrapeFactory,
    upsert: _UpsertOne,
) -> SyncResult:
    """Shared body for the "re-scrape all + idempotent upsert" entities.

    These list endpoints have no ``since`` param, so "incremental" = re-scrape
    everything and rely on the repos' idempotent Core upsert (same PK twice ->
    one refreshed row). The watermark records ``last_synced_at`` only; no
    version/timestamp cursor is advanced because there is no upstream window to
    close. A scraper failure is recorded as a failed SyncRun and a failed
    SyncResult is returned (no re-raise).
    """
    run = await _start_run(
        session, entity_type, channel=channel, scope_id=scope_id
    )

    try:
        recs: list[dict[str, Any]] = [rec async for rec in scrape()]
    except Exception as exc:
        logger.exception("scrape failed for %s", entity_type)
        error = repr(exc)
        await _finish_run(
            session, run, status="failed", fetched=0, upserted=0, error=error
        )
        return SyncResult(
            entity=entity_type,
            status="failed",
            fetched=0,
            upserted=0,
            channel=channel,
            scope_id=scope_id,
            error=error,
        )

    for rec in recs:
        await upsert(session, rec)

    # Stamp the watermark's last_synced_at (no cursor advance — see docstring).
    await set_watermark(
        session, entity_type, channel=channel, scope_id=scope_id
    )

    await _finish_run(
        session, run, status="success", fetched=len(recs), upserted=len(recs)
    )
    return SyncResult(
        entity=entity_type,
        status="success",
        fetched=len(recs),
        upserted=len(recs),
        channel=channel,
        scope_id=scope_id,
    )


async def sync_departments(
    session: AsyncSession, client: YunkeClient
) -> SyncResult:
    """Full re-scrape of the department tree and idempotent upsert."""
    return await _sync_full_rescrape(
        session,
        entity_type=ENTITY_DEPARTMENT,
        channel=None,
        scope_id=None,
        scrape=lambda: scrape_departments(client),
        upsert=upsert_department,
    )


async def sync_wechat_accounts(
    session: AsyncSession, client: YunkeClient
) -> SyncResult:
    """Full re-scrape of all WeChat staff accounts (+ embedded staff rows)."""
    return await _sync_full_rescrape(
        session,
        entity_type=ENTITY_WECHAT_ACCOUNT,
        channel=None,
        scope_id=None,
        scrape=lambda: scrape_wechat_accounts(client),
        upsert=upsert_wechat_account,
    )


async def sync_wechat_friends(
    session: AsyncSession, client: YunkeClient
) -> SyncResult:
    """Full re-scrape of all WeChat friends (customers) and idempotent upsert."""
    return await _sync_full_rescrape(
        session,
        entity_type=ENTITY_WECHAT_FRIEND,
        channel=CHANNEL_WECHAT,
        scope_id=None,
        scrape=lambda: scrape_wechat_friends(client),
        upsert=upsert_customer,
    )


async def sync_wechat_groups(
    session: AsyncSession, client: YunkeClient
) -> SyncResult:
    """Full re-scrape of all WeChat chat-groups (+members) and idempotent upsert."""
    return await _sync_full_rescrape(
        session,
        entity_type=ENTITY_WECHAT_GROUP,
        channel=CHANNEL_WECHAT,
        scope_id=None,
        scrape=lambda: scrape_wechat_groups(client),
        upsert=upsert_chat_group,
    )


async def sync_wechat_voice_calls(
    session: AsyncSession, client: YunkeClient
) -> SyncResult:
    """Full re-scrape of all WeChat voice-call records and idempotent upsert."""
    return await _sync_full_rescrape(
        session,
        entity_type=ENTITY_WECHAT_VOICE_CALL,
        channel=None,
        scope_id=None,
        scrape=lambda: scrape_wechat_voice_calls(client),
        upsert=upsert_voice_call,
    )


# ---------------------------------------------------------------------------
# Coordinator
# ---------------------------------------------------------------------------


async def sync_all_wechat(
    session: AsyncSession,
    client: YunkeClient,
    store: ObjectStore,
    *,
    conversations: list[tuple[str, str]] | None = None,
) -> list[SyncResult]:
    """Run the full WeChat sync sequence: org/entities first, then messages.

    Order: departments → accounts → friends → groups → voice_calls, then the
    per-conversation message sync for every ``(wechat_id, friend_id)`` in
    ``conversations``. Each step returns its own :class:`SyncResult` so a single
    failed entity does NOT abort the batch (the failure is recorded and the
    next entity proceeds). Returns the results in execution order.

    ``conversations`` is the explicit "which direct / room chats to pull
    messages for" list. Full auto-discovery from the talkers endpoint is a
    LATER concern (Task 2.5 wires the scraper; wiring it into the coordinator
    is intentionally deferred to keep this entrypoint bounded and testable).
    When ``conversations`` is ``None`` / empty, only the five full-re-scrape
    entities run.
    """
    results: list[SyncResult] = [
        await sync_departments(session, client),
        await sync_wechat_accounts(session, client),
        await sync_wechat_friends(session, client),
        await sync_wechat_groups(session, client),
        await sync_wechat_voice_calls(session, client),
    ]
    for wechat_id, friend_id in conversations or []:
        results.append(
            await sync_wechat_messages(
                session, client, store, wechat_id, friend_id
            )
        )
    return results