169 lines
9.1 KiB
Python
169 lines
9.1 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""db_query.py - 只读项目数据查询工具（甲类能力，v1 角色 agent / v2 会话 agent 共用）
|
||
|
||
对齐 Hermes 的「用工具查实证据、不靠猜」诊断能力：agent 排查/核对时可直接查
|
||
项目关联表的真实数据（任务/交付件/Bug/审计等），不再只能靠 list_tasks 这类
|
||
固定视图，也不用 run_command 绕道 mysql。
|
||
|
||
安全设计（不靠 LLM 自觉，全部代码层硬门禁）：
|
||
1. 表白名单：只允许查询与项目直接/间接关联的业务表（映射与
|
||
project_capability._dump_project_records 同源），其余表一律拒绝
|
||
（users/permission/rolepermission/llm 等凭据与权限表绝不可查）。
|
||
2. 只读：仅 SELECT，SQL 前缀硬校验；禁多语句（分号截断）。
|
||
3. 强制项目过滤：project_id 由代码注入（参数化），不接受调用方自定义
|
||
WHERE——跨项目/跨租户查询不可能构造出来。
|
||
4. 行数/字符上限 + 写审计（record_audit，append-only）。
|
||
5. 通用会话（generic，无项目）不可用本工具。
|
||
"""
|
||
import json
|
||
import re
|
||
|
||
_MAX_ROWS = 100 # 单次返回行数上限
|
||
_MAX_OUT_CHARS = 20000 # 输出字符上限
|
||
|
||
# 表 → 强制项目过滤 SQL（与 project_capability._dump_project_records 同源）。
|
||
# project_id 参数化注入；间接关联表用子查询回项目。
|
||
TABLE_SCOPES = {
|
||
# 直接关联
|
||
'sd_projects': "SELECT * FROM sd_projects WHERE id=${pid}$",
|
||
'sd_iterations': "SELECT * FROM sd_iterations WHERE project_id=${pid}$",
|
||
'pipeline_tasks': "SELECT * FROM pipeline_tasks WHERE tenant_id=${pid}$",
|
||
'pipeline_deliverables': "SELECT * FROM pipeline_deliverables WHERE project_id COLLATE utf8mb4_unicode_ci=${pid}$",
|
||
'pipeline_project_agents': "SELECT * FROM pipeline_project_agents WHERE project_id COLLATE utf8mb4_unicode_ci=${pid}$",
|
||
'pipeline_agent_questions': "SELECT * FROM pipeline_agent_questions WHERE tenant_id=${pid}$",
|
||
'sd_deploy_envs': "SELECT * FROM sd_deploy_envs WHERE project_id=${pid}$",
|
||
'sd_features': "SELECT * FROM sd_features WHERE project_id=${pid}$",
|
||
'sd_project_repos': "SELECT * FROM sd_project_repos WHERE project_id COLLATE utf8mb4_unicode_ci=${pid}$",
|
||
# 间接关联（子查询回项目）
|
||
'sd_bugs': "SELECT * FROM sd_bugs WHERE iteration_id IN (SELECT id FROM sd_iterations WHERE project_id=${pid}$)",
|
||
'sd_test_plans': "SELECT * FROM sd_test_plans WHERE iteration_id IN (SELECT id FROM sd_iterations WHERE project_id=${pid}$)",
|
||
'sd_test_cases': "SELECT * FROM sd_test_cases WHERE plan_id IN (SELECT id FROM sd_test_plans WHERE iteration_id IN (SELECT id FROM sd_iterations WHERE project_id=${pid}$))",
|
||
'audit_log': "SELECT * FROM audit_log WHERE tenant_id COLLATE utf8mb4_unicode_ci=${pid}$ ORDER BY created_at ASC",
|
||
}
|
||
|
||
# 产线扩展表（投标产线等，同样强制 project_id 过滤；表不存在时运行期容错）
|
||
PIPELINE_TABLE_SCOPES = {
|
||
'bid_chapters': "SELECT * FROM bid_chapters WHERE project_id=${pid}$",
|
||
'bid_doc_requirements': "SELECT * FROM bid_doc_requirements WHERE project_id=${pid}$",
|
||
'bid_qualifications': "SELECT * FROM bid_qualifications WHERE project_id=${pid}$",
|
||
'bid_qc_reviews': "SELECT * FROM bid_qc_reviews WHERE project_id=${pid}$",
|
||
'bid_documents': "SELECT * FROM bid_documents WHERE project_id=${pid}$",
|
||
'bid_kb_docs': "SELECT * FROM bid_kb_docs WHERE project_id=${pid}$",
|
||
'bid_members': "SELECT * FROM bid_members WHERE project_id=${pid}$",
|
||
'bid_reviews': "SELECT * FROM bid_reviews WHERE project_id=${pid}$",
|
||
'bid_analysis_history': "SELECT * FROM bid_analysis_history WHERE project_id=${pid}$",
|
||
'bid_cost_benefit': "SELECT * FROM bid_cost_benefit WHERE project_id=${pid}$",
|
||
'bid_material_gaps': "SELECT * FROM bid_material_gaps WHERE project_id=${pid}$",
|
||
}
|
||
|
||
ALL_SCOPES = dict(TABLE_SCOPES, **PIPELINE_TABLE_SCOPES)
|
||
|
||
|
||
def _rec_to_dict(rec):
|
||
if isinstance(rec, dict):
|
||
return dict(rec)
|
||
try:
|
||
return dict(rec)
|
||
except (TypeError, ValueError):
|
||
return {}
|
||
|
||
|
||
async def tool_query_project_data(sor, table: str, project_id: str,
|
||
where: str = '', order_by: str = '',
|
||
limit: int = _MAX_ROWS,
|
||
who: str = '', agent_id: str = '',
|
||
task_id: str = '') -> str:
|
||
"""查询项目关联表数据（只读，强制项目过滤）。
|
||
|
||
table: 白名单表名（TABLE_SCOPES/PIPELINE_TABLE_SCOPES 键）
|
||
where: 可选附加过滤（AND 拼接；仅允许「列 运算符 字面量」安全形态，代码校验）
|
||
order_by: 可选排序列（仅允许白名单表的列名形态）
|
||
"""
|
||
table = (table or '').strip().lower()
|
||
if not project_id:
|
||
return 'FAIL: 缺少项目上下文（本工具仅项目内可用）'
|
||
sql = ALL_SCOPES.get(table)
|
||
if not sql:
|
||
return ('FAIL: 表 %s 不在可查询白名单。可查询的表: %s'
|
||
% (table or '(空)', ', '.join(sorted(ALL_SCOPES.keys()))))
|
||
|
||
try:
|
||
limit = max(1, min(_MAX_ROWS, int(limit or _MAX_ROWS)))
|
||
except (ValueError, TypeError):
|
||
limit = _MAX_ROWS
|
||
|
||
# 附加 WHERE 安全校验：仅允许「标识符 运算符 字面量」的 AND 组合，
|
||
# 禁子查询/分号/注释/UNION——参数化之外的注入面在这里掐死。
|
||
extra = ''
|
||
if where and str(where).strip():
|
||
w = str(where).strip().rstrip(';')
|
||
# 字面量：单引号串/双引号串/数字/NULL/安全IN列表（引号串或数字，禁嵌套括号）
|
||
_val = (r"('[^']*'|\"[^\"]*\"|[\d.]+|NULL|"
|
||
r"\((?:'[^']*'|[\d.]+)(?:\s*,\s*(?:'[^']*'|[\d.]+))*\))")
|
||
_col = r"[A-Za-z_][A-Za-z0-9_]*(\.[A-Za-z_][A-Za-z0-9_]*)?"
|
||
# 两种条件形态：比较运算符带值 / IS [NOT] NULL 不带值
|
||
_cmp = _col + r"\s*(=|!=|<>|>=|<=|>|<|LIKE|IN)\s*" + _val
|
||
_isnull = _col + r"\s+IS\s+(NOT\s+)?NULL"
|
||
_cond = r"(?:" + _cmp + r"|" + _isnull + r")"
|
||
safe_cond = r"^" + _cond + r"(\s+AND\s+" + _cond + r")*$"
|
||
if not re.match(safe_cond, w, re.I):
|
||
return ('FAIL: where 附加过滤只允许「列名 运算符 字面量」的 AND 组合'
|
||
"（禁子查询/UNION/分号/注释）。示例: state='running'")
|
||
# 2026-09-16 甘肃项目实测 bug：LIKE '%节能%' 的 % 会被 aiomysql 的
|
||
# query % args 参数替换当成格式占位符 → TypeError: not enough arguments。
|
||
# 字面量里的 % 必须转义为 %%（替换后还原为字面 %）。
|
||
extra = ' AND (' + w.replace('%', '%%') + ')'
|
||
order = ''
|
||
if order_by and str(order_by).strip():
|
||
o = str(order_by).strip()
|
||
if not re.match(r"^[A-Za-z_][A-Za-z0-9_]*(\s+(ASC|DESC))?$", o, re.I):
|
||
return 'FAIL: order_by 只允许单个列名（可带 ASC/DESC）'
|
||
order = ' ORDER BY ' + o
|
||
|
||
full_sql = sql + extra + order + (' LIMIT %d' % limit)
|
||
try:
|
||
recs = await sor.sqlExe(full_sql, {"pid": project_id})
|
||
except Exception as e:
|
||
msg = str(e)[:200]
|
||
# 2026-09-15：LLM 常把 params JSON 键当列传（如 stage='analysis_redo'）→ 1054。
|
||
# 回带该表真实列清单，让调用方下一轮自恢复，不再白烧轮次。
|
||
if 'Unknown column' in msg:
|
||
try:
|
||
cols = await sor.sqlExe("SHOW COLUMNS FROM " + table, {})
|
||
names = ', '.join(getattr(c, 'Field', '') or '' for c in (cols or []))
|
||
msg += '。表 %s 可用列: %s（stage/task_kind 等在 params JSON 内，不是列）' % (table, names)
|
||
except Exception:
|
||
pass
|
||
# 产线扩展表在其他产线库可能不存在——容错但报实情
|
||
return 'FAIL: 查询异常（表可能不存在于当前产线库）: ' + msg
|
||
|
||
rows = [_rec_to_dict(r) for r in (recs or [])]
|
||
# 审计（append-only）：谁在查什么
|
||
try:
|
||
from .audit import record_audit
|
||
await record_audit(project_id, 'db_query', table, 'query',
|
||
who=who or 'agent', agent_id=agent_id or '',
|
||
detail='rows=%d where=%s' % (len(rows), (where or '')[:100]),
|
||
sor=sor)
|
||
except Exception:
|
||
pass
|
||
|
||
if not rows:
|
||
return '（表 %s 在项目 %s 下无记录）' % (table, project_id[:8])
|
||
out = json.dumps(rows, ensure_ascii=False, default=str)
|
||
if len(out) > _MAX_OUT_CHARS:
|
||
# 超限：截断并告知（与 read_file 同一「显式截断」纪律）
|
||
kept = []
|
||
acc = 2
|
||
for r in rows:
|
||
s = json.dumps(r, ensure_ascii=False, default=str)
|
||
if acc + len(s) + 1 > _MAX_OUT_CHARS:
|
||
break
|
||
kept.append(s)
|
||
acc += len(s) + 1
|
||
out = '[' + ','.join(kept) + ']'
|
||
out += ('\n\n[⚠️ 截断提示：结果共 %d 行，以上仅展示前 %d 行（输出上限 %d 字符）。'
|
||
'缩小范围请加 where/limit 条件重查。]' % (len(rows), len(kept), _MAX_OUT_CHARS))
|
||
return '表 %s 共 %d 行（limit=%d）:\n%s' % (table, len(rows), limit, out)
|
||
return '表 %s 共 %d 行（limit=%d）:\n%s' % (table, len(rows), limit, out)
|