fastapi-toolsets/src/fastapi_toolsets/fixtures/utils.py

"""Fixture loading utilities for database seeding."""

from collections.abc import Iterator
from enum import Enum
from typing import Any, cast

from sqlalchemy import Table, select
from sqlalchemy import inspect as sa_inspect
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import DeclarativeBase, selectinload
from sqlalchemy.orm.interfaces import ExecutableOption, ORMOption

from ..db import transaction
from ..logger import get_logger
from .enum import LoadStrategy
from .registry import FixtureRegistry, _normalize_contexts

logger = get_logger()


def _instance_to_dict(instance: DeclarativeBase) -> dict[str, Any]:
    """Extract column values from a model instance, skipping unset server-default columns."""
    state = sa_inspect(instance)
    state_dict = state.dict
    result: dict[str, Any] = {}
    for prop in state.mapper.column_attrs:
        if prop.key not in state_dict:
            continue
        val = state_dict[prop.key]
        if val is None:
            col = prop.columns[0]

            if (
                col.server_default is not None
                or (col.default is not None and col.default.is_callable)
                or col.autoincrement is True
            ):
                continue
        result[prop.key] = val
    return result


def _get_table_chain(model_cls: type[DeclarativeBase]) -> list[type[DeclarativeBase]]:
    """Return [root, ..., model_cls] for joined-table inheritance, or [model_cls]."""
    chain: list[type[DeclarativeBase]] = []
    current = sa_inspect(model_cls)
    while current is not None:
        chain.append(current.class_)
        current = current.inherits
    chain.reverse()
    seen: set[int] = set()
    result: list[type[DeclarativeBase]] = []
    for cls in chain:
        tid = id(cls.__table__)
        if tid not in seen:  # pragma: no branch
            seen.add(tid)
            result.append(cls)
    return result


def _instance_to_dict_for_cls(
    instance: DeclarativeBase, cls: type[DeclarativeBase]
) -> dict[str, Any]:
    """Like _instance_to_dict but limited to columns belonging to cls's own table."""
    own_cols = {col.key for col in cls.__table__.columns}
    return {k: v for k, v in _instance_to_dict(instance).items() if k in own_cols}


def _group_by_type(
    instances: list[DeclarativeBase],
) -> list[tuple[type[DeclarativeBase], list[DeclarativeBase]]]:
    """Group instances by their concrete model class, preserving insertion order."""
    groups: dict[type[DeclarativeBase], list[DeclarativeBase]] = {}
    for instance in instances:
        groups.setdefault(type(instance), []).append(instance)
    return list(groups.items())


def _group_by_column_set(
    dicts: list[dict[str, Any]],
    instances: list[DeclarativeBase],
) -> list[tuple[list[dict[str, Any]], list[DeclarativeBase]]]:
    """Group (dict, instance) pairs by their dict key sets."""
    groups: dict[
        frozenset[str], tuple[list[dict[str, Any]], list[DeclarativeBase]]
    ] = {}
    for d, inst in zip(dicts, instances):
        key = frozenset(d)
        if key not in groups:
            groups[key] = ([], [])
        groups[key][0].append(d)
        groups[key][1].append(inst)
    return list(groups.values())


def _grouped_table_dicts(
    model_cls: type[DeclarativeBase], instances: list[DeclarativeBase]
) -> Iterator[
    tuple[type[DeclarativeBase], list[dict[str, Any]], list[DeclarativeBase]]
]:
    """Yield (cls, group_dicts, group_instances) per table in the inheritance
    chain and per column-set group, skipping empty groups.
    """
    for cls in _get_table_chain(model_cls):
        dicts = [_instance_to_dict_for_cls(i, cls) for i in instances]
        for group_dicts, group_instances in _group_by_column_set(dicts, instances):
            if group_dicts and group_dicts[0]:  # pragma: no branch
                yield cls, group_dicts, group_instances


async def _batch_insert(
    session: AsyncSession,
    model_cls: type[DeclarativeBase],
    instances: list[DeclarativeBase],
) -> None:
    """INSERT all instances, raises on conflict."""
    for cls, group_dicts, group_instances in _grouped_table_dicts(model_cls, instances):
        table = cast(Table, cls.__table__)
        missing_pk_cols = [
            col for col in table.primary_key.columns if col.key not in group_dicts[0]
        ]
        if not missing_pk_cols:
            await session.execute(pg_insert(table), group_dicts)
            continue
        stmt = pg_insert(table).returning(
            *missing_pk_cols, sort_by_parameter_order=True
        )
        result = await session.execute(stmt, group_dicts)
        for inst, row in zip(group_instances, result):
            for col, val in zip(missing_pk_cols, row):
                setattr(inst, col.key, val)


