mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 00:41:59 +00:00
审批质询复核修正（第 2 轮，后端）：和 * 写在一起、子查询里没起别名的遮罩列表达式认得出；SQL 只换在找某一个值的字面量，范围比较的边界、LIMIT 照原样；报告口径卡里的指标、SQL 里留着的阈值不当遮罩值；带前缀的字符串、名字、列名里的遮罩值也换
复核提出的问题和改法： - 遮罩列的表达式和 SELECT * 写在同一支里、又没起别名（SELECT *, SUBSTR(buyer_phone, 4)），原值的一截出现在回答正文和质询工件里： 按位置认，* 前面的从头数、最后一个 * 后面的从结果末尾倒着数（结果的列一个不少、按顺序，重名也照留）。同一个根子的几种也一起补上： 子查询、WITH 里没起别名的表达式（外层的 * 原样带出来，列名各数据库起法不同）、两个 * 之间的项认不出落在哪一列，整份快照遮掉； 派生来源里 UNION 的另一支、带列名表的 UNION 按位置对上第一支或列名表的名字；LATERAL、多套几层括号照样认 - 和遮罩列做范围比较时的阈值、5 位以上的数不看上下文就被换成「已遮罩」，封存的 SQL 不是实际执行的那条：SQL 只换在找某一个值的 字面量（等值、不等、IN、LIKE、SQLite 的 IS）——和遮罩列比较的是遮罩值就换，和别的列比较的够长的遮罩值也换；范围比较的边界、 BETWEEN、LIMIT / OFFSET、选出列和算式里的数照原样（字符串的范围边界同样） - 回答正文里复述 SQL 里照原样留着的阈值（「金额不低于 30000 元」）被遮：这些值本来就随 SQL 摆在读者眼前，正文、快照按值遮时免掉 （只免按值认的，SQL 里和遮罩列做等值比较的照旧换） - 报告口径卡里的指标恰好是遮罩值时（单笔目标 20000），引用它的回答整份交不出去、换个问法也绕不过：报告列出的指标和运行输入的值 像问题里的数一样不当遮罩值，回答时和取质询工件时一样 - 带前缀的字符串（N''、E''、U&''、$$…$$）、带引号和不带引号的名字（AS "13900000001 的单"、张伟的单数）里的遮罩值原样封存：所有 字符串类的词只换引号里的内容，名字按值换；快照的列名、回答文档目录里的列名同样换，引用原列名的标记里的列名跟着换，正好对上 - 回答原文的标记里写着遮罩值时，单元格引用「Q1.r0.」后面的长号码会被当成小数的一截认不出来：标记按段换，单元格引用只换列名那一截； 遮罩值在组装之前就换掉，不再组装两遍，组装出来还认得出遮罩值的照旧不交出去 - 新加的接线没有测试守住：补上问题里写了遮罩列的短数时 SQL 原样保留、质询工件按报告快照里的遮罩值遮别的列，以及上面每一种情形 的单元测试和端到端测试。每处改动都做过去掉实现、测试变红的验证
This commit is contained in:
ThearchyHelios
+336
-118
@@ -13,9 +13,11 @@
|
||||
自己查询的格子，或者报告里已有的指标、运行输入；回答文档由 compose_doc 组装，裸数字被标出来（提问人自己写的数除外）。
|
||||
封存之后的追问查的是数据源现在的数据，任务里给报告生成的时间，回答末尾注明；触顶、查询失败、引用解析不了也在末尾注明
|
||||
3. 遮罩（_plan）：哪些列要遮照证据面板的规则（查询当时记下的，加上数据源现在设的），另按 SQL 认出改了名字的遮罩列
|
||||
（起别名、套函数、经子查询转一道，_derived_hidden）；快照按列遮，再按值遮（别的列里写着遮罩值的格子，_masked_view）。
|
||||
回答正文、查询清单里的遮罩值换成「已遮罩」（_Secrets、_scrub_sql）：按值认有门槛，SQL 只换字面量，阈值、日期、LIMIT
|
||||
照原样，封存的是实际执行的语句。组装两遍之后正文里还认得出遮罩值的，这份回答不交出去
|
||||
（起别名、套函数、经子查询转一道、和 * 写在同一支里，_derived_hidden；认不出落在哪一列的，整份快照遮掉）；快照按列
|
||||
遮，再按值遮（别的列里写着遮罩值的格子，_masked_view）。回答正文、查询清单里的遮罩值换成「已遮罩」（_Secrets、
|
||||
_scrub_sql）：按值认有门槛；SQL 只换在找某一个值的字面量（等值、IN、LIKE），范围比较的边界、LIMIT 照原样，封存的是
|
||||
实际执行的语句。已经摆在读者眼前的值不换：问题和选中原文里的、报告列出的指标和运行输入的值、SQL 里照原样留着的
|
||||
阈值。组装之后正文里还认得出遮罩值的，这份回答不交出去
|
||||
4. 写入（_record）：一条 approval.inquiry 事件，node_id 为空，用量写在数据的 usage 里（计进运行用量，见 runner 的
|
||||
_usage_of_events）。质询的查询快照只放在 inquiry_artifacts / queries 里，绝不放在 query_artifact 下：证据接口按那个键
|
||||
收集封存的查询，会把质询的查询当成报告的证据。回答期间运行状态变了（被批准、被放弃、继续运行之后重新封存），
|
||||
@@ -492,6 +494,17 @@ def _report_refs(catalog: dict[str, Any]) -> dict[str, dict[str, Any]]:
|
||||
if isinstance(e, dict) and e.get("kind") in ("metric", "input") and e.get("value") not in (None, "")}
|
||||
|
||||
|
||||
def _ref_values(refs: dict[str, dict[str, Any]]) -> list[str]:
|
||||
"""报告里列出的指标、运行输入的值（原样和渲染出来的写法）：报告里就看得到，碰巧和遮罩列里的某个值一样，回答引用
|
||||
它也不算露出遮罩值（_Secrets.protect）。不免的话，引用「单笔目标」这样的指标整份回答都交不出去，换个问法也一样。"""
|
||||
out: list[str] = []
|
||||
for entry in refs.values():
|
||||
out += _forms(entry.get("value"))
|
||||
if isinstance(entry.get("rendered"), str):
|
||||
out.append(entry["rendered"])
|
||||
return out
|
||||
|
||||
|
||||
def _run_sources(events: Events) -> list[str]:
|
||||
"""这次运行用过的查询工具对应的数据源，按第一次用到的先后。"""
|
||||
out: list[str] = []
|
||||
@@ -640,7 +653,7 @@ def _is_star(item: Any) -> bool:
|
||||
|
||||
|
||||
def _branches(root: Any) -> list[list[Any]]:
|
||||
"""最外层选出的列：一条 SELECT 一支，UNION 这类的每一支各一支（结果的第 i 列来自每一支的第 i 项）。"""
|
||||
"""一块查询选出的列：一条 SELECT 一支，UNION 这类的每一支各一支（结果的第 i 列来自每一支的第 i 项）。"""
|
||||
from sqlglot import exp
|
||||
|
||||
out: list[list[Any]] = []
|
||||
@@ -656,12 +669,108 @@ def _branches(root: Any) -> list[list[Any]]:
|
||||
return out
|
||||
|
||||
|
||||
def _tainted(sql: str, kind: str | None, masked: set[str]) -> tuple[set[str], set[int]] | None:
|
||||
"""SQL 里指代遮罩列的名字（小写），和最外层选出的列里值取自遮罩列的那几列的位置。解析不了返回 None。
|
||||
def _output_name(item: Any) -> str | None:
|
||||
"""一个选出项在结果里的列名（小写）：起了别名的是别名，直接选一列的是列名。别的表达式（函数、运算、CAST、括号、
|
||||
子查询）的列名各数据库起法不同（SQLite 照写法原样记「SUBSTR(buyer_phone, 4)」，Postgres 记函数名或 ?column?），
|
||||
认不准，返回 None。"""
|
||||
from sqlglot import exp
|
||||
|
||||
if isinstance(item, exp.Alias):
|
||||
return item.alias.lower() or None
|
||||
if isinstance(item, exp.Column) and not isinstance(item.this, exp.Star):
|
||||
return item.name.lower() or None
|
||||
return None
|
||||
|
||||
|
||||
def _places(branch: list[Any], width: int | None) -> list[int | None]:
|
||||
"""一支里每一项落在结果的第几列。没有 * 的一一对上；有 * 的，第一个 * 前面的从头数，最后一个 * 后面的从结果末尾
|
||||
倒着数（width 是结果的列数，不知道就数不了）；两个 * 之间的、* 本身认不出位置（None）。"""
|
||||
stars = [i for i, item in enumerate(branch) if _is_star(item)]
|
||||
if not stars:
|
||||
return list(range(len(branch)))
|
||||
first, last = stars[0], stars[-1]
|
||||
tail = len(branch) - 1 - last
|
||||
# * 至少展开成 0 列：倒着数要是数到了 * 前面那几项的位置上，列数对不上，后面的一项都不认
|
||||
countable = width is not None and width - tail >= first
|
||||
return [i if i < first else width - (len(branch) - i) if i > last and countable else None
|
||||
for i in range(len(branch))]
|
||||
|
||||
|
||||
def _blocks(root: Any) -> list[tuple[Any, str, list[str] | None]]:
|
||||
"""SQL 里的每一块查询（一条 SELECT，或者 UNION 这类连起来的一串）和它的结果流到哪：
|
||||
- "result"：最外层，结果就是快照的列；
|
||||
- "source"：当成一张表用的——WITH 里的一段、FROM / JOIN 里的子查询、LATERAL，外层按列名取它的列，* 会把它的列
|
||||
全带出去。写了列名表的（t(p)）带上列名表；
|
||||
- "other"：标量子查询、IN / EXISTS 里的，值经外层的表达式带出（_value_columns 已经算上），列名无关紧要。
|
||||
套在外面的括号（可能好几层）不改变结果流到哪：先越过它们再看。UNION 里的一支不另算，整串按一块。"""
|
||||
from sqlglot import exp
|
||||
|
||||
out: list[tuple[Any, str, list[str] | None]] = []
|
||||
for node in root.find_all(exp.Select, exp.SetOperation):
|
||||
wrappers: list[Any] = []
|
||||
outer = node
|
||||
while isinstance(outer.parent, exp.Subquery):
|
||||
outer = outer.parent
|
||||
wrappers.append(outer)
|
||||
parent = outer.parent
|
||||
if isinstance(parent, exp.SetOperation):
|
||||
continue
|
||||
if parent is None:
|
||||
out.append((node, "result", None))
|
||||
elif isinstance(parent, (exp.CTE, exp.Lateral)) or (wrappers and isinstance(parent, (exp.From, exp.Join))):
|
||||
# 别名（和列名表）记在 WITH 的那一段、LATERAL 上，FROM 里的子查询记在带别名的那层括号上
|
||||
alias = next((h.args["alias"] for h in [parent, *reversed(wrappers)] if h.args.get("alias")), None)
|
||||
listed = [c.name.lower() for c in alias.columns] if alias is not None and alias.columns else None
|
||||
out.append((node, "source", listed))
|
||||
else:
|
||||
out.append((node, "other", None))
|
||||
return out
|
||||
|
||||
|
||||
def _source_names(top: Any, listed: list[str] | None, names: set[str]) -> tuple[set[str], bool]:
|
||||
"""派生来源的结果里值取自遮罩列的那几列的名字，以及有没有认不出名字的（没起别名的表达式、两个 * 之间的项）。
|
||||
写了列名表的按位置对上列名表；没写的，单独一条 SELECT 取各项自己的列名，UNION 这类按位置对上第一支的列名。"""
|
||||
branches = _branches(top)
|
||||
found: set[str] = set()
|
||||
unknown = False
|
||||
|
||||
def take(name: str | None) -> None:
|
||||
nonlocal unknown
|
||||
if name is None:
|
||||
unknown = True
|
||||
else:
|
||||
found.add(name)
|
||||
|
||||
if listed is not None:
|
||||
for branch in branches:
|
||||
for item, place in zip(branch, _places(branch, len(listed))):
|
||||
if names & _value_columns(item):
|
||||
take(listed[place] if place is not None and place < len(listed) else None)
|
||||
return found, unknown
|
||||
first = branches[0] if branches else []
|
||||
for item in first:
|
||||
if names & _value_columns(item):
|
||||
take(_output_name(item))
|
||||
if len(branches) > 1:
|
||||
# 每一支一样宽：拿一支没有 * 的定宽度，* 后面的项才数得出位置
|
||||
width = next((len(b) for b in branches if not any(_is_star(i) for i in b)), None)
|
||||
heads = {place: _output_name(item) for item, place in zip(first, _places(first, width)) if place is not None}
|
||||
for branch in branches[1:]:
|
||||
for item, place in zip(branch, _places(branch, width)):
|
||||
if names & _value_columns(item):
|
||||
take(heads.get(place) if place is not None else None)
|
||||
return found, unknown
|
||||
|
||||
|
||||
def _tainted(sql: str, kind: str | None, masked: set[str], width: int) -> tuple[set[str], set[int]] | None:
|
||||
"""SQL 里指代遮罩列的名字（小写），和结果（width 列）里值取自遮罩列的那几列的位置。解析不了返回 None。
|
||||
|
||||
名字从遮罩列本身起，一层层传下去：选出列的值取自这些名字的（buyer_phone AS 手机号、UPPER(buyer_phone) AS x），它的
|
||||
别名也算；派生来源写了列名表的（WITH t(p) AS (SELECT buyer_phone …)），按位置对上的列名也算，直到不再多出名字。
|
||||
同名不同层的也一并算上（宁可多遮）。位置是给没起别名的表达式用的：UPPER(buyer_phone) 的结果列名对不上任何名字。"""
|
||||
名字也算；派生来源（子查询、WITH 里的一段）按位置对上列名表、UNION 第一支的列名，也算，直到不再多出名字。同名
|
||||
不同层的也一并算上（宁可多遮）。
|
||||
位置是给没起别名的表达式用的（UPPER(buyer_phone) 的结果列名对不上任何名字）：和 * 写在同一支里的，* 前面的从头数、
|
||||
后面的从末尾倒着数。认不出的——两个 * 之间的项、派生来源里没起别名的表达式（外层的 * 会把它原样带出来，列名
|
||||
又认不准）——整份快照的列都算（同解析不了）。"""
|
||||
from sqlglot import exp
|
||||
|
||||
from app.data.sql_semantics import _parse, _read
|
||||
@@ -675,28 +784,32 @@ def _tainted(sql: str, kind: str | None, masked: set[str]) -> tuple[set[str], se
|
||||
try:
|
||||
names = set(masked)
|
||||
selects = list(root.find_all(exp.Select))
|
||||
lists = [a for a in root.find_all(exp.TableAlias) if a.columns
|
||||
and isinstance(a.parent, (exp.CTE, exp.Subquery)) and isinstance(a.parent.this, exp.Select)]
|
||||
for _ in range(len(selects) + len(lists) + 1):
|
||||
blocks = _blocks(root)
|
||||
sources = [(top, listed) for top, role, listed in blocks if role == "source"]
|
||||
for _ in range(len(selects) + len(sources) + 1):
|
||||
before = len(names)
|
||||
for select_ in selects:
|
||||
for item in select_.expressions:
|
||||
name = item.alias_or_name.lower()
|
||||
if name and name not in names and names & _value_columns(item):
|
||||
names.add(name)
|
||||
for alias in lists:
|
||||
for column, item in zip(alias.columns, alias.parent.this.expressions):
|
||||
name = column.name.lower()
|
||||
if name and name not in names and (names & _value_columns(item)
|
||||
or item.alias_or_name.lower() in names):
|
||||
names.add(name)
|
||||
for top, listed in sources:
|
||||
names |= _source_names(top, listed, names)[0]
|
||||
if len(names) == before:
|
||||
break
|
||||
everything = any(_source_names(top, listed, names)[1] for top, listed in sources)
|
||||
positions: set[int] = set()
|
||||
for branch in _branches(root):
|
||||
if not any(_is_star(item) for item in branch):
|
||||
positions |= {i for i, item in enumerate(branch) if names & _value_columns(item)}
|
||||
return names, positions
|
||||
for top, role, _listed in blocks:
|
||||
if role != "result":
|
||||
continue
|
||||
for branch in _branches(top):
|
||||
for item, place in zip(branch, _places(branch, width)):
|
||||
if names & _value_columns(item):
|
||||
if place is None or place >= width:
|
||||
everything = True
|
||||
else:
|
||||
positions.add(place)
|
||||
return names, set(range(width)) if everything else positions
|
||||
except Exception: # noqa: BLE001 - 分析本身的缺陷不能让回答出错：按解析不了处理（宁可多遮）
|
||||
logger.warning("质询的 SQL 遮罩分析失败", exc_info=True)
|
||||
return None
|
||||
@@ -717,7 +830,7 @@ def _derived_hidden(snapshot: dict[str, Any], hidden: set[str], kind: str | None
|
||||
if not base or not isinstance(sql, str) or not sql.strip() or snapshot.get("source") == MERGE_SOURCE:
|
||||
return base, base
|
||||
columns = [str(c) for c in snapshot.get("columns") or []]
|
||||
found = _tainted(sql, kind, base)
|
||||
found = _tainted(sql, kind, base, len(columns))
|
||||
if found is None:
|
||||
return (base | {c.lower() for c in columns}, base) if _mentions(sql, base) else (base, base)
|
||||
names, positions = found
|
||||
@@ -791,9 +904,9 @@ class _Secrets:
|
||||
"""遮罩列里出现过的值。模型看得到原始的查询结果（工具交回的是原样），回答里照着复述、或者把值写进查询条件，
|
||||
都会把它带进回答正文和查询清单；按值找出来换成「已遮罩」。这只减少暴露，不是安全边界（同证据面板的遮罩）。
|
||||
|
||||
strong 是按值认得出的写法（够长的数、两个字以上的文字），正文、SQL 的字符串和注释、别的列里的格子都按它换；weak 是
|
||||
遮罩列里的全部写法（含短的数），只用在 SQL 里和遮罩列直接比较的字面量上。换的时候先按词、按开头两个字筛出正文里
|
||||
真有的那几个值，再拿它们现编正则：遮罩值上万个时，现编一个上万分支的正则要一秒上下，每次都卡住整个进程。"""
|
||||
strong 是按值认得出的写法（够长的数、两个字以上的文字），正文、SQL 里的字符串、名字和注释、别的列里的格子、列名都按
|
||||
它换；weak 是遮罩列里的全部写法（含短的数），只用在 SQL 里和遮罩列做等值比较的字面量上。换的时候先按词、按开头两个
|
||||
字筛出正文里真有的那几个值，再拿它们现编正则：遮罩值上万个时，现编一个上万分支的正则要一秒上下，每次都卡住整个进程。"""
|
||||
|
||||
strong: set[str] = field(default_factory=set)
|
||||
weak: set[str] = field(default_factory=set)
|
||||
@@ -823,9 +936,12 @@ class _Secrets:
|
||||
self.strong.add(form)
|
||||
self._index = None
|
||||
|
||||
def protect(self, texts: Iterable[str | None]) -> None:
|
||||
"""提问人自己写进问题、选中原文里的值不换（「13900000001 这位客户」「排除 300 以下的」）：问题本身就带着它
|
||||
（事件里原样记着），换了只会让回答、查询对不上问题。"""
|
||||
def protect(self, texts: Iterable[str | None], *, strong_only: bool = False) -> None:
|
||||
"""已经摆在读者眼前的值不换：提问人自己写进问题、选中原文里的（「13900000001 这位客户」「排除 300 以下的」），
|
||||
报告里列出的指标、运行输入的值——换了只会让回答、查询对不上问题和报告。
|
||||
|
||||
strong_only：只免正文、快照里按值认的（SQL 里照原样留着的阈值、LIMIT，见 _shown_literals）；SQL 里和遮罩列
|
||||
做等值比较的照旧换（weak 不动）。"""
|
||||
given = " ".join(t for t in texts if t)
|
||||
if not given:
|
||||
return
|
||||
@@ -835,7 +951,8 @@ class _Secrets:
|
||||
return _digits_of(form) in numbers if _NUMERIC.fullmatch(form) else form in given
|
||||
|
||||
self.strong = {f for f in self.strong if not asked(f)}
|
||||
self.weak = {f for f in self.weak if not asked(f)}
|
||||
if not strong_only:
|
||||
self.weak = {f for f in self.weak if not asked(f)}
|
||||
self._index, self._compiled = None, {}
|
||||
|
||||
def hit(self, value: Any) -> bool:
|
||||
@@ -911,19 +1028,36 @@ class _Secrets:
|
||||
return self._index
|
||||
|
||||
|
||||
def _outside_markers(text: str, secrets: _Secrets) -> str:
|
||||
"""只换引用标记之外的文字：标记里是编号、行号、列名，换了会把一个好好的引用弄坏。"""
|
||||
pattern = secrets.pattern(text)
|
||||
if pattern is None:
|
||||
#: 单元格引用的前一截（Q1.r0.）：后面才是列名
|
||||
_CELL_HEAD = re.compile(r"^\s*[QK]\d+\.r\d+\.")
|
||||
|
||||
|
||||
def _scrubbed_source(text: str, secrets: _Secrets) -> str:
|
||||
"""回答原文里的遮罩值换成「已遮罩」：标记之外按值换；标记里的按段换（| 隔开的几段各自换），单元格引用只换列名那
|
||||
一截——编号、行号换不出东西，整段换的话「r0.」会让列名开头的长号码看着像小数的一截，反而认不出来。
|
||||
标记里写着遮罩值的引用换完多半解析不了（列名里写着遮罩值的除外：快照的列名也换过了，正好对上），宁可弄坏一个
|
||||
引用，也不让值留在存下的原文、片段里，或者从解析不了的占位里露出来。"""
|
||||
if not text or not secrets.strong:
|
||||
return text
|
||||
spans = [m.span() for m in MARKER_RE.finditer(text)]
|
||||
word = _masked_word()
|
||||
return pattern.sub(lambda m: m.group(0) if any(s < m.end() and m.start() < e for s, e in spans) else word, text)
|
||||
out: list[str] = []
|
||||
last = 0
|
||||
for m in MARKER_RE.finditer(text):
|
||||
out.append(secrets.scrub(text[last:m.start()]))
|
||||
start, end = m.span("body")
|
||||
parts = text[start:end].split("|")
|
||||
head = _CELL_HEAD.match(parts[0]) if m.group("kind").strip() == "v" else None
|
||||
cut = head.end() if head else 0
|
||||
parts = [parts[0][:cut] + secrets.scrub(parts[0][cut:]), *(secrets.scrub(p) for p in parts[1:])]
|
||||
out.append(text[m.start():start] + "|".join(parts) + text[end:m.end()])
|
||||
last = m.end()
|
||||
out.append(secrets.scrub(text[last:]))
|
||||
return "".join(out)
|
||||
|
||||
|
||||
def _masked_view(snapshot: dict[str, Any], hidden: set[str], secrets: _Secrets) -> tuple[dict[str, Any], list[str]]:
|
||||
"""一份查询快照遮好的样子和遮掉的列：要遮的列（hidden，小写）每一格都换成「已遮罩」、类型记成文本（渲染时不当数）；
|
||||
别的列里写着按值认得出的遮罩值的格子也换（另一张表里没设遮罩的同一个手机号、改了名字没认出来的遮罩列）。
|
||||
列名里写着遮罩值的（agent 把遮罩值写进了别名，「13900000001 的单」），列名也换（同 SQL 里的名字，_scrub_sql）。
|
||||
原快照不动；没有要遮的返回原快照。"""
|
||||
columns = [str(c) for c in snapshot.get("columns") or []]
|
||||
index = {i for i, c in enumerate(columns) if c.lower() in hidden}
|
||||
@@ -933,6 +1067,7 @@ def _masked_view(snapshot: dict[str, Any], hidden: set[str], secrets: _Secrets)
|
||||
word = _masked_word()
|
||||
touched = set(index)
|
||||
position = {c: i for i, c in enumerate(columns)}
|
||||
names = [secrets.scrub(c) for c in columns] if by_value else columns
|
||||
|
||||
def cell(i: int | None, value: Any) -> Any:
|
||||
if i in index:
|
||||
@@ -943,82 +1078,132 @@ def _masked_view(snapshot: dict[str, Any], hidden: set[str], secrets: _Secrets)
|
||||
return word
|
||||
return value
|
||||
|
||||
def key(k: Any) -> Any:
|
||||
i = position.get(str(k))
|
||||
return names[i] if i is not None else k
|
||||
|
||||
rows: list[Any] = []
|
||||
for row in snapshot["rows"]:
|
||||
if isinstance(row, list):
|
||||
rows.append([cell(i, v) for i, v in enumerate(row)])
|
||||
elif isinstance(row, dict):
|
||||
rows.append({k: word if str(k).lower() in hidden else cell(position.get(str(k)), v)
|
||||
rows.append({key(k): word if str(k).lower() in hidden else cell(position.get(str(k)), v)
|
||||
for k, v in row.items()})
|
||||
else:
|
||||
rows.append(row)
|
||||
if not touched:
|
||||
renamed = names != columns
|
||||
if not touched and not renamed:
|
||||
return snapshot, []
|
||||
types = dict(snapshot["column_types"]) if isinstance(snapshot.get("column_types"), dict) else {}
|
||||
types.update({columns[i]: "text" for i in index})
|
||||
return {**snapshot, "rows": rows, "column_types": types}, [columns[i] for i in sorted(touched)]
|
||||
types = {key(k): v for k, v in snapshot["column_types"].items()} \
|
||||
if isinstance(snapshot.get("column_types"), dict) else {}
|
||||
types.update({names[i]: "text" for i in index})
|
||||
view = {**snapshot, "rows": rows, "column_types": types, **({"columns": names} if renamed else {})}
|
||||
return view, [names[i] for i in sorted(touched)]
|
||||
|
||||
|
||||
_COMPARE_TOKENS = frozenset({"EQ", "NEQ", "GT", "GTE", "LT", "LTE", "LIKE", "ILIKE", "GLOB", "RLIKE", "NULLSAFE_EQ"})
|
||||
_LITERAL_TOKENS = frozenset({"STRING", "NUMBER"})
|
||||
#: 在找某一个值的比较：col = 'x'、col <> 'x'、col LIKE 'x'、SQLite 的 col IS 'x'（IN 列表另外认）
|
||||
_EQUAL_TOKENS = frozenset({"EQ", "NEQ", "NULLSAFE_EQ", "LIKE", "ILIKE", "GLOB", "RLIKE", "IS"})
|
||||
#: 范围比较：边界是阈值，不是在找某一个值（BETWEEN 另外认）
|
||||
_RANGE_TOKENS = frozenset({"GT", "GTE", "LT", "LTE"})
|
||||
_NAME_TOKENS = frozenset({"VAR", "IDENTIFIER"})
|
||||
#: 字面量前面可能连着的：正负号、MySQL 的字符集前缀（_utf8mb4'x'）
|
||||
_LEADS = frozenset({"DASH", "PLUS", "INTRODUCER"})
|
||||
#: 能结束一个值的词：减号跟在它们后面是两数相减，不是负号
|
||||
_VALUE_ENDS = frozenset({"NUMBER", "VAR", "IDENTIFIER", "R_PAREN", "R_BRACKET", "NULL", "TRUE", "FALSE"})
|
||||
|
||||
|
||||
def _against(tokens: list[Any], i: int, compared: set[str]) -> bool:
|
||||
"""第 i 个字面量是不是直接和遮罩列比较：col = 'x'、col IN ('x', …)、col NOT LIKE 'x'、col BETWEEN a AND b，
|
||||
以及 'x' = col。列可以带表名（s.buyer_phone），可以是遮罩列的别名（compared 里已经有）。"""
|
||||
kinds = [t.token_type.name for t in tokens]
|
||||
j = i - 1
|
||||
def _is_literal(kind: str) -> bool:
|
||||
"""字符串（含 N'…'、E'…'、$$…$$、U&'…' 这些带前缀的变体）和数字。"""
|
||||
return kind == "NUMBER" or kind.endswith("STRING")
|
||||
|
||||
|
||||
def _operand_start(kinds: list[str], k: int) -> int:
|
||||
"""第 k 个字面量连同它前面的负号、字符集前缀从哪个词开始（-25、_utf8mb4'x'）。"""
|
||||
while k >= 1 and kinds[k - 1] in _LEADS and (
|
||||
kinds[k - 1] == "INTRODUCER" or k < 2 or not (kinds[k - 2] in _VALUE_ENDS or _is_literal(kinds[k - 2]))):
|
||||
k -= 1
|
||||
return k
|
||||
|
||||
|
||||
def _name_at(tokens: list[Any], kinds: list[str], k: int) -> str | None:
|
||||
return tokens[k].text.lower() if 0 <= k < len(kinds) and kinds[k] in _NAME_TOKENS else None
|
||||
|
||||
|
||||
def _context(tokens: list[Any], kinds: list[str], i: int) -> tuple[str, str | None]:
|
||||
"""第 i 个字面量用在哪：("equal", 列名) 在找某一个值——col = 'x'、col NOT LIKE 'x'、col IS NOT 25、
|
||||
col IN ('x', …)、'x' = col；("range", None) 范围比较的边界——col > 25、col BETWEEN a AND b；("other", None) 别的
|
||||
（LIMIT / OFFSET、选出列、函数参数、算式）。列名是比较另一边的那个名字（s.buyer_phone 取 buyer_phone，小写），
|
||||
另一边不是单独一个名字（函数、算式）时为 None。"""
|
||||
j = _operand_start(kinds, i) - 1
|
||||
if j >= 1 and kinds[j] == "NOT" and kinds[j - 1] == "IS":
|
||||
return "equal", _name_at(tokens, kinds, j - 2)
|
||||
if j >= 0 and kinds[j] in _EQUAL_TOKENS:
|
||||
return "equal", _name_at(tokens, kinds, j - 2 if j >= 1 and kinds[j - 1] == "NOT" else j - 1)
|
||||
if j >= 0 and (kinds[j] in _RANGE_TOKENS or kinds[j] == "BETWEEN"):
|
||||
return "range", None
|
||||
if j >= 1 and kinds[j] == "AND" and _is_literal(kinds[j - 1]):
|
||||
# BETWEEN a AND b 的 b
|
||||
lower = _operand_start(kinds, j - 1) - 1
|
||||
if lower >= 0 and kinds[lower] == "BETWEEN":
|
||||
return "range", None
|
||||
# IN 列表：往回越过逗号和别的字面量，找到「IN (」
|
||||
k = j
|
||||
while k >= 0 and kinds[k] in ("COMMA", "STRING", "NUMBER", "DASH"):
|
||||
while k >= 0 and (kinds[k] == "COMMA" or kinds[k] in _LEADS or _is_literal(kinds[k])):
|
||||
k -= 1
|
||||
if k >= 1 and kinds[k] == "L_PAREN" and kinds[k - 1] == "IN":
|
||||
j = k - 2
|
||||
elif j >= 0 and kinds[j] in _COMPARE_TOKENS:
|
||||
j -= 1
|
||||
elif j >= 2 and kinds[j] == "AND" and kinds[j - 1] in _LITERAL_TOKENS and kinds[j - 2] == "BETWEEN":
|
||||
j -= 3
|
||||
elif j >= 0 and kinds[j] == "BETWEEN":
|
||||
j -= 1
|
||||
else:
|
||||
# 字面量在左边：'x' = col、'x' = t.col
|
||||
if i + 2 < len(tokens) and kinds[i + 1] in _COMPARE_TOKENS and kinds[i + 2] in _NAME_TOKENS:
|
||||
last = i + 2
|
||||
while last + 2 < len(tokens) and kinds[last + 1] == "DOT" and kinds[last + 2] in _NAME_TOKENS:
|
||||
last += 2
|
||||
return tokens[last].text.lower() in compared
|
||||
return False
|
||||
if j >= 0 and kinds[j] == "NOT":
|
||||
j -= 1
|
||||
return j >= 0 and kinds[j] in _NAME_TOKENS and tokens[j].text.lower() in compared
|
||||
return "equal", _name_at(tokens, kinds, k - 3 if k >= 2 and kinds[k - 2] == "NOT" else k - 2)
|
||||
# 字面量在左边：'x' = col、25 < t.col
|
||||
after = i + 1
|
||||
if after < len(kinds) and kinds[after] in _RANGE_TOKENS:
|
||||
return "range", None
|
||||
if after < len(kinds) and kinds[after] in _EQUAL_TOKENS:
|
||||
last = after + 1
|
||||
while last + 2 < len(kinds) and kinds[last + 1] == "DOT" and kinds[last + 2] in _NAME_TOKENS:
|
||||
last += 2
|
||||
return "equal", _name_at(tokens, kinds, last)
|
||||
return "other", None
|
||||
|
||||
|
||||
def _replace_literal(raw: str, kind: str, change: Callable[[str], str]) -> str:
|
||||
"""字面量的原文 → 换过的原文：字符串只换引号里面的（引号、N'' 这类前缀照留），数整个换。"""
|
||||
if kind == "NUMBER":
|
||||
return change(raw)
|
||||
for quote in ("'", '"'):
|
||||
start = raw.find(quote)
|
||||
if start >= 0 and len(raw) - start >= 2 and raw.endswith(quote):
|
||||
return raw[:start + 1] + change(raw[start + 1:-1]) + quote
|
||||
def _requote(raw: str, change: Callable[[str], str]) -> str:
|
||||
"""一个字面量或名字的原文 → 只换里面的字：引号、N'' E'' U&'' 这类前缀、$tag$ 照留；不带引号的整个交给 change。"""
|
||||
if raw.startswith("$") and (end := raw.find("$", 1)) > 0:
|
||||
tag = raw[:end + 1]
|
||||
if len(raw) >= 2 * len(tag) and raw.endswith(tag):
|
||||
return tag + change(raw[len(tag):-len(tag)]) + tag
|
||||
for opening, closing in (("'", "'"), ('"', '"'), ("`", "`"), ("[", "]")):
|
||||
start = raw.find(opening)
|
||||
if start >= 0 and len(raw) - start >= 2 and raw.endswith(closing):
|
||||
return raw[:start + 1] + change(raw[start + 1:-1]) + closing
|
||||
return change(raw)
|
||||
|
||||
|
||||
def _scrub_sql(sql: str, secrets: _Secrets, compared: set[str], kind: str | None) -> str:
|
||||
"""SQL 里遮罩列的值换成「已遮罩」，别的一个字不动：查询清单随运行封存，它得是实际执行的那条，审批人看 what-if 的
|
||||
条件（阈值、日期、LIMIT、别的列的条件）看的就是它。
|
||||
|
||||
- 字符串字面量、注释里按值换（同正文：够长的数、两个字以上的文字）；数字字面量够长、又是遮罩值的整个换；
|
||||
- 直接和遮罩列（或它的别名，compared）比较的字面量是遮罩列里出现过的值就换，短的数也换（age = 25）。
|
||||
分词出错（引号没闭合这类）就整条按正文的办法换。"""
|
||||
if not sql or not (secrets.strong or secrets.weak):
|
||||
return sql
|
||||
def _sql_tokens(sql: str, kind: str | None) -> tuple[list[Any], list[str]] | None:
|
||||
from app.data.sql_semantics import _read, _tokens
|
||||
|
||||
try:
|
||||
tokens = _tokens(sql, _read(kind))
|
||||
except Exception: # noqa: BLE001 - 分词器认不了的写法：按正文的办法换
|
||||
except Exception: # noqa: BLE001 - 分词器认不了的写法（引号没闭合这类）
|
||||
return None
|
||||
return tokens, [t.token_type.name for t in tokens]
|
||||
|
||||
|
||||
def _scrub_sql(sql: str, secrets: _Secrets, compared: set[str], kind: str | None) -> str:
|
||||
"""SQL 里遮罩列的值换成「已遮罩」，别的一个字不动：查询清单随运行封存，它得是实际执行的那条，审批人看 what-if 的
|
||||
条件（阈值、日期、LIMIT、别的列的条件）看的就是它。只换在找某一个值的字面量（_context 的 equal）：
|
||||
|
||||
- 和遮罩列（或它的别名，compared）做等值、IN、LIKE 比较的，是遮罩列里出现过的值就换，短的数也换（age = 25）；
|
||||
和别的列做等值比较的，够长的遮罩值也换（另一张表里没设遮罩的同一个手机号）；
|
||||
- 范围比较的边界（age > 31、BETWEEN 10000 AND 50000）、LIMIT / OFFSET、选出列和算式里的数照原样：它们是 agent
|
||||
自己定的阈值，不是在找某一个人，换了封存的就不是实际执行的语句；
|
||||
- 别处的字符串（LIKE 别的列、函数参数）、注释、名字（AS "13900000001 的单"、张伟的单数）按值换（同正文：够长的数、
|
||||
两个字以上的文字）。
|
||||
分词出错就整条按正文的办法换。"""
|
||||
if not sql or not (secrets.strong or secrets.weak):
|
||||
return sql
|
||||
split = _sql_tokens(sql, kind)
|
||||
if split is None:
|
||||
return secrets.scrub(sql)
|
||||
tokens, kinds = split
|
||||
word = _masked_word()
|
||||
out: list[str] = []
|
||||
last = 0
|
||||
@@ -1029,14 +1214,14 @@ def _scrub_sql(sql: str, secrets: _Secrets, compared: set[str], kind: str | None
|
||||
gap = sql[last:start]
|
||||
out.append(secrets.scrub(gap) if gap.strip() else gap)
|
||||
raw = sql[start:end]
|
||||
kind_of = token.token_type.name
|
||||
if kind_of in _LITERAL_TOKENS:
|
||||
if compared and secrets.weak_hit(token.text) and _against(tokens, i, compared):
|
||||
raw = _replace_literal(raw, kind_of, lambda _: word)
|
||||
elif kind_of == "STRING":
|
||||
raw = _replace_literal(raw, kind_of, secrets.scrub)
|
||||
elif secrets.hit(token.text):
|
||||
raw = word
|
||||
if _is_literal(kinds[i]):
|
||||
use, column = _context(tokens, kinds, i)
|
||||
if use == "equal" and ((column in compared and secrets.weak_hit(token.text)) or secrets.hit(token.text)):
|
||||
raw = _requote(raw, lambda _: word)
|
||||
elif use != "range" and kinds[i] != "NUMBER":
|
||||
raw = _requote(raw, secrets.scrub)
|
||||
elif kinds[i] in _NAME_TOKENS:
|
||||
raw = _requote(raw, secrets.scrub)
|
||||
out.append(raw)
|
||||
last = end
|
||||
rest = sql[last:]
|
||||
@@ -1044,6 +1229,22 @@ def _scrub_sql(sql: str, secrets: _Secrets, compared: set[str], kind: str | None
|
||||
return "".join(out)
|
||||
|
||||
|
||||
def _shown_literals(sql: str, kind: str | None) -> list[str]:
|
||||
"""SQL 里不管是什么值都照原样留着的字面量（_scrub_sql 不换的两种）：范围比较的边界，LIMIT、选出列这类地方的数。
|
||||
它们随 SQL 摆在读者眼前，回答正文、快照里同样的值再遮也遮不住什么，只会让回答和 SQL 对不上。"""
|
||||
split = _sql_tokens(sql, kind) if sql else None
|
||||
if split is None:
|
||||
return []
|
||||
tokens, kinds = split
|
||||
out: list[str] = []
|
||||
for i, token in enumerate(tokens):
|
||||
if _is_literal(kinds[i]):
|
||||
use = _context(tokens, kinds, i)[0]
|
||||
if use == "range" or (use == "other" and kinds[i] == "NUMBER"):
|
||||
out.append(token.text)
|
||||
return out
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Plan:
|
||||
"""一次回答（或一次取质询工件）的遮罩：agent 每份查询快照遮好的样子（views）和遮掉的列（masked）、SQL 里指代
|
||||
@@ -1068,9 +1269,11 @@ async def _source_kinds(names: Iterable[str]) -> dict[str, str | None]:
|
||||
return {str(name): kind for name, kind in rows}
|
||||
|
||||
|
||||
async def _plan(own: dict[str, dict[str, Any]], others: list[dict[str, Any]], given: list[str | None]) -> _Plan:
|
||||
async def _plan(own: dict[str, dict[str, Any]], others: list[dict[str, Any]], given: list[str | None],
|
||||
shown: list[tuple[str, str]]) -> _Plan:
|
||||
"""own：agent 自己的查询快照（工件 id → 内容）；others：报告自己的查询快照（被问的那句可能直接写着遮罩列的值，
|
||||
agent 会照着复述、写进查询条件）；given：问题和选中的原文（里面的值不换）。
|
||||
agent 会照着复述、写进查询条件）；given：问题、选中的原文、报告里列出的指标和运行输入的值（都已经摆在读者眼前，
|
||||
里面的值不换）；shown：agent 执行的 SQL 和它查的数据源名（照原样留在 SQL 里的阈值、LIMIT，正文里同样不换）。
|
||||
|
||||
查数据源的遮罩要读库，在这里读；解析 SQL、收上万个值、逐格比对是纯计算，放进线程：放在事件循环里会挡住同一进程
|
||||
里别的运行（同 tools/datasource 的 _sql_checks）。"""
|
||||
@@ -1079,13 +1282,14 @@ async def _plan(own: dict[str, dict[str, Any]], others: list[dict[str, Any]], gi
|
||||
masks = _Masks()
|
||||
named = {artifact: await _hidden_columns(snapshot, masks) for artifact, snapshot in own.items()}
|
||||
based = [(snapshot, (await _hidden_columns(snapshot, masks))[0]) for snapshot in others]
|
||||
kinds = await _source_kinds(str(s.get("source") or "") for s in [*own.values(), *others])
|
||||
return await asyncio.to_thread(_settle_plan, own, named, based, kinds, given)
|
||||
kinds = await _source_kinds([*(str(s.get("source") or "") for s in [*own.values(), *others]),
|
||||
*(name for _, name in shown)])
|
||||
return await asyncio.to_thread(_settle_plan, own, named, based, kinds, given, shown)
|
||||
|
||||
|
||||
def _settle_plan(own: dict[str, dict[str, Any]], named: dict[str, tuple[set[str], list[str]]],
|
||||
based: list[tuple[dict[str, Any], set[str]]], kinds: dict[str, str | None],
|
||||
given: list[str | None]) -> _Plan:
|
||||
given: list[str | None], shown: list[tuple[str, str]]) -> _Plan:
|
||||
secrets = _Secrets()
|
||||
hidden: dict[str, set[str]] = {}
|
||||
tainted: dict[str, set[str]] = {}
|
||||
@@ -1096,6 +1300,8 @@ def _settle_plan(own: dict[str, dict[str, Any]], named: dict[str, tuple[set[str]
|
||||
for snapshot, base in based:
|
||||
secrets.add(snapshot, _derived_hidden(snapshot, base, kinds.get(str(snapshot.get("source") or "")))[0])
|
||||
secrets.protect(given)
|
||||
secrets.protect([value for sql, name in shown for value in _shown_literals(sql, kinds.get(name))],
|
||||
strong_only=True)
|
||||
secrets.prepare()
|
||||
views: dict[str, dict[str, Any]] = {}
|
||||
masked: dict[str, list[str]] = {}
|
||||
@@ -1124,7 +1330,7 @@ def _masking_loader(views: dict[str, dict[str, Any]]) -> Callable[[str], Any]:
|
||||
|
||||
|
||||
class _Leak(Exception):
|
||||
"""组装了两遍，正文里还认得出遮罩值（渲染出来才露的，比如恰好是遮罩值的运行输入）：这份回答不交出去。"""
|
||||
"""组装出来的正文里还认得出遮罩值（渲染出来才露的，比如别的列里一格文字夹着手机号）：这份回答不交出去。"""
|
||||
|
||||
|
||||
def _limit(outcome: VerifyOutcome) -> str | None:
|
||||
@@ -1195,7 +1401,8 @@ async def _answer(op: _Opening, ask: Ask, target: _Target | None, settings: dict
|
||||
own[str(query["artifact"])] = snapshot
|
||||
others = [snapshot for entry in catalog.values() if isinstance(entry, dict) and entry.get("kind") == "query"
|
||||
and _query_shaped(snapshot := _quiet_load(op.loader, entry.get("artifact")))]
|
||||
plan = await _plan(own, others, [ask.question, ask.quote])
|
||||
plan = await _plan(own, others, [ask.question, ask.quote, *_ref_values(refs)],
|
||||
[(str(q.get("sql") or ""), str(q.get("source") or "")) for q in outcome.queries])
|
||||
|
||||
text, doc_id, failed = "", None, None
|
||||
source = (outcome.rationale or {}).get("source") if outcome.result is not None else None
|
||||
@@ -1249,21 +1456,28 @@ def _quiet_load(loader: Callable[[str], Any], artifact: Any) -> Any:
|
||||
def _compose(source: str, catalog: dict[str, Any], *, run_id: str, views: dict[str, dict[str, Any]],
|
||||
secrets: _Secrets, allow: list[str]) -> dict[str, Any]:
|
||||
"""agent 写的回答原文（带引用标记）→ 回答文档：目录是 agent 自己的查询，加上报告里的指标和运行输入；快照按遮罩
|
||||
遮好再渲染。遮罩值先在标记之外的文字里换掉；换完正文里还认得出遮罩值（标记里写了值，解析不了，原样出现在正文里），
|
||||
就连标记一起换了再组装一遍——宁可弄坏一个引用，也不让值露出来；还认得出（渲染出来才露的），抛 _Leak，不交出去。
|
||||
allow 是提问人自己写的数（_given_numbers），回答里复述它们不标「无出处」。"""
|
||||
遮好再渲染。遮罩值在组装之前换掉（_scrubbed_source）；组装出来正文里还认得出遮罩值（渲染出来才露的，比如别的列
|
||||
里一格文字夹着手机号），抛 _Leak，不交出去。allow 是提问人自己写的数（_given_numbers），回答里复述它们不标
|
||||
「无出处」。"""
|
||||
loader = _masking_loader(views)
|
||||
allowed = list(allow) or None
|
||||
doc = compose_doc(_outside_markers(source, secrets), catalog, run_id=run_id, loader=loader, allow_numbers=allowed)
|
||||
# 列名里写着遮罩值的查询：目录里的列名照快照换（_masked_view），整表展开、表头核对、存下的目录都按换过的名字
|
||||
catalog = {alias: _renamed_entry(entry, secrets) for alias, entry in catalog.items()}
|
||||
doc = compose_doc(_scrubbed_source(source, secrets), catalog, run_id=run_id, loader=loader,
|
||||
allow_numbers=allowed)
|
||||
markdown = str(doc.get("markdown") or "")
|
||||
if secrets.scrub(markdown) != markdown:
|
||||
doc = compose_doc(secrets.scrub(source), catalog, run_id=run_id, loader=loader, allow_numbers=allowed)
|
||||
markdown = str(doc.get("markdown") or "")
|
||||
if secrets.scrub(markdown) != markdown:
|
||||
raise _Leak()
|
||||
raise _Leak()
|
||||
return doc
|
||||
|
||||
|
||||
def _renamed_entry(entry: Any, secrets: _Secrets) -> Any:
|
||||
if not (isinstance(entry, dict) and entry.get("kind") == "query" and isinstance(entry.get("columns"), list)):
|
||||
return entry
|
||||
names = [secrets.scrub(str(c)) for c in entry["columns"]]
|
||||
return {**entry, "columns": names} if names != [str(c) for c in entry["columns"]] else entry
|
||||
|
||||
|
||||
def _with_notes(answer: dict[str, Any], op: _Opening, outcome: VerifyOutcome, part: dict[str, Any], lost: list[str],
|
||||
own: dict[str, dict[str, Any]]) -> dict[str, Any]:
|
||||
"""回答末尾的系统说明：用不了的数据源、模型的轮数用完了、有查询执行失败、有引用解析不了、封存之后查的是现在的数据。
|
||||
@@ -1393,11 +1607,12 @@ def _cited_cards(doc_id: Any) -> set[str]:
|
||||
if isinstance(e, dict) and e.get("kind") in ("metric", "input") and e.get("artifact")}
|
||||
|
||||
|
||||
async def _report_snapshots(run_id: str, reports: Iterable[Any]) -> list[dict[str, Any]]:
|
||||
"""这几份报告（报告节点 id）的文档目录里的查询快照：和回答时收遮罩值的是同一批（报告最后一次 report.checked 的文档）。"""
|
||||
async def _report_material(run_id: str, reports: Iterable[Any]) -> tuple[list[dict[str, Any]], list[str]]:
|
||||
"""这几份报告（报告节点 id）的文档目录里的查询快照，和列出的指标、运行输入的值：和回答时收遮罩值、免遮的是同一批
|
||||
（报告最后一次 report.checked 的文档）。"""
|
||||
wanted = sorted({str(r) for r in reports if isinstance(r, str) and r})
|
||||
if not wanted:
|
||||
return []
|
||||
return [], []
|
||||
async with SessionLocal() as session:
|
||||
rows = await session.execute(
|
||||
select(RunEvent.node_id, RunEvent.data)
|
||||
@@ -1405,16 +1620,18 @@ async def _report_snapshots(run_id: str, reports: Iterable[Any]) -> list[dict[st
|
||||
.order_by(RunEvent.seq))
|
||||
docs = {node_id: str(data["doc_artifact"]) for node_id, data in rows
|
||||
if isinstance(data, dict) and data.get("doc_artifact")}
|
||||
out: list[dict[str, Any]] = []
|
||||
snapshots: list[dict[str, Any]] = []
|
||||
values: list[str] = []
|
||||
for artifact in docs.values():
|
||||
doc = _quiet_load(artifact_store.load, artifact)
|
||||
catalog = doc.get("catalog") if isinstance(doc, dict) and isinstance(doc.get("catalog"), dict) else {}
|
||||
values += _ref_values(_report_refs(catalog))
|
||||
for entry in catalog.values():
|
||||
if isinstance(entry, dict) and entry.get("kind") == "query":
|
||||
snapshot = _quiet_load(artifact_store.load, entry.get("artifact"))
|
||||
if _query_shaped(snapshot):
|
||||
out.append(snapshot)
|
||||
return out
|
||||
snapshots.append(snapshot)
|
||||
return snapshots, values
|
||||
|
||||
|
||||
async def inquiry_artifact(run_id: str, artifact_id: str) -> dict[str, Any]:
|
||||
@@ -1440,15 +1657,16 @@ async def inquiry_artifact(run_id: str, artifact_id: str) -> dict[str, Any]:
|
||||
if not _query_shaped(content):
|
||||
return {"id": artifact_id, "content": content, "masked": []}
|
||||
# 遮罩值照回答时收：引用它的那几条质询各自查到的全部快照（查询条件里写着的值，不一定在这一份的结果里），加上报告
|
||||
# 自己的查询快照；提问人自己写的值不换
|
||||
# 自己的查询快照；免遮的也照回答时：提问人自己写的、报告里列出的指标和运行输入的值、SQL 里照原样留着的阈值
|
||||
own: dict[str, dict[str, Any]] = {artifact_id: content}
|
||||
for data in referenced:
|
||||
for other in dict.fromkeys(str(a) for a in data.get("inquiry_artifacts") or []):
|
||||
if other not in own and _query_shaped(snapshot := _quiet_load(artifact_store.load, other)):
|
||||
own[other] = snapshot
|
||||
others = await _report_snapshots(run_id, (d.get("report") for d in referenced))
|
||||
plan = await _plan(own, others, [t for d in referenced for t in (d.get("question"), d.get("quote"))
|
||||
if isinstance(t, str)])
|
||||
others, values = await _report_material(run_id, (d.get("report") for d in referenced))
|
||||
asked = [t for d in referenced for t in (d.get("question"), d.get("quote")) if isinstance(t, str)]
|
||||
plan = await _plan(own, others, [*asked, *values],
|
||||
[(str(s.get("sql") or ""), str(s.get("source") or "")) for s in own.values()])
|
||||
shown = plan.views[artifact_id]
|
||||
recorded = next((q["sql"] for d in referenced for q in d.get("queries") or []
|
||||
if isinstance(q, dict) and q.get("artifact") == artifact_id and isinstance(q.get("sql"), str)),
|
||||
|
||||
@@ -11,8 +11,9 @@
|
||||
设置里关了 409；没停在那张审批卡上、还在执行（按情形说下一步）、封存核对不通过、开放途中运行被继续了都 409；
|
||||
事件里记的报告文档不是这次运行的 422；输入不合法 422，指明是哪一项
|
||||
- 回答：先给结论，数字只能引用自己查询的格子或报告已有的指标，裸数字被标出（提问人自己写的数不标）；遮罩列（起了
|
||||
别名的也算）的格子渲染成「已遮罩」，回答正文、查询清单、事件、质询工件里都不出现遮罩列的值；短的遮罩值不误伤阈值、
|
||||
天数、LIMIT
|
||||
别名的、和 * 写在一起没起别名的表达式也算）的格子渲染成「已遮罩」，回答正文、查询清单、事件、质询工件里都不出现
|
||||
遮罩列的值；遮罩值不误伤阈值、天数、LIMIT（长的数也一样）、提问人自己写的数、报告口径卡里的指标；质询工件按值遮
|
||||
只出现在报告快照里的遮罩值
|
||||
- 接线：设置里的上限、运行固定的数据版本交到 agent 手里（上传表格重传之后查的仍是运行固定的那一版）；同名重建的
|
||||
数据源、没固定过版本的上传表格不查
|
||||
- 触顶、失败照样记一条事件，说法给提问人看得懂、说下一步；轮数用完不算查询次数到顶，查询失败、引用解析不了在末尾注明
|
||||
@@ -162,6 +163,14 @@ def gated(source: str, *, gate: bool = True) -> dict:
|
||||
)
|
||||
|
||||
|
||||
def fetching(source: str, sql: str) -> dict:
|
||||
"""取数（给定的 SQL）→ 报告 → 人工审批 → 成果。"""
|
||||
return chain(node("start", "input"), node("fetch", "tool", tool=f"db_query__{source}", args={"sql": sql}),
|
||||
node("write", "report", "写周报", instructions="写周报"),
|
||||
node("gate", "human", mode="approve", title="周报可以出具吗"),
|
||||
node("out", "output", fields=[{"name": "周报", "value": "{{ nodes.write.text }}"}]))
|
||||
|
||||
|
||||
@dataclass
|
||||
class Play:
|
||||
"""质询 agent 对某个问题的剧本：先看几次表结构（schema），再依次执行的查询，最后交的回答。final 给了就照它交
|
||||
@@ -979,6 +988,154 @@ async def test_numbers_that_merely_match_a_short_masked_value_are_left_alone(mon
|
||||
await engines.invalidate(source_id)
|
||||
|
||||
|
||||
STAR_TAIL = "SELECT *, SUBSTR(buyer_phone, 4) FROM sales ORDER BY amount DESC LIMIT 1"
|
||||
|
||||
|
||||
async def test_an_unnamed_expression_next_to_a_star_is_masked_everywhere(monkeypatch, masked_shop, client):
|
||||
"""和 * 写在同一支里、没起别名的表达式（取手机号的后几位）：结果的列名对不上任何名字，按位置认（从末尾倒着数）。
|
||||
回答、事件、回答文档、质询工件里都不露出原值的任何一截。"""
|
||||
plays = {"买得最多": Play([STAR_TAIL], "金额最高的一笔来自尾号 [[v:Q1.r0.SUBSTR(buyer_phone, 4)]] 的客户，"
|
||||
"金额 [[v:Q1.r0.amount]]。[[see:Q1]]")}
|
||||
Models(monkeypatch, masked_shop, plays=plays)
|
||||
paused = await pause(masked_shop)
|
||||
[approval] = await pending(paused.id)
|
||||
r = await ask(client, paused.id, report="write", unit=None, question="买得最多的是哪位客户？",
|
||||
approval_id=approval.id)
|
||||
assert r.status_code == 200, r.text
|
||||
item = r.json()["inquiry"]
|
||||
assert item["answer"]["text"] == "金额最高的一笔来自尾号 已遮罩 的客户，金额 1,200。"
|
||||
[query] = item["queries"]
|
||||
assert query["sql"] == STAR_TAIL
|
||||
[event] = await events(paused.id, "approval.inquiry")
|
||||
doc = artifact_store.load(item["answer"]["doc_artifact"])
|
||||
got = await client.get(f"/api/runs/{paused.id}/inquiries/artifacts/{query['artifact']}")
|
||||
assert got.status_code == 200, got.text
|
||||
for shown in (r.text, json.dumps(event.data, ensure_ascii=False), json.dumps(doc, ensure_ascii=False), got.text):
|
||||
assert "00000001" not in shown and "1390000000" not in shown
|
||||
assert got.json()["content"]["rows"] == [[1, "华东", "已遮罩", 1200, 0, "已遮罩"]]
|
||||
assert got.json()["masked"] == ["buyer_phone", "SUBSTR(buyer_phone, 4)"]
|
||||
|
||||
|
||||
#: 金额设了遮罩的大单：金额都是五位数，和阈值、LIMIT、访问次数的条件撞在一起
|
||||
BIG = ("CREATE TABLE sales (id INTEGER PRIMARY KEY, region TEXT, amount INTEGER, visits INTEGER);"
|
||||
"INSERT INTO sales VALUES (1, '华东', 10000, 25000), (2, '华东', 12000, 18000), (3, '华南', 20000, 30000),"
|
||||
" (4, '华南', 30000, 22000), (5, '华北', 50000, 9000);")
|
||||
#: 报告取数选了金额：金额列里的值都算遮罩值
|
||||
AMOUNTS = "SELECT region AS 区域, amount AS 金额 FROM sales ORDER BY id"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def big_shop(tmp_path):
|
||||
name, source_id = await _source(tmp_path, options={"mask_columns": "amount"}, script=BIG)
|
||||
yield name
|
||||
await engines.invalidate(source_id)
|
||||
|
||||
|
||||
async def test_thresholds_and_limits_are_sealed_as_written_even_when_they_match_masked_values(
|
||||
monkeypatch, big_shop, client):
|
||||
"""遮罩列的值都是五位数：范围比较的边界（和金额比的、和访问次数比的）、LIMIT 撞上遮罩值也照原样封存，回答正文里
|
||||
复述它们也不遮——它们就摆在 SQL 里。在找某一笔的（amount = 50000）照旧遮。"""
|
||||
large = "SELECT COUNT(*) AS 笔数 FROM sales WHERE amount >= 30000 AND visits > 20000 LIMIT 10000"
|
||||
exact = "SELECT COUNT(*) AS 笔数 FROM sales WHERE amount = 50000"
|
||||
plays = {"大单": Play([large, exact], "金额不低于 30000 元、访问超过 20000 次的大单有 [[v:Q1.r0.笔数]] 笔，"
|
||||
"恰好 50000 元的有 [[v:Q2.r0.笔数]] 笔。[[see:Q1]]")}
|
||||
Models(monkeypatch, big_shop, writer="华东首单金额 [[v:Q1.r0.金额]]。[[see:Q1]]", plays=plays)
|
||||
paused = await pause(big_shop, graph=fetching(big_shop, AMOUNTS))
|
||||
[approval] = await pending(paused.id)
|
||||
r = await ask(client, paused.id, report="write", unit=None, question="3 万元以上的大单有几笔？",
|
||||
approval_id=approval.id)
|
||||
assert r.status_code == 200, r.text
|
||||
item = r.json()["inquiry"]
|
||||
assert item["answer"]["text"] == "金额不低于 30000 元、访问超过 20000 次的大单有 1 笔，恰好 已遮罩 元的有 1 笔。"
|
||||
sealed = [large, "SELECT COUNT(*) AS 笔数 FROM sales WHERE amount = 已遮罩"]
|
||||
assert [q["sql"] for q in item["queries"]] == sealed
|
||||
[event] = await events(paused.id, "approval.inquiry")
|
||||
assert [q["sql"] for q in event.data["queries"]] == sealed
|
||||
got = await client.get(f"/api/runs/{paused.id}/inquiries/artifacts/{item['queries'][0]['artifact']}")
|
||||
assert got.status_code == 200 and got.json()["content"]["sql"] == large
|
||||
doc = artifact_store.load(item["answer"]["doc_artifact"])
|
||||
assert [s["text"] for s in iter_segments(doc) if s.get("issue") == "uncited_number"] == ["30000", "20000"], \
|
||||
"不遮不等于有出处：agent 自己定的阈值照样标「无出处」"
|
||||
|
||||
|
||||
async def test_a_metric_from_the_report_that_equals_a_masked_value_can_be_cited(monkeypatch, big_shop, client):
|
||||
"""报告口径卡里的指标（单笔目标 20000）恰好和遮罩列里的某一笔一样：报告里本来就看得到，引用它的回答照常交出去，
|
||||
换个问法也不会被拦。"""
|
||||
metrics = [{"id": "target", "name": "单笔目标", "unit": "元", "expression": "vars.kpi.target"}]
|
||||
graph = chain(node("start", "input"),
|
||||
node("fetch", "tool", tool=f"db_query__{big_shop}", args={"sql": AMOUNTS}),
|
||||
node("calc", "transform", mode="expression", expression="{'target': 20000}", assign_to="kpi"),
|
||||
node("caliber", "metrics", caliber="大单口径", caliber_version="v1", metrics=metrics),
|
||||
node("write", "report", "写周报", instructions="写周报"),
|
||||
node("gate", "human", mode="approve", title="周报可以出具吗"),
|
||||
node("out", "output", fields=[{"name": "周报", "value": "{{ nodes.write.text }}"}]))
|
||||
smallest = "SELECT amount AS 金额 FROM sales WHERE amount > 15000 ORDER BY amount LIMIT 1"
|
||||
plays = {"定得高不高": Play([smallest], "单笔目标为 [[m:target]]，达标的最小一笔是 [[v:Q1.r0.金额]]。[[see:Q1]]")}
|
||||
Models(monkeypatch, big_shop, writer="单笔目标 [[m:target]]。[[see:m:target]]", plays=plays)
|
||||
paused = await pause(big_shop, graph=graph)
|
||||
[approval] = await pending(paused.id)
|
||||
r = await ask(client, paused.id, report="write", unit=None, question="单笔目标定得高不高？",
|
||||
approval_id=approval.id)
|
||||
assert r.status_code == 200, r.text
|
||||
answer = r.json()["inquiry"]["answer"]
|
||||
assert (answer["ok"], answer["error"]) == (True, None)
|
||||
assert answer["text"] == "单笔目标为 20,000元，达标的最小一笔是 已遮罩。"
|
||||
|
||||
|
||||
async def test_a_masked_value_the_asker_wrote_stays_in_the_sql(monkeypatch, tmp_path, client):
|
||||
"""提问人自己写的数恰好是遮罩列里的值（金额设了遮罩，问的就是 300 那一笔）：问题里原样写着，SQL 里和金额做等值
|
||||
比较也不换，事件、质询工件里封存的就是实际执行的那条。"""
|
||||
name, source_id = await _source(tmp_path, options={"mask_columns": "amount"})
|
||||
without = "SELECT COUNT(*) AS 笔数 FROM sales WHERE amount <> 300"
|
||||
plays = {"300 的那笔": Play([without], "去掉金额为 300 的那笔后还剩 [[v:Q1.r0.笔数]] 笔。[[see:Q1]]")}
|
||||
Models(monkeypatch, name, writer="华东首单金额 [[v:Q1.r0.金额]]。[[see:Q1]]", plays=plays)
|
||||
try:
|
||||
paused = await pause(name, graph=fetching(name, AMOUNTS))
|
||||
[approval] = await pending(paused.id)
|
||||
r = await ask(client, paused.id, report="write", unit=None, question="去掉金额为 300 的那笔会怎样？",
|
||||
approval_id=approval.id)
|
||||
assert r.status_code == 200, r.text
|
||||
item = r.json()["inquiry"]
|
||||
assert item["answer"]["text"] == "去掉金额为 300 的那笔后还剩 3 笔。"
|
||||
[query] = item["queries"]
|
||||
assert query["sql"] == without
|
||||
[event] = await events(paused.id, "approval.inquiry")
|
||||
assert event.data["queries"][0]["sql"] == without
|
||||
got = await client.get(f"/api/runs/{paused.id}/inquiries/artifacts/{query['artifact']}")
|
||||
assert got.status_code == 200 and got.json()["content"]["sql"] == without
|
||||
finally:
|
||||
await engines.invalidate(source_id)
|
||||
|
||||
|
||||
#: 另有一张联系方式表：同一个手机号写在没设遮罩的列里
|
||||
CONTACTS = SALES + ("CREATE TABLE contacts (id INTEGER PRIMARY KEY, phone_num TEXT, visits INTEGER);"
|
||||
"INSERT INTO contacts VALUES (1, '13900000001', 5), (2, '13900000003', 2);")
|
||||
|
||||
|
||||
async def test_a_value_seen_only_in_the_reports_snapshot_is_masked_in_the_artifact_too(monkeypatch, tmp_path, client):
|
||||
"""agent 查的表没设遮罩，里面却写着报告快照里遮罩列的值（另一张表里的同一个手机号）：质询工件照回答时的办法收
|
||||
遮罩值，报告自己的查询快照也算，按值遮掉。"""
|
||||
name, source_id = await _source(tmp_path, options={"mask_columns": "buyer_phone"}, script=CONTACTS)
|
||||
visits = "SELECT phone_num AS 号码, visits AS 次数 FROM contacts ORDER BY id"
|
||||
plays = {"常来": Play([visits], "最常来的客户来过 [[v:Q1.r0.次数]] 次。[[see:Q1]]")}
|
||||
Models(monkeypatch, name, writer="金额最高的一笔 [[v:Q1.r0.金额]]。[[see:Q1]]", plays=plays)
|
||||
try:
|
||||
report_sql = "SELECT buyer_phone, amount AS 金额 FROM sales ORDER BY amount DESC"
|
||||
paused = await pause(name, graph=fetching(name, report_sql))
|
||||
[approval] = await pending(paused.id)
|
||||
r = await ask(client, paused.id, report="write", unit=None, question="常来的客户来过几次？",
|
||||
approval_id=approval.id)
|
||||
assert r.status_code == 200, r.text
|
||||
assert r.json()["inquiry"]["answer"]["text"] == "最常来的客户来过 5 次。"
|
||||
[query] = r.json()["inquiry"]["queries"]
|
||||
got = await client.get(f"/api/runs/{paused.id}/inquiries/artifacts/{query['artifact']}")
|
||||
assert got.status_code == 200, got.text
|
||||
assert got.json()["content"]["rows"] == [["已遮罩", 5], ["已遮罩", 2]]
|
||||
assert got.json()["masked"] == ["号码"] and "1390000000" not in got.text
|
||||
finally:
|
||||
await engines.invalidate(source_id)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 接线：两道审批的竞态、上传表格的版本、设置、开放之后状态又变了、张冠李戴的报告文档
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
@@ -1,19 +1,21 @@
|
||||
"""审批质询的遮罩（engine/inquiry.py）：回答正文、查询清单、质询工件里都不露出数据源遮罩列的值，也不误伤别的数。
|
||||
|
||||
- 改了名字的遮罩列（起别名、套函数、经子查询或 CTE 转一道、UNION 的另一支）照样遮；COUNT 只数个数、窗口的分区键
|
||||
只决定次序，不遮；解析不了、SQL 里又写着遮罩列的，整份快照遮掉
|
||||
- 改了名字的遮罩列（起别名、套函数、经子查询或 CTE 转一道、UNION 的另一支、和 * 写在同一支里没起别名的表达式）
|
||||
照样遮；COUNT 只数个数、窗口的分区键只决定次序，不遮；认不出落在哪一列的（两个 * 之间、子查询里没起别名的表达式）、
|
||||
解析不了又写着遮罩列的，整份快照遮掉
|
||||
- 按值认的遮罩值有门槛：纯数字够长（SECRET_DIGITS 位）、文字至少两个字；日期、小数、千分位里的一截不算；提问人自己
|
||||
写进问题、选中原文里的不换
|
||||
- SQL 只换字面量（和注释）里的遮罩值：直接和遮罩列比较的字面量是遮罩列出现过的值就换，短的数也换；LIMIT、日期、
|
||||
别的列的条件照原样
|
||||
- 快照按列遮，另按值遮：别的列里写着遮罩值的格子也换成「已遮罩」
|
||||
- 组装回答两遍之后正文里还认得出遮罩值的，不交出去
|
||||
写进问题、选中原文里的，报告列出的指标和运行输入的值，SQL 里照原样留着的阈值，都不换
|
||||
- SQL 只换在找某一个值的字面量（等值、IN、LIKE，带前缀的字符串也算）和注释、名字里的遮罩值：和遮罩列比较的字面量是
|
||||
遮罩列出现过的值就换，短的数也换；范围比较的边界、BETWEEN、LIMIT、日期、别的列的范围条件照原样
|
||||
- 快照按列遮，另按值遮：别的列里写着遮罩值的格子也换成「已遮罩」；列名里写着遮罩值的，列名也换
|
||||
- 组装回答之后正文里还认得出遮罩值的，不交出去
|
||||
- 遮罩值上万个时换一遍也不卡：不现编上万个分支的正则
|
||||
|
||||
数据都是虚构的：手机号、年龄、区域、金额。
|
||||
数据都是虚构的：手机号、姓名、年龄、区域、金额。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
|
||||
import pytest
|
||||
@@ -34,9 +36,10 @@ def secrets_of(snapshot: dict, hidden: set[str]) -> inquiry_mod._Secrets:
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def hidden_of(sql: str, columns: list[str], masked: set[str] = frozenset({"buyer_phone"})) -> set[str]:
|
||||
def hidden_of(sql: str, columns: list[str], masked: set[str] = frozenset({"buyer_phone"}), kind: str = "sqlite"
|
||||
) -> set[str]:
|
||||
snapshot = {"source": "shop", "sql": sql, "columns": columns, "rows": []}
|
||||
return inquiry_mod._derived_hidden(snapshot, set(masked), "sqlite")[0]
|
||||
return inquiry_mod._derived_hidden(snapshot, set(masked), kind)[0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql, columns, hidden", [
|
||||
@@ -54,12 +57,70 @@ def test_a_masked_column_under_another_name_is_masked(sql, columns, hidden):
|
||||
assert hidden_of(sql, columns) == hidden
|
||||
|
||||
|
||||
SALES_COLUMNS = ["id", "region", "buyer_phone", "amount", "is_test"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql, columns, hidden", [
|
||||
# * 后面没起别名的表达式：从结果末尾倒着数（SQLite 照写法原样给列名，起不出能往下传的名字）
|
||||
("SELECT *, SUBSTR(buyer_phone, 4) FROM sales ORDER BY amount DESC LIMIT 1",
|
||||
[*SALES_COLUMNS, "SUBSTR(buyer_phone, 4)"], {"buyer_phone", "substr(buyer_phone, 4)"}),
|
||||
("SELECT s.*, buyer_phone || '-x' FROM sales s", [*SALES_COLUMNS, "buyer_phone || '-x'"],
|
||||
{"buyer_phone", "buyer_phone || '-x'"}),
|
||||
# sqlglot 给 CAST 记的名字是里面的列名，结果的列名却是整个写法：照样按位置认
|
||||
("SELECT *, CAST(buyer_phone AS TEXT) FROM sales", [*SALES_COLUMNS, "CAST(buyer_phone AS TEXT)"],
|
||||
{"buyer_phone", "cast(buyer_phone as text)"}),
|
||||
# * 前面的从头数，* 后面不是遮罩列的不遮
|
||||
("SELECT UPPER(buyer_phone), s.*, amount * 2 FROM sales s", ["UPPER(buyer_phone)", *SALES_COLUMNS, "amount * 2"],
|
||||
{"buyer_phone", "upper(buyer_phone)"}),
|
||||
# 派生来源里 UNION 的另一支：结果的列名取第一支的
|
||||
("SELECT * FROM (SELECT region FROM sales UNION ALL SELECT buyer_phone FROM sales) t", ["region"],
|
||||
{"buyer_phone", "region"}),
|
||||
("WITH t(p) AS (SELECT region FROM sales UNION ALL SELECT buyer_phone FROM sales) SELECT p FROM t", ["p"],
|
||||
{"buyer_phone", "p"}),
|
||||
("SELECT * FROM (SELECT *, buyer_phone AS p FROM sales) t", [*SALES_COLUMNS, "p"], {"buyer_phone", "p"}),
|
||||
# 多套几层括号不改变结果流到哪
|
||||
("((SELECT UPPER(buyer_phone) FROM sales))", ["UPPER(buyer_phone)"], {"buyer_phone", "upper(buyer_phone)"}),
|
||||
("SELECT * FROM ((SELECT buyer_phone AS p FROM sales)) t", ["p"], {"buyer_phone", "p"}),
|
||||
])
|
||||
def test_an_unnamed_expression_next_to_a_star_or_under_a_union_is_masked(sql, columns, hidden):
|
||||
assert hidden_of(sql, columns) == hidden
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql, columns", [
|
||||
# 两个 * 之间：前后都数不出位置
|
||||
("SELECT a.*, SUBSTR(a.buyer_phone, 4), b.* FROM sales a JOIN sales b ON a.id = b.id",
|
||||
[*SALES_COLUMNS, "SUBSTR(a.buyer_phone, 4)", *SALES_COLUMNS]),
|
||||
# 子查询里没起别名的表达式：外层的 * 原样带出来，列名各数据库起法不同，认不准
|
||||
("SELECT * FROM (SELECT SUBSTR(buyer_phone, 4), amount FROM sales) t", ["SUBSTR(buyer_phone, 4)", "amount"]),
|
||||
("SELECT * FROM (SELECT SUBSTR(buyer_phone, 4) FROM sales)", ["SUBSTR(buyer_phone, 4)"]),
|
||||
("WITH t AS ((SELECT SUBSTR(buyer_phone, 4) FROM sales)) SELECT * FROM t", ["SUBSTR(buyer_phone, 4)"]),
|
||||
# UNION 第一支是 *：另一支的遮罩列落在哪个名字上说不清
|
||||
("SELECT * FROM (SELECT * FROM sales UNION ALL SELECT id, region, SUBSTR(buyer_phone, 4), amount, is_test "
|
||||
"FROM sales) t", SALES_COLUMNS),
|
||||
])
|
||||
def test_a_masked_expression_that_cannot_be_placed_hides_the_whole_snapshot(sql, columns):
|
||||
assert hidden_of(sql, columns) == {"buyer_phone", *(c.lower() for c in columns)}
|
||||
|
||||
|
||||
def test_a_lateral_subquery_is_read_like_a_derived_table():
|
||||
"""Postgres 的 LATERAL：别名记在 LATERAL 上，起了别名的照名字传，没起别名的整份遮。"""
|
||||
assert hidden_of("SELECT x.* FROM sales s CROSS JOIN LATERAL (SELECT s.buyer_phone AS p) x", ["p"],
|
||||
kind="postgresql") == {"buyer_phone", "p"}
|
||||
columns = [*SALES_COLUMNS, "substr"]
|
||||
assert hidden_of("SELECT * FROM sales s CROSS JOIN LATERAL (SELECT SUBSTR(s.buyer_phone, 4)) x", columns,
|
||||
kind="postgresql") == {*columns}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql, columns", [
|
||||
("SELECT region AS 区域, COUNT(DISTINCT buyer_phone) AS 买家数 FROM sales GROUP BY region", ["区域", "买家数"]),
|
||||
("SELECT region AS 区域, SUM(amount) AS 销售额 FROM sales WHERE buyer_phone = '13900000001' GROUP BY region",
|
||||
["区域", "销售额"]),
|
||||
("SELECT ROW_NUMBER() OVER (PARTITION BY buyer_phone ORDER BY amount) AS 序号 FROM sales", ["序号"]),
|
||||
("SELECT (SELECT COUNT(*) FROM sales s2 WHERE s2.buyer_phone = s.buyer_phone) AS 笔数 FROM sales s", ["笔数"]),
|
||||
# 子查询里只数个数、条件里的子查询：值带不出来，外层的 * 也不怕
|
||||
("SELECT * FROM (SELECT region, COUNT(DISTINCT buyer_phone) FROM sales GROUP BY region) t",
|
||||
["region", "COUNT(DISTINCT buyer_phone)"]),
|
||||
("SELECT region FROM sales WHERE amount > (SELECT AVG(LENGTH(buyer_phone)) FROM sales)", ["region"]),
|
||||
])
|
||||
def test_counting_filtering_and_ordering_by_a_masked_column_shows_nothing(sql, columns):
|
||||
assert hidden_of(sql, columns) == {"buyer_phone"}
|
||||
@@ -144,7 +205,7 @@ def people() -> inquiry_mod._Secrets:
|
||||
("SELECT COUNT(*) AS 笔数 FROM sales WHERE buyer_phone = '13900000001'",
|
||||
"SELECT COUNT(*) AS 笔数 FROM sales WHERE buyer_phone = '已遮罩'"),
|
||||
("SELECT COUNT(*) FROM sales s WHERE s.buyer_phone IN ('13900000001', '13900000003') AND age BETWEEN 25 AND 31",
|
||||
"SELECT COUNT(*) FROM sales s WHERE s.buyer_phone IN ('已遮罩', '已遮罩') AND age BETWEEN 已遮罩 AND 已遮罩"),
|
||||
"SELECT COUNT(*) FROM sales s WHERE s.buyer_phone IN ('已遮罩', '已遮罩') AND age BETWEEN 25 AND 31"),
|
||||
("SELECT COUNT(*) FROM sales WHERE age = 25 AND amount > 25 AND paid_on >= '2026-01-25' LIMIT 25",
|
||||
"SELECT COUNT(*) FROM sales WHERE age = 已遮罩 AND amount > 25 AND paid_on >= '2026-01-25' LIMIT 25"),
|
||||
("SELECT COUNT(*) FROM sales WHERE 25 = age AND note LIKE '%13900000003%' -- 13900000001 的单\nLIMIT 20",
|
||||
@@ -152,11 +213,76 @@ def people() -> inquiry_mod._Secrets:
|
||||
("SELECT region FROM sales WHERE buyer_phone NOT IN ('13900000009') AND age <> 40",
|
||||
"SELECT region FROM sales WHERE buyer_phone NOT IN ('13900000009') AND age <> 40"),
|
||||
("SELECT amount FROM sales WHERE amount = 13900000001", "SELECT amount FROM sales WHERE amount = 已遮罩"),
|
||||
# 和遮罩列做范围比较：边界是阈值，不是在找某个人，照原样
|
||||
("SELECT COUNT(*) FROM sales WHERE age > 31 AND age <= 25 AND 25 < age AND age NOT BETWEEN 25 AND 31",
|
||||
"SELECT COUNT(*) FROM sales WHERE age > 31 AND age <= 25 AND 25 < age AND age NOT BETWEEN 25 AND 31"),
|
||||
# 不等、SQLite 的 IS、NOT LIKE 也是在找某一个值
|
||||
("SELECT COUNT(*) FROM sales WHERE age != 25 AND age IS NOT 31 AND buyer_phone NOT LIKE '13900000001'",
|
||||
"SELECT COUNT(*) FROM sales WHERE age != 已遮罩 AND age IS NOT 已遮罩 AND buyer_phone NOT LIKE '已遮罩'"),
|
||||
])
|
||||
def test_sql_keeps_everything_but_the_masked_values(people, sql, shown):
|
||||
assert inquiry_mod._scrub_sql(sql, people, {"buyer_phone", "age"}, "sqlite") == shown
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def amounts() -> inquiry_mod._Secrets:
|
||||
"""金额设了遮罩，金额都是五位数：和阈值、LIMIT、别的列的条件撞在一起。"""
|
||||
return secrets_of({"columns": ["amount"], "rows": [[10000], [12000], [20000], [30000], [50000]]}, {"amount"})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql, shown", [
|
||||
# 范围比较、LIMIT / OFFSET、选出列里的数：长的数也照原样
|
||||
("SELECT COUNT(*) FROM sales WHERE visits > 20000 AND amount >= 30000 AND amount BETWEEN 10000 AND 50000 "
|
||||
"LIMIT 10000 OFFSET 12000",
|
||||
"SELECT COUNT(*) FROM sales WHERE visits > 20000 AND amount >= 30000 AND amount BETWEEN 10000 AND 50000 "
|
||||
"LIMIT 10000 OFFSET 12000"),
|
||||
("SELECT 20000 AS 门槛, amount - 10000 AS 差额 FROM sales", "SELECT 20000 AS 门槛, amount - 10000 AS 差额 FROM sales"),
|
||||
# 等值、IN：和遮罩列比较的换，和别的列比较的、够长的遮罩值也换
|
||||
("SELECT COUNT(*) FROM sales WHERE amount = 30000 OR amount IN (10000, -12000) OR visits = 20000",
|
||||
"SELECT COUNT(*) FROM sales WHERE amount = 已遮罩 OR amount IN (已遮罩, -已遮罩) OR visits = 已遮罩"),
|
||||
])
|
||||
def test_long_masked_numbers_are_only_masked_where_a_value_is_looked_up(amounts, sql, shown):
|
||||
assert inquiry_mod._scrub_sql(sql, amounts, {"amount"}, "sqlite") == shown
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind, sql, shown", [
|
||||
("oracle", "SELECT COUNT(*) FROM sales WHERE buyer_phone = N'13900000001'",
|
||||
"SELECT COUNT(*) FROM sales WHERE buyer_phone = N'已遮罩'"),
|
||||
("postgresql", "SELECT COUNT(*) FROM sales WHERE buyer_phone = E'13900000001' OR note = $$13900000003$$ "
|
||||
"OR note = $t$13900000001$t$ OR buyer_phone = U&'13900000003'",
|
||||
"SELECT COUNT(*) FROM sales WHERE buyer_phone = E'已遮罩' OR note = $$已遮罩$$ "
|
||||
"OR note = $t$已遮罩$t$ OR buyer_phone = U&'已遮罩'"),
|
||||
# 名字里写着遮罩值：带引号的、不带引号的（SQLite 的中文名字可以不加引号）
|
||||
("postgresql", 'SELECT COUNT(*) AS "13900000001 的单" FROM sales',
|
||||
'SELECT COUNT(*) AS "已遮罩 的单" FROM sales'),
|
||||
("mysql", "SELECT COUNT(*) AS `13900000001 的单` FROM sales", "SELECT COUNT(*) AS `已遮罩 的单` FROM sales"),
|
||||
])
|
||||
def test_prefixed_strings_and_names_are_scrubbed_too(people, kind, sql, shown):
|
||||
assert inquiry_mod._scrub_sql(sql, people, {"buyer_phone", "age"}, kind) == shown
|
||||
|
||||
|
||||
def test_a_name_written_into_an_unquoted_alias_is_scrubbed():
|
||||
secrets = secrets_of({"columns": ["buyer_name"], "rows": [["张伟"], ["李娜"]]}, {"buyer_name"})
|
||||
assert inquiry_mod._scrub_sql("SELECT COUNT(*) AS 张伟的单数 FROM sales WHERE buyer_name = '张伟'", secrets,
|
||||
{"buyer_name"}, "sqlite") == (
|
||||
"SELECT COUNT(*) AS 已遮罩的单数 FROM sales WHERE buyer_name = '已遮罩'")
|
||||
|
||||
|
||||
def test_numbers_left_as_written_in_the_sql_are_not_masked_in_the_answer(amounts):
|
||||
"""范围比较的边界、LIMIT 照原样留在 SQL 里，回答正文里同样的数再遮也遮不住什么，只会和 SQL 对不上；等值比较
|
||||
里换掉的照旧遮。"""
|
||||
sql = "SELECT COUNT(*) FROM sales WHERE amount >= 30000 AND amount <> 50000 LIMIT 10000"
|
||||
assert inquiry_mod._shown_literals(sql, "sqlite") == ["30000", "10000"]
|
||||
amounts.protect(inquiry_mod._shown_literals(sql, "sqlite"), strong_only=True)
|
||||
assert amounts.scrub("金额不低于 30000 元、不是 50000 元的，取前 10000 条") == (
|
||||
"金额不低于 30000 元、不是 已遮罩 元的，取前 10000 条")
|
||||
assert inquiry_mod._scrub_sql(sql, amounts, {"amount"}, "sqlite") == (
|
||||
"SELECT COUNT(*) FROM sales WHERE amount >= 30000 AND amount <> 已遮罩 LIMIT 10000")
|
||||
assert inquiry_mod._scrub_sql("SELECT COUNT(*) FROM sales WHERE amount = 30000", amounts, {"amount"},
|
||||
"sqlite") == "SELECT COUNT(*) FROM sales WHERE amount = 已遮罩", \
|
||||
"只免了正文里的：和遮罩列做等值比较的照旧换"
|
||||
|
||||
|
||||
def test_a_comparison_with_a_renamed_masked_column_is_masked_too(people):
|
||||
sql = "SELECT * FROM (SELECT buyer_phone AS p FROM sales) t WHERE p = '13900000003' AND 1 = 1"
|
||||
assert inquiry_mod._scrub_sql(sql, people, {"buyer_phone", "p"}, "sqlite") == (
|
||||
@@ -173,18 +299,59 @@ def test_sql_that_cannot_be_split_into_words_is_scrubbed_like_text(people):
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def query_entry(columns: list[str]) -> dict:
|
||||
return {"alias": "Q1", "kind": "query", "artifact": "a1", "source": "shop", "columns": columns, "rows": 1,
|
||||
"truncated": False}
|
||||
|
||||
|
||||
def test_an_answer_still_showing_a_masked_value_is_not_handed_out():
|
||||
"""运行输入里恰好是一个遮罩值（报告目录里的 i: 条目）：标记里看不出来，渲染出来才露，连标记一起换也换不掉——
|
||||
宁可不交这份回答。"""
|
||||
"""别的列里一格文字夹着遮罩值（备注里写着手机号）：按整格比认不出来，标记里也看不出来，渲染出来才露——宁可不交
|
||||
这份回答。"""
|
||||
secrets = secrets_of({"columns": ["phone"], "rows": [["13900000001"]]}, {"phone"})
|
||||
catalog = {"i:客户": {"alias": "i:客户", "kind": "input", "eid": "input:客户", "locator": {"field": "客户"},
|
||||
"value": "13900000001", "label": "客户 = 13900000001"}}
|
||||
snapshot = {"columns": ["备注"], "rows": [["客户13900000001来电改址"]]}
|
||||
views = {"a1": inquiry_mod._masked_view(snapshot, set(), secrets)[0]}
|
||||
assert views["a1"]["rows"] == [["客户13900000001来电改址"]]
|
||||
catalog = {"Q1": query_entry(["备注"])}
|
||||
with pytest.raises(inquiry_mod._Leak):
|
||||
inquiry_mod._compose("这位客户是 [[i:客户]]。", catalog, run_id="r1", views={}, secrets=secrets, allow=[])
|
||||
doc = inquiry_mod._compose("没有遮罩值。", catalog, run_id="r1", views={}, secrets=secrets, allow=[])
|
||||
inquiry_mod._compose("备注是 [[v:Q1.r0.备注]]。", catalog, run_id="r1", views=views, secrets=secrets, allow=[])
|
||||
doc = inquiry_mod._compose("没有遮罩值。", catalog, run_id="r1", views=views, secrets=secrets, allow=[])
|
||||
assert doc["markdown"] == "没有遮罩值。"
|
||||
|
||||
|
||||
def test_report_values_that_happen_to_be_masked_values_are_shown():
|
||||
"""报告口径卡里的指标恰好和遮罩列里的某个值一样（单笔目标 20000，金额也有 20000 的单）：报告里本来就看得到，
|
||||
引用它的回答照常交出去。"""
|
||||
secrets = secrets_of({"columns": ["amount"], "rows": [[20000], [35000]]}, {"amount"})
|
||||
catalog = {"m:target": {"alias": "m:target", "kind": "metric", "eid": "metric:c1:target", "artifact": "c1",
|
||||
"locator": {"metric": "target"}, "value": 20000, "rendered": "20,000",
|
||||
"label": "单笔目标 = 20,000"}}
|
||||
with pytest.raises(inquiry_mod._Leak):
|
||||
inquiry_mod._compose("单笔目标为 [[m:target]]。", catalog, run_id="r1", views={}, secrets=secrets, allow=[])
|
||||
secrets.protect(inquiry_mod._ref_values(inquiry_mod._report_refs(catalog)))
|
||||
doc = inquiry_mod._compose("单笔目标为 [[m:target]]，有一笔 35000 元。", catalog, run_id="r1", views={},
|
||||
secrets=secrets, allow=[])
|
||||
assert doc["markdown"] == "单笔目标为 20,000，有一笔 已遮罩 元。"
|
||||
|
||||
|
||||
def test_a_column_named_after_a_masked_value_is_renamed():
|
||||
"""agent 把遮罩值写进了别名（「13900000001 的单」）：快照、回答文档里的列名同样换掉。回答引用了原列名的，标记里
|
||||
的列名照样换，正好对上换过的列名，引用照样解析得了。"""
|
||||
secrets = secrets_of({"columns": ["phone"], "rows": [["13900000001"]]}, {"phone"})
|
||||
snapshot = {"columns": ["13900000001 的单", "金额"], "column_types": {"13900000001 的单": "number"},
|
||||
"rows": [[2, 1200]]}
|
||||
shown, masked = inquiry_mod._masked_view(snapshot, set(), secrets)
|
||||
assert shown["columns"] == ["已遮罩 的单", "金额"] and shown["rows"] == [[2, 1200]]
|
||||
assert shown["column_types"] == {"已遮罩 的单": "number"} and masked == []
|
||||
assert snapshot["columns"][0] == "13900000001 的单", "原快照不动"
|
||||
rows_as_dicts = {**snapshot, "rows": [{"13900000001 的单": 2, "金额": 1200}]}
|
||||
assert inquiry_mod._masked_view(rows_as_dicts, set(), secrets)[0]["rows"] == [{"已遮罩 的单": 2, "金额": 1200}]
|
||||
|
||||
doc = inquiry_mod._compose("这位客户共 [[v:Q1.r0.13900000001 的单]] 笔。", {"Q1": query_entry(snapshot["columns"])},
|
||||
run_id="r1", views={"a1": shown}, secrets=secrets, allow=[])
|
||||
assert doc["markdown"] == "这位客户共 2 笔。"
|
||||
assert "13900000001" not in json.dumps(doc, ensure_ascii=False)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 规模
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user