Добавлен поллер накладной

app/repositories/order/repository.py

@@ -1,12 +1,15 @@
 """PostgreSQL order repository."""
 
+from collections.abc import Sequence
 from contextlib import AbstractAsyncContextManager
 from dataclasses import dataclass
+from datetime import datetime
 from typing import Any
 
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
 
+from app.domain.cdek_polling import TERMINAL_ORDER_STATUSES
 from app.repositories.order.models import Order
 
 
@@ -72,15 +75,72 @@ class OrderRepository:
         session: AsyncSession,
         order_uuid: str,
         cdek_order_uuid: str,
-        cdek_waybill_uuid: str | None = None,
-        cdek_waybill_url: str | None = None,
     ) -> Order | None:
         order = await self.get_order_by_order_uuid(session, order_uuid)
         if order is None:
             return None
 
         order.cdek_order_uuid = cdek_order_uuid
-        order.cdek_waybill_uuid = cdek_waybill_uuid
-        order.cdek_waybill_url = cdek_waybill_url
+        await session.flush()
+        return order
+
+    async def list_orders_pending_waybill(
+        self,
+        session: AsyncSession,
+        *,
+        limit: int,
+    ) -> Sequence[Order]:
+        statement = (
+            select(Order)
+            .where(
+                Order.cdek_order_uuid.is_not(None),
+                Order.cdek_waybill_url.is_(None),
+                (
+                    Order.cdek_order_status.is_(None)
+                    | Order.cdek_order_status.not_in(TERMINAL_ORDER_STATUSES)
+                ),
+            )
+            .order_by(Order.cdek_polled_at.asc().nulls_first())
+            .limit(limit)
+            .with_for_update(skip_locked=True)
+        )
+        result = await session.execute(statement)
+        return result.scalars().all()
+
+    async def record_order_poll(
+        self,
+        session: AsyncSession,
+        *,
+        order_uuid: str,
+        order_status: str | None,
+        waybill_uuid: str | None,
+        polled_at: datetime,
+    ) -> Order | None:
+        order = await self.get_order_by_order_uuid(session, order_uuid)
+        if order is None:
+            return None
+
+        order.cdek_order_status = order_status
+        if waybill_uuid is not None and order.cdek_waybill_uuid is None:
+            order.cdek_waybill_uuid = waybill_uuid
+        order.cdek_polled_at = polled_at
+        await session.flush()
+        return order
+
+    async def record_waybill_poll(
+        self,
+        session: AsyncSession,
+        *,
+        order_uuid: str,
+        waybill_url: str | None,
+        polled_at: datetime,
+    ) -> Order | None:
+        order = await self.get_order_by_order_uuid(session, order_uuid)
+        if order is None:
+            return None
+
+        if waybill_url is not None and order.cdek_waybill_url is None:
+            order.cdek_waybill_url = waybill_url
+        order.cdek_polled_at = polled_at
         await session.flush()
         return order