group_xinghuo_jinrong/scripts/demo/subscribe_alerts.py

#!/usr/bin/env python3
"""risk:pub:alert 预警推送订阅演示（B9a · PRD FR-4 通知 / A-3 验收第三条）。

演示 SOP：
  终端 1  python scripts/demo/subscribe_alerts.py
  终端 2  uvicorn app.main:app --reload → Swagger 发演示交易
          （A-3：CUST-3001 申购 50 万 PROD-510300，RISK-001/002 命中）
  终端 1 应实时打印预警推送行；无推送时每 30s 打印心跳提示仍在监听。

payload 口径（PRD §4 FR-4 · 02-redis-keys.md §2.4）：
  {alert_id, alert_type, customer_id_mask, risk_score, trace_id, notify_role}

--duration N：监听 N 秒后自动退出（默认 0 = 不限，Ctrl+C 退出）。
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT))

from app.config.settings import settings  # noqa: E402

CHANNEL = "risk:pub:alert"


def format_alert(payload: dict[str, Any]) -> str:
    """推送 payload → 单行可读文本（演示展示用，字段缺失容错）。"""
    if "raw" in payload:
        return f"[raw] {payload['raw']}"
    roles = ",".join(payload.get("notify_role") or [])
    return (
        f"[{payload.get('alert_type', '?'):>12}] {payload.get('alert_id', '?')}"
        f" score={payload.get('risk_score', '?')}"
        f" customer={payload.get('customer_id_mask', '?')}"
        f" trace={payload.get('trace_id', '?')}"
        f" notify=[{roles}]"
    )


def run(duration: float) -> None:
    try:
        import redis
    except ImportError:
        print("缺少 redis 包：请执行 python -m pip install -r requirements.txt", file=sys.stderr)
        raise SystemExit(1)

    client = redis.Redis.from_url(settings.redis_url, decode_responses=True)
    try:
        client.ping()
    except Exception as exc:
        print(f"Redis 不可达（{settings.redis_url}）：{exc}", file=sys.stderr)
        print("请先启动本机 Redis 服务（见 FLOW §0 本机状态）。", file=sys.stderr)
        raise SystemExit(1)

    pubsub = client.pubsub(ignore_subscribe_messages=True)
    pubsub.subscribe(CHANNEL)
    print(f"已订阅 {CHANNEL} @ {settings.redis_url}（Ctrl+C 退出）")
    print("等待预警推送…（另开终端发演示交易，SOP 见脚本头注释）")

    deadline = time.monotonic() + duration if duration > 0 else None
    received = 0
    last_heartbeat = time.monotonic()
    try:
        while True:
            msg = pubsub.get_message(timeout=1.0)
            if msg and msg.get("type") == "message":
                received += 1
                try:
                    payload = json.loads(msg["data"])
                    if not isinstance(payload, dict):
                        payload = {"raw": str(msg["data"])}
                except (TypeError, ValueError):
                    payload = {"raw": str(msg["data"])}
                print(f"{time.strftime('%H:%M:%S')}  #{received}  {format_alert(payload)}")
            elif deadline is not None and time.monotonic() >= deadline:
                print(f"已监听 {duration:g}s，共收到 {received} 条推送，退出。")
                return
            elif deadline is None and time.monotonic() - last_heartbeat >= 30:
                last_heartbeat = time.monotonic()
                print(f"…监听中（已收到 {received} 条）")
    except KeyboardInterrupt:
        print(f"\n退出（共收到 {received} 条推送）。")
    finally:
        pubsub.close()
        client.close()


def main() -> None:
    parser = argparse.ArgumentParser(description="订阅 risk:pub:alert（风控预警推送演示）")
    parser.add_argument("--duration", type=float, default=0.0, help="监听秒数；0=不限（Ctrl+C 退出）")
    args = parser.parse_args()
    run(args.duration)


if __name__ == "__main__":
    main()