"""Helpers for filtering and sorting extra fields.""" from __future__ import annotations import json from dataclasses import dataclass from typing import TYPE_CHECKING import sqlalchemy from sqlalchemy import Alias, ColumnElement, Select from sqlalchemy.ext.compiler import compiles from sqlalchemy.orm import aliased from sqlalchemy.sql.expression import FunctionElement from spoolman.database import models from spoolman.database.utils import ( LIKE_ESCAPE, SortOrder, escape_like, order_by_clauses, split_datetime_range_filter, ) from spoolman.extra_field_registry import EntityType, ExtraField, ExtraFieldType, get_extra_fields if TYPE_CHECKING: from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm.attributes import InstrumentedAttribute class _JsonArrayFirstElement(FunctionElement): """Cross-database helper: return the first element of a JSON array stored as text.""" name = "json_array_first_element" inherit_cache = True @compiles(_JsonArrayFirstElement, "postgresql") @compiles(_JsonArrayFirstElement, "cockroachdb") def _compile_json_array_first_pg(element: _JsonArrayFirstElement, compiler: object, **kw: object) -> str: # type: ignore[misc] """PostgreSQL/CockroachDB: CAST(value AS JSON)->>0 returns the first element as TEXT.""" (col_expr,) = element.clauses col_sql = compiler.process(col_expr, **kw) # type: ignore[union-attr] return f"(CAST({col_sql} AS JSON)->>0)" @compiles(_JsonArrayFirstElement) def _compile_json_array_first_default(element: _JsonArrayFirstElement, compiler: object, **kw: object) -> str: # type: ignore[misc] """SQLite/MariaDB: json_extract(value, '$[0]') returns the first element as a scalar.""" (col_expr,) = element.clauses col_sql = compiler.process(col_expr, **kw) # type: ignore[union-attr] return f"JSON_EXTRACT({col_sql}, '$[0]')" class _JsonScalarText(FunctionElement): """Cross-database helper: decode a top-level JSON scalar and return it as unquoted TEXT. Used to compare against the *decoded* value rather than reconstructing the exact JSON serialization the client wrote. This decouples filtering from json.dumps/JSON.stringify encoding quirks (non-ASCII escaping, surrounding quotes, etc.). Only valid for JSON string scalars (text/choice/datetime) - booleans and numbers are handled with numeric casts instead, since scalar decoding is not consistent for them across dialects. """ name = "json_scalar_text" inherit_cache = True @compiles(_JsonScalarText, "postgresql") @compiles(_JsonScalarText, "cockroachdb") def _compile_json_scalar_text_pg(element: _JsonScalarText, compiler: object, **kw: object) -> str: # type: ignore[misc] """PostgreSQL/CockroachDB: CAST(value AS JSON) #>> '{}' extracts the root scalar as TEXT.""" (col_expr,) = element.clauses col_sql = compiler.process(col_expr, **kw) # type: ignore[union-attr] return f"(CAST({col_sql} AS JSON) #>> '{{}}')" @compiles(_JsonScalarText, "mysql") def _compile_json_scalar_text_mysql(element: _JsonScalarText, compiler: object, **kw: object) -> str: # type: ignore[misc] """MySQL/MariaDB: JSON_EXTRACT keeps the surrounding quotes, so JSON_UNQUOTE is required.""" (col_expr,) = element.clauses col_sql = compiler.process(col_expr, **kw) # type: ignore[union-attr] return f"JSON_UNQUOTE(JSON_EXTRACT({col_sql}, '$'))" @compiles(_JsonScalarText) def _compile_json_scalar_text_default(element: _JsonScalarText, compiler: object, **kw: object) -> str: # type: ignore[misc] """SQLite: json_extract(value, '$') already returns string scalars unquoted.""" (col_expr,) = element.clauses col_sql = compiler.process(col_expr, **kw) # type: ignore[union-attr] return f"JSON_EXTRACT({col_sql}, '$')" class _JsonArraySecondElement(FunctionElement): """Cross-database helper: return the second element of a JSON array stored as text.""" name = "json_array_second_element" inherit_cache = True @compiles(_JsonArraySecondElement, "postgresql") @compiles(_JsonArraySecondElement, "cockroachdb") def _compile_json_array_second_pg(element: _JsonArraySecondElement, compiler: object, **kw: object) -> str: # type: ignore[misc] """PostgreSQL/CockroachDB: CAST(value AS JSON)->>1 returns the second element as TEXT.""" (col_expr,) = element.clauses col_sql = compiler.process(col_expr, **kw) # type: ignore[union-attr] return f"(CAST({col_sql} AS JSON)->>1)" @compiles(_JsonArraySecondElement) def _compile_json_array_second_default(element: _JsonArraySecondElement, compiler: object, **kw: object) -> str: # type: ignore[misc] """SQLite/MariaDB: json_extract(value, '$[1]') returns the second element as a scalar.""" (col_expr,) = element.clauses col_sql = compiler.process(col_expr, **kw) # type: ignore[union-attr] return f"JSON_EXTRACT({col_sql}, '$[1]')" @dataclass class ExtraFieldJoin: """One extra field pulled onto a query as a plain column, ready to group or order by. Built by :func:`extra_field_join`, applied with :meth:`apply`. The two steps are separate because the value expression has to exist before the SELECT is constructed, while the join can only be added afterwards. """ alias: Alias field_key: str id_column_name: str #: The DB-decoded scalar, i.e. the string the user typed rather than its JSON encoding. value: ColumnElement[str] def apply(self, stmt: Select, base_id_column: InstrumentedAttribute[int]) -> Select: """Outer-join the field's row onto `stmt`. OUTER, so entities with no row for this key survive with a NULL value — the same "unset" bucket a nullable built-in column produces. """ return stmt.join( self.alias, sqlalchemy.and_( self.alias.c[self.id_column_name] == base_id_column, self.alias.c.key == self.field_key, ), isouter=True, ) def extra_field_value_text(column: ColumnElement[str]) -> ColumnElement[str]: """Decode a stored extra-field value to its scalar, as plain TEXT. Compare against this rather than a reconstructed JSON string, so matching doesn't depend on how the value was encoded when it was written (quoting, non-ASCII escaping). """ return _JsonScalarText(column) def extra_field_join(entity_type: EntityType, field_key: str) -> ExtraFieldJoin: """Prepare an extra field to be selected as a column of an entity query. The field table is aliased so it stays a distinct FROM element: the extra-field *filters* reference the un-aliased table inside correlated subqueries, and without the alias those subqueries would correlate to this join instead of standing on their own. """ field_table = _get_field_table_for_entity(entity_type) id_column_name = _get_entity_id_column(field_table).key alias = field_table.__table__.alias() return ExtraFieldJoin( alias=alias, field_key=field_key, id_column_name=id_column_name, value=extra_field_value_text(alias.c.value), ) def _get_field_table_for_entity(entity_type: EntityType) -> type[models.Base]: """Map an entity type to its extra-field table.""" if entity_type == EntityType.spool: return models.SpoolField if entity_type == EntityType.filament: return models.FilamentField if entity_type == EntityType.vendor: return models.VendorField raise ValueError(f"Unknown entity type: {entity_type}") def _get_entity_id_column(field_table: type[models.Base]) -> InstrumentedAttribute[int]: """Map an extra-field table to its owning entity id column.""" if field_table == models.SpoolField: return models.SpoolField.spool_id if field_table == models.FilamentField: return models.FilamentField.filament_id if field_table == models.VendorField: return models.VendorField.vendor_id raise ValueError(f"Unknown field table: {field_table}") def _parse_boolean_filter(value: str) -> bool: """Parse a boolean filter using explicit true/false tokens only.""" normalized = value.strip().lower() if normalized == "true": return True if normalized == "false": return False raise ValueError(f"Invalid boolean filter value: {value!r}") async def apply_extra_field_filters_and_sort( *, db: AsyncSession, stmt: Select, base_obj: type[models.Base], entity_type: EntityType, extra_field_filters: dict[str, str] | None, sort_by: dict[str, SortOrder] | None, ) -> Select: """Apply extra-field filtering and sorting to a query.""" if not extra_field_filters and not (sort_by is not None and any(field.startswith("extra.") for field in sort_by)): return stmt extra_fields = await get_extra_fields(db, entity_type) extra_fields_dict: dict[str, ExtraField] = {field.key: field for field in extra_fields} if extra_field_filters: field_table = _get_field_table_for_entity(entity_type) id_select = sqlalchemy.select(_get_entity_id_column(field_table)) for field_key, value in extra_field_filters.items(): field = extra_fields_dict.get(field_key) if field is None: continue stmt = add_where_clause_extra_field( stmt=stmt, link_column=base_obj.id, id_select=id_select, field_table=field_table, field_key=field_key, field_type=field.field_type, value=value, multi_choice=field.multi_choice if field.field_type == ExtraFieldType.choice else None, ) if sort_by is not None: for field_name, order in sort_by.items(): if not field_name.startswith("extra."): continue field_key = field_name[6:] extra_field = extra_fields_dict.get(field_key) if extra_field is None: continue stmt = add_order_by_extra_field( stmt=stmt, base_obj=base_obj, entity_type=entity_type, field_key=field_key, field_type=extra_field.field_type, order=order, ) return stmt async def apply_spool_related_extra_filters( *, db: AsyncSession, stmt: Select, filament_filters: dict[str, str] | None, vendor_filters: dict[str, str] | None, link_column: InstrumentedAttribute[int] | None = None, ) -> Select: """Filter a spool query by extra fields that live on the spool's filament or its vendor. Spool extra fields are handled by apply_extra_field_filters_and_sort; this handles the two related entities, linking through Spool.filament_id (a spool matches when its filament — or that filament's vendor — matches the extra-field condition). Both conditions are ultimately expressed as a set of filament ids, so `link_column` can name any column holding one. It defaults to Spool.filament_id — the spool query this was written for — and is Filament.id for the query that lists filaments themselves, empty ones included (see find_groups' include_empty). """ link = models.Spool.filament_id if link_column is None else link_column if filament_filters: fields = {f.key: f for f in await get_extra_fields(db, EntityType.filament)} for field_key, value in filament_filters.items(): field = fields.get(field_key) if field is None: continue stmt = add_where_clause_extra_field( stmt=stmt, link_column=link, id_select=sqlalchemy.select(models.FilamentField.filament_id), field_table=models.FilamentField, field_key=field_key, field_type=field.field_type, value=value, multi_choice=field.multi_choice if field.field_type == ExtraFieldType.choice else None, ) if vendor_filters: fields = {f.key: f for f in await get_extra_fields(db, EntityType.vendor)} for field_key, value in vendor_filters.items(): field = fields.get(field_key) if field is None: continue # Map matching vendors back to filament ids: a spool matches when its filament's vendor # matches. Alias Filament so this subquery doesn't collide with the outer query's join. fil = aliased(models.Filament) id_select = sqlalchemy.select(fil.id).join( models.VendorField, models.VendorField.vendor_id == fil.vendor_id, ) stmt = add_where_clause_extra_field( stmt=stmt, link_column=link, id_select=id_select, field_table=models.VendorField, field_key=field_key, field_type=field.field_type, value=value, multi_choice=field.multi_choice if field.field_type == ExtraFieldType.choice else None, ) return stmt def add_where_clause_extra_field( # noqa: C901, PLR0912, PLR0915 stmt: Select, *, link_column: InstrumentedAttribute[int], id_select: Select, field_table: type[models.Base], field_key: str, field_type: ExtraFieldType, value: str, multi_choice: bool | None = None, ) -> Select: """Add a where clause to a select statement for an extra field. `link_column` is the column on the outer query to constrain (e.g. Spool.id, or Spool.filament_id when filtering spools by a filament/vendor extra field). `id_select` is a SELECT of ids in that same space (with any join to `field_table` already applied); this function narrows a copy of it by the field key/value conditions. This indirection is what lets one query filter by extra fields that live on a related entity. """ conditions = [] for value_part in value.split(","): # Empty-string filters follow the existing string-query API semantics. if len(value_part) == 0: empty_conditions = [ field_table.value.is_(None), field_table.value == "null", # A stored empty string is unset too, the same way `location=` matches both NULL # and '' on a built-in string column (add_where_clause_str_opt). Clearing a text # box writes one of these, and without this they would be counted in the "no # value" group by group_by but missing from the listing of that same group. field_table.value == json.dumps(""), ] if field_type == ExtraFieldType.boolean: empty_conditions.append(field_table.value == json.dumps(bool(0))) field_has_empty_value = id_select.where( sqlalchemy.and_(field_table.key == field_key, sqlalchemy.or_(*empty_conditions)) ) # The linked entity either stores an empty value, or has no row for this key at all. conditions.append(link_column.in_(field_has_empty_value)) conditions.append(link_column.not_in(id_select.where(field_table.key == field_key))) continue exact_match = value_part.startswith('"') and value_part.endswith('"') parsed_value = value_part[1:-1] if exact_match else value_part if field_type == ExtraFieldType.text: # Compare against the DB-decoded scalar rather than a reconstructed JSON string, so # matching is independent of how the value was encoded (quote handling, non-ASCII # escaping). Substring search escapes LIKE wildcards to match user input literally. decoded = _JsonScalarText(field_table.value) field_condition = ( decoded == parsed_value if exact_match else decoded.ilike(f"%{escape_like(parsed_value)}%", escape=LIKE_ESCAPE) ) elif field_type == ExtraFieldType.integer: if ":" in parsed_value: min_val_str, max_val_str = parsed_value.split(":", 1) int_conditions = [] try: stored = sqlalchemy.cast(field_table.value, sqlalchemy.Integer) if min_val_str: int_conditions.append(stored >= int(min_val_str)) if max_val_str: int_conditions.append(stored <= int(max_val_str)) except (ValueError, TypeError) as exc: raise ValueError(f"Invalid integer range filter value for '{field_key}': {parsed_value}") from exc if not int_conditions: raise ValueError(f"Invalid integer range filter value for '{field_key}': {parsed_value}") field_condition = sqlalchemy.and_(*int_conditions) else: try: field_condition = field_table.value == json.dumps(int(parsed_value)) except ValueError as exc: raise ValueError(f"Invalid integer filter value for '{field_key}': {parsed_value}") from exc elif field_type == ExtraFieldType.float: if ":" in parsed_value: min_val_str, max_val_str = parsed_value.split(":", 1) float_conditions = [] try: stored = sqlalchemy.cast(field_table.value, sqlalchemy.Float) if min_val_str: float_conditions.append(stored >= float(min_val_str)) if max_val_str: float_conditions.append(stored <= float(max_val_str)) except (ValueError, TypeError) as exc: raise ValueError(f"Invalid float range filter value for '{field_key}': {parsed_value}") from exc if not float_conditions: raise ValueError(f"Invalid float range filter value for '{field_key}': {parsed_value}") field_condition = sqlalchemy.and_(*float_conditions) else: try: # Compare numerically rather than by JSON string so that int-typed storage # (e.g. "2") and non-canonical decimals (e.g. "2.50") match a "2.0" filter. field_condition = sqlalchemy.cast(field_table.value, sqlalchemy.Float) == float(parsed_value) except ValueError as exc: raise ValueError(f"Invalid float filter value for '{field_key}': {parsed_value}") from exc elif field_type == ExtraFieldType.boolean: field_condition = field_table.value == json.dumps(_parse_boolean_filter(parsed_value)) elif field_type == ExtraFieldType.choice: if multi_choice: # Multi-choice is stored as a JSON array; match the JSON-encoded token as a # substring. json.dumps gives the exact quoted, escaped form the array element # is stored as, and LIKE wildcards in the token are escaped. token = json.dumps(parsed_value, ensure_ascii=False) field_condition = field_table.value.like(f"%{escape_like(token)}%", escape=LIKE_ESCAPE) else: # Compare against the DB-decoded scalar, independent of JSON encoding. field_condition = _JsonScalarText(field_table.value) == parsed_value elif field_type == ExtraFieldType.datetime: # Compare decoded ISO-8601 strings. Both bounds and stored values are the frontend's # canonical toISOString() output, so lexicographic comparison is chronological. That # assumption is what makes this differ from the built-in datetime columns, which parse # the bound and compare real datetimes: here an equivalent spelling of the same instant # (a UTC offset, or a missing '.000') compares as a different string. # # The *grammar* is shared even though the comparison is not, so the split comes from # the same helper the built-in columns use (see add_where_clause_datetime_opt). decoded = _JsonScalarText(field_table.value) ends = split_datetime_range_filter(parsed_value, field_key) if ends is None: field_condition = decoded == parsed_value else: start_str, end_str = ends dt_conditions = [] if start_str: dt_conditions.append(decoded >= start_str) if end_str: dt_conditions.append(decoded <= end_str) field_condition = sqlalchemy.and_(*dt_conditions) elif field_type in (ExtraFieldType.integer_range, ExtraFieldType.float_range): if ":" not in parsed_value: raise ValueError( f"Invalid range filter value for '{field_key}': {parsed_value}. Expected ':'." ) min_val_str, max_val_str = parsed_value.split(":", 1) converter = int if field_type == ExtraFieldType.integer_range else float range_conditions = [] try: cast_type = sqlalchemy.Integer if field_type == ExtraFieldType.integer_range else sqlalchemy.Float if min_val_str: # stored_min >= filter_min: the range starts at or after the requested minimum. stored_min = sqlalchemy.cast(_JsonArrayFirstElement(field_table.value), cast_type) range_conditions.append(stored_min >= converter(min_val_str)) if max_val_str: # stored_max <= filter_max: the range ends at or before the requested maximum. stored_max = sqlalchemy.cast(_JsonArraySecondElement(field_table.value), cast_type) range_conditions.append(stored_max <= converter(max_val_str)) except (ValueError, TypeError) as exc: range_kind = "integer" if field_type == ExtraFieldType.integer_range else "float" raise ValueError(f"Invalid {range_kind} range filter value for '{field_key}': {parsed_value}") from exc if not range_conditions: raise ValueError( f"Invalid range filter value for '{field_key}': {parsed_value}. Expected ':'." ) field_condition = sqlalchemy.and_(*range_conditions) else: raise ValueError(f"Unsupported extra field type for '{field_key}': {field_type}") matching_entities = id_select.where(sqlalchemy.and_(field_table.key == field_key, field_condition)) conditions.append(link_column.in_(matching_entities)) if not conditions: return stmt return stmt.where(sqlalchemy.or_(*conditions)) async def find_extra_field_values( *, db: AsyncSession, entity_type: EntityType, field_key: str, ) -> list[str]: """Find all distinct values currently stored for a scalar extra field. Intended for text/single-choice/datetime fields, whose values are stored as JSON string scalars. Values are DB-decoded (see _JsonScalarText) so the result is independent of JSON encoding, mirroring the built-in distinct-value endpoints (materials, locations, ...). Empty and null values are omitted; the result is sorted for a stable option order. """ field_table = _get_field_table_for_entity(entity_type) decoded = _JsonScalarText(field_table.value) stmt = sqlalchemy.select(decoded).where(field_table.key == field_key).distinct() rows = await db.execute(stmt) values = {row[0] for row in rows.all() if row[0] is not None and row[0] != ""} return sorted(values) def add_order_by_extra_field( stmt: Select, base_obj: type[models.Base], entity_type: EntityType, field_key: str, field_type: ExtraFieldType, order: SortOrder, ) -> Select: """Add an order-by clause to a select statement for an extra field.""" field_table = _get_field_table_for_entity(entity_type) entity_id_column = _get_entity_id_column(field_table) value_subq = ( sqlalchemy.select(field_table.value) .where( sqlalchemy.and_( field_table.key == field_key, entity_id_column == base_obj.id, ) ) .scalar_subquery() .correlate(base_obj) ) if field_type == ExtraFieldType.integer: sort_expr = sqlalchemy.cast(value_subq, sqlalchemy.Integer) elif field_type == ExtraFieldType.float: sort_expr = sqlalchemy.cast(value_subq, sqlalchemy.Float) elif field_type in (ExtraFieldType.integer_range, ExtraFieldType.float_range): cast_type = sqlalchemy.Integer if field_type == ExtraFieldType.integer_range else sqlalchemy.Float # Use dialect-specific JSON first-element extraction, then cast to numeric. sort_expr = sqlalchemy.cast(_JsonArrayFirstElement(value_subq), cast_type) else: sort_expr = value_subq # A spool that has no value for the field yields NULL here, and "not filled in" belongs at # the bottom of the list in both directions -- see order_by_clauses. return stmt.order_by(*order_by_clauses([sort_expr], order))