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

"""APScheduler-based periodic sync for the yunke data mirror (M6, Part A).

Wraps the orchestrator's ``sync_*`` functions behind a long-running scheduler
so the data mirror stays fresh without manual kicks. Two paces:

* **Daily full sync** at a configurable hour (default ``02:00`` local time,
  off-hours) — re-scrapes the org tree, accounts, friends, groups, voice
  calls, plus phone calls and SMS (all entities that have no upstream
  ``since`` cursor and therefore rely on idempotent re-upsert).
* **Incremental message sync** more frequently (default every **4 hours**).
  Messages DO have a real ``[last_timestamp, now]`` window (see
  :func:`app.sync.orchestrator.sync_wechat_messages`), so a fast cadence
  catches new chats promptly without re-pulling history. The conversation
  list is discovered from the already-mirrored ``customer`` table.

Auth: the scheduler builds a :class:`YunkeClient` per run from env (honoring
``.env`` via :func:`scripts.smoke_wechat._push_settings_into_environ`).
Session-cookie mode (``YK_USERID`` + ``YK_COMPANY`` [+ ``YK_COOKIE``]) is
preferred and is the only mode exercised in production today; pcLogin mode
(``YUNKE_USERNAME`` + ``YUNKE_PASSWORD``) is wired but deferred (see README
"Current limitations").

Failure policy: every run is wrapped in a ``try/except`` that logs the error
and lets the scheduler keep its next appointment. A single broken entity
(scraper 500, network blip) does NOT bring down the loop. Per-entity
failures are also captured by the orchestrator as ``status='failed'``
``SyncRun`` rows.

Threading: :class:`SyncScheduler` is built on :class:`AsyncIOScheduler` and
is intended to be ``await``-ed inside a long-running asyncio process (the
CLI's ``scheduler`` command — see :mod:`app.sync.cli`). ``run_once`` is
also awaitable for manual / API-triggered syncs.
"""

from __future__ import annotations

import asyncio
import datetime as dt
import json
import logging
from collections.abc import Awaitable, Callable
from typing import Any

import sqlalchemy as sa
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from apscheduler.triggers.cron import CronTrigger
from apscheduler.triggers.interval import IntervalTrigger
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine

from app.config import get_settings
from app.db import models  # noqa: F401  (register tables on Base.metadata)
from app.db.base import Base
from app.session.yunke_client import YunkeClient
from app.storage.object_store import ObjectStore
from app.sync.orchestrator import (
    SyncResult,
    sync_all_wechat,
    sync_departments,
    sync_phone_calls,
    sync_sms,
    sync_wechat_accounts,
    sync_wechat_friends,
    sync_wechat_groups,
    sync_wechat_messages,
    sync_wechat_voice_calls,
)

logger = logging.getLogger(__name__)

__all__ = ["SyncScheduler", "ENTITY_FUNCS"]

# Cap on the number of conversations one message-sync run will pull. Bounds
# runtime + upstream request volume per run (the orchestrator paces each
# request, so 200 conversations × N pages each can still take a while).
DEFAULT_MAX_CONVERSATIONS_PER_RUN = 200


# ---------------------------------------------------------------------------
# Conversation discovery (drives the incremental message sync)
# ---------------------------------------------------------------------------


async def _discover_conversations(
    session: AsyncSession,
    *,
    limit: int = DEFAULT_MAX_CONVERSATIONS_PER_RUN,
) -> list[tuple[str, str]]:
    """Return distinct ``(owner_account_id, friend_id)`` WeChat pairs to sync.

    Source: the local ``customer`` mirror (where
    ``channel='wechat'`` and ``owner_account_id`` is non-empty). Reading the
    already-mirrored table — instead of the upstream ``talkerList`` endpoint
    — keeps the discovery call off the risk-control budget and makes the
    incremental run purely a function of the customers we already know
    about. Order is unspecified; ``limit`` caps work per run.
    """
    from app.db.models import Customer

    stmt = (
        sa.select(Customer.owner_account_id, Customer.friend_id)
        .where(Customer.channel == "wechat")
        .where(Customer.owner_account_id.is_not(None))
        .where(Customer.owner_account_id != "")
        .distinct()
        .limit(limit)
    )
    rows = (await session.execute(stmt)).all()
    return [
        (owner, friend_id)
        for owner, friend_id in rows
        if owner and friend_id
    ]


