| import asyncio |
| import logging |
| from typing import Any |
|
|
| from sqlalchemy import select |
| from sqlalchemy.ext.asyncio import AsyncSession |
|
|
| from app.database import AsyncSessionLocal |
| from app.models import Job |
| from app.schemas import job_to_status_response |
|
|
| logger = logging.getLogger(__name__) |
|
|
| _queues: dict[str, list[asyncio.Queue[dict[str, Any]]]] = {} |
|
|
|
|
| def _qs(job_id: str) -> list[asyncio.Queue[dict[str, Any]]]: |
| return _queues.setdefault(job_id, []) |
|
|
|
|
| def subscribe(job_id: str) -> asyncio.Queue[dict[str, Any]]: |
| q: asyncio.Queue[dict[str, Any]] = asyncio.Queue(maxsize=64) |
| _qs(job_id).append(q) |
| return q |
|
|
|
|
| def unsubscribe(job_id: str, q: asyncio.Queue[dict[str, Any]]) -> None: |
| lst = _qs(job_id) |
| if q in lst: |
| lst.remove(q) |
| if not lst and job_id in _queues: |
| del _queues[job_id] |
|
|
|
|
| async def broadcast_job(job_id: str) -> None: |
| try: |
| async with AsyncSessionLocal() as session: |
| result = await session.execute(select(Job).where(Job.id == job_id)) |
| job = result.scalar_one_or_none() |
| if not job: |
| return |
| payload = job_to_status_response(job).model_dump(mode="json") |
| message = {"type": "job_status", "data": payload} |
| for q in list(_qs(job_id)): |
| try: |
| q.put_nowait(message) |
| except asyncio.QueueFull: |
| try: |
| q.get_nowait() |
| except asyncio.QueueEmpty: |
| pass |
| try: |
| q.put_nowait(message) |
| except asyncio.QueueFull: |
| logger.debug("Progress queue still full for job %s", job_id) |
| except Exception: |
| logger.exception("broadcast_job failed for %s", job_id) |
|
|
