"""Storage location catalog — single write path for spool location fields (#1004).""" from __future__ import annotations import logging import re import time from dataclasses import dataclass import httpx from sqlalchemy import func, select, update from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from backend.app.models.location import Location from backend.app.models.spool import Spool logger = logging.getLogger(__name__) DUPLICATE_LOCATION_NAME = "A location with this name already exists" # AMS residency markers, not storage locations. Bambuddy used to write the slot # a spool was loaded into -- " - AMS A1", the shape # `SpoolmanClient.convert_ams_slot_to_location` still produces -- straight into # Spoolman's `location` field. That writer went away when Storage Location # became a place the user chooses (#1114), but the strings survive on people's # Spoolman spools, and importing them offers a printer slot as somewhere to put # a spool away. A slot is where a spool is loaded, not where it is stored, and # Bambuddy tracks that separately through slot assignments. _AMS_SLOT_LOCATION_RE = re.compile( r"^(?:.+\s-\s)?(?:AMS[- ]HT [A-Z]\d+|AMS [A-Z]\d+|External Spool)$", re.IGNORECASE, ) def is_ams_slot_location(name: str) -> bool: """True when a location string names a printer slot rather than a storage place.""" return bool(_AMS_SLOT_LOCATION_RE.match(name.strip())) def normalize_location_name(name: str) -> str: trimmed = name.strip() if not trimmed: raise ValueError("name must not be empty") return trimmed def location_name_key(name: str) -> str: """Case-insensitive lookup key stored on Location.name_key.""" return normalize_location_name(name).lower() def assign_location_name(location: Location, name: str) -> None: normalized = normalize_location_name(name) location.name = normalized location.name_key = location_name_key(normalized) @dataclass(frozen=True) class SpoolLocationFields: """Canonical spool location state: FK + denormalized string for Spoolman/display.""" location_id: int | None storage_location: str | None async def get_location_by_id(db: AsyncSession, location_id: int) -> Location | None: result = await db.execute(select(Location).where(Location.id == location_id)) return result.scalar_one_or_none() async def get_location_by_name(db: AsyncSession, name: str) -> Location | None: key = location_name_key(name) result = await db.execute(select(Location).where(Location.name_key == key)) return result.scalar_one_or_none() async def get_locations_by_name_keys(db: AsyncSession, keys: set[str]) -> dict[str, Location]: if not keys: return {} result = await db.execute(select(Location).where(Location.name_key.in_(keys))) return {loc.name_key: loc for loc in result.scalars().all()} async def _create_location_or_get_existing(db: AsyncSession, normalized: str) -> Location: """Insert a location row, returning the winner on concurrent name_key collision.""" existing = await get_location_by_name(db, normalized) if existing: return existing location = Location() assign_location_name(location, normalized) try: async with db.begin_nested(): db.add(location) await db.flush() return location except IntegrityError as exc: winner = await get_location_by_name(db, normalized) if winner: return winner raise ValueError(DUPLICATE_LOCATION_NAME) from exc async def _insert_location_if_absent(db: AsyncSession, name: str) -> bool: """Stage a new location row when absent. Returns True when one was added.""" normalized = normalize_location_name(name) if await get_location_by_name(db, normalized): return False location = Location() assign_location_name(location, normalized) try: async with db.begin_nested(): db.add(location) await db.flush() return True except IntegrityError: # Race: another writer inserted the same name between our check and # flush. The row already exists by definition — surface as "not added" # rather than re-raising. Anything else (NULL constraint, FK, check # constraint) would be a programming bug — re-fetch to verify so we # don't silently drop unrelated IntegrityErrors. if await get_location_by_name(db, normalized): return False logger.warning("IntegrityError on insert of location %r without surviving row", normalized) raise async def resolve_location_by_name(db: AsyncSession, name: str, *, create: bool = True) -> Location | None: """Find a location by name (case-insensitive), optionally creating it.""" normalized = normalize_location_name(name) existing = await get_location_by_name(db, normalized) if existing: return existing if not create: return None return await _create_location_or_get_existing(db, normalized) async def resolve_spool_location_fields( db: AsyncSession, *, location_id: int | None = None, storage_location: str | None = None, fields_set: set[str], ) -> SpoolLocationFields | None: """Resolve location_id + storage_location from API input. ``location_id`` wins when both fields appear in ``fields_set``. Returns ``None`` when neither location field was provided. """ if "location_id" in fields_set: if location_id is None: return SpoolLocationFields(location_id=None, storage_location=None) loc = await get_location_by_id(db, location_id) if not loc: raise ValueError(f"Location {location_id} not found") return SpoolLocationFields(location_id=loc.id, storage_location=loc.name) if "storage_location" in fields_set: if not storage_location: return SpoolLocationFields(location_id=None, storage_location=None) loc = await resolve_location_by_name(db, storage_location) if not loc: return SpoolLocationFields(location_id=None, storage_location=None) return SpoolLocationFields(location_id=loc.id, storage_location=loc.name) return None async def prepare_internal_spool_payload(db: AsyncSession, data: dict, fields_set: set[str]) -> dict: """Apply resolved location fields before creating or updating an internal spool.""" payload = dict(data) resolved = await resolve_spool_location_fields( db, location_id=payload.get("location_id"), storage_location=payload.get("storage_location"), fields_set=fields_set, ) if resolved is not None: payload["location_id"] = resolved.location_id payload["storage_location"] = resolved.storage_location return payload async def resolve_spoolman_location_string( db: AsyncSession, *, location_id: int | None = None, storage_location: str | None = None, fields_set: set[str], ) -> tuple[str | None, bool]: """Return (Spoolman location string, changed) for proxy writes.""" resolved = await resolve_spool_location_fields( db, location_id=location_id, storage_location=storage_location, fields_set=fields_set, ) if resolved is None: return None, False return resolved.storage_location, True async def count_internal_spools_at_location(db: AsyncSession, location_id: int) -> int: result = await db.execute( select(func.count()) .select_from(Spool) .where( Spool.location_id == location_id, Spool.archived_at.is_(None), ) ) return int(result.scalar() or 0) async def count_spools_at_location_by_name(db: AsyncSession, name: str) -> int: normalized = name.strip() if not normalized: return 0 result = await db.execute( select(func.count()) .select_from(Spool) .where( Spool.archived_at.is_(None), func.lower(func.trim(Spool.storage_location)) == normalized.lower(), ) ) return int(result.scalar() or 0) async def enrich_spool_dicts_with_location_id(db: AsyncSession, spools: list[dict]) -> None: """Attach location_id to mapped Spoolman-style spool dicts in place.""" keys = {location_name_key(s["storage_location"]) for s in spools if (s.get("storage_location") or "").strip()} if not keys: for s in spools: s["location_id"] = None return by_key = await get_locations_by_name_keys(db, keys) for s in spools: raw = (s.get("storage_location") or "").strip() if not raw: s["location_id"] = None continue loc = by_key.get(location_name_key(raw)) s["location_id"] = loc.id if loc else None async def rename_location(db: AsyncSession, location: Location, new_name: str) -> Location: normalized = normalize_location_name(new_name) existing = await get_location_by_name(db, normalized) if existing and existing.id != location.id: raise ValueError(DUPLICATE_LOCATION_NAME) old_name = location.name # Mirror the SQL TRIM on the Python side so a legacy row whose # `storage_location` has trailing whitespace still matches against the # `old_name` we just lifted off the Location row. Without `.strip()` the # equality is asymmetric (SQL strips the column; Python doesn't) and # legacy rows quietly fall out of the rename cascade. old_name_key = old_name.strip().lower() assign_location_name(location, normalized) await db.execute(update(Spool).where(Spool.location_id == location.id).values(storage_location=normalized)) # Keep legacy rows in sync when only storage_location was set. await db.execute( update(Spool) .where( Spool.location_id.is_(None), func.lower(func.trim(Spool.storage_location)) == old_name_key, ) .values(storage_location=normalized, location_id=location.id) ) try: await db.flush() except IntegrityError as exc: raise ValueError(DUPLICATE_LOCATION_NAME) from exc return location async def sync_locations_from_spoolman(db: AsyncSession, client) -> bool: """Import distinct Spoolman location strings into the local catalog. Returns True when new rows were staged (caller must commit). Logs and returns False on Spoolman fetch failures so the calling read path keeps serving the local catalog instead of 500ing; bare-Exception swallow used to be the shape here and hid both transport errors and shape regressions. """ from backend.app.services.spoolman import SpoolmanClientError, SpoolmanUnavailableError try: names = await client.get_distinct_locations() except (SpoolmanUnavailableError, SpoolmanClientError, httpx.HTTPError) as exc: logger.warning("location sync from Spoolman failed: %s", exc) return False # Collapse case variants before insert — Spoolman may return both # "Drybox 1" and "DRYBOX 1" in the same payload. by_key: dict[str, str] = {} for raw in names: name = (raw or "").strip() if not name: continue if is_ams_slot_location(name): logger.debug("Skipping AMS slot marker %r from the Spoolman location import", name) continue key = location_name_key(name) if key not in by_key: by_key[key] = name changed = False for name in by_key.values(): if await _insert_location_if_absent(db, name): changed = True return changed # Per-URL last-sync timestamp guard. Calling list_spools runs the sync, so on # a polling UI without this guard every refetch round-trips to Spoolman and # opens a write transaction — measurable latency and SQLite write contention. # 60s is long enough to absorb dashboard polling, short enough that a manual # spool rename in Spoolman shows up on the next minute's refresh. _SPOOLMAN_LOCATION_SYNC_TTL_SECONDS = 60.0 _spoolman_location_sync_last_run: dict[str, float] = {} def _spoolman_location_sync_cache_clear() -> None: """Test hook: drop the TTL cache so each test starts from a clean slate.""" _spoolman_location_sync_last_run.clear() async def maybe_sync_spoolman_locations(db: AsyncSession, *, client=None) -> bool: """Sync Spoolman location names into the local catalog when integration is enabled. Pass ``client`` when the caller has already resolved one (the GET /spools route does); otherwise the function falls back to ``init_spoolman_client``. Passing the route's client keeps test fixtures honest — without it, the fall-back path imports from ``backend.app.services.spoolman`` directly and bypasses any patch that targets the route module's alias, which causes real TCP connects to whatever ``spoolman_url`` happens to point at. """ from backend.app.api.routes._spoolman_helpers import assert_safe_spoolman_url from backend.app.models.settings import Settings result = await db.execute(select(Settings)) settings = {s.key: s.value for s in result.scalars().all()} if settings.get("spoolman_enabled", "false").lower() != "true": return False url = settings.get("spoolman_url", "").strip() if not url: return False # Debounce: skip the round-trip when we synced this URL recently. cache_key = url.rstrip("/") last_run = _spoolman_location_sync_last_run.get(cache_key, 0.0) now = time.monotonic() if now - last_run < _SPOOLMAN_LOCATION_SYNC_TTL_SECONDS: return False try: assert_safe_spoolman_url(url) except ValueError as exc: logger.warning("Spoolman URL rejected by SSRF guard during location sync: %s", exc) return False if client is None: from backend.app.services.spoolman import get_spoolman_client, init_spoolman_client client = await get_spoolman_client() if not client or client.base_url != cache_key: client = await init_spoolman_client(url) if not client: return False changed = await sync_locations_from_spoolman(db, client) _spoolman_location_sync_last_run[cache_key] = now return changed