# ---------------------------------------------------------------------------
# Entity registry — name -> async callable(session, client, store) -> list[SyncResult]
# ---------------------------------------------------------------------------
#
# One entry per logical "thing the scheduler can run". Each callable returns a
# list of SyncResult (one for the entity, or one per conversation for the
# message entry). Keeping this as a registry means run_once("friends") and the
# scheduler's daily job can share the same dispatch.


async def _run_departments(
    session: AsyncSession, client: YunkeClient, store: ObjectStore
) -> list[SyncResult]:
    return [await sync_departments(session, client)]


async def _run_accounts(
    session: AsyncSession, client: YunkeClient, store: ObjectStore
) -> list[SyncResult]:
    return [await sync_wechat_accounts(session, client)]


async def _run_friends(
    session: AsyncSession, client: YunkeClient, store: ObjectStore
) -> list[SyncResult]:
    return [await sync_wechat_friends(session, client)]


async def _run_groups(
    session: AsyncSession, client: YunkeClient, store: ObjectStore
) -> list[SyncResult]:
    return [await sync_wechat_groups(session, client)]


async def _run_voice_calls(
    session: AsyncSession, client: YunkeClient, store: ObjectStore
) -> list[SyncResult]:
    return [await sync_wechat_voice_calls(session, client)]


async def _run_phone_calls(
    session: AsyncSession, client: YunkeClient, store: ObjectStore
) -> list[SyncResult]:
    return [await sync_phone_calls(session, client)]


async def _run_sms(
    session: AsyncSession, client: YunkeClient, store: ObjectStore
) -> list[SyncResult]:
    return [await sync_sms(session, client)]


async def _run_messages(
    session: AsyncSession, client: YunkeClient, store: ObjectStore
) -> list[SyncResult]:
    """Incremental message sync across all locally-known WeChat conversations."""
    conversations = await _discover_conversations(session)
    results: list[SyncResult] = []
    for wechat_id, friend_id in conversations:
        results.append(
            await sync_wechat_messages(
                session, client, store, wechat_id, friend_id
            )
        )
    return results


async def _run_all_wechat(
    session: AsyncSession, client: YunkeClient, store: ObjectStore
) -> list[SyncResult]:
    """Full WeChat sequence: org/entities first, then per-conversation messages.

    Conversations are discovered from the local ``customer`` mirror so the
    daily job is self-contained (no separate message kick needed).
    """
    conversations = await _discover_conversations(session)
    return await sync_all_wechat(
        session, client, store, conversations=conversations
    )


async def _run_all(
    session: AsyncSession, client: YunkeClient, store: ObjectStore
) -> list[SyncResult]:
    """Everything: full WeChat sequence + phone calls + SMS.

    This is the entity the scheduler's daily job runs.
    """
    results = await _run_all_wechat(session, client, store)
    results.append(await sync_phone_calls(session, client))
    results.append(await sync_sms(session, client))
    return results


# Logical name -> runner. ``messages`` is the fast-cadence entry; ``all`` is
# the daily full-sync entry. The single-entity names are exposed for manual
# ``run_once("<entity>")`` triggers from the CLI / API.
ENTITY_FUNCS: dict[
    str, Callable[[AsyncSession, YunkeClient, ObjectStore], Awaitable[list[SyncResult]]]
] = {
    "departments": _run_departments,
    "accounts": _run_accounts,
    "friends": _run_friends,
    "groups": _run_groups,
    "voice_calls": _run_voice_calls,
    "phone_calls": _run_phone_calls,
    "sms": _run_sms,
    "messages": _run_messages,
    "all_wechat": _run_all_wechat,
    "all": _run_all,
}


# ---------------------------------------------------------------------------
# JSON-line logger — one line per run, easy to ship to any aggregator
# ---------------------------------------------------------------------------


