g2s-aggregator/app/services/waybill_poller.py

"""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.info(
                    "waybill_poll_tick "
                    f"processed={summary.processed} "
                    f"succeeded={summary.succeeded} "
                    f"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,
        )