Source code for fsh_lib.loading

"""Eager-load application for list endpoints.

:func:`apply_eager_loads` is the companion of
:func:`fsh_lib.ordering.apply_ordering`: it turns a representation's
:class:`EagerLoad` plan into ``select(...).options(...)`` loader
options.  The one piece of coordination is ``joined`` -- the set of
relationships ``apply_ordering`` already LEFT-joined for an
``ORDER BY``.  A top-level relationship in that set is loaded with
``contains_eager`` (reuse the existing join) instead of being
fetched a second time by ``selectinload`` / ``joinedload``.

Generated list handlers wire the two together explicitly::

    stmt = select(Task)
    stmt, joined = apply_ordering(stmt=stmt, sort_clauses=body.sort, ...)
    stmt = apply_filters(stmt=stmt, node=body.filter, model=Task)
    stmt = apply_eager_loads(
        stmt, [EagerLoad(Task.project, "joined")], joined=joined
    )
"""

from dataclasses import dataclass
from enum import StrEnum
from typing import TYPE_CHECKING

from sqlalchemy.orm import (
    contains_eager,
    joinedload,
    selectinload,
    subqueryload,
)

if TYPE_CHECKING:
    from collections.abc import Callable, Iterable

    from sqlalchemy import Select
    from sqlalchemy.orm import InstrumentedAttribute, RelationshipProperty
    from sqlalchemy.orm.strategy_options import _AbstractLoad




class EagerStrategy(StrEnum):
    """Eager-loading strategy for an :class:`EagerLoad`.

    The three relationship loader strategies that make sense for a
    representation dump.  Each member's value is the string token
    SQLAlchemy's ``relationship(lazy=...)`` accepts -- a precise
    subset of SQLAlchemy's (private) ``_LazyLoadArgumentType``.

    ``contains_eager`` is deliberately not a member: it isn't a
    per-field choice but is applied automatically by
    :func:`apply_eager_loads` when a sort already joined the
    relationship.
    """

    SELECTIN = "selectin"
    """Separate ``WHERE ... IN`` query -- safe for collections."""
    JOINED = "joined"
    """LEFT JOIN -- right for scalar to-one relationships."""
    SUBQUERY = "subquery"
    """Separate query correlated by subquery."""



# Strategy -> the ``sqlalchemy.orm`` loader function that *starts*
# a loader chain, e.g. ``selectinload(Task.project)``.
_HEAD_LOADER: dict[EagerStrategy, Callable[..., _AbstractLoad]] = {
    EagerStrategy.SELECTIN: selectinload,
    EagerStrategy.JOINED: joinedload,
    EagerStrategy.SUBQUERY: subqueryload,
}
# Strategy -> the ``Load`` *method name* that extends an existing
# loader to a nested relationship, e.g.
# ``head_load.selectinload(Project.owner)``.  A standalone function
# can't chain -- only a method on the parent ``Load`` can.
_CHAIN_METHOD: dict[EagerStrategy, str] = {
    EagerStrategy.SELECTIN: "selectinload",
    EagerStrategy.JOINED: "joinedload",
    EagerStrategy.SUBQUERY: "subqueryload",
}




@dataclass(frozen=True)
class EagerLoad:
    """One relationship to eager-load, with its strategy.

    Attributes:
        attr: The relationship attribute, e.g. ``Task.project``.
        strategy: The :class:`EagerStrategy` to load it with.
            Overridden by ``contains_eager`` for a top-level
            relationship a sort already joined.
        children: Nested eager loads one level deeper, e.g. a
            ``project``'s own ``owner``.  Their strategy is always
            honoured -- ``contains_eager`` reuse only applies to a
            top-level relationship a sort actually joined.

    """

    attr: InstrumentedAttribute
    strategy: EagerStrategy = EagerStrategy.SELECTIN
    children: tuple[EagerLoad, ...] = ()





def apply_eager_loads(
    stmt: Select,
    loads: Iterable[EagerLoad],
    joined: set[RelationshipProperty] | None = None,
) -> Select:
    """Apply eager-load options to a SELECT, reusing ordering's joins.

    Args:
        stmt: The SQLAlchemy SELECT to add loader options to.
        loads: The representation's :class:`EagerLoad` plan.
        joined: Relationships already LEFT-joined by
            :func:`fsh_lib.ordering.apply_ordering` (its second
            return value).  A top-level load whose relationship is
            in this set is loaded with ``contains_eager`` rather
            than re-fetched.  ``None`` means nothing was joined.

    Returns:
        The statement with ``.options(...)`` applied.

    """
    joined = joined or set()

    for load in loads:
        stmt = stmt.options(_loader(load, joined, head=True))

    return stmt



def _loader(
    load: EagerLoad,
    joined: set[RelationshipProperty],
    *,
    head: bool,
) -> _AbstractLoad:
    """Build the loader option for *load* and its nested children.

    A *head* relationship already joined for an ORDER BY is loaded
    with ``contains_eager`` -- it reuses that join's columns
    instead of issuing a second fetch.  Otherwise its configured
    strategy applies.  Children always use their own strategy.
    """
    if head and load.attr.property in joined:
        option = contains_eager(load.attr)

    else:
        option = _HEAD_LOADER[load.strategy](load.attr)

    for child in load.children:
        option = _extend(option, child)

    return option


def _extend(option: _AbstractLoad, child: EagerLoad) -> _AbstractLoad:
    """Chain a nested *child* eager load onto an existing loader."""
    sub: _AbstractLoad = getattr(option, _CHAIN_METHOD[child.strategy])(
        child.attr
    )

    for grandchild in child.children:
        sub = _extend(sub, grandchild)

    return sub