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

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

Three invocations (plus analytics backfills, see below)::

    # 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

    # One-shot analytics backfills / self-heal (all idempotent):
    python -m app.sync.cli daily-stats
    python -m app.sync.cli staff-daily-stats
    python -m app.sync.cli sensitive-scan --full
    python -m app.sync.cli first-reply-facts
    python -m app.sync.cli employee-metrics

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."
        ),
    )

    # ``staff-daily-stats`` ---------------------------------------------------
    p_staff_stats = sub.add_parser(
        "staff-daily-stats",
        help="Rebuild stat_daily_staff (staff x day rollup).",
        description=(
            "Full rebuild of the staff daily rollup (stat_daily_staff) from "
            "msg_daily_stats. Used as the nightly self-heal step in "
            "scripts/full_sync_all.sh and as the one-off deploy backfill. "
            "Idempotent."
        ),
    )

    # ``sensitive-scan`` ------------------------------------------------------
    p_scan = sub.add_parser(
        "sensitive-scan",
        help="Sensitive-word scan (incremental; --full for backfill).",
        description=(
            "Incremental sensitive-word scan by fetched_at watermark. "
            "First run on a populated DB MUST use --full to build the "
            "watermark (run_incremental refuses no-watermark on big tables). "
            "Adding a new word does NOT backfill history (watermark only "
            "sees new rows) — rerun --full to scan history for it."
        ),
    )
    p_scan.add_argument(
        "--full",
        action="store_true",
        help="Full backfill scan (first deploy / new-word history).",
    )

    # ``first-reply-facts`` ----------------------------------------------
    p_fr = sub.add_parser(
        "first-reply-facts",
        help="Rebuild first_reply_fact (cross-day exact threads).",
        description=(
            "Full rebuild of the first-reply fact table from the message "
            "table (cross-day exact thread algorithm — overnight samples "
            "not lost, unlike msg_daily_stats day-cut). Deploy backfill + "
            "nightly self-heal step. Idempotent."
        ),
    )

    # ``employee-metrics`` ------------------------------------------------
    p_emp = sub.add_parser(
        "employee-metrics",
        help="Rebuild employee_metric_daily + employee_call_daily.",
        description=(
            "Full rebuild of the management wide table (per-day x channel "
            "x staff, 1:1 caliber) and the call daily table. Must run "
            "AFTER first-reply-facts (the fact table feeds Batch-B "
            "scoring and drilldown; both recompute threads from messages "
            "via one shared algorithm). Deploy backfill + nightly "
            "self-heal. 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()


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

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


async def _cmd_sensitive_scan(args: argparse.Namespace) -> int:
    """敏感词扫描:默认增量(水位后);--full 全量回填(首跑/新词历史)。"""
    settings = get_settings()
    engine, session_maker = await _init_db(settings.db_url)
    try:
        from app.analytics.sensitive import scanner
        from app.analytics.sensitive import words as words_svc

        async with session_maker() as session:
            seeded = await words_svc.ensure_starter_words(session)  # 幂等补种起步词库
            if args.full:
                out = await scanner.run_full(session)
            else:
                out = await scanner.run_incremental(session)
            await session.commit()
        print(f"[sensitive-scan] seeded={seeded} {out}")
        return 0
    finally:
        await engine.dispose()


async def _cmd_first_reply_facts(args: argparse.Namespace) -> int:
    """全量重建 first_reply_fact(跨日精确串;部署回填/夜间自愈,幂等)。"""
    settings = get_settings()
    engine, session_maker = await _init_db(settings.db_url)
    try:
        from app.storage.first_reply_facts import recompute_all

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


async def _cmd_employee_metrics(args: argparse.Namespace) -> int:
    """全量重建 employee_metric_daily + employee_call_daily(幂等)。"""
    settings = get_settings()
    engine, session_maker = await _init_db(settings.db_url)
    try:
        import sqlalchemy as sa

        from app.db.models import MsgDailyStats, PhoneCall
        from app.storage.employee_metrics import (
            recompute_call_days,
            recompute_metric_days,
        )
        from app.storage.summary import day_key

        async with session_maker() as session:
            msg_days = (
                await session.execute(sa.select(MsgDailyStats.day).distinct())
            ).scalars().all()
            # 宽表全走消息派生源,只算消息日。
            rows = await recompute_metric_days(session, list(msg_days))
            # 通话日 = 消息日 ∪ 通话日:消息空窗日(节假日/同步缺口)的通话
            # 也要进小表。start_call_time→上海日必须在 Python 侧走 day_key
            # ——MySQL 会话时区是 UTC,SQL 日期函数会切错日。
            call_ts = (
                await session.execute(
                    sa.select(PhoneCall.start_call_time)
                    .where(PhoneCall.start_call_time.is_not(None))
                    .distinct()
                )
            ).scalars().all()
            call_days = sorted(set(msg_days) | {day_key(ts) for ts in call_ts})
            calls = await recompute_call_days(session, call_days)
            await session.commit()
        print(
            f"[employee-metrics] rebuilt {rows} wide rows, {calls} call rows "
            f"({len(call_days)} days)"
        )
        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))
        if args.command == "staff-daily-stats":
            return asyncio.run(_cmd_staff_daily_stats(args))
        if args.command == "sensitive-scan":
            return asyncio.run(_cmd_sensitive_scan(args))
        if args.command == "first-reply-facts":
            return asyncio.run(_cmd_first_reply_facts(args))
        if args.command == "employee-metrics":
            return asyncio.run(_cmd_employee_metrics(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())