g2s-aggregator/app/services/waybill_poller.py

"""Background service that polls providers for waybill updates."""

import asyncio
from collections.abc import Callable, Mapping, Sequence
from contextlib import AbstractAsyncContextManager
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Protocol

import structlog

logger = structlog.get_logger(__name__)


class OrderInfoAdapterProtocol(Protocol):
    async def get_order(self, provider_order_id: str) -> object: ...


class WaybillInfoAdapterProtocol(Protocol):
    async def get_waybill(self, provider_waybill_id: str) -> object: ...


class OrderRecord(Protocol):
    order_uuid: str
    provider: str
    provider_order_id: str | None
    provider_waybill_id: 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_adapters: Mapping[str, OrderInfoAdapterProtocol],
        waybill_info_adapters: Mapping[str, WaybillInfoAdapterProtocol],
        batch_size: int,
        datetime_now: Callable[[], datetime] = lambda: datetime.now(timezone.utc),
    ) -> None:
        self._repository = order_repository
        self._order_info_adapters = dict(order_info_adapters)
        self._waybill_info_adapters = dict(waybill_info_adapters)
        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,
                        provider=order.provider,
                        provider_order_id=order.provider_order_id,
                        provider_waybill_id=order.provider_waybill_id,
                    )
            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.provider_waybill_id is None:
            provider_order_id = order.provider_order_id
            if provider_order_id is None:
                return
            adapter = self._resolve_order_info_adapter(order.provider)
            info = await adapter.get_order(provider_order_id)
            status_code = _extract_status_code(info)
            waybill_id = _extract_waybill_id(info)
            await self._repository.record_order_poll(
                session,
                order_uuid=order.order_uuid,
                order_status=status_code,
                waybill_uuid=waybill_id,
                polled_at=polled_at,
            )
            logger.info(
                "waybill_poll_order_result",
                order_uuid=order.order_uuid,
                provider=order.provider,
                provider_order_id=provider_order_id,
                provider_order_status=status_code,
                provider_waybill_id=waybill_id,
            )
            return

        adapter = self._resolve_waybill_info_adapter(order.provider)
        waybill = await adapter.get_waybill(order.provider_waybill_id)
        waybill_url = _extract_waybill_url(waybill)
        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,
            provider=order.provider,
            provider_waybill_id=order.provider_waybill_id,
            provider_waybill_url=waybill_url,
        )

    def _resolve_order_info_adapter(self, provider: str) -> OrderInfoAdapterProtocol:
        adapter = self._order_info_adapters.get(provider)
        if adapter is None:
            raise RuntimeError(f"Order info adapter is not configured for {provider}.")
        return adapter

    def _resolve_waybill_info_adapter(
        self, provider: str
    ) -> WaybillInfoAdapterProtocol:
        adapter = self._waybill_info_adapters.get(provider)
        if adapter is None:
            raise RuntimeError(
                f"Waybill info adapter is not configured for {provider}."
            )
        return adapter


def _extract_status_code(info: object) -> str | None:
    value = getattr(info, "status_code", None)
    return value if isinstance(value, str) and value else None


def _extract_waybill_id(info: object) -> str | None:
    for attr in ("waybill_id", "waybill_uuid", "waybill_number"):
        value = getattr(info, attr, None)
        if isinstance(value, str) and value:
            return value
    return None


def _extract_waybill_url(info: object) -> str | None:
    value = getattr(info, "url", None)
    return value if isinstance(value, str) and value else None