monitor-shared-utils/src/monitor_utils/database.py

"""Database helpers for upserting monitored products and tracking price history.

Provides the shared Motor client, collection access, and the upsert path that
rewrites current product fields while preserving prior prices in a pastPrices
array deduplicated by price, stamping lastUpdated on every write, and emitting
alert references for new products and price changes.
"""

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 BulkWriteError, PyMongoError

from monitor_utils.constants import MONGO_DATABASE_NAME, MONGO_URL
from monitor_utils.models import Product

LOGGER = logging.getLogger(__name__)

# Drain rank for products with no MSRP. Their percentage is 0 by construction, so
# without this they sort dead last and starve behind any backlog. 6 puts them
# directly under the 70%+ band, tying with 60-69% and broken by createdAt.
NO_MSRP_PRIORITY = 6
# Floor rank for new listings, which send regardless of discount and would
# otherwise sort near 0 and starve the same way MSRP-less products did.
NEW_LISTING_PRIORITY = 6


def get_collection(
    collection_name: str,
    mongo_client: AsyncIOMotorClient = None,
    database_name=MONGO_DATABASE_NAME,
) -> AsyncIOMotorCollection:
    # rename mongo_client to client and add optional database name arg
    """Return the named collection from the shared client's database.

    Args:
        collection_name: Name of the collection to return.
        mongo_client: optional client for projects with already configured client
        database_name: optional name to use besides the default mongo database name

    Returns:
        The requested collection.
    """

    if mongo_client is None:
        mongo_client = get_mongo_client()

    return mongo_client[database_name][collection_name]


def resolve_alerts_collection(
    product_collection: AsyncIOMotorCollection,
) -> AsyncIOMotorCollection:
    """Return the alerts collection beside a product collection.

    Alerts co-locate with the products they reference: the same database and
    the same development mode. The mode is read from the product collection's
    own name (setup_config appends a _testing suffix outside production), so a
    monitor writing products to <name>_testing also writes alerts to
    alerts_testing and the development bot draining alerts_testing sees them.

    Args:
        product_collection: The mode-resolved product collection the monitor
            already writes to.

    Returns:
        The matching alerts collection in the same database and mode.
    """

    suffix = "_testing" if product_collection.name.endswith("_testing") else ""
    return product_collection.database[f"alerts{suffix}"]


@cache
def get_mongo_client(mongo_url: str = MONGO_URL) -> AsyncIOMotorClient:
    """Return a cached Motor client configured for the shared cluster.

    Args:
        mongo_url: The mongo url to use for connections.

    Returns:
        The pooled, timezone-aware Motor client.
    """

    return AsyncIOMotorClient(
        mongo_url,
        minPoolSize=10,
        connectTimeoutMS=20_000,
        socketTimeoutMS=60_000,
        serverSelectionTimeoutMS=30_000,
        retryWrites=True,
        retryReads=True,
        tz_aware=True,
        tzinfo=timezone.utc,
    )


async def get_product_from_database(
    collection: AsyncIOMotorCollection,
    product_id: str,
) -> dict:
    """Fetch a single product document by its _id."""

    return await collection.find_one({"_id": product_id})


