Source code for fsh_lib.ordering

"""ORDER BY application from typed Pydantic sort clauses."""

from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, cast

from fsh_lib._relationships import resolve_sort_target, sort_expr

if TYPE_CHECKING:
    from collections.abc import Sequence

    from pydantic import BaseModel
    from sqlalchemy import Select
    from sqlalchemy.orm import RelationshipProperty

    from fsh_lib._relationships import SortDirection

__all__ = ["OrderKey", "apply_ordering"]




@dataclass(frozen=True)
class OrderKey:
    """One resolved ORDER BY key, in statement order.

    Carries the SQL expression actually sorted on (for a
    relationship sort that's the related table's representative
    column, not an attribute of the listed model) so keyset
    pagination can build a continuation predicate and select the
    boundary values without re-resolving the sort field.
    """

    field: str
    direction: SortDirection
    expr: Any
    """SQLAlchemy column element the ORDER BY uses."""





def apply_ordering(
    stmt: Select,
    sort_clauses: Sequence[BaseModel] | None,
    model: type,
    default_field: str,
    default_dir: SortDirection = "asc",
) -> tuple[Select, set[RelationshipProperty], list[OrderKey]]:
    """Apply one or more sort clauses to a SELECT statement.

    Each clause is a Pydantic model with ``field`` (an enum
    whose ``.value`` is the sort field) and ``dir``
    (:data:`~fsh_lib._relationships.SortDirection`).

    A sort field naming a column sorts by that column.  A field
    naming a (to-one) *relationship* -- ``"vendor"`` on a model
    with a ``vendor`` many-to-one -- is intuited as the related
    row's representative column (``vendor.name``, falling back to
    its primary key), with the related table LEFT-joined in so
    rows with a null relationship still appear.  The sort field
    stays the bare relationship name end to end -- callers never
    see or send ``"vendor.name"``.

    When *sort_clauses* is ``None`` or empty, the default
    field and direction are used.

    Args:
        stmt: The SQLAlchemy SELECT statement to sort.
        sort_clauses: List of sort clause models, or ``None``.
        model: The SQLAlchemy model class providing columns.
        default_field: Field to sort by when no clauses are
            provided.
        default_dir: Direction for the default sort.

    Returns:
        ``(stmt, joined, keys)`` -- the statement with ORDER BY
        applied, the set of relationships LEFT-joined to resolve a
        relationship sort field, and the resolved
        :class:`OrderKey` list in ORDER BY order.  Hand the
        *joined* set to :func:`fsh_lib.loading.apply_eager_loads`
        so an eager load of one of those relationships reuses the
        join (``contains_eager``) instead of fetching it again;
        hand *keys* to
        :func:`fsh_lib.pagination.run_keyset_query` so the keyset
        continuation predicate matches the applied order.

    """
    # Relationships already joined -- a second sort clause through
    # the same relationship must not join it twice.
    joined: set[RelationshipProperty] = set()
    keys: list[OrderKey] = []

    def order_by_field(
        stmt: Select,
        field: str,
        direction: SortDirection,
    ) -> Select:
        col, join_attr = resolve_sort_target(model, field)

        if join_attr is not None:
            rel = cast("RelationshipProperty", join_attr.property)

            if rel not in joined:
                joined.add(rel)
                stmt = stmt.join(join_attr, isouter=True)

        keys.append(OrderKey(field=field, direction=direction, expr=col))
        return stmt.order_by(sort_expr(col, direction))

    if not sort_clauses:
        return (
            order_by_field(stmt, default_field, default_dir),
            joined,
            keys,
        )

    for clause in sort_clauses:
        field_enum = getattr(clause, "field", None)

        if field_enum is None:
            continue

        direction: SortDirection = getattr(clause, "direction", "asc")
        stmt = order_by_field(stmt, field_enum.value, direction)

    return stmt, joined, keys