60 lines
2.3 KiB
Python
60 lines
2.3 KiB
Python
import asyncio
|
|
from collections import defaultdict
|
|
from enum import unique, IntEnum
|
|
from typing import List
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
from sqlalchemy.orm import joinedload
|
|
|
|
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]):
|
|
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)
|
|
for update in updates:
|
|
# Working with marketplaces
|
|
stmt = (
|
|
select(
|
|
MarketplaceProduct.marketplace_id.distinct()
|
|
)
|
|
.where(
|
|
MarketplaceProduct.product_id == update.product_id,
|
|
MarketplaceProduct.marketplace_id.in_([9, 41])
|
|
)
|
|
)
|
|
stmt_result = await self.session.execute(stmt)
|
|
marketplace_ids = stmt_result.scalars().all()
|
|
if not marketplace_ids:
|
|
continue
|
|
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: len(x[1]))
|
|
|
|
tasks = []
|
|
for marketplace_id, marketplace_updates in updates_list:
|
|
tasks.append(self.update_marketplace(marketplace_id, marketplace_updates))
|
|
await asyncio.gather(*tasks)
|