perf: CTE任务树加载器 + MOM跨库查询缓存
消除两个核心N+1性能瓶颈： 1. CTE任务树加载器 (task_tree_loader.py) - PostgreSQL Recursive CTE一次性加载完整任务树 - 无论树深度多大，仅2条SQL（CTE + records selectinload） - set_committed_value安全注入，避免Session脏数据 - 修复add_task_record双重加载问题 - 移除get_all_tasks中冗余的selectinload(child_tasks) 2. MOM跨库查询缓存 (mom_cache.py) - 零依赖TTL内存缓存（threading.RLock + time.monotonic） - 参数化ANY(:user_ids)替代OR拼接LIKE（防注入） - get_all_products中3次调用共享缓存，2h TTL内零跨库查询
This commit is contained in:
154
backend/app/services/mom_cache.py
Normal file
154
backend/app/services/mom_cache.py
Normal file
@ -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
|
||||
@ -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(
|
||||
|
||||
@ -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)
|
||||
|
||||
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