async def _batch_merge(
    session: AsyncSession,
    model_cls: type[DeclarativeBase],
    instances: list[DeclarativeBase],
) -> None:
    """UPSERT: insert new rows, update existing ones with the provided values."""
    for cls, group_dicts, _ in _grouped_table_dicts(model_cls, instances):
        table = cast(Table, cls.__table__)
        pk_names = [col.name for col in table.primary_key]
        pk_names_set = set(pk_names)
        own_col_keys = {col.key for col in table.columns}
        non_pk_cols = [k for k in own_col_keys if k not in pk_names_set]

        stmt = pg_insert(table).values(group_dicts)

        inserted_keys = set(group_dicts[0])
        update_cols = [col for col in non_pk_cols if col in inserted_keys]

        if update_cols:
            stmt = stmt.on_conflict_do_update(
                index_elements=pk_names,
                set_={col: stmt.excluded[col] for col in update_cols},
            )
        else:
            stmt = stmt.on_conflict_do_nothing(index_elements=pk_names)

        await session.execute(stmt)


async def _batch_skip_existing(
    session: AsyncSession,
    model_cls: type[DeclarativeBase],
    instances: list[DeclarativeBase],
) -> list[DeclarativeBase]:
    """INSERT only rows that do not already exist; return the inserted ones."""
    if len(_get_table_chain(model_cls)) > 1:
        loaded: list[DeclarativeBase] = []
        for inst in instances:
            pk = _get_primary_key(inst)
            if pk is None or not await session.get(model_cls, pk):
                session.add(inst)
                loaded.append(inst)
        await session.flush()
        return loaded

    mapper = model_cls.__mapper__
    pk_names = [col.name for col in mapper.primary_key]

    no_pk: list[DeclarativeBase] = []
    with_pk_pairs: list[tuple[DeclarativeBase, Any]] = []
    for inst in instances:
        pk = _get_primary_key(inst)
        if pk is None:
            no_pk.append(inst)
        else:
            with_pk_pairs.append((inst, pk))

    loaded = list(no_pk)
    if no_pk:
        no_pk_dicts = [_instance_to_dict(i) for i in no_pk]
        for group_dicts, group_instances in _group_by_column_set(no_pk_dicts, no_pk):
            stmt = pg_insert(cast(Table, model_cls.__table__)).returning(
                *mapper.primary_key, sort_by_parameter_order=True
            )
            result = await session.execute(stmt, group_dicts)
            for inst, row in zip(group_instances, result):
                for col, val in zip(mapper.primary_key, row):
                    setattr(inst, cast(str, col.key), val)

    if with_pk_pairs:
        with_pk = [i for i, _ in with_pk_pairs]
        with_pk_dicts = [_instance_to_dict(i) for i in with_pk]
        for group_dicts, group_insts in _group_by_column_set(with_pk_dicts, with_pk):
            stmt = (
                pg_insert(model_cls)
                .values(group_dicts)
                .on_conflict_do_nothing(index_elements=pk_names)
            )
            result = await session.execute(stmt.returning(*mapper.primary_key))
            inserted_pks = {
                row[0] if len(pk_names) == 1 else tuple(row) for row in result
            }
            loaded.extend(
                inst
                for inst, pk in zip(
                    group_insts, [_get_primary_key(i) for i in group_insts]
                )
                if pk in inserted_pks
            )

    return loaded


def _relationship_load_options(model: type[DeclarativeBase]) -> list[ExecutableOption]:
    """Build selectinload options for all direct relationships on a model."""
    return [
        selectinload(getattr(model, rel.key)) for rel in model.__mapper__.relationships
    ]


async def _reload_with_relationships(
    session: AsyncSession,
    instances: list[DeclarativeBase],
    load_options: list[ExecutableOption],
) -> list[DeclarativeBase]:
    """Reload instances in a single bulk query with relationship eager-loading."""
    model = type(instances[0])
    mapper = model.__mapper__
    pk_cols = mapper.primary_key

    if len(pk_cols) == 1:
        pk_key = cast(str, pk_cols[0].key)
        pk_attr = getattr(model, pk_key)
        pks = [getattr(inst, pk_key) for inst in instances]
        result = await session.execute(
            select(model).where(pk_attr.in_(pks)).options(*load_options)
        )
        by_pk = {getattr(row, pk_key): row for row in result.unique().scalars()}
        return [by_pk[pk] for pk in pks]

    # Composite PK: fall back to per-instance reload
    reloaded: list[DeclarativeBase] = []
    for instance in instances:
        pk = _get_primary_key(instance)
        refreshed = await session.get(
            model,
            pk,
            options=cast(list[ORMOption], load_options),
            populate_existing=True,
        )
        if refreshed is not None:  # pragma: no branch
            reloaded.append(refreshed)
    return reloaded


