"""Background service that polls CDEK for waybill updates.""" import asyncio from collections.abc import Callable, Sequence from contextlib import AbstractAsyncContextManager from dataclasses import dataclass from datetime import datetime, timezone from typing import Protocol import structlog from app.adapters.delivery_providers.cdek.order_mapper import ( CDEKOrderInfo, CDEKWaybillInfo, ) logger = structlog.get_logger(__name__) class CDEKOrderInfoAdapterProtocol(Protocol): async def get_order(self, cdek_order_uuid: str) -> CDEKOrderInfo: ... class CDEKWaybillInfoAdapterProtocol(Protocol): async def get_waybill(self, cdek_waybill_uuid: str) -> CDEKWaybillInfo: ... class OrderRecord(Protocol): order_uuid: str cdek_order_uuid: str | None cdek_waybill_uuid: str | None class WaybillPollerRepositoryProtocol(Protocol): def session(self) -> AbstractAsyncContextManager[object]: ... async def list_orders_pending_waybill( self, session: object, *, limit: int ) -> Sequence[OrderRecord]: ... async def record_order_poll( self, session: object, *, order_uuid: str, order_status: str | None, waybill_uuid: str | None, polled_at: datetime, ) -> object | None: ... async def record_waybill_poll( self, session: object, *, order_uuid: str, waybill_url: str | None, polled_at: datetime, ) -> object | None: ... @dataclass(frozen=True) class PollBatchSummary: processed: int succeeded: int failed: int class WaybillPollerService: def __init__( self, *, order_repository: WaybillPollerRepositoryProtocol, order_info_adapter: CDEKOrderInfoAdapterProtocol, waybill_info_adapter: CDEKWaybillInfoAdapterProtocol, batch_size: int, datetime_now: Callable[[], datetime] = lambda: datetime.now(timezone.utc), ) -> None: self._repository = order_repository self._order_info_adapter = order_info_adapter self._waybill_info_adapter = waybill_info_adapter self._batch_size = batch_size self._datetime_now = datetime_now async def poll_once(self) -> PollBatchSummary: async with self._repository.session() as session: orders = await self._repository.list_orders_pending_waybill( session, limit=self._batch_size ) succeeded = 0 failed = 0 for order in orders: try: await self._handle_order(session, order) succeeded += 1 except Exception: failed += 1 logger.exception( "waybill_poll_order_failed", order_uuid=order.order_uuid, cdek_order_uuid=order.cdek_order_uuid, cdek_waybill_uuid=order.cdek_waybill_uuid, ) return PollBatchSummary( processed=len(orders), succeeded=succeeded, failed=failed, ) async def run_forever( self, *, interval_seconds: float, stop_event: asyncio.Event, ) -> None: while not stop_event.is_set(): try: summary = await self.poll_once() logger.debug( "waybill_poll_tick", processed=summary.processed, succeeded=summary.succeeded, failed=summary.failed, ) except Exception: logger.exception("waybill_poll_tick_failed") try: await asyncio.wait_for( stop_event.wait(), timeout=interval_seconds ) except asyncio.TimeoutError: continue async def _handle_order(self, session: object, order: OrderRecord) -> None: polled_at = self._datetime_now() if order.cdek_waybill_uuid is None: cdek_order_uuid = order.cdek_order_uuid if cdek_order_uuid is None: return info = await self._order_info_adapter.get_order(cdek_order_uuid) await self._repository.record_order_poll( session, order_uuid=order.order_uuid, order_status=info.status_code, waybill_uuid=info.waybill_uuid, polled_at=polled_at, ) logger.info( "waybill_poll_order_result", order_uuid=order.order_uuid, cdek_order_uuid=cdek_order_uuid, cdek_order_status=info.status_code, cdek_waybill_uuid=info.waybill_uuid, ) return waybill = await self._waybill_info_adapter.get_waybill(order.cdek_waybill_uuid) await self._repository.record_waybill_poll( session, order_uuid=order.order_uuid, waybill_url=waybill.url, polled_at=polled_at, ) logger.info( "waybill_poll_waybill_result", order_uuid=order.order_uuid, cdek_waybill_uuid=order.cdek_waybill_uuid, cdek_waybill_url=waybill.url, )