guardia-itsm/core/events.py

"""경량 인메모리 이벤트 버스 — SSE 브로드캐스트용.

단일 워커 환경(개발/온프레미스 데모)에서 모든 SSE 구독자에게
실시간 이벤트를 전달합니다.

사용법:
    from core.events import broadcast, subscribe, unsubscribe

    # 이벤트 발행 (router 내 async 함수에서)
    await broadcast("sr_updated", {"sr_id": "SR-...", "new_status": "APPROVED"})

    # SSE 엔드포인트에서 구독
    queue = subscribe()
    try:
        msg = await asyncio.wait_for(queue.get(), timeout=20.0)
    finally:
        unsubscribe(queue)
"""
import asyncio
import json
from typing import Any

# 구독 중인 클라이언트 큐 목록
_subscribers: list[asyncio.Queue] = []


async def broadcast(event_type: str, data: dict[str, Any] | None = None) -> None:
    """모든 SSE 구독자에게 이벤트를 전송합니다.

    큐가 가득 찬(오프라인) 구독자는 자동으로 제거됩니다.
    """
    msg = json.dumps({"type": event_type, **(data or {})}, ensure_ascii=False)
    dead: list[asyncio.Queue] = []
    for q in list(_subscribers):
        try:
            q.put_nowait(msg)
        except asyncio.QueueFull:
            dead.append(q)
    for q in dead:
        unsubscribe(q)


def subscribe() -> asyncio.Queue:
    """새 구독 큐를 생성하고 등록합니다."""
    q: asyncio.Queue = asyncio.Queue(maxsize=100)
    _subscribers.append(q)
    return q


def unsubscribe(q: asyncio.Queue) -> None:
    """구독을 해제합니다."""
    try:
        _subscribers.remove(q)
    except ValueError:
        pass


def subscriber_count() -> int:
    """현재 활성 구독자 수를 반환합니다."""
    return len(_subscribers)


# ── 배치 실행 로그 SSE 스트리밍 ────────────────────────────────────────────────

async def stream_batch_log(run_id: int):
    """
    배치 실행 로그를 SSE 형식으로 스트리밍합니다.

    BatchRun.stdout_tail 을 주기적으로 폴링하여 새 내용을 yield.
    result 가 RUNNING 이 아닌 상태가 되면 마지막 로그를 전송 후 종료.

    Usage (FastAPI StreamingResponse):
        return StreamingResponse(
            stream_batch_log(run_id),
            media_type="text/event-stream",
        )

    Yields:
        SSE 형식 문자열: "data: {json}\\n\\n"
    """
    import asyncio
    from database import SessionLocal
    from models import BatchRun, BatchRunResult

    POLL_INTERVAL = 2.0    # 2초 폴링
    MAX_WAIT_SEC  = 3600   # 최대 1시간 대기
    elapsed       = 0.0
    sent_lines    = 0      # 이미 전송한 줄 수 (중복 방지)

    while elapsed < MAX_WAIT_SEC:
        async with SessionLocal() as db:
            from sqlalchemy import select
            row = (await db.execute(
                select(BatchRun).where(BatchRun.id == run_id)
            )).scalars().first()

        if not row:
            yield f"data: {json.dumps({'type': 'error', 'message': '실행 이력을 찾을 수 없습니다.'}, ensure_ascii=False)}\n\n"
            return

        # stdout_tail 에서 새 줄 추출
        stdout = row.stdout_tail or ""
        all_lines = stdout.splitlines()
        new_lines = all_lines[sent_lines:]

        if new_lines:
            for line in new_lines:
                payload = json.dumps(
                    {"type": "log", "run_id": run_id, "line": line},
                    ensure_ascii=False,
                )
                yield f"data: {payload}\n\n"
            sent_lines += len(new_lines)

        # 완료/실패/타임아웃 → 종료 이벤트 전송
        if row.result != BatchRunResult.RUNNING:
            result_payload = json.dumps({
                "type":      "done",
                "run_id":    run_id,
                "result":    row.result,
                "exit_code": row.exit_code,
                "ended_at":  row.ended_at.isoformat() if row.ended_at else None,
                "error_msg": row.error_msg,
            }, ensure_ascii=False)
            yield f"data: {result_payload}\n\n"
            return

        # 아직 실행 중 — heartbeat 전송 후 대기
        heartbeat = json.dumps(
            {"type": "heartbeat", "run_id": run_id, "elapsed": int(elapsed)},
            ensure_ascii=False,
        )
        yield f"data: {heartbeat}\n\n"

        await asyncio.sleep(POLL_INTERVAL)
        elapsed += POLL_INTERVAL

    # 타임아웃 — 스트리밍 강제 종료
    timeout_payload = json.dumps(
        {"type": "timeout", "run_id": run_id, "message": "스트리밍 최대 대기 시간 초과"},
        ensure_ascii=False,
    )
    yield f"data: {timeout_payload}\n\n"