v1.0

updaters/stocks_updater.py

@@ -1,13 +1,16 @@
 import asyncio
+import time
 from collections import defaultdict
 from enum import unique, IntEnum
-from typing import List
+from typing import List, Union
 
-from sqlalchemy import select
+from sqlalchemy import select, or_
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import joinedload
+from sqlalchemy.orm import joinedload, selectinload
 
+from backend.session import get_session, session_factory
 from database import Marketplace, MarketplaceProduct, Warehouse, Company
+from database.sipro.enums.general import BaseMarketplace
 from schemas.general import StockUpdate
 from updaters.factory import UpdaterFactory
 
@@ -24,24 +27,77 @@ class StocksUpdater:
         ])
         return marketplace
 
-    async def update_marketplace(self, marketplace_id: int, updates: List[StockUpdate]):
+    async def get_marketplaces(self, marketplace_ids: Union[list[int], None] = None) -> List[Marketplace]:
+        if not marketplace_ids:
+            marketplace_ids = []
+        stmt = (
+            select(
+                Marketplace
+            )
+            .join(
+                Company
+            )
+            .options(
+                selectinload(Marketplace.warehouses).selectinload(Warehouse.suppliers),
+                selectinload(Marketplace.warehouses).selectinload(Warehouse.company_warehouses),
+                joinedload(Marketplace.company).joinedload(Company.warehouse)
+            )
+            .where(
+                Company.is_deleted == False,
+                Company.is_archived == False,
+                Marketplace.is_deleted == False,
+                Marketplace.base_marketplace.in_([
+                    BaseMarketplace.OZON,
+                    BaseMarketplace.WILDBERRIES,
+                    BaseMarketplace.YANDEX_MARKET
+                ]),
+                or_(
+                    marketplace_ids == [],
+                    Marketplace.id.in_(marketplace_ids)
+                )
+            )
+        )
+        query_result = await self.session.scalars(stmt)
+        return query_result.all()
+
+    async def full_update_marketplace(self, marketplace_id: int):
         marketplace = await self.get_marketplace(marketplace_id)
         updater = UpdaterFactory.get_updater(self.session, marketplace)
-        if not updater:
-            return
-        await updater.update(updates)
+        await updater.update_all()
+
+    async def full_update_all_marketplaces(self, marketplace_ids: Union[List[int], None] = None):
+        marketplaces = await self.get_marketplaces(marketplace_ids)
+
+        async def update_marketplace(marketplace):
+            async with session_factory() as session:
+                updater = UpdaterFactory.get_updater(session, marketplace)
+                await updater.update_all()
+
+        tasks = [update_marketplace(marketplace) for marketplace in marketplaces]
+        await asyncio.gather(*tasks)
+
+    async def update_marketplace(self, marketplace_id: int, updates: List[StockUpdate]):
+        marketplace = await self.get_marketplace(marketplace_id)
+        async with session_factory() as session:
+            updater = UpdaterFactory.get_updater(session, marketplace)
+            if not updater:
+                return
+            await updater.update(updates)
 
     async def update(self, updates: list[StockUpdate]):
         updates_dict = defaultdict(list)
         for update in updates:
-            # Working with marketplaces
             stmt = (
                 select(
                     MarketplaceProduct.marketplace_id.distinct()
                 )
+                .join(Marketplace)
+                .join(Company)
                 .where(
                     MarketplaceProduct.product_id == update.product_id,
-                    MarketplaceProduct.marketplace_id.in_([9, 41])
+                    Marketplace.is_deleted == False,
+                    Company.is_deleted == False,
+                    Company.is_archived == False
                 )
             )
             stmt_result = await self.session.execute(stmt)
@@ -52,7 +108,8 @@ class StocksUpdater:
                 updates_dict[marketplace_id].append(update)
         updates_list = list(updates_dict.items())
         updates_list = sorted(updates_list, key=lambda x: len(x[1]))
-
+        print(updates_list)
+        return
         tasks = []
         for marketplace_id, marketplace_updates in updates_list:
             tasks.append(self.update_marketplace(marketplace_id, marketplace_updates))