"""Utility functions for the database module.""" from collections.abc import Sequence from enum import Enum from typing import Any, TypeVar import sqlalchemy from sqlalchemy import Select from sqlalchemy.orm import attributes from spoolman.database import models class SortOrder(Enum): ASC = 1 DESC = 2 def parse_sort(sort: str | None) -> dict[str, SortOrder]: """Parse the ``sort`` query parameter into field/direction pairs. Shared by the spool, filament and vendor endpoints, which each used to do this inline with a bare ``split(":")`` -- so ``?sort=name`` raised "not enough values to unpack" and ``?sort=name:sideways`` a KeyError, both surfacing as a 500. Args: sort: The raw parameter, e.g. ``name:asc,filament.material:desc``. May be None. Returns: dict[str, SortOrder]: Field name to direction, empty if nothing was requested. Raises: ValueError: If an entry has no direction, or a direction that is not asc/desc. """ sort_by: dict[str, SortOrder] = {} if sort is None: return sort_by for sort_item in sort.split(","): item = sort_item.strip() if not item: continue field, separator, direction = item.partition(":") if not separator: raise ValueError(f"Invalid sort '{item}', expected the form 'field:asc' or 'field:desc'.") if not field: raise ValueError(f"Invalid sort '{item}', no field name was given.") try: sort_by[field] = SortOrder[direction.strip().upper()] except KeyError: raise ValueError( f"Invalid sort direction '{direction}' for field '{field}', expected 'asc' or 'desc'.", ) from None return sort_by def parse_nested_field(base_obj: type[models.Base], field: str) -> attributes.InstrumentedAttribute[Any]: """Parse a nested field string into a sqlalchemy field object.""" fields = field.split(".") if fields[0] == "filament" and hasattr(base_obj, "filament"): if len(fields) == 1: raise ValueError("No field specified for filament") return parse_nested_field(models.Filament, ".".join(fields[1:])) if fields[0] == "vendor" and hasattr(base_obj, "vendor"): if len(fields) == 1: raise ValueError("No field specified for vendor") return parse_nested_field(models.Vendor, ".".join(fields[1:])) # Only mapped columns, not any attribute that happens to exist. `hasattr` also accepted # `metadata`, `registry` and relationships, which then reached order_by() and blew up there -- # `?sort=metadata:asc` raised "'MetaData' object has no attribute 'asc'" as a 500. if fields[0] not in sqlalchemy.inspect(base_obj).columns: raise ValueError(f"Invalid field name '{field}', '{fields[0]}' is not a valid field on '{base_obj.__name__}'.") if len(fields) > 1: raise ValueError(f"Field '{fields[0]}' does not have any nested fields") return getattr(base_obj, fields[0]) def add_where_clause_str_opt( stmt: Select, field: attributes.InstrumentedAttribute[str | None], value: str | None, ) -> Select: """Add a where clause to a select statement for an optional string field.""" if value is not None: conditions = [] for value_part in value.split(","): # If part is empty, search for empty fields if len(value_part) == 0: conditions.append(field.is_(None)) conditions.append(field == "") # Do exact match if value_part is surrounded by quotes elif value_part[0] == '"' and value_part[-1] == '"': conditions.append(field == value_part[1:-1]) # Do fuzzy match if value_part is not surrounded by quotes else: conditions.append(field.ilike(f"%{value_part}%")) stmt = stmt.where(sqlalchemy.or_(*conditions)) return stmt def add_where_clause_str( stmt: Select, field: attributes.InstrumentedAttribute[str], value: str | None, ) -> Select: """Add a where clause to a select statement for a string field.""" if value is not None: conditions = [] for value_part in value.split(","): # If part is empty, search for empty fields if len(value_part) == 0: conditions.append(field == "") # Do exact match if value_part is surrounded by quotes elif value_part[0] == '"' and value_part[-1] == '"': conditions.append(field == value_part[1:-1]) # Do fuzzy match if value_part is not surrounded by quotes else: conditions.append(field.ilike(f"%{value_part}%")) stmt = stmt.where(sqlalchemy.or_(*conditions)) return stmt def add_where_clause_int( stmt: Select, field: attributes.InstrumentedAttribute[int], value: int | Sequence[int] | None, ) -> Select: """Add a where clause to a select statement for a field.""" if value is not None: if isinstance(value, int): value = [value] stmt = stmt.where(field.in_(value)) return stmt def add_where_clause_int_opt( stmt: Select, field: attributes.InstrumentedAttribute[int | None], value: int | Sequence[int] | None, ) -> Select: """Add a where clause to a select statement for a field.""" if value is not None: if isinstance(value, int): value = [value] statements = [] for value_part in value: if value_part == -1: statements.append(field.is_(None)) else: statements.append(field == value_part) stmt = stmt.where(sqlalchemy.or_(*statements)) return stmt T = TypeVar("T") def add_where_clause_int_in( stmt: Select, field: attributes.InstrumentedAttribute[T], value: Sequence[T] | None, ) -> Select: """Add a where clause to a select statement for a field.""" if value is not None: stmt = stmt.where(field.in_(value)) return stmt