Spoolman/spoolman/database/vendor.py

"""Helper functions for interacting with vendor database objects."""

import logging
from datetime import datetime

import sqlalchemy
from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession

from spoolman.api.v1.models import EventType, Vendor, VendorEvent
from spoolman.database import models
from spoolman.database.extra_field_query import apply_extra_field_filters_and_sort
from spoolman.database.utils import (
    SortOrder,
    add_where_clause_str,
    add_where_clause_str_opt,
    parse_nested_field,
)
from spoolman.exceptions import ItemNotFoundError
from spoolman.extra_field_registry import EntityType
from spoolman.ws import websocket_manager

logger = logging.getLogger(__name__)


async def create(
    *,
    db: AsyncSession,
    name: str | None = None,
    comment: str | None = None,
    empty_spool_weight: float | None = None,
    external_id: str | None = None,
    extra: dict[str, str | None] | None = None,
) -> models.Vendor:
    """Add a new vendor to the database."""
    vendor = models.Vendor(
        name=name,
        registered=datetime.utcnow().replace(microsecond=0),
        comment=comment,
        empty_spool_weight=empty_spool_weight,
        external_id=external_id,
        extra=[models.VendorField(key=k, value=v) for k, v in (extra or {}).items() if v is not None],
    )
    db.add(vendor)
    await db.commit()
    await vendor_changed(vendor, EventType.ADDED)
    return vendor


async def get_by_id(db: AsyncSession, vendor_id: int) -> models.Vendor:
    """Get a vendor object from the database by the unique ID."""
    vendor = await db.get(models.Vendor, vendor_id)
    if vendor is None:
        raise ItemNotFoundError(f"No vendor with ID {vendor_id} found.")
    return vendor


async def find(
    *,
    db: AsyncSession,
    name: str | None = None,
    external_id: str | None = None,
    extra_field_filters: dict[str, str] | None = None,
    sort_by: dict[str, SortOrder] | None = None,
    limit: int | None = None,
    offset: int = 0,
) -> tuple[list[models.Vendor], int]:
    """Find a list of vendor objects by search criteria.

    Returns a tuple containing the list of items and the total count of matching items.
    """
    stmt = select(models.Vendor)

    stmt = add_where_clause_str(stmt, models.Vendor.name, name)
    stmt = add_where_clause_str_opt(stmt, models.Vendor.external_id, external_id)

    total_count = None

    stmt = await apply_extra_field_filters_and_sort(
        db=db,
        stmt=stmt,
        base_obj=models.Vendor,
        entity_type=EntityType.vendor,
        extra_field_filters=extra_field_filters,
        sort_by=sort_by,
    )

    if sort_by is not None:
        for fieldstr, order in sort_by.items():
            # Check if this is a custom field sort
            if fieldstr.startswith("extra."):
                continue

            field = parse_nested_field(models.Vendor, fieldstr)
            if order == SortOrder.ASC:
                stmt = stmt.order_by(field.asc())
            elif order == SortOrder.DESC:
                stmt = stmt.order_by(field.desc())

    if limit is not None:
        total_count_stmt = stmt.with_only_columns(func.count(), maintain_column_froms=True).order_by(None)
        total_count = (await db.execute(total_count_stmt)).scalar()
        stmt = stmt.offset(offset).limit(limit)

    rows = await db.execute(
        stmt,
        execution_options={"populate_existing": True},
    )
    result = list(rows.unique().scalars().all())
    if total_count is None:
        total_count = len(result)

    return result, total_count


async def update(
    *,
    db: AsyncSession,
    vendor_id: int,
    data: dict,
) -> models.Vendor:
    """Update the fields of a vendor object."""
    vendor = await get_by_id(db, vendor_id)
    for k, v in data.items():
        if k == "extra":
            # Merged per key, the same as a spool's and a filament's: only the keys present
            # in the patch are touched, and a null value means the vendor has no value for
            # that field, so its row is dropped and not re-added — that is how a value that
            # has been set gets cleared.
            vendor.extra = [f for f in vendor.extra if f.key not in v]
            vendor.extra.extend([models.VendorField(key=k, value=v) for k, v in v.items() if v is not None])
        else:
            setattr(vendor, k, v)
    await db.commit()
    await vendor_changed(vendor, EventType.UPDATED)
    return vendor


async def delete(db: AsyncSession, vendor_id: int) -> None:
    """Delete a vendor object."""
    vendor = await get_by_id(db, vendor_id)
    await db.delete(vendor)
    # Commit before notifying so the deletion is durable and visible to subsequent
    # requests; post-commit notification must be the last, infallible step.
    await db.commit()
    await vendor_changed(vendor, EventType.DELETED)


async def clear_extra_field(db: AsyncSession, key: str) -> None:
    """Delete all extra fields with a specific key."""
    await db.execute(
        sqlalchemy.delete(models.VendorField).where(models.VendorField.key == key),
    )
    await db.commit()


async def vendor_changed(vendor: models.Vendor, typ: EventType) -> None:
    """Notify websocket clients that a vendor has changed."""
    try:
        await websocket_manager.send(
            ("vendor", str(vendor.id)),
            VendorEvent(
                type=typ,
                resource="vendor",
                date=datetime.utcnow(),
                payload=Vendor.from_db(vendor),
            ),
        )
    except Exception:
        # Important to have a catch-all here since we don't want to stop the call if this fails.
        logger.exception("Failed to send websocket message")