location_service.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354
  1. """Storage location catalog — single write path for spool location fields (#1004)."""
  2. from __future__ import annotations
  3. import logging
  4. import time
  5. from dataclasses import dataclass
  6. import httpx
  7. from sqlalchemy import func, select, update
  8. from sqlalchemy.exc import IntegrityError
  9. from sqlalchemy.ext.asyncio import AsyncSession
  10. from backend.app.models.location import Location
  11. from backend.app.models.spool import Spool
  12. logger = logging.getLogger(__name__)
  13. DUPLICATE_LOCATION_NAME = "A location with this name already exists"
  14. def normalize_location_name(name: str) -> str:
  15. trimmed = name.strip()
  16. if not trimmed:
  17. raise ValueError("name must not be empty")
  18. return trimmed
  19. def location_name_key(name: str) -> str:
  20. """Case-insensitive lookup key stored on Location.name_key."""
  21. return normalize_location_name(name).lower()
  22. def assign_location_name(location: Location, name: str) -> None:
  23. normalized = normalize_location_name(name)
  24. location.name = normalized
  25. location.name_key = location_name_key(normalized)
  26. @dataclass(frozen=True)
  27. class SpoolLocationFields:
  28. """Canonical spool location state: FK + denormalized string for Spoolman/display."""
  29. location_id: int | None
  30. storage_location: str | None
  31. async def get_location_by_id(db: AsyncSession, location_id: int) -> Location | None:
  32. result = await db.execute(select(Location).where(Location.id == location_id))
  33. return result.scalar_one_or_none()
  34. async def get_location_by_name(db: AsyncSession, name: str) -> Location | None:
  35. key = location_name_key(name)
  36. result = await db.execute(select(Location).where(Location.name_key == key))
  37. return result.scalar_one_or_none()
  38. async def get_locations_by_name_keys(db: AsyncSession, keys: set[str]) -> dict[str, Location]:
  39. if not keys:
  40. return {}
  41. result = await db.execute(select(Location).where(Location.name_key.in_(keys)))
  42. return {loc.name_key: loc for loc in result.scalars().all()}
  43. async def _create_location_or_get_existing(db: AsyncSession, normalized: str) -> Location:
  44. """Insert a location row, returning the winner on concurrent name_key collision."""
  45. existing = await get_location_by_name(db, normalized)
  46. if existing:
  47. return existing
  48. location = Location()
  49. assign_location_name(location, normalized)
  50. try:
  51. async with db.begin_nested():
  52. db.add(location)
  53. await db.flush()
  54. return location
  55. except IntegrityError as exc:
  56. winner = await get_location_by_name(db, normalized)
  57. if winner:
  58. return winner
  59. raise ValueError(DUPLICATE_LOCATION_NAME) from exc
  60. async def _insert_location_if_absent(db: AsyncSession, name: str) -> bool:
  61. """Stage a new location row when absent. Returns True when one was added."""
  62. normalized = normalize_location_name(name)
  63. if await get_location_by_name(db, normalized):
  64. return False
  65. location = Location()
  66. assign_location_name(location, normalized)
  67. try:
  68. async with db.begin_nested():
  69. db.add(location)
  70. await db.flush()
  71. return True
  72. except IntegrityError:
  73. # Race: another writer inserted the same name between our check and
  74. # flush. The row already exists by definition — surface as "not added"
  75. # rather than re-raising. Anything else (NULL constraint, FK, check
  76. # constraint) would be a programming bug — re-fetch to verify so we
  77. # don't silently drop unrelated IntegrityErrors.
  78. if await get_location_by_name(db, normalized):
  79. return False
  80. logger.warning("IntegrityError on insert of location %r without surviving row", normalized)
  81. raise
  82. async def resolve_location_by_name(db: AsyncSession, name: str, *, create: bool = True) -> Location | None:
  83. """Find a location by name (case-insensitive), optionally creating it."""
  84. normalized = normalize_location_name(name)
  85. existing = await get_location_by_name(db, normalized)
  86. if existing:
  87. return existing
  88. if not create:
  89. return None
  90. return await _create_location_or_get_existing(db, normalized)
  91. async def resolve_spool_location_fields(
  92. db: AsyncSession,
  93. *,
  94. location_id: int | None = None,
  95. storage_location: str | None = None,
  96. fields_set: set[str],
  97. ) -> SpoolLocationFields | None:
  98. """Resolve location_id + storage_location from API input.
  99. ``location_id`` wins when both fields appear in ``fields_set``.
  100. Returns ``None`` when neither location field was provided.
  101. """
  102. if "location_id" in fields_set:
  103. if location_id is None:
  104. return SpoolLocationFields(location_id=None, storage_location=None)
  105. loc = await get_location_by_id(db, location_id)
  106. if not loc:
  107. raise ValueError(f"Location {location_id} not found")
  108. return SpoolLocationFields(location_id=loc.id, storage_location=loc.name)
  109. if "storage_location" in fields_set:
  110. if not storage_location:
  111. return SpoolLocationFields(location_id=None, storage_location=None)
  112. loc = await resolve_location_by_name(db, storage_location)
  113. if not loc:
  114. return SpoolLocationFields(location_id=None, storage_location=None)
  115. return SpoolLocationFields(location_id=loc.id, storage_location=loc.name)
  116. return None
  117. async def prepare_internal_spool_payload(db: AsyncSession, data: dict, fields_set: set[str]) -> dict:
  118. """Apply resolved location fields before creating or updating an internal spool."""
  119. payload = dict(data)
  120. resolved = await resolve_spool_location_fields(
  121. db,
  122. location_id=payload.get("location_id"),
  123. storage_location=payload.get("storage_location"),
  124. fields_set=fields_set,
  125. )
  126. if resolved is not None:
  127. payload["location_id"] = resolved.location_id
  128. payload["storage_location"] = resolved.storage_location
  129. return payload
  130. async def resolve_spoolman_location_string(
  131. db: AsyncSession,
  132. *,
  133. location_id: int | None = None,
  134. storage_location: str | None = None,
  135. fields_set: set[str],
  136. ) -> tuple[str | None, bool]:
  137. """Return (Spoolman location string, changed) for proxy writes."""
  138. resolved = await resolve_spool_location_fields(
  139. db,
  140. location_id=location_id,
  141. storage_location=storage_location,
  142. fields_set=fields_set,
  143. )
  144. if resolved is None:
  145. return None, False
  146. return resolved.storage_location, True
  147. async def count_internal_spools_at_location(db: AsyncSession, location_id: int) -> int:
  148. result = await db.execute(
  149. select(func.count())
  150. .select_from(Spool)
  151. .where(
  152. Spool.location_id == location_id,
  153. Spool.archived_at.is_(None),
  154. )
  155. )
  156. return int(result.scalar() or 0)
  157. async def count_spools_at_location_by_name(db: AsyncSession, name: str) -> int:
  158. normalized = name.strip()
  159. if not normalized:
  160. return 0
  161. result = await db.execute(
  162. select(func.count())
  163. .select_from(Spool)
  164. .where(
  165. Spool.archived_at.is_(None),
  166. func.lower(func.trim(Spool.storage_location)) == normalized.lower(),
  167. )
  168. )
  169. return int(result.scalar() or 0)
  170. async def enrich_spool_dicts_with_location_id(db: AsyncSession, spools: list[dict]) -> None:
  171. """Attach location_id to mapped Spoolman-style spool dicts in place."""
  172. keys = {location_name_key(s["storage_location"]) for s in spools if (s.get("storage_location") or "").strip()}
  173. if not keys:
  174. for s in spools:
  175. s["location_id"] = None
  176. return
  177. by_key = await get_locations_by_name_keys(db, keys)
  178. for s in spools:
  179. raw = (s.get("storage_location") or "").strip()
  180. if not raw:
  181. s["location_id"] = None
  182. continue
  183. loc = by_key.get(location_name_key(raw))
  184. s["location_id"] = loc.id if loc else None
  185. async def rename_location(db: AsyncSession, location: Location, new_name: str) -> Location:
  186. normalized = normalize_location_name(new_name)
  187. existing = await get_location_by_name(db, normalized)
  188. if existing and existing.id != location.id:
  189. raise ValueError(DUPLICATE_LOCATION_NAME)
  190. old_name = location.name
  191. # Mirror the SQL TRIM on the Python side so a legacy row whose
  192. # `storage_location` has trailing whitespace still matches against the
  193. # `old_name` we just lifted off the Location row. Without `.strip()` the
  194. # equality is asymmetric (SQL strips the column; Python doesn't) and
  195. # legacy rows quietly fall out of the rename cascade.
  196. old_name_key = old_name.strip().lower()
  197. assign_location_name(location, normalized)
  198. await db.execute(update(Spool).where(Spool.location_id == location.id).values(storage_location=normalized))
  199. # Keep legacy rows in sync when only storage_location was set.
  200. await db.execute(
  201. update(Spool)
  202. .where(
  203. Spool.location_id.is_(None),
  204. func.lower(func.trim(Spool.storage_location)) == old_name_key,
  205. )
  206. .values(storage_location=normalized, location_id=location.id)
  207. )
  208. try:
  209. await db.flush()
  210. except IntegrityError as exc:
  211. raise ValueError(DUPLICATE_LOCATION_NAME) from exc
  212. return location
  213. async def sync_locations_from_spoolman(db: AsyncSession, client) -> bool:
  214. """Import distinct Spoolman location strings into the local catalog.
  215. Returns True when new rows were staged (caller must commit). Logs and
  216. returns False on Spoolman fetch failures so the calling read path keeps
  217. serving the local catalog instead of 500ing; bare-Exception swallow used
  218. to be the shape here and hid both transport errors and shape regressions.
  219. """
  220. from backend.app.services.spoolman import SpoolmanClientError, SpoolmanUnavailableError
  221. try:
  222. names = await client.get_distinct_locations()
  223. except (SpoolmanUnavailableError, SpoolmanClientError, httpx.HTTPError) as exc:
  224. logger.warning("location sync from Spoolman failed: %s", exc)
  225. return False
  226. # Collapse case variants before insert — Spoolman may return both
  227. # "Drybox 1" and "DRYBOX 1" in the same payload.
  228. by_key: dict[str, str] = {}
  229. for raw in names:
  230. name = (raw or "").strip()
  231. if not name:
  232. continue
  233. key = location_name_key(name)
  234. if key not in by_key:
  235. by_key[key] = name
  236. changed = False
  237. for name in by_key.values():
  238. if await _insert_location_if_absent(db, name):
  239. changed = True
  240. return changed
  241. # Per-URL last-sync timestamp guard. Calling list_spools runs the sync, so on
  242. # a polling UI without this guard every refetch round-trips to Spoolman and
  243. # opens a write transaction — measurable latency and SQLite write contention.
  244. # 60s is long enough to absorb dashboard polling, short enough that a manual
  245. # spool rename in Spoolman shows up on the next minute's refresh.
  246. _SPOOLMAN_LOCATION_SYNC_TTL_SECONDS = 60.0
  247. _spoolman_location_sync_last_run: dict[str, float] = {}
  248. def _spoolman_location_sync_cache_clear() -> None:
  249. """Test hook: drop the TTL cache so each test starts from a clean slate."""
  250. _spoolman_location_sync_last_run.clear()
  251. async def maybe_sync_spoolman_locations(db: AsyncSession, *, client=None) -> bool:
  252. """Sync Spoolman location names into the local catalog when integration is enabled.
  253. Pass ``client`` when the caller has already resolved one (the GET /spools
  254. route does); otherwise the function falls back to ``init_spoolman_client``.
  255. Passing the route's client keeps test fixtures honest — without it, the
  256. fall-back path imports from ``backend.app.services.spoolman`` directly and
  257. bypasses any patch that targets the route module's alias, which causes
  258. real TCP connects to whatever ``spoolman_url`` happens to point at.
  259. """
  260. from backend.app.api.routes._spoolman_helpers import assert_safe_spoolman_url
  261. from backend.app.models.settings import Settings
  262. result = await db.execute(select(Settings))
  263. settings = {s.key: s.value for s in result.scalars().all()}
  264. if settings.get("spoolman_enabled", "false").lower() != "true":
  265. return False
  266. url = settings.get("spoolman_url", "").strip()
  267. if not url:
  268. return False
  269. # Debounce: skip the round-trip when we synced this URL recently.
  270. cache_key = url.rstrip("/")
  271. last_run = _spoolman_location_sync_last_run.get(cache_key, 0.0)
  272. now = time.monotonic()
  273. if now - last_run < _SPOOLMAN_LOCATION_SYNC_TTL_SECONDS:
  274. return False
  275. try:
  276. assert_safe_spoolman_url(url)
  277. except ValueError as exc:
  278. logger.warning("Spoolman URL rejected by SSRF guard during location sync: %s", exc)
  279. return False
  280. if client is None:
  281. from backend.app.services.spoolman import get_spoolman_client, init_spoolman_client
  282. client = await get_spoolman_client()
  283. if not client or client.base_url != cache_key:
  284. client = await init_spoolman_client(url)
  285. if not client:
  286. return False
  287. changed = await sync_locations_from_spoolman(db, client)
  288. _spoolman_location_sync_last_run[cache_key] = now
  289. return changed