Split UME alarm WebSocket consumer into state, log, subscription, and loop modules.

Keep the public facade stable and re-export shared mutable log/connection objects for existing tests.

Co-authored-by: Cursor <cursoragent@cursor.com>

netx_api/ume_alarm_ws_subscription.py

@@ -0,0 +1,356 @@
+"""UME alarm WSS subscription state APIs and manual control."""
+from __future__ import annotations
+
+import json
+import logging
+import time
+from typing import Any, Callable
+
+from sqlalchemy.orm import Session
+
+from . import ume_alarm_ws_state as ws_state
+from .config import settings
+from .ume_alarm_subscription_store import (
+    DEFAULT_SUBSCRIPTION_KEY,
+    clear_subscription,
+    load_subscription,
+    save_subscription,
+)
+from .ume_alarm_ws_log import (
+    append_ws_log,
+    is_ume_subscription_missing_error,
+    parse_subscription_id_from_already_exists_error,
+)
+from .ume_client import UMEClient
+from .ume_sync_service import (
+    apply_alarm_to_current,
+    extract_alarm_from_notification,
+    normalize_yang_alarm,
+    _utc_now_naive,
+)
+
+_ws_log = logging.getLogger("netx.ume.alarm_ws")
+
+
+def mark_ume_subscription_lost(reason: str) -> None:
+    detail = str(reason or "").strip()[:500]
+    with ws_state._ume_lost_lock:
+        if ws_state._ume_subscription_lost and detail == ws_state._ume_subscription_lost_reason:
+            return
+        ws_state._ume_subscription_lost = True
+        ws_state._ume_subscription_lost_reason = detail
+    append_ws_log(
+        f"UME 侧订阅已丢失: {detail[:200]}",
+        level="warning",
+        dedup_cooldown_s=ws_state._LOOP_STATUS_COOLDOWN_S,
+    )
+    request_ws_reconnect()
+
+
+def clear_ume_subscription_lost_flag() -> None:
+    with ws_state._ume_lost_lock:
+        ws_state._ume_subscription_lost = False
+        ws_state._ume_subscription_lost_reason = ""
+
+
+def is_ume_subscription_lost() -> bool:
+    with ws_state._ume_lost_lock:
+        return bool(ws_state._ume_subscription_lost)
+
+
+def _server_subscription_lost_fields() -> dict[str, Any]:
+    with ws_state._ume_lost_lock:
+        lost = bool(ws_state._ume_subscription_lost)
+        reason = str(ws_state._ume_subscription_lost_reason or "")
+    return {
+        "server_subscription_lost": lost,
+        "server_subscription_lost_reason": reason,
+    }
+
+
+def _set_ws_connection_state(state: str, *, detail: str = "") -> None:
+    with ws_state._ws_connection_lock:
+        ws_state._ws_connection_state = str(state or "").strip() or "unknown"
+        ws_state._ws_connection_detail = str(detail or "").strip()[:240]
+
+
+def get_ws_connection_status() -> dict[str, Any]:
+    with ws_state._ws_connection_lock:
+        state = str(ws_state._ws_connection_state or "init")
+        detail = str(ws_state._ws_connection_detail or "")
+    return {
+        "state": state,
+        "label": ws_state._WS_CONNECTION_LABELS.get(state, state),
+        "detail": detail,
+    }
+
+
+def _notify_ws_connection(
+    state: str,
+    *,
+    detail: str = "",
+    subscription_id: str = "",
+    on_status: Callable[[str], None] | None = None,
+    log_level: str = "info",
+    log_cooldown_s: float = ws_state._WS_NOTIFY_COOLDOWN_S,
+) -> None:
+    _set_ws_connection_state(state, detail=detail)
+    label = ws_state._WS_CONNECTION_LABELS.get(state, state)
+    log_msg = str(detail or label).strip()
+    notify_key = f"{state}\0{log_msg}\0{subscription_id}"
+    now = time.monotonic()
+    should_log = True
+    with ws_state._WS_NOTIFY_LOCK:
+        if notify_key == ws_state._WS_NOTIFY_LAST_KEY and now - ws_state._WS_NOTIFY_LAST_TS < max(1.0, float(log_cooldown_s)):
+            should_log = False
+        else:
+            ws_state._WS_NOTIFY_LAST_KEY = notify_key
+            ws_state._WS_NOTIFY_LAST_TS = now
+    if should_log:
+        append_ws_log(
+            log_msg,
+            level=log_level,
+            subscription_id=subscription_id,
+            dedup=False,
+        )
+    if on_status is not None:
+        on_status(f"ws:{state}")
+
+
+def _parse_ws_message(raw: str) -> dict[str, Any] | None:
+    text = str(raw or "").strip()
+    if not text:
+        return None
+    try:
+        data = json.loads(text)
+    except json.JSONDecodeError:
+        return None
+    return data if isinstance(data, dict) else None
+
+
+def process_alarm_notification(db: Session, payload: dict[str, Any]) -> tuple[str, bool]:
+    alarm = extract_alarm_from_notification(payload)
+    if alarm is None:
+        return "skipped", False
+    return apply_alarm_to_current(db, alarm, touch_ts=_utc_now_naive())
+
+
+def _delete_subscription_on_ume(client: UMEClient, subscription_id: str) -> None:
+    sub_id = str(subscription_id or "").strip()
+    if not sub_id:
+        return
+    client.delete_alarm_subscription(sub_id)
+
+
+def _wait_for_shared_token(client: UMEClient, *, timeout_s: float = 300.0) -> bool:
+    deadline = time.time() + max(5.0, float(timeout_s))
+    while time.time() < deadline:
+        client._sync_token_from_store()
+        if client.has_valid_token():
+            return True
+        time.sleep(2.0)
+    return False
+
+
+def _set_active_subscription(subscription_id: str, wss_uri: str, *, topic: str = "ALARM") -> None:
+    with ws_state._subscription_lock:
+        ws_state._subscription_id = str(subscription_id or "").strip()
+        ws_state._subscription_uri = str(wss_uri or "").strip()
+        ws_state._subscription_topic = str(topic or "ALARM").strip() or "ALARM"
+
+
+def _clear_active_subscription() -> None:
+    with ws_state._subscription_lock:
+        ws_state._subscription_id = ""
+        ws_state._subscription_uri = ""
+        ws_state._subscription_topic = "ALARM"
+
+
+def get_active_subscription() -> tuple[str, str]:
+    with ws_state._subscription_lock:
+        return str(ws_state._subscription_id or ""), str(ws_state._subscription_uri or "")
+
+
+def load_persisted_subscription() -> bool:
+    """Load subscription from DB into memory (call on process startup)."""
+    db = SessionLocal()
+    try:
+        loaded = load_subscription(db)
+        if loaded is None:
+            _clear_active_subscription()
+            return False
+        sub_id, uri, topic = loaded
+        _set_active_subscription(sub_id, uri, topic=topic)
+        append_ws_log(f"loaded persisted subscription id={sub_id}", subscription_id=sub_id)
+        return True
+    finally:
+        db.close()
+
+
+def get_subscription_status() -> dict[str, Any]:
+    with ws_state._subscription_lock:
+        sub_id = str(ws_state._subscription_id or "")
+        uri = str(ws_state._subscription_uri or "")
+        topic = str(ws_state._subscription_topic or "ALARM")
+    return {
+        "active": bool(sub_id and uri),
+        "subscription_id": sub_id,
+        "wss_uri": uri,
+        "topic": topic,
+        **_server_subscription_lost_fields(),
+    }
+
+
+def is_wss_active_for_current_alarms() -> bool:
+    """True when WSS subscription is active and should own ume_alarms_current updates."""
+    if not bool(getattr(settings, "ume_alarm_ws_enabled", True)):
+        return False
+    if is_ume_subscription_lost():
+        return False
+    return bool(get_subscription_status().get("active"))
+
+
+def get_current_alarms_mode() -> str:
+    return "wss" if is_wss_active_for_current_alarms() else "rest"
+
+
+def get_alarms_coordination_status() -> dict[str, Any]:
+    wss_active = is_wss_active_for_current_alarms()
+    skip_when_ws = bool(getattr(settings, "ume_sync_alarms_current_skip_when_ws", True))
+    return {
+        "current_alarms_mode": get_current_alarms_mode(),
+        "wss_active_for_current_alarms": wss_active,
+        "scheduled_sync_skipped": bool(wss_active and skip_when_ws),
+    }
+
+
+def clear_local_alarm_subscription_manual(db: Session) -> dict[str, Any]:
+    """Drop persisted/in-memory subscription without calling UME delete."""
+    clear_subscription(db, cache_key=DEFAULT_SUBSCRIPTION_KEY)
+    db.commit()
+    _clear_active_subscription()
+    clear_ume_subscription_lost_flag()
+    request_ws_reconnect()
+    append_ws_log("cleared local subscription record (UME delete skipped)")
+    return get_subscription_status()
+
+
+def request_ws_reconnect() -> None:
+    ws_state._ws_wake_event.set()
+
+
+def establish_alarm_subscription_manual(
+    client: UMEClient,
+    db: Session,
+    *,
+    force_reestablish: bool = False,
+) -> dict[str, Any]:
+    """Manually establish ALARM subscription (UME API + persist). Does not open WSS."""
+    if force_reestablish or is_ume_subscription_lost():
+        clear_local_alarm_subscription_manual(db)
+
+    existing = load_subscription(db)
+    if existing is not None:
+        sub_id, uri, topic = existing
+        _set_active_subscription(sub_id, uri, topic=topic)
+        request_ws_reconnect()
+        append_ws_log(f"establish skipped: already in DB id={sub_id}", subscription_id=sub_id)
+        status_out = get_subscription_status()
+        return {**status_out, "already_exists": True}
+
+    mem_id, mem_uri = get_active_subscription()
+    if mem_id and mem_uri:
+        request_ws_reconnect()
+        append_ws_log(f"establish skipped: in-memory id={mem_id}", subscription_id=mem_id)
+        status_out = get_subscription_status()
+        return {**status_out, "already_exists": True}
+
+    if not _wait_for_shared_token(client, timeout_s=120.0):
+        raise RuntimeError("ume_ws_no_valid_token:wait_timeout")
+    client._sync_token_from_store()
+    if not client.has_valid_token():
+        raise RuntimeError("ume_ws_no_valid_token")
+
+    topic = str(getattr(settings, "ume_notification_topic", "ALARM") or "ALARM").strip() or "ALARM"
+    try:
+        sub_id, uri = client.establish_alarm_subscription(topic=topic)
+    except RuntimeError as exc:
+        orphan_id = parse_subscription_id_from_already_exists_error(str(exc))
+        if not orphan_id:
+            raise
+        append_ws_log(
+            f"establish: orphan on UME id={orphan_id}, deleting then retry",
+            level="warning",
+            subscription_id=orphan_id,
+        )
+        try:
+            _delete_subscription_on_ume(client, orphan_id)
+        except RuntimeError as del_exc:
+            raise RuntimeError(
+                f"ume_orphan_subscription_delete_failed:id={orphan_id}:{str(del_exc)[:180]}"
+            ) from del_exc
+        sub_id, uri = client.establish_alarm_subscription(topic=topic)
+    save_subscription(db, subscription_id=sub_id, wss_uri=uri, topic=topic)
+    db.commit()
+
+    with ws_state._shutdown_lock:
+        ws_state._active_client = client
+    _set_active_subscription(sub_id, uri, topic=topic)
+    clear_ume_subscription_lost_flag()
+    request_ws_reconnect()
+    append_ws_log(f"manual establish subscription id={sub_id}", subscription_id=sub_id)
+    return {**get_subscription_status(), "already_exists": False}
+
+
+def cancel_alarm_subscription_manual(
+    client: UMEClient,
+    db: Session,
+    *,
+    force_clear_local: bool = False,
+) -> dict[str, Any]:
+    """Manually cancel subscription (UME delete + clear DB + drop WSS)."""
+    loaded = load_subscription(db)
+    sub_id, _uri = get_active_subscription()
+    if not sub_id and loaded is not None:
+        sub_id = str(loaded[0] or "")
+    ume_already_missing = False
+    if sub_id:
+        client.refresh_if_needed()
+        try:
+            _delete_subscription_on_ume(client, sub_id)
+        except RuntimeError as exc:
+            if is_ume_subscription_missing_error(str(exc)):
+                ume_already_missing = True
+                mark_ume_subscription_lost(str(exc))
+                if not force_clear_local:
+                    return {
+                        **get_subscription_status(),
+                        "ok": False,
+                        "ume_already_missing": True,
+                        "needs_local_cleanup": True,
+                        "message": (
+                            "UME 侧告警订阅已不存在或已过期（服务器订阅已丢失）。"
+                            "请确认是否清除本地订阅记录，然后重新建立订阅。"
+                        ),
+                    }
+            else:
+                raise
+
+    clear_subscription(db, cache_key=DEFAULT_SUBSCRIPTION_KEY)
+    db.commit()
+    _clear_active_subscription()
+    clear_ume_subscription_lost_flag()
+    request_ws_reconnect()
+    append_ws_log(
+        f"manual cancel subscription id={sub_id or '(none)'}",
+        subscription_id=sub_id,
+    )
+    status_out = get_subscription_status()
+    return {
+        **status_out,
+        "ok": True,
+        "cleared_local": True,
+        "ume_already_missing": ume_already_missing,
+    }
+
+