"""Database operations for Location model.""" from typing import Optional from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import selectinload from spoolman.database import models from spoolman.exceptions import ItemNotFoundError async def create( *, db: AsyncSession, name: str, parent_id: Optional[int] = None, location_type: Optional[str] = None, description: Optional[str] = None, ) -> models.Location: """Create a new location.""" # Verify parent exists if specified if parent_id is not None: parent = await db.get(models.Location, parent_id) if parent is None: raise ItemNotFoundError(f"Parent location with ID {parent_id} not found.") db_item = models.Location( name=name, parent_id=parent_id, location_type=location_type, description=description, ) db.add(db_item) await db.commit() return db_item async def get_by_id(db: AsyncSession, location_id: int) -> models.Location: """Get a location by ID.""" stmt = select(models.Location).where(models.Location.id == location_id) result = await db.execute(stmt) db_item = result.scalar_one_or_none() if db_item is None: raise ItemNotFoundError(f"Location with ID {location_id} not found.") return db_item async def find( *, db: AsyncSession, parent_id: Optional[int] = None, ) -> list[models.Location]: """Find locations, optionally filtered by parent.""" stmt = select(models.Location).order_by(models.Location.name) if parent_id is not None: stmt = stmt.where(models.Location.parent_id == parent_id) result = await db.execute(stmt) return list(result.scalars().all()) async def find_all(db: AsyncSession) -> list[models.Location]: """Get all locations.""" stmt = select(models.Location).order_by(models.Location.name) result = await db.execute(stmt) return list(result.scalars().all()) async def update( *, db: AsyncSession, location_id: int, data: dict, ) -> models.Location: """Update a location.""" db_item = await get_by_id(db, location_id) # Verify parent exists if changing parent if "parent_id" in data and data["parent_id"] is not None: # Prevent setting parent to self if data["parent_id"] == location_id: raise ValueError("Location cannot be its own parent.") parent = await db.get(models.Location, data["parent_id"]) if parent is None: raise ItemNotFoundError(f"Parent location with ID {data['parent_id']} not found.") for key, value in data.items(): setattr(db_item, key, value) await db.commit() return db_item async def delete(db: AsyncSession, location_id: int) -> None: """Delete a location.""" db_item = await get_by_id(db, location_id) await db.delete(db_item) await db.commit() async def get_full_path(db: AsyncSession, location_id: int) -> str: """Get the full path of a location as a string (e.g., 'Room > Shelf > Bin').""" parts = [] current_id: Optional[int] = location_id while current_id is not None: location = await get_by_id(db, current_id) parts.insert(0, location.name) current_id = location.parent_id return " > ".join(parts) async def get_tree(db: AsyncSession) -> list[dict]: """Get all locations as a nested tree structure.""" all_locations = await find_all(db) # Build lookup dict by_id = {loc.id: loc for loc in all_locations} # Build tree def build_node(loc: models.Location) -> dict: return { "id": loc.id, "name": loc.name, "location_type": loc.location_type, "description": loc.description, "children": [ build_node(by_id[child.id]) for child in all_locations if child.parent_id == loc.id ], } # Get root nodes (no parent) roots = [loc for loc in all_locations if loc.parent_id is None] return [build_node(loc) for loc in roots] async def get_spool_count(db: AsyncSession, location_id: int) -> int: """Get the number of spools at this location (not including children).""" stmt = ( select(func.count(models.Spool.id)) .where(models.Spool.location_id == location_id) .where((models.Spool.archived == False) | (models.Spool.archived == None)) # noqa: E712 ) result = await db.execute(stmt) return result.scalar() or 0