chore: fork from IRIS track 供 LICA 部门独立运行
- 复制来源: /home/yueli/track @ 192c8ee (feature/ai-audit-update) - 组织隔离目标: LICA - 端口规划: 前端 8030 / 后端 8031 / 数据库 8032 - 已排除 deploy.sh、deploy_full.sh、docker-compose.prod.yml（IRIS 生产发布脚本） - 已排除工作区未提交改动，取干净的 192c8ee 状态
This commit is contained in:
176
backend/app/services/task_tree_loader.py
Normal file
176
backend/app/services/task_tree_loader.py
Normal file
@ -0,0 +1,176 @@
|
||||
"""
|
||||
共享 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
|
||||
Reference in New Issue
Block a user