diff --git a/backend/app/services/mom_cache.py b/backend/app/services/mom_cache.py new file mode 100644 index 0000000..cdbbbe4 --- /dev/null +++ b/backend/app/services/mom_cache.py @@ -0,0 +1,154 @@ +""" +MOM 跨库查询缓存模块 — 使用本地 TTL 缓存消除冗余跨库请求 + +解决的问题: + 1. _lookup_display_names 在 get_all_products 中被调用 3 次,每次都打开/关闭 + MOM 数据库连接,150 条产品的列表页 = 3 根管线查询。 + 2. 同一批 username 在短时间内(用户翻页、多人同时访问)被反复查询。 + 3. 旧实现用 OR 拼接 LIKE 条件,存在注入风险。 + +方案:python -m 内置模块(零依赖)实现线程安全 TTL 缓存 + 参数化 ANY 查询。 + +TTL: 2 小时(人员姓名不会频繁变动,可调)。 +""" +from __future__ import annotations + +import threading +import time + +from app.core.mom_database import MomSessionLocal + + +# ============================================================ +# 零依赖 TTL 缓存(线程安全) +# ============================================================ + +class _TTLCache: + """线程安全的内存 TTL 缓存,用于 MOM 只读查询结果""" + + def __init__(self, ttl_seconds: int = 7200) -> None: + self._store: dict[str, str] = {} + self._expiry: dict[str, float] = {} + self._ttl = ttl_seconds + self._lock = threading.RLock() + + def get_many(self, keys: list[str]) -> tuple[dict[str, str], list[str]]: + """ + 批量获取 → (命中字典, 未命中 key 列表)。 + + 内部自动清理过期条目。 + """ + hits: dict[str, str] = {} + missed: list[str] = [] + now = time.monotonic() + + with self._lock: + for k in keys: + exp = self._expiry.get(k) + if exp is not None and now < exp: + hits[k] = self._store[k] + else: + missed.append(k) + # 清理过期残留 + if k in self._store: + del self._store[k] + del self._expiry[k] + + return hits, missed + + def set_many(self, mapping: dict[str, str]) -> None: + """批量写入,所有 key 共享同一过期时间""" + expiry = time.monotonic() + self._ttl + with self._lock: + for k, v in mapping.items(): + self._store[k] = v + self._expiry[k] = expiry + + +# ============================================================ +# 全局缓存实例(2h TTL) +# ============================================================ + +_user_name_cache = _TTLCache(ttl_seconds=7200) + + +# ============================================================ +# 公开 API +# ============================================================ + +def get_display_names(user_ids: list[str]) -> dict[str, str]: + """ + 批量查询 MOM sys_user,将 username 映射为真实姓名(带 2h TTL 缓存)。 + + 缓存穿透流程: + 1. 去重 → 从缓存批量读取 + 2. 计算 miss 差集 + 3. miss 非空时,用参数化 ANY(:user_ids) 查 MOM(1 条 SQL) + 4. 写回缓存 + 5. 合并 hits + fresh 返回 + + 参数: + user_ids: 短用户名列表,如 ["zhangsan01", "lisi02"] + + 返回: + {"zhangsan01": "张三", "lisi02": "李四"} + 不存在的 key 不会出现在返回字典中。 + + SQL 安全: + 使用 SPLIT_PART(username, '/', 2) = ANY(:user_ids) 参数化查询, + 杜绝旧实现中 OR 拼接 LIKE 的注入风险。 + """ + if not user_ids: + return {} + + # 过滤特殊值 + 去重保序 + seen: set[str] = set() + real_ids: list[str] = [] + for uid in user_ids: + if uid and uid != "virtual_warehouse" and uid not in seen: + seen.add(uid) + real_ids.append(uid) + + if not real_ids: + return {} + + # ── Step 1: 批量查缓存 ── + hits, missed = _user_name_cache.get_many(real_ids) + + # ── Step 2: 仅对 miss 查 MOM ── + if missed: + db = MomSessionLocal() + try: + from sqlalchemy import text + + # 参数化 ANY 查询 — 安全防注入 + # SPLIT_PART('张三/zhangsan01', '/', 2) = 'zhangsan01' + # OR username = ANY(...) 兜底无斜杠的用户名(如 admin) + sql = text(""" + SELECT username, + SPLIT_PART(username, '/', 1) AS display_name + FROM sys_user + WHERE SPLIT_PART(username, '/', 2) = ANY(:user_ids) + OR username = ANY(:user_ids) + """) + result = db.execute(sql, {"user_ids": missed}) + rows = result.fetchall() + finally: + db.close() + + # ── Step 3: 解析结果 + 写回缓存 ── + fresh: dict[str, str] = {} + for row in rows: + full_username: str = row[0] + display_name: str = row[1] + # "张三/zhangsan01" → short="zhangsan01" + short = full_username.split("/")[-1] if "/" in full_username else full_username + fresh[short] = display_name + + if fresh: + _user_name_cache.set_many(fresh) + + # ── Step 4: 合并 ── + hits.update(fresh) + + return hits diff --git a/backend/app/services/product_service.py b/backend/app/services/product_service.py index a30275b..71207d3 100644 --- a/backend/app/services/product_service.py +++ b/backend/app/services/product_service.py @@ -46,35 +46,10 @@ def _task_to_response(task: Task) -> TaskResponse: async def _load_task_tree(db: AsyncSession, product_id: uuid.UUID) -> list[TaskResponse]: - """递归加载产品下的完整任务树""" - # 先取顶层任务 - result = await db.execute( - select(Task) - .options(selectinload(Task.child_tasks), selectinload(Task.records)) - .where( - Task.product_id == product_id, - Task.parent_task_id.is_(None), - ) - .order_by(Task.created_at) - ) - top_tasks = result.scalars().all() - - # 递归加载每层子任务 - async def _load_children(t: Task): - for child in t.child_tasks: - child_result = await db.execute( - select(Task) - .options(selectinload(Task.child_tasks), selectinload(Task.records)) - .where(Task.id == child.id) - ) - refreshed = child_result.scalar_one() - t.child_tasks[t.child_tasks.index(child)] = refreshed - await _load_children(refreshed) - - for task in top_tasks: - await _load_children(task) - - return [_task_to_response(t) for t in top_tasks] + """使用 PostgreSQL Recursive CTE 一次性加载产品下完整任务树(消除 N+1)""" + from app.services.task_tree_loader import load_task_trees_by_product + tasks = await load_task_trees_by_product(db, product_id) + return [_task_to_response(t) for t in tasks] async def get_product_by_serial(db: AsyncSession, serial_number: str) -> ProductScanResponse: @@ -344,32 +319,9 @@ async def update_overall_status( def _lookup_display_names(location_ids: list[str]) -> dict[str, str]: - """批量查询 MOM sys_user,将 username 映射为真实姓名""" - if not location_ids: - return {} - from app.core.mom_database import MomSessionLocal - from sqlalchemy import text - db = MomSessionLocal() - try: - # 过滤掉特殊值 - real_ids = [uid for uid in location_ids if uid and uid != "virtual_warehouse"] - if not real_ids: - return {} - # 用 LIKE 模糊匹配批量查出 - conditions = " OR ".join([f"username LIKE '%/{uid}'" for uid in real_ids]) - result = db.execute( - text(f"SELECT username, SPLIT_PART(username, '/', 1) as display_name FROM sys_user WHERE {conditions}") - ) - mapping = {} - for row in result: - full_username = row[0] - display_name = row[1] - # 从 full_username 末尾提取短用户名: "张三/zhangsan01" → "zhangsan01" - short = full_username.split("/")[-1] if "/" in full_username else full_username - mapping[short] = display_name - return mapping - finally: - db.close() + """批量查询 MOM sys_user,将 username 映射为真实姓名(带 2h TTL 缓存)""" + from app.services.mom_cache import get_display_names + return get_display_names(location_ids) async def get_all_products( diff --git a/backend/app/services/task_service.py b/backend/app/services/task_service.py index 91803ff..cdc88df 100644 --- a/backend/app/services/task_service.py +++ b/backend/app/services/task_service.py @@ -125,39 +125,9 @@ async def _get_task_or_404(db: AsyncSession, task_id: uuid.UUID) -> Task: async def _get_task_with_children_recursive(db: AsyncSession, task_id: uuid.UUID) -> Task: - """递归加载任务及其所有子孙任务""" - result = await db.execute( - select(Task) - .options( - selectinload(Task.child_tasks), - selectinload(Task.records), - ) - .where(Task.id == task_id) - ) - task = result.scalar_one_or_none() - if not task: - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail=f"任务不存在: {task_id}", - ) - - # 递归加载每一层子任务 - async def _load_children(t: Task): - for child in t.child_tasks: - child_result = await db.execute( - select(Task) - .options( - selectinload(Task.child_tasks), - selectinload(Task.records), - ) - .where(Task.id == child.id) - ) - refreshed_child = child_result.scalar_one() - t.child_tasks[t.child_tasks.index(child)] = refreshed_child - await _load_children(refreshed_child) - - await _load_children(task) - return task + """使用 PostgreSQL Recursive CTE 一次性加载任务及其所有子孙任务(消除 N+1)""" + from app.services.task_tree_loader import load_task_tree_by_root + return await load_task_tree_by_root(db, task_id) def _to_flat_response(task: Task) -> TaskResponse: @@ -326,7 +296,6 @@ async def get_all_tasks( ) -> TaskListResponse: """获取任务列表,可按产品/负责人筛选""" stmt = select(Task).options( - selectinload(Task.child_tasks), selectinload(Task.records), selectinload(Task.product), ) @@ -1085,5 +1054,6 @@ async def add_task_record( await db.commit() await db.refresh(record) - # 重新加载 task 带上新 record - return await get_task(db, task_id) + # 手动追加新 record,避免二次加载整棵树(task 已在 L1051 由 CTE 完整加载) + task.records.append(record) + return _to_response(task) diff --git a/backend/app/services/task_tree_loader.py b/backend/app/services/task_tree_loader.py new file mode 100644 index 0000000..8ad0493 --- /dev/null +++ b/backend/app/services/task_tree_loader.py @@ -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