Source code for fsh_lib.ordering

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:
    field: str
    direction: SortDirection
    expr: Any



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]]:
    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