skeleton-api/app/crud/augmented_part.py

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

from app.models.associations import augmented_part_imported_parts
from app.models.augmented_part import AugmentedPart
from app.models.imported_part import ImportedPart

_RELATIONSHIPS = ["colour", "manufacturer", "scope_model", "part_type", "sub_part_type"]


async def get(db: AsyncSession, augmented_part_id: int) -> AugmentedPart | None:
    return await db.get(AugmentedPart, augmented_part_id)


async def list_all(
    db: AsyncSession, skip: int = 0, limit: int = 100
) -> list[AugmentedPart]:
    result = await db.execute(
        select(AugmentedPart)
        .order_by(AugmentedPart.description)
        .offset(skip)
        .limit(limit)
    )
    return list(result.scalars().all())


async def list_filtered(
    db: AsyncSession,
    description: str = "",
    colour_id: int | None = None,
    manufacturer_id: int | None = None,
    scope_model_id: int | None = None,
    part_type_id: int | None = None,
    sub_part_type_id: int | None = None,
    skip: int = 0,
    limit: int = 100,
) -> list[AugmentedPart]:
    q = select(AugmentedPart)
    if description.strip():
        q = q.where(AugmentedPart.description.ilike(f"%{description.strip()}%"))
    if colour_id is not None:
        q = q.where(AugmentedPart.colour_id == colour_id)
    if manufacturer_id is not None:
        q = q.where(AugmentedPart.manufacturer_id == manufacturer_id)
    if scope_model_id is not None:
        q = q.where(AugmentedPart.scope_model_id == scope_model_id)
    if part_type_id is not None:
        q = q.where(AugmentedPart.part_type_id == part_type_id)
    if sub_part_type_id is not None:
        q = q.where(AugmentedPart.sub_part_type_id == sub_part_type_id)
    result = await db.execute(
        q.order_by(AugmentedPart.description).offset(skip).limit(limit)
    )
    return list(result.scalars().all())


async def count_filtered(
    db: AsyncSession,
    description: str = "",
    colour_id: int | None = None,
    manufacturer_id: int | None = None,
    scope_model_id: int | None = None,
    part_type_id: int | None = None,
    sub_part_type_id: int | None = None,
) -> int:
    from sqlalchemy import func

    q = select(func.count()).select_from(AugmentedPart)
    if description.strip():
        q = q.where(AugmentedPart.description.ilike(f"%{description.strip()}%"))
    if colour_id is not None:
        q = q.where(AugmentedPart.colour_id == colour_id)
    if manufacturer_id is not None:
        q = q.where(AugmentedPart.manufacturer_id == manufacturer_id)
    if scope_model_id is not None:
        q = q.where(AugmentedPart.scope_model_id == scope_model_id)
    if part_type_id is not None:
        q = q.where(AugmentedPart.part_type_id == part_type_id)
    if sub_part_type_id is not None:
        q = q.where(AugmentedPart.sub_part_type_id == sub_part_type_id)
    result = await db.execute(q)
    return result.scalar_one()


async def get_unlinked_imported_parts_filtered(
    db: AsyncSession,
    description: str = "",
    colour: str = "",
    skip: int = 0,
    limit: int = 50,
) -> list[ImportedPart]:
    q = select(ImportedPart).where(
        ImportedPart.InventoryID.not_in(
            select(augmented_part_imported_parts.c.imported_part_id)
        )
    )
    if description.strip():
        q = q.where(ImportedPart.Description.ilike(f"%{description.strip()}%"))
    if colour.strip():
        q = q.where(ImportedPart.TBXColour.ilike(f"%{colour.strip()}%"))
    result = await db.execute(
        q.order_by(ImportedPart.InventoryID).offset(skip).limit(limit)
    )
    return list(result.scalars().all())


async def count_unlinked_filtered(
    db: AsyncSession,
    description: str = "",
    colour: str = "",
) -> int:
    from sqlalchemy import func

    q = (
        select(func.count())
        .select_from(ImportedPart)
        .where(
            ImportedPart.InventoryID.not_in(
                select(augmented_part_imported_parts.c.imported_part_id)
            )
        )
    )
    if description.strip():
        q = q.where(ImportedPart.Description.ilike(f"%{description.strip()}%"))
    if colour.strip():
        q = q.where(ImportedPart.TBXColour.ilike(f"%{colour.strip()}%"))
    result = await db.execute(q)
    return result.scalar_one()


async def distinct_unlinked_colours(db: AsyncSession) -> list[str]:
    q = (
        select(distinct(ImportedPart.TBXColour))
        .where(
            ImportedPart.TBXColour.isnot(None),
            ImportedPart.InventoryID.not_in(
                select(augmented_part_imported_parts.c.imported_part_id)
            ),
        )
        .order_by(ImportedPart.TBXColour)
    )
    result = await db.execute(q)
    return [row[0] for row in result.all()]


async def create(
    db: AsyncSession,
    *,
    sku: str | None,
    description: str,
    colour_id: int,
    manufacturer_id: int,
    scope_model_id: int | None,
    part_type_id: int,
    sub_part_type_id: int | None,
) -> AugmentedPart:
    ap = AugmentedPart(
        sku=sku,
        description=description,
        colour_id=colour_id,
        manufacturer_id=manufacturer_id,
        scope_model_id=scope_model_id,
        part_type_id=part_type_id,
        sub_part_type_id=sub_part_type_id,
    )
    db.add(ap)
    await db.commit()
    await db.refresh(ap, attribute_names=_RELATIONSHIPS)
    return ap


async def update(db: AsyncSession, ap: AugmentedPart, data: dict) -> AugmentedPart:
    for field, value in data.items():
        setattr(ap, field, value)
    await db.commit()
    await db.refresh(ap, attribute_names=_RELATIONSHIPS)
    return ap


async def get_unlinked_imported_parts(db: AsyncSession) -> list[ImportedPart]:
    """ImportedParts that have no corresponding AugmentedPart link."""
    result = await db.execute(
        select(ImportedPart)
        .where(
            ImportedPart.InventoryID.not_in(
                select(augmented_part_imported_parts.c.imported_part_id)
            )
        )
        .order_by(ImportedPart.InventoryID)
    )
    return list(result.scalars().all())


async def delete(db: AsyncSession, ap: AugmentedPart) -> None:
    await db.delete(ap)
    await db.commit()