chore: rework DB module (#324)

src/fastapi_toolsets/db/m2m.py

@@ -0,0 +1,170 @@
+"""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)