async def add_products_to_database(
    products: list[Product],
    collection: AsyncIOMotorCollection,
    category_identifier: str,
    monitor: str,
    use_composite: bool = False,
):
    """Upsert products, tracking deduplicated price history and emitting alerts.

    Rewrites current fields on each upsert. priceBeforeDrop is derived here,
    not by monitors: it is set to the previous stored currentPrice when the new
    price is lower (a drop) and cleared otherwise, so any price_before_drop on
    the incoming product is ignored. newListing is likewise stamped here, not
    by monitors: True only on the cycle whose write inserts the document, and
    copied onto the discovery alert, whose copy never flips. Prices are recorded in a pastPrices array
    of price/date entries; a price already present in the array is never
    recorded again, so each distinct price keeps only its first-seen timestamp.
    New products and never-before-recorded prices
    upsert an alert referencing the product; a price returning to a value
    already in pastPrices emits nothing, so per-source price oscillation
    (member vs non-member quotes for the same SKU) alerts once per distinct
    price instead of every run. Alerts are deduplicated so a product holds
    at most one unclaimed alert regardless of concurrent category or page
    tasks detecting it, and write only after the product write succeeds so
    a failed upsert cannot produce an orphaned alert.

    Args:
        products: Products to upsert.
        collection: Target MongoDB collection.
        category_identifier: Category suffix appended to composite ids.
        monitor: Monitor name stamped onto emitted alerts for consumer
            routing.
        use_composite: Build composite ids so multi-variant monitors (Woot)
            keep one doc per variant; other monitors collapse to one doc per
            product_id.
    """

    # Callers are expected to filter empty batches, but guard here too: an empty
    # products list would make collection.bulk_write([]) raise InvalidOperation,
    # which callers would then log as a spurious category failure every cycle.
    if not products:
        return

    doc_ids = [
        build_doc_id(p.product_id, category_identifier, use_composite) for p in products
    ]
    cursor = collection.find(
        {"_id": {"$in": doc_ids}}, {"currentPrice": 1, "pastPrices": 1}
    )
    existing = {str(doc["_id"]): doc async for doc in cursor}

    ops = []
    alert_ops = []
    now = datetime.now(timezone.utc)
    for product, doc_id in zip(products, doc_ids):
        document = convert_product_to_mongo(product)
        document["_id"] = doc_id
        existing_doc = existing.get(doc_id)
        # Stamped here, not by monitors: only the database knows whether the
        # doc already exists. Must precede build_alert so the alert's copy of
        # the flag is taken from the stamped value.
        document["newListing"] = existing_doc is None
        new_price = document.get("currentPrice")
        previous_price = existing_doc.get("currentPrice") if existing_doc else None
        document["priceBeforeDrop"] = resolve_price_before_drop(
            previous_price, new_price
        )

        recorded_prices = extract_recorded_prices(existing_doc)
        never_seen_price = new_price is not None and new_price not in recorded_prices
        set_on_insert = {"discoveredDate": now}
        if existing_doc is None or never_seen_price:
            alert = build_alert(monitor, document, now)
            alert_ops.append(
                UpdateOne(
                    {"product_id": doc_id, "claimed": False},
                    {"$setOnInsert": alert},
                    upsert=True,
                )
            )

        update = {"$set": document, "$setOnInsert": set_on_insert}
        if never_seen_price:
            update["$push"] = {"pastPrices": {"price": new_price, "date": now}}

        ops.append(UpdateOne({"_id": doc_id}, update, upsert=True))

    try:
        await collection.bulk_write(ops, ordered=False)
        if alert_ops:
            await write_alert_ops(alert_ops, resolve_alerts_collection(collection))
    except PyMongoError:
        LOGGER.exception(
            "failed to add products to database",
            extra={"category": category_identifier},
        )
        raise


async def write_alert_ops(
    alert_ops: list[UpdateOne], alerts_collection: AsyncIOMotorCollection
):
    """Write alert upserts, swallowing duplicate-key races.

    Concurrent tasks upserting an alert for the same product can both pass
    the upsert's match phase and race the insert; the partial unique index
    on unclaimed product_id makes the loser fail with a duplicate-key
    error, which is the dedupe working as intended and is ignored. Any
    other write error re-raises.

    Args:
        alert_ops: Alert upsert operations to write.
        alerts_collection: The mode-resolved alerts collection to write to,
            supplied by the caller so the writer stays mode-agnostic.
    """

    try:
        await alerts_collection.bulk_write(alert_ops, ordered=False)
    except BulkWriteError as exc:
        real_errors = [
            write_error
            for write_error in exc.details["writeErrors"]
            if write_error.get("code") != 11000
        ]
        if real_errors:
            raise


def resolve_price_before_drop(
    previous_price: float | None, new_price: float | None
) -> float | None:
    """Return the pre-drop price when new_price is a drop from previous_price.

    Owns price-drop detection so monitors do not: a drop is a new price strictly
    below the previously stored one, and its priceBeforeDrop is that prior price.
    Any other case (no prior price, unchanged, or a rise) is not a drop.

    Args:
        previous_price: The product's last stored currentPrice, or None when new.
        new_price: The incoming currentPrice.

    Returns:
        previous_price on a drop, otherwise None.
    """

    if (
        previous_price is not None
        and new_price is not None
        and new_price < previous_price
    ):
        return previous_price

    return None


def extract_recorded_prices(existing_doc: dict | None) -> set[float]:
    """Return the set of prices already recorded in a document's pastPrices.

    Args:
        existing_doc: Existing product document, or None for new products.

    Returns:
        Prices present in the pastPrices array, empty when none exist.
    """

    if not existing_doc:
        return set()

    return {entry["price"] for entry in existing_doc.get("pastPrices", [])}


