fastapi-toolsets/src/fastapi_toolsets/db/m2m.py

"""Many-to-Many association-table helpers (direct, without loading collections)."""

from typing import Any, TypeVar, cast

from sqlalchemy import ColumnElement, Table, delete, tuple_
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import DeclarativeBase, QueryableAttribute
from sqlalchemy.orm.relationships import RelationshipProperty

_M = TypeVar("_M", bound=DeclarativeBase)


def _m2m_prop(rel_attr: QueryableAttribute) -> tuple[RelationshipProperty, Table]:  # type: ignore[type-arg]
    """Return the validated M2M RelationshipProperty and its secondary table.

    Raises TypeError if *rel_attr* is not a Many-to-Many relationship.
    """
    prop = rel_attr.property
    if not isinstance(prop, RelationshipProperty) or prop.secondary is None:
        raise TypeError(
            f"m2m helpers require a Many-to-Many relationship attribute, "
            f"got {rel_attr!r}. Use a relationship with a secondary table."
        )
    return prop, cast(Table, prop.secondary)


def _parent_where(
    prop: RelationshipProperty,  # type: ignore[type-arg]
    instance: DeclarativeBase,
) -> list[ColumnElement[bool]]:
    """Build the WHERE clauses matching the owner side of *instance*."""
    return [
        assoc_col == getattr(instance, cast(str, parent_col.key))
        for parent_col, assoc_col in prop.synchronize_pairs
    ]


async def m2m_add(
    session: AsyncSession,
    instance: DeclarativeBase,
    rel_attr: QueryableAttribute,
    *related: DeclarativeBase,
    ignore_conflicts: bool = False,
) -> None:
    """Insert rows into a Many-to-Many association table without loading the ORM collection.

    Args:
        session: DB async session.
        instance: The "owner" side model instance (e.g. the ``A`` in ``A.b_list``).
        rel_attr: The M2M relationship attribute on the model class (e.g. ``A.b_list``).
        *related: One or more related instances to associate with ``instance``.
        ignore_conflicts: When ``True``, silently skip rows that already exist
            in the association table (``ON CONFLICT DO NOTHING``).

    Raises:
        TypeError: If ``rel_attr`` is not a Many-to-Many relationship.

    Example:
        ```python
        from fastapi_toolsets.db import m2m_add, transaction

        async with transaction(session):
            await m2m_add(session, post, Post.tags, tag1, tag2)
        ```
    """
    prop, secondary = _m2m_prop(rel_attr)
    if not related:
        return

    sync_pairs = prop.secondary_synchronize_pairs
    assert sync_pairs is not None  # set whenever secondary is set

    # synchronize_pairs:           [(parent_col, assoc_col), ...]
    # secondary_synchronize_pairs: [(related_col, assoc_col), ...]
    rows: list[dict[str, Any]] = []
    for rel_instance in related:
        row: dict[str, Any] = {}
        for parent_col, assoc_col in prop.synchronize_pairs:
            row[assoc_col.name] = getattr(instance, cast(str, parent_col.key))
        for related_col, assoc_col in sync_pairs:
            row[assoc_col.name] = getattr(rel_instance, cast(str, related_col.key))
        rows.append(row)

    stmt = pg_insert(secondary).values(rows)
    if ignore_conflicts:
        stmt = stmt.on_conflict_do_nothing()
    await session.execute(stmt)


async def m2m_remove(
    session: AsyncSession,
    instance: DeclarativeBase,
    rel_attr: QueryableAttribute,
    *related: DeclarativeBase,
) -> None:
    """Remove rows from a Many-to-Many association table without loading the ORM collection.

    Args:
        session: DB async session.
        instance: The "owner" side model instance (e.g. the ``A`` in ``A.b_list``).
        rel_attr: The M2M relationship attribute on the model class (e.g. ``A.b_list``).
        *related: One or more related instances to disassociate from ``instance``.

    Raises:
        TypeError: If ``rel_attr`` is not a Many-to-Many relationship.

    Example:
        ```python
        from fastapi_toolsets.db import m2m_remove, transaction

        async with transaction(session):
            await m2m_remove(session, post, Post.tags, tag1)
        ```
    """
    prop, secondary = _m2m_prop(rel_attr)
    if not related:
        return

    related_pairs = prop.secondary_synchronize_pairs
    assert related_pairs is not None  # set whenever secondary is set

    parent_where = _parent_where(prop, instance)

    if len(related_pairs) == 1:
        related_col, assoc_col = related_pairs[0]
        related_values = [getattr(r, cast(str, related_col.key)) for r in related]
        related_where = assoc_col.in_(related_values)
    else:
        assoc_cols = [ac for _, ac in related_pairs]
        rel_cols = [rc for rc, _ in related_pairs]
        related_values_t = [
            tuple(getattr(r, cast(str, rc.key)) for rc in rel_cols) for r in related
        ]
        related_where = tuple_(*assoc_cols).in_(related_values_t)

    await session.execute(delete(secondary).where(*parent_where, related_where))


async def m2m_set(
    session: AsyncSession,
    instance: DeclarativeBase,
    rel_attr: QueryableAttribute,
    *related: DeclarativeBase,
) -> None:
    """Replace the entire Many-to-Many association set atomically.

    Args:
        session: DB async session.
        instance: The "owner" side model instance (e.g. the ``A`` in ``A.b_list``).
        rel_attr: The M2M relationship attribute on the model class (e.g. ``A.b_list``).
        *related: The new complete set of related instances.

    Raises:
        TypeError: If ``rel_attr`` is not a Many-to-Many relationship.

    Example:
        ```python
        from fastapi_toolsets.db import m2m_set, transaction

        async with transaction(session):
            await m2m_set(session, post, Post.tags, tag1, tag2)  # replaces all
        ```
    """
    prop, secondary = _m2m_prop(rel_attr)

    await session.execute(delete(secondary).where(*_parent_where(prop, instance)))

    if related:
        await m2m_add(session, instance, rel_attr, *related)