def _log_run(
    entity: str,
    results: list[SyncResult],
    *,
    started_at: dt.datetime,
    error: str | None = None,
) -> None:
    """Emit one JSON-line log entry summarising a finished run.

    Always emits exactly one line, even on failure, so a log shipper can
    parse the stream uniformly. ``level`` is ``info`` on full success,
    ``warning`` if any sub-result failed, ``error`` if the run itself blew
    up before producing results.
    """
    finished_at = dt.datetime.now(dt.UTC)
    fetched = sum(r.fetched for r in results)
    upserted = sum(r.upserted for r in results)
    failed = [r.entity for r in results if r.status != "success"]

    if error is not None:
        level = "error"
    elif failed:
        level = "warning"
    else:
        level = "info"

    payload: dict[str, Any] = {
        "ts": finished_at.isoformat(),
        "event": "sync_run",
        "entity": entity,
        "started_at": started_at.isoformat(),
        "duration_sec": round((finished_at - started_at).total_seconds(), 3),
        "results": [
            {
                "entity": r.entity,
                "status": r.status,
                "fetched": r.fetched,
                "upserted": r.upserted,
                "channel": r.channel,
                "scope_id": r.scope_id,
                "error": r.error,
            }
            for r in results
        ],
        "totals": {"fetched": fetched, "upserted": upserted, "count": len(results)},
        "failed_entities": failed,
        "level": level,
    }
    if error is not None:
        payload["error"] = error

    line = json.dumps(payload, ensure_ascii=False, default=str)
    # Logger so operators can route with structured-log filters; the JSON body
    # is the whole message (no %-format) so a downstream parser sees one
    # object per line.
    logger.info(line)
    # Also echo to stdout for human/`docker logs` consumers; flush so a
    # piped reader sees it immediately.
    print(line, flush=True)


# ---------------------------------------------------------------------------
# SyncScheduler
# ---------------------------------------------------------------------------


class SyncScheduler:
    """Periodic-sync driver wrapping an :class:`AsyncIOScheduler`.

    Construction is cheap (no I/O); :meth:`start` opens the engines and
    registers jobs, :meth:`stop` tears them down. :meth:`run_once` is a
    standalone manual trigger that does NOT require :meth:`start` — useful
    from the CLI / API.

    Parameters
    ----------
    full_sync_hour:
        Local-time hour at which the daily full sync fires (0..23). Default 2.
    message_interval_hours:
        Hours between incremental message syncs. Default 4.
    session_maker / engine:
        Optional injected async engine + session maker (tests inject an
        in-memory SQLite). When omitted, a production engine + maker are
        built from :attr:`Settings.db_url`.
    timezone:
        APScheduler timezone string for the daily cron. Defaults to the
        local machine timezone (matches operator expectation).
    """

    def __init__(
        self,
        *,
        full_sync_hour: int = 2,
        message_interval_hours: int = 4,
        engine: Any | None = None,
        session_maker: async_sessionmaker[AsyncSession] | None = None,
        timezone: str | None = None,
        max_conversations_per_run: int = DEFAULT_MAX_CONVERSATIONS_PER_RUN,
    ) -> None:
        self.full_sync_hour = full_sync_hour
        self.message_interval_hours = message_interval_hours
        self.timezone = timezone
        self.max_conversations_per_run = max_conversations_per_run

        # Engine ownership flag — if the caller injected an engine, they own
        # its lifecycle (a test will dispose it). Otherwise we own + dispose
        # on :meth:`stop`.
        self._owns_engine = engine is None
        if engine is not None:
            self._engine = engine
            self._session_maker = session_maker or async_sessionmaker(
                engine, expire_on_commit=False
            )
        else:
            settings = get_settings()
            self._engine = create_async_engine(
                settings.db_url, pool_pre_ping=True, future=True
            )
            self._session_maker = async_sessionmaker(
                self._engine, expire_on_commit=False
            )

        self._scheduler: AsyncIOScheduler | None = None

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    async def start(self) -> None:
        """Open the engine, create tables if missing, register jobs, start.

        Idempotent: calling :meth:`start` twice without an intervening
        :meth:`stop` is a no-op (the second call logs and returns).
        """
        if self._scheduler is not None:
            logger.warning("SyncScheduler.start() called twice; ignoring.")
            return

        # Create tables if missing (mirrors scripts.smoke_wechat). Cheap on a
        # warm DB (a no-op CREATE that already-exists). Lets the scheduler
        # boot on a fresh container without a separate migration step.
        async with self._engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)

        sched = AsyncIOScheduler(timezone=self.timezone)
        sched.add_job(
            self._safe_run,
            CronTrigger(hour=self.full_sync_hour, timezone=self.timezone),
            args=["all"],
            id="yunke_full_sync",
            replace_existing=True,
            coalesce=True,
            max_instances=1,
        )
        sched.add_job(
            self._safe_run,
            IntervalTrigger(hours=self.message_interval_hours),
            args=["messages"],
            id="yunke_message_sync",
            replace_existing=True,
            coalesce=True,
            max_instances=1,
        )
        sched.start()
        self._scheduler = sched
        logger.info(
            "SyncScheduler started: full sync daily at hour=%s, messages every %sh",
            self.full_sync_hour,
            self.message_interval_hours,
        )

    async def stop(self, *, wait: bool = True) -> None:
        """Stop the scheduler and (if owned) dispose the engine."""
        if self._scheduler is not None:
            self._scheduler.shutdown(wait=wait)
            self._scheduler = None
        if self._owns_engine:
            await self._engine.dispose()

    async def run_once(self, entity: str = "all") -> list[SyncResult]:
        """Run one sync for ``entity`` (e.g. ``"friends"``, ``"all"``).

        Does NOT require :meth:`start` — opens its own client + session for
        the run and tears them down before returning. Returns the
        orchestrator's :class:`SyncResult` list (one per entity / per
        conversation). Raises ``ValueError`` on an unknown entity name so
        the CLI can surface usage cleanly.
        """
        if entity not in ENTITY_FUNCS:
            raise ValueError(
                f"unknown entity: {entity!r}. "
                f"Known: {sorted(ENTITY_FUNCS)}"
            )
        # _safe_run handles the client build + try/except + logging; surface
        # the results it produced (possibly an empty list on hard failure).
        return await self._safe_run(entity, log=True)

    # ------------------------------------------------------------------
    # Internals
    # ------------------------------------------------------------------

    async def _safe_run(
        self, entity: str, *, log: bool = True
    ) -> list[SyncResult]:
        """Run one entity, swallowing + logging any unexpected exception.

        This is the failure-policy wrapper that protects the scheduler loop:
        a raised exception becomes a logged error and an empty result list,
        and the scheduler keeps its next appointment.
        """
        started_at = dt.datetime.now(dt.UTC)
        client: YunkeClient | None = None
        results: list[SyncResult] = []
        try:
            client = await _build_client_from_env()
            async with self._session_maker() as session:
                runner = ENTITY_FUNCS[entity]
                results = await runner(session, client, ObjectStore())
                await session.commit()
        except Exception as exc:  # noqa: BLE001 — scheduler MUST NOT crash
            logger.exception("sync %s failed", entity)
            if log:
                _log_run(entity, results, started_at=started_at, error=repr(exc))
            return results

        if log:
            _log_run(entity, results, started_at=started_at)
        return results