async def _refresh_loaded(
    session: AsyncSession, instances: list[DeclarativeBase]
) -> list[DeclarativeBase]:
    """Re-select freshly written rows, eager-loading relationships."""
    if not instances:
        return []
    refreshed: list[DeclarativeBase | None] = [None] * len(instances)
    for model_cls, group in _group_by_type(instances):
        positions = [i for i, inst in enumerate(instances) if type(inst) is model_cls]
        load_options = _relationship_load_options(model_cls)
        for pos, new in zip(
            positions, await _reload_with_relationships(session, group, load_options)
        ):
            refreshed[pos] = new
    return cast(list[DeclarativeBase], refreshed)


async def _load_ordered(
    session: AsyncSession,
    registry: FixtureRegistry,
    ordered_names: list[str],
    strategy: LoadStrategy,
    contexts: tuple[str, ...] | None = None,
) -> dict[str, list[DeclarativeBase]]:
    """Load fixtures in order using batch Core INSERT statements."""
    results: dict[str, list[DeclarativeBase]] = {}

    for name in ordered_names:
        variants = (
            registry.get_load_variants(name, *contexts)
            if contexts is not None
            else registry.get_variants(name)
        )

        if not variants:  # pragma: no cover
            results[name] = []
            continue

        instances = [inst for v in variants for inst in v.func()]

        if not instances:
            results[name] = []
            continue

        model_name = type(instances[0]).__name__
        loaded: list[DeclarativeBase] = []

        async with transaction(session):
            for model_cls, group in _group_by_type(instances):
                match strategy:
                    case LoadStrategy.INSERT:
                        await _batch_insert(session, model_cls, group)
                        loaded.extend(group)
                    case LoadStrategy.MERGE:
                        await _batch_merge(session, model_cls, group)
                        loaded.extend(group)
                    case LoadStrategy.SKIP_EXISTING:
                        inserted = await _batch_skip_existing(session, model_cls, group)
                        loaded.extend(inserted)
                    case _:  # pragma: no cover
                        pass

            loaded = await _refresh_loaded(session, loaded)

        results[name] = loaded
        logger.info("Loaded fixture '%s': %d %s(s)", name, len(loaded), model_name)

    return results


def _get_primary_key(instance: DeclarativeBase) -> Any | None:
    """Get the primary key value of a model instance."""
    mapper = instance.__class__.__mapper__
    pk_cols = mapper.primary_key

    if len(pk_cols) == 1:
        return getattr(instance, pk_cols[0].name, None)

    pk_values = tuple(getattr(instance, col.name, None) for col in pk_cols)
    if all(v is not None for v in pk_values):
        return pk_values
    return None


async def load_fixtures(
    session: AsyncSession,
    registry: FixtureRegistry,
    *names: str,
    strategy: LoadStrategy = LoadStrategy.MERGE,
) -> dict[str, list[DeclarativeBase]]:
    """Load specific fixtures by name with dependencies.

    All context variants of each requested fixture are loaded and merged.

    Args:
        session: Database session
        registry: Fixture registry
        *names: Fixture names to load (dependencies auto-resolved)
        strategy: How to handle existing records

    Returns:
        Dict mapping fixture names to loaded instances
    """
    ordered = registry.resolve_dependencies(*names)
    return await _load_ordered(session, registry, ordered, strategy)


async def load_fixtures_by_context(
    session: AsyncSession,
    registry: FixtureRegistry,
    *contexts: str | Enum,
    strategy: LoadStrategy = LoadStrategy.MERGE,
) -> dict[str, list[DeclarativeBase]]:
    """Load all fixtures for specific contexts.

    Args:
        session: Database session
        registry: Fixture registry
        *contexts: Contexts to load (e.g., ``Context.TESTING``, or plain
            strings for custom contexts)
        strategy: How to handle existing records

    Returns:
        Dict mapping fixture names to loaded instances
    """
    context_strings = tuple(_normalize_contexts(contexts))
    ordered = registry.resolve_context_dependencies(*contexts)
    return await _load_ordered(
        session, registry, ordered, strategy, contexts=context_strings
    )