123

updaters/stocks_updater.py

@@ -1,53 +1,47 @@
+import asyncio
 from collections import defaultdict
-from dataclasses import dataclass
 from enum import unique, IntEnum
-from typing import List, Union
+from typing import List
 
 from sqlalchemy import select
-from sqlalchemy.dialects.postgresql import insert
 from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import joinedload
 
-import database
-from database import Marketplace, MarketplaceProduct, DailyStock
-
-
-@unique
-class StockUpdateType(IntEnum):
-    SALE = 0
-    SUPPLIER_UPDATE = 1
-    WAREHOUSE_UPDATE = 2
-
-@dataclass
-class StockUpdate:
-    product_id: int
-    type: StockUpdateType
-    quantity: int
+from database import Marketplace, MarketplaceProduct, Warehouse, Company
+from schemas.general import StockUpdate
+from updaters.factory import UpdaterFactory
 
 
 class StocksUpdater:
     def __init__(self, session: AsyncSession):
         self.session = session
 
+    async def get_marketplace(self, marketplace_id: int):
+        marketplace = await self.session.get(Marketplace, marketplace_id, options=[
+            joinedload(Marketplace.warehouses).joinedload(Warehouse.suppliers),
+            joinedload(Marketplace.warehouses).joinedload(Warehouse.company_warehouses),
+            joinedload(Marketplace.company).joinedload(Company.warehouse)
+        ])
+        return marketplace
+
     async def update_marketplace(self, marketplace_id: int, updates: List[StockUpdate]):
-        pass
+        marketplace = await self.get_marketplace(marketplace_id)
+        updater = UpdaterFactory.get_updater(self.session, marketplace)
+        if not updater:
+            return
+        await updater.update(updates)
 
     async def update(self, updates: list[StockUpdate]):
         updates_dict = defaultdict(list)
-        stock_update_values = []
         for update in updates:
-            # Working with sold today
-            if update.type == StockUpdateType.SALE:
-                stock_update_values.append({
-                    'product_id': update.product_id,
-                    'sold_today': update.quantity
-                })
             # Working with marketplaces
             stmt = (
                 select(
                     MarketplaceProduct.marketplace_id.distinct()
                 )
                 .where(
-                    MarketplaceProduct.product_id == update.product_id
+                    MarketplaceProduct.product_id == update.product_id,
+                    MarketplaceProduct.marketplace_id.in_([9, 41])
                 )
             )
             stmt_result = await self.session.execute(stmt)
@@ -57,27 +51,9 @@ class StocksUpdater:
             for marketplace_id in marketplace_ids:
                 updates_dict[marketplace_id].append(update)
         updates_list = list(updates_dict.items())
-        updates_list = sorted(updates_list, key=lambda x: x[1])
-
-        # Updating DailyStock-s
-        insert_stmt = (
-            insert(
-                DailyStock
-            )
-            .values(
-                stock_update_values
-            )
-        )
-        insert_stmt = (
-            insert_stmt.on_conflict_do_update(
-                index_elements=['product_id'],
-                set_={
-                    'sold_today': DailyStock.sold_today + insert_stmt.excluded.sold_today
-                }
-            )
-        )
-        await self.session.execute(insert_stmt)
-        await self.session.commit()
+        updates_list = sorted(updates_list, key=lambda x: len(x[1]))
 
+        tasks = []
         for marketplace_id, marketplace_updates in updates_list:
-            await self.update_marketplace(marketplace_id, marketplace_updates)
+            tasks.append(self.update_marketplace(marketplace_id, marketplace_updates))
+        await asyncio.gather(*tasks)