# ---------------------------------------------------------------------------
# Client factory (env-driven; session-cookie mode preferred)
# ---------------------------------------------------------------------------


async def _build_client_from_env() -> YunkeClient:
    """Construct + authenticate a :class:`YunkeClient` from env / ``.env``.

    Lifts non-empty auth values from :func:`app.config.get_settings` (which
    loads ``.env``) into ``os.environ`` first, then delegates to
    :func:`scripts.smoke_wechat.build_client_from_env` so the precedence
    (session-cookie mode wins; pcLogin fallback) stays in one place.

    In login mode, performs the actual ``await client.login()`` (real RSA +
    ``pcLogin``); in session-cookie mode the client is already authenticated
    and no network call is made here.

    Raises ``RuntimeError`` if no auth env is configured at all.
    """
    # Imported lazily so a misconfigured ``scripts`` module never breaks
    # import of :mod:`app.sync.scheduler` (e.g. in unit tests that inject
    # their own engine / client).
    from scripts.smoke_wechat import (
        _push_settings_into_environ,
        build_client_from_env,
    )

    _push_settings_into_environ()
    client = build_client_from_env()
    if client is None:
        raise RuntimeError(
            "no yunke auth configured: set YK_USERID+YK_COMPANY (preferred) "
            "or YUNKE_USERNAME+YUNKE_PASSWORD in .env / environment."
        )
    if not client.logged_in:
        # pcLogin mode — actually log in now. May raise; the caller's
        # _safe_run wrapper will catch + log.
        await client.login()
    return client