2daaa8c8a4
Второй CI-гейт после тестов: ruff как статический анализатор всего Python-кода (services + scripts). Раньше ни линта, ни проверки типов в CI не было вовсе. Конфиг ruff.toml: правила E/F/W/I/UP/B/SIM/C4, line-length 100. Осознанно выключены E501 (длину держит форматтер; длинные RU-комментарии — норма), B008 (Depends()/Query() в дефолтах — идиома FastAPI, не баг) и UP042 ((str, Enum)→StrEnum меняет __str__/сериализацию — не трогаем). Починено под ноль находок: - B904 (11): raise ... from exc / from None — читаемые цепочки исключений в Celery-ретраях и HTTPException, ошибки обработки не маскируют исходные. - SIM105 (5): try/except/pass → contextlib.suppress (faiss remove_ids, lsh.remove, сброс кэша, ws-disconnect, парс года). - C416/SIM108/B905/F841/UP035/UP017/F401/I001: dict(rows), тернарник, zip strict, мёртвая переменная, устаревшие импорты, timezone.utc→UTC, чистка/сортировка. Обвязка: scripts/run_lint.sh (ruff в изолированном python:3.11-slim), шаг «Линт» в job test перед юнит-тестами (падаем раньше). make lint / make lint-fix. Все 41 юнит-тест по-прежнему зелёные, изменённые файлы компилируются. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
86 lines
3.3 KiB
Python
86 lines
3.3 KiB
Python
"""WebSocket менеджер для real-time обновлений статуса задач."""
|
||
|
||
import contextlib
|
||
import json
|
||
import logging
|
||
from typing import Any
|
||
|
||
from fastapi import WebSocket
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class ConnectionManager:
|
||
"""Менеджер WebSocket соединений с поддержкой подписки на задачи."""
|
||
|
||
def __init__(self) -> None:
|
||
# task_id -> список активных соединений
|
||
self._connections: dict[str, list[WebSocket]] = {}
|
||
|
||
async def connect(self, task_id: str, websocket: WebSocket) -> None:
|
||
"""Принять WebSocket соединение и зарегистрировать его для задачи."""
|
||
await websocket.accept()
|
||
if task_id not in self._connections:
|
||
self._connections[task_id] = []
|
||
self._connections[task_id].append(websocket)
|
||
logger.info(f"WebSocket подключён к задаче {task_id!r}")
|
||
|
||
def disconnect(self, task_id: str, websocket: WebSocket) -> None:
|
||
"""Удалить соединение из реестра."""
|
||
if task_id in self._connections:
|
||
with contextlib.suppress(ValueError):
|
||
self._connections[task_id].remove(websocket)
|
||
if not self._connections[task_id]:
|
||
del self._connections[task_id]
|
||
logger.info(f"WebSocket отключён от задачи {task_id!r}")
|
||
|
||
async def send_task_update(self, task_id: str, data: dict[str, Any]) -> None:
|
||
"""
|
||
Отправить обновление статуса задачи всем подключённым клиентам.
|
||
|
||
Args:
|
||
task_id: ID задачи
|
||
data: Словарь с обновлением (status, queue_position, result и т.д.)
|
||
"""
|
||
connections = self._connections.get(task_id, [])
|
||
if not connections:
|
||
return
|
||
|
||
message = json.dumps(data, ensure_ascii=False, default=str)
|
||
dead_connections = []
|
||
|
||
for websocket in connections:
|
||
try:
|
||
await websocket.send_text(message)
|
||
except Exception as e:
|
||
logger.warning(f"Ошибка отправки WebSocket сообщения: {e}")
|
||
dead_connections.append(websocket)
|
||
|
||
# Очистить мёртвые соединения
|
||
for ws in dead_connections:
|
||
self.disconnect(task_id, ws)
|
||
|
||
async def broadcast(self, data: dict[str, Any]) -> None:
|
||
"""Отправить сообщение всем подключённым клиентам."""
|
||
message = json.dumps(data, ensure_ascii=False, default=str)
|
||
all_dead = []
|
||
|
||
for task_id, connections in self._connections.items():
|
||
for websocket in connections:
|
||
try:
|
||
await websocket.send_text(message)
|
||
except Exception:
|
||
all_dead.append((task_id, websocket))
|
||
|
||
for task_id, ws in all_dead:
|
||
self.disconnect(task_id, ws)
|
||
|
||
@property
|
||
def active_connections_count(self) -> int:
|
||
"""Количество активных WebSocket соединений."""
|
||
return sum(len(conns) for conns in self._connections.values())
|
||
|
||
|
||
# Глобальный экземпляр менеджера
|
||
ws_manager = ConnectionManager()
|