Files
Sipro-Stocks/updaters/stocks_updater.py
2024-07-01 06:01:50 +03:00

84 lines
2.5 KiB
Python

from collections import defaultdict
from dataclasses import dataclass
from enum import unique, IntEnum
from typing import List, Union
from sqlalchemy import select
from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.ext.asyncio import AsyncSession
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
class StocksUpdater:
def __init__(self, session: AsyncSession):
self.session = session
async def update_marketplace(self, marketplace_id: int, updates: List[StockUpdate]):
pass
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
)
)
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: 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()
for marketplace_id, marketplace_updates in updates_list:
await self.update_marketplace(marketplace_id, marketplace_updates)