Fulfillment-Backend/marketplaces/wildberries/core.py

import time
from collections import defaultdict

from sqlalchemy import select
from sqlalchemy.orm import selectinload

from external.marketplace.base.product_synchronizer import BaseProductSynchronizer
from external.marketplace.wildberries.core import WildberriesMarketplaceApi
from marketplaces.base.core import BaseMarketplaceController
from models import Product, ProductBarcode, ProductImage, WildberriesProduct


class WildberriesProductSynchronizer(BaseProductSynchronizer):
    marketplace_products: list[WildberriesProduct]
    api: WildberriesMarketplaceApi

    async def _write(self):
        instances = self.products + self.marketplace_products + self.barcodes + self.images
        start = time.time()
        self.session.add_all(instances)
        await self.session.commit()
        print(f'Write time: {time.time() - start}')

    async def _get_synchronized_nm_uuids(self):
        marketplace_id = self.marketplace.id
        synchronized_nm_uuids = set(
            (
                await self.session.scalars(
                    select(
                        WildberriesProduct.nm_uuid
                    )
                    .where(
                        WildberriesProduct.marketplace_id == marketplace_id
                    )
                )
            ).all()
        )
        return synchronized_nm_uuids

    async def _create_product(self, card, size_value):
        return Product(
            client_id=self.marketplace.client_id,
            name=card['title'],
            article=card['vendorCode'],
            size=size_value
        )

    async def _create_barcodes(self, product, skus):
        barcodes = []
        for sku in skus:
            barcode = ProductBarcode(
                product=product,
                barcode=sku
            )
            barcodes.append(barcode)
        return barcodes

    async def _create_images(self, product, photos):
        images = []
        for photo in photos[:1]:
            image = ProductImage(
                product=product,
                image_url=photo['big']
            )
            images.append(image)
        return images

    async def _create_wildberries_product(self, product, nm_uuid):
        return WildberriesProduct(
            marketplace_id=self.marketplace.id,
            product=product,
            nm_uuid=nm_uuid
        )

    async def _update_product_info(self, product, card):
        product.name = card['title']
        product.article = card['vendorCode']

    async def _update_barcodes(self, product, skus):
        existing_barcodes = {barcode.barcode for barcode in product.barcodes}
        new_barcodes = []
        for sku in skus:
            if sku not in existing_barcodes:
                barcode = ProductBarcode(
                    product=product,
                    barcode=sku
                )
                new_barcodes.append(barcode)
        return new_barcodes

    async def _update_images(self, product, photos):
        existing_images = {image.image_url for image in product.images}
        new_images = []
        for photo in photos[:1]:
            if photo['big'] not in existing_images:
                image = ProductImage(
                    product=product,
                    image_url=photo['big']
                )
                new_images.append(image)
        return new_images

    async def _process_product(self, card, size_value, nm_uuid, size):
        product = await self._create_product(card, size_value)

        barcodes = await self._create_barcodes(product, size.get('skus') or [])
        images = await self._create_images(product, card.get('photos') or [])
        wildberries_product = await self._create_wildberries_product(product, nm_uuid)

        self.products.append(product)
        self.barcodes.extend(barcodes)
        self.images.extend(images)
        self.marketplace_products.append(wildberries_product)

    async def create_products(self):
        self._clear()
        synchronized_nm_uuids = await self._get_synchronized_nm_uuids()
        async for card in self.api.get_all_products():
            nm_uuid = card['nmUUID']
            if nm_uuid in synchronized_nm_uuids:
                continue
            sizes = card.get('sizes') or []
            for size in sizes:
                tech_size = size.get('techSize')
                wb_size = size.get('wbSize')
                size_value = tech_size or wb_size
                await self._process_product(
                    card,
                    size_value,
                    nm_uuid,
                    size
                )
        await self._write()

    async def synchronize_products(self):
        self._clear()
        synchronized_products_stmt = (
            select(Product)
            .join(
                WildberriesProduct
            )
            .options(
                selectinload(Product.barcodes),
                selectinload(Product.wildberries_products)
            )
            .where(WildberriesProduct.marketplace_id == self.marketplace.id)
        )
        synchronized_products = await self.session.execute(synchronized_products_stmt)
        synchronized_products = synchronized_products.scalars().all()
        synchronized_products_nm_id_dict = defaultdict(list)
        for product in synchronized_products:
            for wb_product in product.wildberries_products:
                synchronized_products_nm_id_dict[wb_product.nm_uuid].append(product)
        synchronized_nm_uuids = list(synchronized_products_nm_id_dict.keys())
        async for card in self.api.get_all_products():
            nm_uuid = card['nmUUID']
            if nm_uuid not in synchronized_nm_uuids:
                continue
            products = synchronized_products_nm_id_dict[nm_uuid]
            existing_sizes = {product.size for product in products}
            size_product_dict = {product.size: product for product in products}

            sizes = card.get('sizes') or []
            for size in sizes:
                tech_size = size.get('techSize')
                wb_size = size.get('wbSize')
                size_value = tech_size or wb_size
                if size_value in existing_sizes:
                    product = size_product_dict[size_value]
                    await self._update_product_info(product, card)
                    self.barcodes.extend(
                        await self._update_barcodes(product, size.get('skus') or [])
                    )
                    self.images.extend(
                        await self._update_images(product, card.get('photos') or [])
                    )
                    continue
                await self._process_product(
                    card,
                    size_value,
                    nm_uuid,
                    size
                )
        await self._write()


class WildberriesController(
    BaseMarketplaceController
):
    api: WildberriesMarketplaceApi

    def __init__(self, session, marketplace):
        super().__init__(session, marketplace)
        self.product_synchronizer = WildberriesProductSynchronizer(session, marketplace, self.api)

    async def create_products(self):
        await self.product_synchronizer.create_products()

    async def synchronize_products(self):
        await self.product_synchronizer.synchronize_products()