spoolman2/spoolman/database/spool.py

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

from datetime import datetime, timezone
from typing import Optional

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import contains_eager, joinedload

from spoolman.database import filament, models
from spoolman.exceptions import ItemCreateError, ItemNotFoundError
from spoolman.math import weight_from_length


async def create(
    *,
    db: AsyncSession,
    filament_id: int,
    remaining_weight: Optional[float] = None,
    used_weight: Optional[float] = None,
    first_used: Optional[datetime] = None,
    last_used: Optional[datetime] = None,
    location: Optional[str] = None,
    lot_nr: Optional[str] = None,
    comment: Optional[str] = None,
) -> models.Spool:
    """Add a new spool to the database. Leave weight empty to assume full spool."""
    filament_item = await filament.get_by_id(db, filament_id)
    if used_weight is None:
        if remaining_weight is not None:
            if filament_item.weight is None:
                raise ItemCreateError("remaining_weight can only be used if the filament type has a weight set.")
            used_weight = max(filament_item.weight - remaining_weight, 0)
        else:
            used_weight = 0

    db_item = models.Spool(
        filament=filament_item,
        used_weight=used_weight,
        first_used=first_used,
        last_used=last_used,
        location=location,
        lot_nr=lot_nr,
        comment=comment,
    )
    db.add(db_item)
    await db.flush()
    return db_item


async def get_by_id(db: AsyncSession, spool_id: int, with_for_update: Optional[bool] = None) -> models.Spool:
    """Get a spool object from the database by the unique ID."""
    spool = await db.get(
        models.Spool,
        spool_id,
        with_for_update=with_for_update,  # type: ignore  # noqa: PGH003
        options=[joinedload("*")],  # Load all nested objects as well
    )
    if spool is None:
        raise ItemNotFoundError(f"No spool with ID {spool_id} found.")
    return spool


async def find(
    *,
    db: AsyncSession,
    filament_name: Optional[int] = None,
    filament_id: Optional[int] = None,
    filament_material: Optional[str] = None,
    vendor_name: Optional[str] = None,
    vendor_id: Optional[int] = None,
    location: Optional[str] = None,
    lot_nr: Optional[str] = None,
) -> list[models.Spool]:
    """Find a list of spool objects by search criteria."""
    stmt = (
        select(models.Spool)
        .join(models.Spool.filament)
        .join(models.Filament.vendor)
        .options(contains_eager(models.Spool.filament).contains_eager(models.Filament.vendor))
    )
    if filament_name is not None:
        stmt = stmt.where(models.Filament.name.ilike(f"%{filament_name}%"))
    if filament_id is not None:
        stmt = stmt.where(models.Spool.filament_id == filament_id)
    if filament_material is not None:
        stmt = stmt.where(models.Filament.material.ilike(f"%{filament_material}%"))
    if vendor_name is not None:
        stmt = stmt.where(models.Vendor.name.ilike(f"%{vendor_name}%"))
    if vendor_id is not None:
        stmt = stmt.where(models.Filament.vendor_id == vendor_id)
    if location is not None:
        stmt = stmt.where(models.Spool.location.ilike(f"%{location}%"))
    if lot_nr is not None:
        stmt = stmt.where(models.Spool.lot_nr.ilike(f"%{lot_nr}%"))

    rows = await db.execute(stmt)
    return list(rows.scalars().all())


async def update(
    *,
    db: AsyncSession,
    spool_id: int,
    data: dict,
) -> models.Spool:
    """Update the fields of a spool object."""
    spool = await get_by_id(db, spool_id)
    for k, v in data.items():
        if k == "filament_id":
            spool.filament = await filament.get_by_id(db, v)
        elif k == "remaining_weight":
            if spool.filament.weight is None:
                raise ItemCreateError("remaining_weight can only be used if the filament type has a weight set.")
            spool.used_weight = max(spool.filament.weight - v, 0)
        else:
            setattr(spool, k, v)
    await db.flush()
    return spool


async def delete(db: AsyncSession, spool_id: int) -> None:
    """Delete a spool object."""
    spool = await get_by_id(db, spool_id)
    await db.delete(spool)


# TODO: Make unit tests for race conditions on these
async def use_weight(db: AsyncSession, spool_id: int, weight: float) -> models.Spool:
    """Consume filament from a spool by weight.

    Increases the used_weight attribute of the spool.
    Updates the first_used and last_used attributes where appropriate.

    Args:
        db (AsyncSession): Database session
        spool_id (int): Spool ID
        weight (float): Filament weight to consume, in grams

    Returns:
        models.Spool: Updated spool object
    """
    spool = await get_by_id(db, spool_id, with_for_update=True)
    spool.used_weight += weight

    if spool.first_used is None:
        spool.first_used = datetime.now(tz=timezone.utc)
    spool.last_used = datetime.now(tz=timezone.utc)

    await db.flush()
    return spool


async def use_length(db: AsyncSession, spool_id: int, length: float) -> models.Spool:
    """Consume filament from a spool by length.

    Increases the used_weight attribute of the spool.
    Updates the first_used and last_used attributes where appropriate.

    Args:
        db (AsyncSession): Database session
        spool_id (int): Spool ID
        length (float): Length of filament to consume, in mm

    Returns:
        models.Spool: Updated spool object
    """
    spool = await get_by_id(db, spool_id, with_for_update=True)

    filament = spool.filament

    spool.used_weight += weight_from_length(
        length=length,
        radius=filament.diameter / 2,
        density=filament.density,
    )

    if spool.first_used is None:
        spool.first_used = datetime.now(tz=timezone.utc)
    spool.last_used = datetime.now(tz=timezone.utc)

    await db.flush()
    return spool