track-LICA/backend/app/services/task_tree_loader.py

"""
共享 CTE 任务树加载器 — 使用 PostgreSQL Recursive CTE 一次性拉取完整任务树

解决问题：原 _load_task_tree / _get_task_with_children_recursive 使用
Python 递归逐层 SELECT，N 个节点产生 N+1 次数据库查询。
现在无论树深度多大，仅执行 2 条查询（CTE + records selectinload）。
"""
from __future__ import annotations

import uuid

from fastapi import HTTPException, status
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import noload, selectinload
from sqlalchemy.orm.attributes import set_committed_value

from app.models.task import Task


# ============================================================
# 内存树组装（O(N) 时间 / O(N) 空间）
# ============================================================

def _build_tree_in_memory(tasks: list[Task]) -> dict[uuid.UUID, Task]:
    """
    给定扁平 Task ORM 列表，在内存中通过哈希表组装嵌套树结构。

    关键安全设计：
    - 使用临时字典 temp_children_map 暂存父子关系，绝对不直接操作 ORM 的 child_tasks。
    - 通过 set_committed_value 注入最终列表，告诉 SQLAlchemy 这是"已提交数据"，
      避免 add_task_record 等场景中 db.commit() 时触发级联 UPDATE 污染数据库。

    时间复杂度: O(N)，空间复杂度: O(N)。
    """
    if not tasks:
        return {}

    # ── Pass 1: 临时字典存储关系（不触碰 ORM 属性）──
    temp_children_map: dict[uuid.UUID, list[Task]] = {t.id: [] for t in tasks}
    task_map: dict[uuid.UUID, Task] = {t.id: t for t in tasks}

    # ── Pass 2: 挂载到临时字典 ──
    for t in tasks:
        pid = t.parent_task_id
        if pid is not None and pid in temp_children_map:
            temp_children_map[pid].append(t)

    # ── Pass 3: 排序 + set_committed_value 安全注入 ──
    for t in tasks:
        children = temp_children_map[t.id]
        if children:
            children.sort(key=lambda x: x.created_at)
        # 关键：标记为已提交数据，SQLAlchemy 不会对其生成 UPDATE
        set_committed_value(t, 'child_tasks', children)

    return task_map


# ============================================================
# 公开 API：按单一任务 ID 加载子树
# ============================================================

async def load_task_tree_by_root(
    db: AsyncSession, task_id: uuid.UUID
) -> Task:
    """
    使用 Recursive CTE 加载以 task_id 为根的完整任务子树。

    返回: 根 Task ORM 对象（child_tasks 已递归填充）。

    Raises:
        HTTPException(404): 根任务不存在。
    """
    # ── Step 1: Recursive CTE — 收集所有子孙节点 ID ──
    #   WITH RECURSIVE task_tree AS (
    #       SELECT tasks.* FROM tasks WHERE tasks.id = :tid
    #       UNION ALL
    #       SELECT tasks.* FROM tasks
    #         JOIN task_tree ON tasks.parent_task_id = task_tree.id
    #   )
    anchor = (
        select(Task)
        .where(Task.id == task_id)
        .cte(name="task_tree", recursive=True)
    )
    task_tree_cte = anchor.union_all(
        select(Task).join(anchor, Task.parent_task_id == anchor.c.id)
    )

    # ── Step 2: 批量加载所有任务 + 关联数据 ──
    stmt = (
        select(Task)
        .options(
            noload(Task.child_tasks),       # 禁掉模型默认 selectinload，由内存树接管
            noload(Task.parent_task),       # 组装树不需要 parent 引用
            selectinload(Task.records),     # 🔥 一次性预加载所有进度记录
            selectinload(Task.product),     # 🔥 一次性预加载产品引用
        )
        .where(Task.id.in_(select(task_tree_cte.c.id)))
    )

    result = await db.execute(stmt)
    all_tasks = result.unique().scalars().all()

    if not all_tasks:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"任务不存在: {task_id}",
        )

    # ── Step 3: 内存组装 ──
    task_map = _build_tree_in_memory(all_tasks)

    # 根任务一定在 map 中（CTE anchor 保证了这一点）
    return task_map[task_id]


# ============================================================
# 公开 API：按产品 ID 加载所有任务树
# ============================================================

async def load_task_trees_by_product(
    db: AsyncSession, product_id: uuid.UUID
) -> list[Task]:
    """
    使用 Recursive CTE 加载指定产品下的所有任务树。

    返回: 顶层任务列表（parent_task_id IS NULL），每项的 child_tasks 已递归填充。
    若无任务则返回空列表。
    """
    # ── Step 1: Recursive CTE ──
    #   WITH RECURSIVE product_task_tree AS (
    #       SELECT tasks.* FROM tasks
    #         WHERE tasks.product_id = :pid AND tasks.parent_task_id IS NULL
    #       UNION ALL
    #       SELECT tasks.* FROM tasks
    #         JOIN product_task_tree ON tasks.parent_task_id = product_task_tree.id
    #   )
    anchor = (
        select(Task)
        .where(
            Task.product_id == product_id,
            Task.parent_task_id.is_(None),
        )
        .cte(name="product_task_tree", recursive=True)
    )
    task_tree_cte = anchor.union_all(
        select(Task).join(anchor, Task.parent_task_id == anchor.c.id)
    )

    # ── Step 2: 批量加载 ──
    stmt = (
        select(Task)
        .options(
            noload(Task.child_tasks),
            noload(Task.parent_task),
            selectinload(Task.records),
            selectinload(Task.product),
        )
        .where(Task.id.in_(select(task_tree_cte.c.id)))
    )

    result = await db.execute(stmt)
    all_tasks = result.unique().scalars().all()

    if not all_tasks:
        return []

    # ── Step 3: 内存组装 ──
    _build_tree_in_memory(all_tasks)

    # ── Step 4: 返回排序后的顶层任务 ──
    roots = [t for t in all_tasks if t.parent_task_id is None]
    roots.sort(key=lambda t: t.created_at)
    return roots