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

"""Command-line entry point for running syncs + the scheduler (M6, Part B).

Three invocations::

    # One-shot full sync (every entity) into the configured DB:
    python -m app.sync.cli sync

    # One-shot single-entity sync (departments / friends / messages / ...):
    python -m app.sync.cli sync --entity friends

    # Long-running scheduler (APScheduler; full sync daily + messages every 4h):
    python -m app.sync.cli scheduler

Auth + DB
---------
Auth comes from env / ``.env`` via the same helpers the smoke script uses
(:func:`scripts.smoke_wechat._push_settings_into_environ` +
:func:`scripts.smoke_wechat.build_client_from_env`) — session-cookie mode
(``YK_USERID`` + ``YK_COMPANY``) is preferred; pcLogin mode is wired but
deferred.

The DB is taken from :attr:`Settings.db_url` (MySQL in prod, SQLite in dev /
tests). Tables are created if missing (``Base.metadata.create_all``) on every
command so a fresh checkout works without a separate migration step — the
CLI is intentionally idempotent / bootstrappy that way. Real schema evolution
goes through Alembic; this is a safety net for the operator.

Output
------
Progress is printed as JSON lines (one per entity) so the output is both
human-skimmable and machine-parseable. ``--quiet`` suppresses the per-run
JSON lines; ``--help`` lists the entities.

Exit codes: ``0`` success, ``1`` exception during run, ``2`` bad CLI usage /
missing auth.
"""

from __future__ import annotations

import argparse
import asyncio
import logging
import signal
import sys
from typing import Any

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.sync.scheduler import ENTITY_FUNCS, SyncScheduler

logger = logging.getLogger("app.sync.cli")


# ---------------------------------------------------------------------------
# argparse setup
# ---------------------------------------------------------------------------


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="python -m app.sync.cli",
        description=(
            "Run a one-shot yunke sync or start the periodic scheduler. "
            "Auth + DB come from .env (session-cookie mode preferred)."
        ),
    )
    sub = parser.add_subparsers(dest="command", required=True, metavar="<command>")

    # ``sync`` ----------------------------------------------------------------
    p_sync = sub.add_parser(
        "sync",
        help="Run a one-shot sync (all entities by default).",
        description=(
            "Run a one-shot sync into the configured DB. Without --entity, "
            "runs the full set (org tree + messages + phone calls + SMS). "
            "With --entity <name>, runs just that one."
        ),
    )
    p_sync.add_argument(
        "--entity",
        choices=sorted(ENTITY_FUNCS),
        default="all",
        help="Entity to sync (default: all).",
    )
    p_sync.add_argument(
        "--quiet",
        action="store_true",
        help="Suppress per-run JSON progress lines (errors still logged).",
    )

    # ``daily-stats`` ---------------------------------------------------------
    p_stats = sub.add_parser(
        "daily-stats",
        help="Rebuild msg_daily_stats from the message table.",
        description=(
            "Full rebuild of the daily conversation-summary table "
            "(msg_daily_stats). Used as the nightly self-heal step in "
            "scripts/full_sync_all.sh and as the one-off deploy backfill "
            "(see docs/DEPLOY.md). Idempotent."
        ),
    )

    # ``scheduler`` -----------------------------------------------------------
    p_sched = sub.add_parser(
        "scheduler",
        help="Start the APScheduler (long-running).",
        description=(
            "Start the APScheduler loop: full sync daily at FULL_SYNC_HOUR "
            "(default 02:00 local), incremental message sync every "
            "MESSAGE_INTERVAL_HOURS (default 4). Runs forever until "
            "interrupted (Ctrl-C / SIGTERM)."
        ),
    )
    p_sched.add_argument(
        "--full-sync-hour",
        type=int,
        default=None,
        help=(
            "Local-time hour (0-23) for the daily full sync. "
            "Defaults to FULL_SYNC_HOUR env var or 2."
        ),
    )
    p_sched.add_argument(
        "--message-interval-hours",
        type=int,
        default=None,
        help=(
            "Hours between incremental message syncs. "
            "Defaults to MESSAGE_INTERVAL_HOURS env var or 4."
        ),
    )
    p_sched.add_argument(
        "--quiet",
        action="store_true",
        help="Suppress per-run JSON progress lines (errors still logged).",
    )

    return parser


# ---------------------------------------------------------------------------
# Command implementations
# ---------------------------------------------------------------------------


async def _cmd_sync(args: argparse.Namespace) -> int:
    """Run a one-shot sync and print a per-entity JSON-line summary."""
    settings = get_settings()

    # Create tables if missing (idempotent). Lets a fresh container boot
    # without a separate migration step.
    engine, session_maker = await _init_db(settings.db_url)
    try:
        scheduler = SyncScheduler(
            engine=engine, session_maker=session_maker
        )
        if args.quiet:
            # The --quiet flag only governs *our* explicit log echo; the
            # orchestrator's own per-entity logger.exception stays on so
            # failures are never silent.
            logging.getLogger("app.sync").setLevel(logging.WARNING)
        results = await scheduler.run_once(args.entity)
        _print_sync_summary(args.entity, results)
        # Exit non-zero only if every result failed; a partial failure is
        # still useful work delivered and should be exit 0 (the JSON line
        # surfaces which entities failed).
        any_success = any(r.status == "success" for r in results)
        return 0 if any_success or not results else 1
    finally:
        await engine.dispose()