def build_alert(monitor: str, document: dict, now: datetime) -> dict:
    """Build an alert referencing a product for consumer delivery.

    The alert carries its own copy of the document's newListing flag and
    priceBeforeDrop: the product doc is rewritten every cycle, so its flag
    flips False one cycle after discovery and its priceBeforeDrop is
    cleared by the next cycle that sees no drop, while an alert may sit in
    a backed-up queue longer than that. The alert's copies are written once
    and cannot rot.

    Args:
        monitor: Monitor that produced the product, used by the
            consumer to resolve the owning collection and thresholds.
        document: Product document supplying the id, pricing for the
            priority calculation, and the stamped newListing flag and
            priceBeforeDrop.
        now: Detection timestamp stored as createdAt.

    Returns:
        The alert document for upsert via $setOnInsert.
    """

    return {
        "monitor": monitor,
        "product_id": document["_id"],
        "priority": calculate_priority(document),
        "claimed": False,
        "attempts": 0,
        "createdAt": now,
        "newListing": document.get("newListing", False),
        "priceBeforeDrop": document.get("priceBeforeDrop"),
    }


def calculate_priority(document: dict) -> int:
    """Derive an alert's priority bucket from its discount percentage.

    Always recalculated from current and original price so the priority
    reflects the prices at alert time rather than any stale stored
    percentage. Missing prices, zero original prices, and price increases
    resolve to priority 0. A product with no MSRP has no percentage to rank
    on and takes NO_MSRP_PRIORITY instead, since a real 0 would leave it
    permanently behind every other alert in the claim query's sort. A new
    listing is floored at NEW_LISTING_PRIORITY for the same reason: it may
    carry any discount at all, and its percentage rank still wins when higher.

    Args:
        document: Product document supplying the prices.

    Returns:
        Priority from 0 to 10, one bucket per 10 percent of discount.
    """

    current_price = document.get("currentPrice")
    original_price = document.get("originalPrice")
    if current_price is None or not original_price:
        return 0

    if not document.get("hasMsrp", True):
        return NO_MSRP_PRIORITY

    percentage = int(((original_price - current_price) / original_price) * 100)
    priority = min(max(percentage // 10, 0), 10)
    if document.get("newListing"):
        return max(priority, NEW_LISTING_PRIORITY)

    return priority


def build_doc_id(
    product_id: str, category_identifier: str, use_composite: bool = False
) -> str:
    """Return the composite id when use_composite is set, else the bare product id.

    Args:
        product_id: Retailer-specific product identifier.
        category_identifier: Category suffix for multi-variant monitors.
        use_composite: Append category_identifier to disambiguate variants.

    Returns:
        The composite id when use_composite is set, otherwise product_id.
    """

    return f"{product_id}_{category_identifier}" if use_composite else product_id


def convert_product_to_mongo(product: Product) -> dict:
    """Flatten a Product into its MongoDB document representation.

    Stamps lastUpdated with the current UTC time. Prices are already floats via
    the Product's pydantic validation (numeric strings coerced, comma strings
    rejected at construction), so no coercion happens here. The _id is assigned
    by the caller via build_doc_id.

    Args:
        product: Product to convert.

    Returns:
        The MongoDB document dict for the product.
    """

    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"]

    return {
        "title": data["title"],
        "url": data["url"],
        "currentPrice": current_price,
        "originalPrice": original_price,
        "imageUrl": metadata["image_url"],
        "customizationIndex": metadata["customization_index"],
        "memberNonMemberComparison": metadata["member_non_member_comparison"],
        "overstockText": metadata["overstock_text"],
        "affiliateUrl": metadata["affiliate_url"],
        "promoCode": metadata["promo_code"],
        "promoPercentage": metadata["promo_percentage"],
        "promoDollarOff": metadata["promo_dollar_off"],
        "promoMinSpend": metadata["promo_min_spend"],
        "asin": metadata["asin"],
        "startDate": timing["start_date"],
        "endDate": timing["end_date"],
        "priceBeforeDrop": pricing["price_before_drop"],
        "eduPrice": pricing["edu_price"],
        "hasMsrp": pricing["has_msrp"],
        "available": availability["available"],
        "availabilityLastChecked": availability["availability_last_checked"],
        "condition": availability["condition"],
        "appExclusive": availability["app_exclusive"],
        "memberPick": availability["member_pick"],
        "memberDeal": availability["member_deal"],
        "lastUpdated": datetime.now(timezone.utc),
    }