"""Соединение scan-gateway с NATS: core NATS без JetStream (design D-2, D-7, D-8; spec image-recognition-gateway, «Соединение с NATS без JetStream»; docs/context/03-architecture.md, «Контракт обмена через NATS»). - Аутентификация как в шаблоне: заданы `NATS_SEED` и `NATS_JWT` — nkeys/JWT; оба пусты — без учётных данных; задан только один или seed не читается — `NatsConfigError` (подключиться «как-нибудь» без учётных данных нельзя). - `NATS_ENABLED=false` — `connect()` ничего не делает и возвращает `False` (dev без NATS). - Переподключение — встроенное в nats-py (`max_reconnect_attempts=-1`): клиент сам повторяет и первое подключение, поэтому `connect()` возвращается только при успехе (вызывающий оборачивает его в задачу или `wait_for`). После переподключения nats-py сам отправляет SUB для всех подписок с прежней queue group — подписку повторно оформлять не нужно. - `subscribe_requests(cb)` — подписка на `files.img_uploaded` в queue group `scan-gateway`: из нескольких инстансов сообщение получает ровно один. `cb` получает исходное сообщение nats-py (`msg.data`, `msg.headers`); исключения `cb` логируются, подписка живёт. - `publish(subject, payload, msg_id)` — публикация с заголовком `Nats-Msg-Id` и `flush` (подтверждение сервером) с повторами на временных ошибках (первая попытка + до `PUBLISH_RETRIES` повторов); исчерпаны — `PublishFailed`. Воркер логирует error и идёт дальше: ядро закроет задание по `reply_deadline` (D-7). - Stream'ов и consumer'ов JetStream нет: `jetstream()` не вызывается. """ import asyncio import base64 import json from collections.abc import Awaitable, Callable from contextlib import suppress from typing import Any import nkeys from loguru import logger from nats import errors as nats_errors from nats.aio.client import Client as NATS from nats.aio.msg import Msg from src.contract import ( IMAGE_REQUEST_SUBJECT, NATS_MSG_ID_HEADER, RESULT_SUBJECT, MessageStatus, RecognitionResultMessage, ack_msg_id, result_msg_id, ) from src.core.config import NATSSettings, Settings REQUEST_QUEUE_GROUP = "scan-gateway" MAX_RECONNECT_ATTEMPTS = -1 # переподключаться бесконечно RECONNECT_TIME_WAIT = 2 # секунд между попытками переподключения nats-py CONNECT_TIMEOUT = 2 # секунд на одну попытку TCP-подключения FLUSH_TIMEOUT = 5 # секунд на подтверждение публикации сервером (PING/PONG) RETRY_DELAY = 0.5 # секунд между попытками публикации MessageHandler = Callable[[Msg], Awaitable[None]] SleepFn = Callable[[float], Awaitable[None]] # Временные ошибки публикации: соединение закрыто/переподключается/переполнен буфер на время переподключения, # сервер не подтвердил flush. Повтор имеет смысл. Остальное (MaxPayloadError, BadSubjectError, ошибки программы) # повтором не лечится — `PublishFailed` сразу. TRANSIENT_PUBLISH_ERRORS: tuple[type[BaseException], ...] = ( nats_errors.ConnectionClosedError, nats_errors.ConnectionReconnectingError, nats_errors.ConnectionDrainingError, nats_errors.OutboundBufferLimitError, nats_errors.StaleConnectionError, nats_errors.NoServersError, TimeoutError, # nats.errors.TimeoutError / FlushTimeoutError — его подклассы ConnectionError, ) class NatsError(Exception): """Базовая ошибка слоя NATS сервиса.""" class NatsConfigError(NatsError, RuntimeError): """Учётные данные NATS заданы некорректно: подключение не выполняется.""" class NatsNotConnectedError(NatsError): """Операция требует соединения с NATS.""" class PublishFailed(NatsError): """Публикация не подтверждена: повторы исчерпаны или ошибка не временная. Причина последней попытки — в `__cause__`; `attempts` — сколько попыток фактически сделано (0 — соединения не было, при исчерпанных повторах — `retries + 1`).""" def __init__(self, message: str, *, subject: str, msg_id: str, attempts: int) -> None: super().__init__(message) self.subject = subject self.msg_id = msg_id self.attempts = attempts def build_nats_auth(settings: NATSSettings) -> dict[str, Any]: """Параметры аутентификации для `Client.connect`: `{}` без учётных данных или колбэки nkeys/JWT шаблона.""" seed = settings.NATS_SEED.strip() jwt = settings.NATS_JWT.strip() if not seed and not jwt: return {} if not seed or not jwt: raise NatsConfigError( "Неполные учётные данные NATS: NATS_SEED и NATS_JWT задаются вместе или не задаются вовсе" ) try: keypair = nkeys.from_seed(seed.encode("utf-8")) except Exception as exc: raise NatsConfigError(f"Некорректный NATS_SEED: {type(exc).__name__}") from None def user_jwt_cb() -> bytes: return jwt.encode("utf-8") def signature_cb(nonce: str | bytes) -> bytes: raw = nonce.encode("utf-8") if isinstance(nonce, str) else nonce signature: bytes = keypair.sign(raw) return base64.b64encode(signature) return {"user_jwt_cb": user_jwt_cb, "signature_cb": signature_cb} def message_msg_id(message: RecognitionResultMessage, attempt: int) -> str: """`Nats-Msg-Id` конверта: `ack::` для `processing`, иначе `result::`.""" if message.status is MessageStatus.PROCESSING: return ack_msg_id(message.job_id, attempt) return result_msg_id(message.job_id, attempt) class NatsGateway: """Соединение сервиса с NATS: подключение, подписка на запросы, публикация с `Nats-Msg-Id` и повторами. `client_factory`, `sleep`, `retry_delay`, `flush_timeout` — подмены для тестов и тайминги. """ def __init__( # noqa: PLR0913 - все параметры, кроме settings, именованные: подмены для тестов и тайминги self, settings: Settings, *, client_factory: Callable[[], Any] = NATS, sleep: SleepFn = asyncio.sleep, retry_delay: float = RETRY_DELAY, flush_timeout: int = FLUSH_TIMEOUT, ) -> None: self._settings = settings self._client_factory = client_factory self._sleep = sleep self._retry_delay = retry_delay self._flush_timeout = flush_timeout self._client: Any = None self._closing = False # --- состояние и жизненный цикл --------------------------------------------------------------------------- @property def is_connected(self) -> bool: """Фактическое состояние: `False` до подключения, при `NATS_ENABLED=false`, во время переподключения и после `close()`. Источник поля `nats` в `/healthcheck`.""" return self._client is not None and bool(self._client.is_connected) async def connect(self) -> bool: """Подключается к `NATS_URL`; `False` без подключения, если `NATS_ENABLED=false`. Ждёт сервер сколь угодно долго (nats-py повторяет и первое подключение). Ошибка учётных данных — `NatsConfigError` до попытки подключения. """ if not self._settings.NATS_ENABLED: logger.info("NATS выключен (NATS_ENABLED=false): подключение не выполняется") return False nats_settings = self._settings.OPENBAO.NATS auth = build_nats_auth(nats_settings) client = self._client_factory() self._closing = False try: await client.connect( nats_settings.NATS_URL, name=self._settings.SERVICE_NAME, error_cb=self._on_error, disconnected_cb=self._on_disconnected, reconnected_cb=self._on_reconnected, closed_cb=self._on_closed, max_reconnect_attempts=MAX_RECONNECT_ATTEMPTS, reconnect_time_wait=RECONNECT_TIME_WAIT, connect_timeout=CONNECT_TIMEOUT, **auth, ) except BaseException: with suppress(Exception): await client.close() raise self._client = client logger.info("Подключение к NATS установлено. url={url}", url=nats_settings.NATS_URL) return True async def close(self) -> None: """Подключённый клиент — `drain` (дообработать полученное и отправить буфер), иначе `close`; без исключений.""" client = self._client if client is None or client.is_closed: return self._closing = True try: if client.is_connected: await client.drain() else: await client.close() except Exception as exc: logger.warning( "Ошибка при закрытии соединения NATS. error_type={error_type}, error={error}", error_type=type(exc).__name__, error=repr(exc), ) with suppress(Exception): await client.close() # --- подписка --------------------------------------------------------------------------------------------- async def subscribe_requests(self, cb: MessageHandler, queue: str = REQUEST_QUEUE_GROUP) -> Any: """Подписка на `files.img_uploaded` в queue group `queue`; возвращает подписку nats-py после подтверждения сервером (`flush`). `cb` получает исходное `Msg` (`data`, `headers`); его исключения логируются, подписка не падает. `cb` должен возвращаться быстро: nats-py вызывает колбэки одной подписки последовательно, долгая обработка (ожидание ИС) задерживает следующие запросы — её выносят в отдельную задачу. """ if self._client is None: raise NatsNotConnectedError(f"Нет соединения с NATS: подписка на {IMAGE_REQUEST_SUBJECT} невозможна") async def safe_cb(msg: Msg) -> None: try: await cb(msg) except Exception as exc: logger.exception( "Ошибка обработки сообщения NATS. subject={subject}, error_type={error_type}, error={error}", subject=msg.subject, error_type=type(exc).__name__, error=repr(exc), ) subscription = await self._client.subscribe(IMAGE_REQUEST_SUBJECT, queue=queue, cb=safe_cb) await self._client.flush(timeout=self._flush_timeout) # SUB принят сервером: подписка действует по возврату logger.info( "Подписка NATS активна. subject={subject}, queue={queue}", subject=IMAGE_REQUEST_SUBJECT, queue=queue ) return subscription # --- публикация ------------------------------------------------------------------------------------------- async def publish(self, subject: str, payload: bytes, msg_id: str, *, retries: int | None = None) -> None: """Публикует `payload` с заголовком `Nats-Msg-Id: ` и ждёт подтверждения сервера (`flush`). `retries` — число **повторов** после первой попытки (спека: «повторить до трёх раз»; по умолчанию `PUBLISH_RETRIES` = 3, т.е. не более 4 попыток; `retries = 0` — ровно одна попытка), перед каждым повтором — пауза `retry_delay`. Повторяются только временные ошибки (`TRANSIENT_PUBLISH_ERRORS`). Повторы исчерпаны, ошибка не временная или соединения не было вовсе (`attempts = 0`) — `PublishFailed`. Дубликаты при повторе безопасны: ядро дедуплицирует по `Nats-Msg-Id`. `PublishFailed` не означает «не доставлено»: сообщения, буферизованные nats-py во время переподключения, могут уйти после восстановления соединения; ядро отбросит копии по `Nats-Msg-Id`. """ retry_count = self._settings.PUBLISH_RETRIES if retries is None else retries if retry_count < 0: raise ValueError(f"retries должно быть ≥ 0, получено {retry_count}") attempts_total = retry_count + 1 if self._client is None: raise PublishFailed( f"Нет соединения с NATS: {msg_id} в {subject} не опубликовано", subject=subject, msg_id=msg_id, attempts=0, ) headers = {NATS_MSG_ID_HEADER: msg_id} for attempt in range(1, attempts_total + 1): try: await self._client.publish(subject, payload, headers=headers) await self._client.flush(timeout=self._flush_timeout) return except TRANSIENT_PUBLISH_ERRORS as exc: if attempt == attempts_total: raise PublishFailed( f"Публикация {msg_id} в {subject} не удалась, попыток: {attempt} ({type(exc).__name__})", subject=subject, msg_id=msg_id, attempts=attempt, ) from exc logger.warning( "Публикация в NATS не удалась, повтор. subject={subject}, msg_id={msg_id}, attempt={attempt}, " "error_type={error_type}", subject=subject, msg_id=msg_id, attempt=attempt, error_type=type(exc).__name__, ) await self._sleep(self._retry_delay) except Exception as exc: raise PublishFailed( f"Публикация {msg_id} в {subject} не удалась, ошибка не временная ({type(exc).__name__})", subject=subject, msg_id=msg_id, attempts=attempt, ) from exc async def publish_result(self, message: RecognitionResultMessage, attempt: int) -> None: """Публикует конверт в `files.tabular_parsed`: тело — `message.to_wire()` (JSON, UTF-8), `Nats-Msg-Id` — `ack:…` для `processing`, `result:…` для итога; `attempt` — из запроса. Ошибки — как у `publish`.""" payload = json.dumps(message.to_wire(), ensure_ascii=False).encode("utf-8") await self.publish(RESULT_SUBJECT, payload, message_msg_id(message, attempt)) # --- колбэки nats-py -------------------------------------------------------------------------------------- async def _on_error(self, exc: Exception) -> None: logger.warning( "Ошибка соединения NATS. error_type={error_type}, error={error}", error_type=type(exc).__name__, error=repr(exc), ) async def _on_disconnected(self) -> None: if self._closing: return logger.warning("Соединение с NATS потеряно, nats-py переподключается") async def _on_reconnected(self) -> None: logger.info("Соединение с NATS восстановлено, подписки восстановлены клиентом") async def _on_closed(self) -> None: logger.info("Соединение с NATS закрыто")