async def _cmd_scheduler(args: argparse.Namespace) -> int:
    """Start the long-running scheduler and block until interrupted."""
    # Env-var fallbacks so docker-compose / .env can drive the cadence
    # without re-editing the compose command.
    full_sync_hour = args.full_sync_hour
    if full_sync_hour is None:
        full_sync_hour = _env_int("FULL_SYNC_HOUR", default=2)
    message_interval_hours = args.message_interval_hours
    if message_interval_hours is None:
        message_interval_hours = _env_int("MESSAGE_INTERVAL_HOURS", default=4)

    if not (0 <= full_sync_hour <= 23):
        print(
            f"full-sync-hour must be 0..23, got {full_sync_hour}",
            file=sys.stderr,
        )
        return 2

    scheduler = SyncScheduler(
        full_sync_hour=full_sync_hour,
        message_interval_hours=message_interval_hours,
    )
    await scheduler.start()
    print(
        f"[scheduler] running: full sync daily at hour {full_sync_hour}, "
        f"messages every {message_interval_hours}h. Ctrl-C to stop.",
        flush=True,
    )

    # Block forever until SIGINT / SIGTERM. AsyncIOScheduler runs jobs on
    # its own event loop thread; we just need to keep the process alive.
    stop_event = asyncio.Event()

    def _signal_handler(*_: Any) -> None:
        print("\n[scheduler] stop signal received, shutting down...", flush=True)
        stop_event.set()

    # KeyboardInterrupt is delivered as SIGINT on POSIX, but on Windows
    # the convention is a Ctrl-C event; loop.add_signal_handler isn't
    # implemented there. Install both: signal.signal for Windows + the
    # asyncio version where available.
    for sig in (signal.SIGINT, signal.SIGTERM):
        try:
            asyncio.get_running_loop().add_signal_handler(sig, _signal_handler)
        except (NotImplementedError, RuntimeError):
            # Windows / non-main thread — fall back to the stdlib handler.
            signal.signal(sig, _signal_handler)

    try:
        await stop_event.wait()
    finally:
        await scheduler.stop()
    return 0


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


async def _init_db(db_url: str) -> tuple[Any, Any]:
    """Create the async engine + session maker and ensure tables exist.

    Returns ``(engine, session_maker)``. The caller owns disposal. Tables
    are created with ``Base.metadata.create_all`` — idempotent, so this is
    safe to call on a warm DB.
    """
    from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

    engine = create_async_engine(
        db_url,
        pool_pre_ping=True,
        pool_recycle=3600,  # recycle before MySQL wait_timeout closes us
        future=True,
    )
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    session_maker = async_sessionmaker(engine, expire_on_commit=False)
    return engine, session_maker


def _env_int(name: str, *, default: int) -> int:
    """Read an int env var; fall back to ``default`` on missing / unparseable."""
    import os

    raw = os.environ.get(name, "").strip()
    if not raw:
        return default
    try:
        return int(raw)
    except ValueError:
        print(
            f"[cli] {name}={raw!r} is not an int; using default {default}",
            file=sys.stderr,
        )
        return default


def _print_sync_summary(entity: str, results: list[Any]) -> None:
    """Print a compact human summary of a finished run (after the JSON lines)."""
    if not results:
        print(f"[sync] {entity}: no results (run failed before any entity completed)")
        return
    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"]
    line = (
        f"[sync] {entity}: {len(results)} sub-run(s), "
        f"fetched={fetched}, upserted={upserted}"
    )
    if failed:
        line += f", FAILED={failed}"
    print(line, flush=True)


async def _cmd_daily_stats(args: argparse.Namespace) -> int:
    """全量重建 msg_daily_stats(夜间自愈 / 部署回填,幂等)。"""
    settings = get_settings()
    engine, session_maker = await _init_db(settings.db_url)
    try:
        from app.storage.summary import recompute_all

        async with session_maker() as session:
            pairs = await recompute_all(session)
            await session.commit()
        print(f"[daily-stats] rebuilt {pairs} conversation pairs")
        return 0
    finally:
        await engine.dispose()


def main(argv: list[str] | None = None) -> int:
    """CLI entry point. Returns a process exit code."""
    # Configure logging so the scheduler's JSON line + the orchestrator's
    # exception tracebacks are visible. Quiet flag is applied per-command
    # after parsing.
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    )

    parser = _build_parser()
    args = parser.parse_args(argv)

    try:
        if args.command == "sync":
            return asyncio.run(_cmd_sync(args))
        if args.command == "scheduler":
            return asyncio.run(_cmd_scheduler(args))
        if args.command == "daily-stats":
            return asyncio.run(_cmd_daily_stats(args))
    except KeyboardInterrupt:
        # So `python -m app.sync.cli scheduler` exits cleanly on Ctrl-C.
        return 0
    except ValueError as exc:
        # Bad entity name or auth misconfig surfaced as a ValueError —
        # print a clear message instead of a traceback.
        print(f"[cli] {exc}", file=sys.stderr)
        return 2
    except Exception as exc:  # noqa: BLE001
        logger.exception("cli command failed")
        print(f"[cli] {exc}", file=sys.stderr)
        return 1

    # argparse with required=True subparsers won't get here, but keep the
    # lint happy.
    parser.print_help(sys.stderr)
    return 2


if __name__ == "__main__":
    sys.exit(main())