small bug fix

src/monitor_utils/database.py

@@ -0,0 +1,118 @@
+import logging
+from dataclasses import asdict
+from datetime import datetime, timezone
+from functools import cache
+
+from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorCollection
+from pymongo import UpdateOne
+from pymongo.errors import PyMongoError
+
+from monitor_utils.src.models import Product
+from monitor_utils.src.constants import STATUS_PENDING, MONGO_DATABASE_NAME, MONGO_URL
+
+LOGGER = logging.getLogger(__name__)
+
+
+def get_collection(collection_name: str) -> AsyncIOMotorCollection:
+    return get_mongo_client()[MONGO_DATABASE_NAME][collection_name]
+
+
+@cache
+def get_mongo_client() -> AsyncIOMotorClient:
+    return AsyncIOMotorClient(
+        MONGO_URL,
+        maxPoolSize=200,
+        minPoolSize=10,
+        retryWrites=True,
+        retryReads=True,
+        tz_aware=True,
+        tzinfo=timezone.utc,
+    )
+
+
+async def get_product_from_database(
+    collection: AsyncIOMotorCollection,
+    product_id: str,
+) -> dict | None:
+    return await collection.find_one({"_id": product_id})
+
+
+async def add_products_to_database(
+    products: list[Product],
+    collection: AsyncIOMotorCollection,
+    category_identifier: str,
+) -> None:
+    if not products:
+        return
+
+    composite_ids = [f"{p.product_id}_{category_identifier}" for p in products]
+    cursor = collection.find({"_id": {"$in": composite_ids}}, {"currentPrice": 1})
+    existing_prices = {str(doc["_id"]): doc.get("currentPrice") async for doc in cursor}
+
+    ops = []
+    for product in products:
+        document = convert_product_to_mongo(product, category_identifier)
+        doc_id = document["_id"]
+        old_price = existing_prices.get(doc_id)
+        new_price = document.get("currentPrice")
+
+        price_dropped = old_price is None or (
+            new_price is not None and new_price < old_price
+        )
+
+        if price_dropped:
+            document["status"] = STATUS_PENDING
+            document["priceBeforeDrop"] = old_price
+
+        set_on_insert = {"discoveredDate": datetime.now(timezone.utc)}
+        if not price_dropped:
+            set_on_insert["status"] = STATUS_PENDING
+
+        ops.append(
+            UpdateOne(
+                {"_id": doc_id},
+                {"$set": document, "$setOnInsert": set_on_insert},
+                upsert=True,
+            )
+        )
+
+    try:
+        await collection.bulk_write(ops, ordered=False)
+    except PyMongoError:
+        LOGGER.exception(f"{category_identifier} - failed to add products to database")
+        raise
+
+
+def convert_product_to_mongo(product: Product, category_identifier: str) -> dict:
+    data = asdict(product)
+    metadata = data["metadata"]
+    timing = data["timing"]
+    pricing = data["pricing"]
+    availability = data["availability"]
+    current_price = pricing["current_price"]
+    original_price = pricing["original_price"]
+    if isinstance(current_price, str):
+        current_price = float(current_price.replace(",", ""))
+        original_price = float(original_price.replace(",", ""))
+
+    return {
+        "_id": f"{data['product_id']}_{category_identifier}",
+        "title": data["title"],
+        "url": data["url"],
+        "currentPrice": current_price,
+        "originalPrice": original_price,
+        "percentage": metadata["percentage"],
+        "imageUrl": metadata["image_url"],
+        "customizationIndex": metadata["customization_index"],
+        "memberNonMemberComparison": metadata["member_non_member_comparison"],
+        "overstockText": metadata["overstock_text"],
+        "startDate": timing["start_date"],
+        "endDate": timing["end_date"],
+        "priceBeforeDrop": pricing["price_before_drop"],
+        "eduPrice": pricing["edu_price"],
+        "available": availability["available"],
+        "condition": availability["condition"],
+        "appExclusive": availability["app_exclusive"],
+        "memberPick": availability["member_pick"],
+        "memberDeal": availability["member_deal"],
+    }