mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 02:51:58 +00:00
K3 复核修正（第 2 轮）：* 把遮罩列带进别的列名时按展开落到的位置遮；函数参数里的 JSON 箭头认作遮罩列；agent 查到的遮罩值和报告里的分开记，上限收满也不丢
- * 改名（复核点名的 blocker）：WITH t(a, b, …) AS (SELECT * …) 这类列名表、UNION 后一支的 *、FROM 子查询或 LATERAL 的列名表、标量子查询和 ARRAY 子查询里的 *，以前 * 本身不算取自哪一列，原值原样交给核查 agent、落库，第二道防线也 认不出。现在 * 带出遮罩列的值时，它展开落到的那几列都算（masking._star_places）：只展开一个来源、每一列都数得清的 按位置只遮取自遮罩列的那几列——子查询按选出项数，数据源里的表按查询当时冻结的表结构数（_schema_columns；工具边界 用数据源的 schema_cache，事后按快照的 schema_artifact，是同一份），几张列一样的月度表 SELECT * 连起来汇总，金额照常； 数不清的（没有表结构、列数对不上、几个来源）整段遮。表结构只用来数位置，不用来断定表里没有遮罩列。标量子查询里的 * 按 _value_columns 展开（展开数据源里的表一律算），EXISTS 里的 * 不算；递归的 CTE 不会绕不出来；带 EXCLUDE / REPLACE 的 * 不按位置数 - JSON 箭头：函数参数里没写表名的 col->'path' 被 sqlglot 解析成以 col 为参数的 lambda，_value_columns 把 lambda 的 参数名照列名收。SQLite json_each(profile->'$.tels')、Postgres jsonb_array_elements_text(profile->'tels')、MySQL JSON_UNQUOTE(profile->'$.tel') 这几种拆开、取出来的值在工具边界就遮 - 上限：run_verify_agent 以前先把调用方交进来的报告快照值并进 seen，收满 SECRETS_MAX 之后 agent 查到的值并不进去，后面 的查询、回答时的第二道防线都认不出。现在两份分开记（masking.SeenValues：given / found），查询工具按值遮、交回调用方 （VerifyOutcome.masked_values）时合起来，agent 查到的在前；反方审稿理由的按值遮也改成 agent 查到的先收 - 已知局限写进 inquiry 模块说明：表结构过期、列数没变却换了次序或改了名时可能数错；EXCEPT 后一支照 UNION 算（多遮） 测试：新增 test_masking_star_renames（真跑 SQLite 的各种 * 改名、按表结构数位置、表结构对不上、快照记着的表结构）， 工具边界（交给模型的、落库的、通用工件接口）和审批质询端到端的 * 改名、JSON 箭头、上限收满；每处实现都做过去掉之后 测试变红的验证。
This commit is contained in:
ThearchyHelios
@@ -56,7 +56,14 @@
|
||||
IS NOT NULL THEN amount ELSE 0 END)）同写在 WHERE 里，不算
|
||||
- FROM、JOIN、LATERAL 里的表值函数（json_each、json_tree、unnest、jsonb_each_text、JSON_TABLE、LATERAL VALUES）拆开
|
||||
遮罩列的，按出列名认它出的列（常见函数的出列名、写的列名表、别名、函数名，_function_columns），同名的别的列一并遮
|
||||
（多遮）；不认得的函数又没写列名表的，整份快照都遮
|
||||
（多遮）；不认得的函数又没写列名表的，整份快照都遮。函数参数里没写表名的 col->'path'（JSON 箭头，解析器当成以 col
|
||||
为参数的 lambda）按 col 认，真的 lambda（transform(arr, x -> …)）的参数和遮罩列同名时也算（多遮）
|
||||
- * 把遮罩列带进别的列名（WITH t(a, b, …) AS (SELECT * …) 这类列名表、UNION 后一支的 *、标量子查询里的 *）：只展开一个
|
||||
来源、那里每一列都数得清的，按位置只遮取自遮罩列的那几列——子查询按它的选出项数，数据源里的表按查询当时冻结的表结构
|
||||
数（几张列一样的月度表 SELECT * 连起来汇总，金额照常）。表结构里查不到这张表、数出来的列数和实际展开的对不上，或者
|
||||
* 连着几个来源的，* 展开落到的那一段整段遮，同一段里的金额这类别的列一并遮（多遮）；表结构过期、列数没变却换了次序
|
||||
或改了名的，可能数错（同按表结构认标签列）。标量子查询里的 * 展开数据源里的表一律算取自遮罩列。EXCEPT 后一支只用来
|
||||
去掉行，也照 UNION 的后一支算（多遮）
|
||||
- 按遮罩列分档的标签（CASE、IIF 的结果全是写死的常量）：标签文字是 agent 写在 SQL 里的，不当遮罩值，SQL、回答正文里
|
||||
照原样；标签列当分组键（每一档一行）照常显示，不分组的逐行标签列整列遮（格子里看不出某一笔落在哪一档），可 agent
|
||||
在正文里写出「这一笔是大单」遮不住，按 id 分组再按档数个数也看得出（同按阈值查出某几行）。只有只装着标签的列不收：
|
||||
@@ -93,7 +100,8 @@
|
||||
报告快照里的遮罩值、agent 前面的查询里遮掉的值都看不到。只有一种照原样：先查了写着号码的备注（那时还不知道这个号码
|
||||
是遮罩值），后来才查出号码——备注快照落库时原样，从通用接口打开照原样，质询工件接口照事件里记下的位置换。报告自己
|
||||
的查询快照（写作时的查询）仍原样存着，从那里打开不经这里的遮罩
|
||||
- 遮罩值最多收 SECRETS_MAX 种写法，再多的按值认不出（按列遮照旧）。agent 自己查到的值先收，报告快照里的后收
|
||||
- 遮罩值最多收 SECRETS_MAX 种写法，再多的按值认不出（按列遮照旧）。agent 自己查到的值先收，报告快照里的后收：核查
|
||||
agent 把两份分开记（masking.SeenValues），查询时按值遮、交回调用方、回答时合起来，都是 agent 查到的在前
|
||||
- 换了写法的号码只认数字之间夹着一到四个分隔符（空格、制表符、全角空格、不换行空格、窄空格这类空白，连字符、减号、
|
||||
点、括号，含全角），带 +86 / 0086 / 86 前缀、全角数字：写成汉字（一三九……）、中间夹着别的字、拆成几处写的、跨行
|
||||
写的认不出；和别的数字只隔着一个连字符或点连着的（13900000001-2、19914.74）不算号码；SQL 里不带引号的数不这样认
|
||||
|
||||
+323
-41
@@ -11,16 +11,19 @@
|
||||
审稿理由里的按值遮降为第二道防线。查询失败时报错原话里的值也换（mask_error_text）
|
||||
|
||||
- 哪些列（_derived_hidden）：从遮罩列本身起，按 SQL 一层层认出值取自遮罩列的结果列（别名、函数、子查询、WITH、UNION 的
|
||||
另一支、和 * 写在同一支里没起别名的表达式、FROM 里拆开遮罩列的表值函数）；COUNT 只数个数、窗口的分区键只决定次序、
|
||||
按遮罩列分档又当分组键或套在汇总里数个数的、汇总里的 CASE 只在条件里用到的不算（_value_columns、_band_only）。认不出
|
||||
落在哪一列的、SQL 解析不了又写着遮罩列的，整份快照的列都算。双盲复算按同一套认法在派盲算之前挡下取自遮罩列的指标
|
||||
（derived_columns）
|
||||
另一支、和 * 写在同一支里没起别名的表达式、FROM 里拆开遮罩列的表值函数、函数参数里的 JSON 箭头 col->'path'）；* 把
|
||||
遮罩列带进别的列名的（写了列名表的派生来源、UNION 后一支的 *、标量子查询里的 *）按它展开落到的位置算（_star_places，
|
||||
展开数据源里的表的按查询当时冻结的表结构数位置，_schema_columns）。
|
||||
COUNT 只数个数、窗口的分区键只决定次序、按遮罩列分档又当分组键或套在汇总里数个数的、汇总里的 CASE 只在条件里用到的
|
||||
不算（_value_columns、_band_only）。认不出落在哪一列的、SQL 解析不了又写着遮罩列的，整份快照的列都算。双盲复算按同一套
|
||||
认法在派盲算之前挡下取自遮罩列的指标（derived_columns）
|
||||
- 只装着分档标签的列（_label_columns）：标签是 agent 写在 SQL 里的字，不当遮罩值；* 展开数据源里的表时按冻结的表结构
|
||||
认它带不带同名的列（_schema_tables）
|
||||
- 哪些值（_Secrets）：遮罩列里出现过的值，按值认有门槛（SECRET_MIN、SECRET_DIGITS），换的时候先按词、按开头几个字筛出
|
||||
正文里真有的那几个值，不现编上万个分支的正则；正文、文字格、SQL 的字符串和注释里换了写法的长号码（199-1474-2785、
|
||||
+86 …、全角数字）按纯数字写法在集合里查（loose_spans）；已经摆在读者眼前的值免遮，之后并进来的照样免（protected）
|
||||
- 快照遮好的样子（_masked_view），回答时才认得出、事后要照位置重放的地方（view_edits、replay_edits）
|
||||
- 一次核查里认出的遮罩值（SeenValues）：调用方交进来的和 agent 查询时遮掉的分开记，合起来时 agent 查到的在前
|
||||
|
||||
除 mask_query_result、换了写法的号码和这一轮新加的几样之外，这里的函数原来都在 engine/inquiry.py 里，名字不变，
|
||||
inquiry 照旧导出（别处按 inquiry 的名字引用）。
|
||||
@@ -32,7 +35,7 @@ import functools
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from collections.abc import Iterable
|
||||
from collections.abc import Callable, Iterable
|
||||
from dataclasses import dataclass, field
|
||||
from decimal import Decimal
|
||||
from typing import Any
|
||||
@@ -72,16 +75,26 @@ def _masked_word() -> str:
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
#: _value_columns 里「值可能取自数据源里一张表的任何一列」：套在里面的子查询里的 * 展开了表，SQL 里看不出这张表有哪几列。
|
||||
#: 按它认（_reads）一律算取自遮罩列（宁可多遮）；真实的列名不会是这个
|
||||
_ANY_COLUMN = "*"
|
||||
|
||||
|
||||
def _value_columns(node: Any) -> set[str]:
|
||||
"""一个选出列的值取自哪些列（列名小写）。COUNT 里的只数个数，窗口的分区键、排序键只决定次序，结果全是写死常量的
|
||||
分档（_band_only：当分组键，或者套在汇总里按条件数个数）只决定一行落在哪一档，都带不出值，不算；套在里面的子查询
|
||||
只看它选出来的列（它的条件、连接里写的列带不出值）。汇总里的 CASE / IIF 只看各个结果：条件只决定这一行算不算进来，
|
||||
同写在 WHERE 里（SUM(CASE WHEN buyer_phone IS NOT NULL THEN amount ELSE 0 END) 的值全部来自金额）；不汇总的逐行
|
||||
CASE 照旧连条件一起算（每一行落在哪一支本身就是这一行的信息）。"""
|
||||
CASE 照旧连条件一起算（每一行落在哪一支本身就是这一行的信息）。
|
||||
套在里面的子查询选的是 *（标量子查询、ARRAY(SELECT * …)）：值取自 * 展开的来源（_star_inputs），派生来源看它的各个
|
||||
选出项，表值函数看参数，展开了数据源里的表的记成 _ANY_COLUMN；EXISTS 里的 * 只说有没有，不算。函数参数里的
|
||||
col->'path'（JSON 箭头）被解析器当成以 col 为参数的 lambda，参数名照列名收（真的 lambda 的参数多算一列无害）。"""
|
||||
from sqlglot import exp
|
||||
|
||||
out: set[str] = set()
|
||||
stack = [node]
|
||||
# 已经展开过的派生来源：递归的 CTE 里 * 展开它自己时不再绕回去
|
||||
expanded = {id(node)}
|
||||
while stack:
|
||||
item = stack.pop()
|
||||
if isinstance(item, exp.Column):
|
||||
@@ -94,13 +107,32 @@ def _value_columns(node: Any) -> set[str]:
|
||||
stack.extend(r for r in results if r is not None)
|
||||
elif isinstance(item, exp.Window):
|
||||
stack.append(item.this)
|
||||
elif isinstance(item, exp.Lambda):
|
||||
out.update(p.name.lower() for p in item.expressions if isinstance(p, exp.Identifier) and p.name)
|
||||
stack.append(item.this)
|
||||
elif isinstance(item, exp.Select):
|
||||
stack.extend(item.expressions)
|
||||
for part in item.expressions:
|
||||
if not _is_star(part):
|
||||
stack.append(part)
|
||||
elif not isinstance(item.parent, exp.Exists):
|
||||
inputs = _star_inputs(part, item)
|
||||
if inputs is None:
|
||||
out.add(_ANY_COLUMN)
|
||||
continue
|
||||
fresh = [n for n in inputs if id(n) not in expanded]
|
||||
expanded.update(id(n) for n in fresh)
|
||||
stack.extend(fresh)
|
||||
else:
|
||||
stack.extend(item.iter_expressions())
|
||||
return out
|
||||
|
||||
|
||||
def _reads(names: set[str], columns: set[str]) -> bool:
|
||||
"""值取自 columns（_value_columns）的一项是不是取自 names：有一列对上，或者 * 展开了数据源里的表（_ANY_COLUMN，说不清
|
||||
带出的是哪一列）。"""
|
||||
return bool(names & columns) or _ANY_COLUMN in columns
|
||||
|
||||
|
||||
def _aggregated(node: Any) -> bool:
|
||||
"""node 套在同一层 SELECT 的汇总函数里（SUM(…)、AVG(…) OVER (…) 也算）。"""
|
||||
from sqlglot import exp
|
||||
@@ -230,6 +262,26 @@ def _schema_tables(schema: Any) -> dict[str, set[str]] | None:
|
||||
return out
|
||||
|
||||
|
||||
def _schema_columns(schema: Any) -> dict[str, list[str]] | None:
|
||||
"""冻结的表结构 → {表名（小写）: 按表里次序的列名（小写）}，名字的记法同 _schema_tables；几个库里同名、列又不一样的表
|
||||
不记（说不清 * 展开的是哪一张）。* 展开数据源里的表时按它数位置（_star_columns）；整份拿不到的为 None。"""
|
||||
tables = schema.get("tables") if isinstance(schema, dict) else None
|
||||
if not isinstance(tables, dict):
|
||||
return None
|
||||
out: dict[str, list[str]] = {}
|
||||
clashed: set[str] = set()
|
||||
for key, meta in tables.items():
|
||||
columns = meta.get("columns") if isinstance(meta, dict) else None
|
||||
if not isinstance(columns, list) or not all(isinstance(c, dict) and c.get("name") is not None for c in columns):
|
||||
continue
|
||||
names = [str(c["name"]).lower() for c in columns]
|
||||
for name in {str(key).lower(), str(key).lower().rsplit(".", 1)[-1],
|
||||
str(meta.get("qualified") or key).lower().rsplit(".", 1)[-1]}:
|
||||
if out.setdefault(name, names) != names:
|
||||
clashed.add(name)
|
||||
return {name: names for name, names in out.items() if name not in clashed}
|
||||
|
||||
|
||||
def _table_lacks(source: Any, name: str, memo: _LabelMemo) -> bool:
|
||||
"""FROM 里的这个来源是数据源里的一张表，冻结的表结构里查得到它、又没有 name 这一列：* 展开它带不出 name。"""
|
||||
from sqlglot import exp
|
||||
@@ -577,6 +629,165 @@ def _places(branch: list[Any], width: int | None) -> list[int | None]:
|
||||
for i in range(len(branch))]
|
||||
|
||||
|
||||
#: 按位置数 * 展开成的列最多往下数几层：再深的（递归的 CTE 绕回自己）按数不清算
|
||||
_STAR_DEPTH = 16
|
||||
#: * 上的改动（DuckDB 一类的 EXCLUDE / REPLACE / RENAME / ILIKE）：展开成的列不再是来源那几列的原样，不按位置数
|
||||
_STAR_EDITS = ("except_", "replace", "rename", "ilike")
|
||||
|
||||
|
||||
def _star_sources(star: Any, select: Any) -> list[Any] | None:
|
||||
"""select 里的这个 * 展开的是哪几个来源：没写表名的是 FROM、JOIN、LATERAL VIEW 里的全部，写了的（t.*）只有那一个；
|
||||
写的表名对不上任何来源（外层的表）为 None。"""
|
||||
from sqlglot import exp
|
||||
|
||||
sources = _from_sources(select) + [((lateral.alias_or_name or "").lower(), lateral)
|
||||
for lateral in select.args.get("laterals") or []]
|
||||
qualifier = (star.table.lower() or None) if isinstance(star, exp.Column) else None
|
||||
if qualifier is None:
|
||||
return [source for _, source in sources]
|
||||
found = [source for alias, source in sources if alias == qualifier]
|
||||
return found or None
|
||||
|
||||
|
||||
def _star_inputs(star: Any, select: Any) -> list[Any] | None:
|
||||
"""* 带出的值取自哪里：要接着看的节点——派生来源的每一支（值取自它的各个选出项），表值函数、VALUES 这类来源本身
|
||||
（值取自参数）。展开了数据源里的表、写的表名对不上任何来源的为 None：SQL 里看不出这张表有哪几列，哪一列都可能带出来。"""
|
||||
from sqlglot import exp
|
||||
|
||||
sources = _star_sources(star, select)
|
||||
if sources is None:
|
||||
return None
|
||||
out: list[Any] = []
|
||||
for source in sources:
|
||||
derived = _derived_block(source)
|
||||
if derived is not None:
|
||||
out += _leaves(derived[0])
|
||||
elif isinstance(source, exp.Table) and isinstance(source.this, exp.Identifier):
|
||||
return None
|
||||
else:
|
||||
out.append(source)
|
||||
return out
|
||||
|
||||
|
||||
def _star_feeds(star: Any, select: Any, names: set[str]) -> bool:
|
||||
"""select 里的这个 * 带出的值里有没有取自 names 的（展开了数据源里的表的一律算，见 _star_inputs）。"""
|
||||
inputs = _star_inputs(star, select)
|
||||
return inputs is None or any(_reads(names, _value_columns(node)) for node in inputs)
|
||||
|
||||
|
||||
def _table_columns(source: Any, tables: dict[str, list[str]] | None) -> list[str] | None:
|
||||
"""FROM 里的这个来源是数据源里的一张表、冻结的表结构里查得到它时，它按次序的列名（_schema_columns）；别的为 None。"""
|
||||
from sqlglot import exp
|
||||
|
||||
if not tables or not isinstance(source, exp.Table) or not isinstance(source.this, exp.Identifier):
|
||||
return None
|
||||
return tables.get(source.name.lower())
|
||||
|
||||
|
||||
def _single_source(star: Any, select: Any, depth: int) -> Any:
|
||||
"""* 只展开一个来源、展开成的就是那个来源的列原样（没带 EXCLUDE 这类改动）时，那个来源；别的为 None。"""
|
||||
from sqlglot import exp
|
||||
|
||||
node = star.this if isinstance(star, exp.Column) else star
|
||||
if depth > _STAR_DEPTH or any(node.args.get(k) for k in _STAR_EDITS):
|
||||
return None
|
||||
sources = _star_sources(star, select)
|
||||
return sources[0] if sources is not None and len(sources) == 1 else None
|
||||
|
||||
|
||||
def _star_columns(star: Any, select: Any, names: set[str], memo: dict[int, Any],
|
||||
tables: dict[str, list[str]] | None, depth: int = 0) -> list[bool] | None:
|
||||
"""* 展开成的每一列（按位置）取不取自 names：只展开一个来源时才数得清——派生来源按它的每一列（_block_columns），数据源
|
||||
里的表按查询当时冻结的表结构（tables，_schema_columns）。表结构里查不到的表、几个来源（按同名列连起来的会合成一列）、
|
||||
表值函数，或者 * 带着 EXCLUDE / REPLACE 这类改动的，为 None。"""
|
||||
source = _single_source(star, select, depth)
|
||||
if source is None:
|
||||
return None
|
||||
if (derived := _derived_block(source)) is not None:
|
||||
return _block_columns(derived[0], names, memo, tables, depth + 1)
|
||||
columns = _table_columns(source, tables)
|
||||
return [column in names for column in columns] if columns is not None else None
|
||||
|
||||
|
||||
def _block_columns(block: Any, names: set[str], memo: dict[int, Any], tables: dict[str, list[str]] | None,
|
||||
depth: int = 0) -> list[bool] | None:
|
||||
"""一块查询（派生来源）的每一列（按位置）取不取自 names：每一支一项项数，* 按它展开成的几列数（_star_columns），任何
|
||||
一支的那一列取自 names 就算；有一支数不清、几支不一样宽的为 None。memo 记下数过的（id(那块查询) → 结果）。"""
|
||||
from sqlglot import exp
|
||||
|
||||
if id(block) in memo:
|
||||
return memo[id(block)]
|
||||
memo[id(block)] = None
|
||||
rows: list[list[bool]] = []
|
||||
for leaf in _leaves(block):
|
||||
if not isinstance(leaf, exp.Select):
|
||||
return None
|
||||
row: list[bool] = []
|
||||
for item in leaf.expressions:
|
||||
if not _is_star(item):
|
||||
row.append(_reads(names, _value_columns(item)))
|
||||
elif (spread := _star_columns(item, leaf, names, memo, tables, depth)) is not None:
|
||||
row += spread
|
||||
else:
|
||||
return None
|
||||
rows.append(row)
|
||||
if not rows or len({len(row) for row in rows}) != 1:
|
||||
return None
|
||||
memo[id(block)] = [any(column) for column in zip(*rows)]
|
||||
return memo[id(block)]
|
||||
|
||||
|
||||
def _branch_names(select: Any, tables: dict[str, list[str]] | None, depth: int = 0) -> list[str | None] | None:
|
||||
"""一支（一条 SELECT）结果里按次序的列名：没起别名的表达式为 None（_output_name 认不准的）；* 按它展开成的列名
|
||||
（只展开一个来源时：数据源里的表按表结构，派生来源按它写的列名表、第一支的列名）。有 * 说不清展开成哪几列的为 None。
|
||||
UNION 第一支是 * 时，后一支的值按位置落进的就是这几个名字（_source_names）。"""
|
||||
from sqlglot import exp
|
||||
|
||||
out: list[str | None] = []
|
||||
for item in select.expressions:
|
||||
if not _is_star(item):
|
||||
out.append(_output_name(item))
|
||||
continue
|
||||
source = _single_source(item, select, depth)
|
||||
if source is None:
|
||||
return None
|
||||
if (derived := _derived_block(source)) is not None:
|
||||
block, listed = derived
|
||||
first = _leaves(block)[0]
|
||||
names = list(listed) if listed is not None else (
|
||||
_branch_names(first, tables, depth + 1) if isinstance(first, exp.Select) else None)
|
||||
else:
|
||||
names = _table_columns(source, tables)
|
||||
if names is None:
|
||||
return None
|
||||
out += names
|
||||
return out
|
||||
|
||||
|
||||
def _star_places(select: Any, width: int | None, names: set[str], tables: dict[str, list[str]] | None = None
|
||||
) -> set[int] | None:
|
||||
"""select 这一支里的 * 把取自 names 的值带到了结果（width 列）的哪几列。* 展开的列落在第一个 * 起、最后一个 * 后面
|
||||
那几项之前的那一段：只有一个 *、它展开成的每一列都数得清的（_star_columns：派生来源，表结构里查得到的表）只算取自
|
||||
names 的那几列，别的整段都算。数出来的列数和这一段对不上的（表结构过期了）也整段算。没有 * 带出这样的值为空；不知道
|
||||
结果有几列、* 前后的几项就比结果还宽的为 None。
|
||||
用在 * 展开的列换了名字的地方：写了列名表的派生来源（WITH t(a, b) AS (SELECT * FROM sales)）、UNION 后一支的 *
|
||||
（结果的列名取第一支的）。"""
|
||||
branch = list(select.expressions)
|
||||
stars = [i for i, item in enumerate(branch) if _is_star(item)]
|
||||
# 展开数据源里的表的一律算带出了遮罩列的值：表结构只用来数位置，不用来断定这张表没有遮罩列（表结构可能过期）
|
||||
if not any(_star_feeds(branch[i], select, names) for i in stars):
|
||||
return set()
|
||||
if width is None:
|
||||
return None
|
||||
first, end = stars[0], width - (len(branch) - 1 - stars[-1])
|
||||
if end < first:
|
||||
return None
|
||||
spread = _star_columns(branch[first], select, names, {}, tables) if len(stars) == 1 else None
|
||||
if spread is not None and len(spread) == end - first:
|
||||
return {first + k for k, hit in enumerate(spread) if hit}
|
||||
return set(range(first, end))
|
||||
|
||||
|
||||
def _blocks(root: Any) -> list[tuple[Any, str, list[str] | None]]:
|
||||
"""SQL 里的每一块查询（一条 SELECT，或者 UNION 这类连起来的一串）和它的结果流到哪：
|
||||
- "result"：最外层，结果就是快照的列；
|
||||
@@ -608,10 +819,16 @@ def _blocks(root: Any) -> list[tuple[Any, str, list[str] | None]]:
|
||||
return out
|
||||
|
||||
|
||||
def _source_names(top: Any, listed: list[str] | None, names: set[str]) -> tuple[set[str], bool]:
|
||||
def _source_names(top: Any, listed: list[str] | None, names: set[str],
|
||||
tables: dict[str, list[str]] | None = None) -> tuple[set[str], bool]:
|
||||
"""派生来源的结果里值取自遮罩列的那几列的名字，以及有没有认不出名字的（没起别名的表达式、两个 * 之间的项）。
|
||||
写了列名表的按位置对上列名表；没写的，单独一条 SELECT 取各项自己的列名，UNION 这类按位置对上第一支的列名。"""
|
||||
branches = _branches(top)
|
||||
写了列名表的按位置对上列名表；没写的，单独一条 SELECT 取各项自己的列名，UNION 这类按位置对上第一支的列名。
|
||||
* 带出遮罩列的值、展开的列又换了名字的（写了列名表、UNION 后一支的 *），它展开落到的那几列照位置对上（_star_places）；
|
||||
单独一条 SELECT、第一支的 * 不改名，遮罩列照原名传出去，按名字就认得出。第一支的 * 展开成哪几列数得清的（_branch_names），
|
||||
后一支落进去的位置按它对上列名。tables 是查询当时冻结的表结构（_schema_columns），* 展开数据源里的表时按它数位置。"""
|
||||
from sqlglot import exp
|
||||
|
||||
selects = [leaf for leaf in _leaves(top) if isinstance(leaf, exp.Select)]
|
||||
found: set[str] = set()
|
||||
unknown = False
|
||||
|
||||
@@ -622,28 +839,39 @@ def _source_names(top: Any, listed: list[str] | None, names: set[str]) -> tuple[
|
||||
else:
|
||||
found.add(name)
|
||||
|
||||
def taken(select: Any, width: int | None, head: Callable[[int | None], str | None]) -> None:
|
||||
branch = list(select.expressions)
|
||||
for item, place in zip(branch, _places(branch, width)):
|
||||
if not _is_star(item) and _reads(names, _value_columns(item)):
|
||||
take(head(place))
|
||||
spread = _star_places(select, width, names, tables)
|
||||
for place in sorted(spread) if spread is not None else [None]:
|
||||
take(head(place))
|
||||
|
||||
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)
|
||||
for select in selects:
|
||||
taken(select, len(listed), lambda p: listed[p] if p is not None and p < len(listed) else None)
|
||||
return found, unknown
|
||||
first = branches[0] if branches else []
|
||||
first = list(selects[0].expressions) if selects else []
|
||||
for item in first:
|
||||
if names & _value_columns(item):
|
||||
if not _is_star(item) and _reads(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)
|
||||
if len(selects) > 1:
|
||||
header = _branch_names(selects[0], tables)
|
||||
if header is not None:
|
||||
width, heads = len(header), dict(enumerate(header))
|
||||
else:
|
||||
# 每一支一样宽：拿一支没有 * 的定宽度，* 后面的项才数得出位置
|
||||
width = next((len(s.expressions) for s in selects if not any(_is_star(i) for i in s.expressions)), None)
|
||||
heads = {place: _output_name(item) for item, place in zip(first, _places(first, width))
|
||||
if place is not None}
|
||||
for select in selects[1:]:
|
||||
taken(select, width, lambda p: heads.get(p) if p 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:
|
||||
def _tainted(sql: str, kind: str | None, masked: set[str], width: int, tables: dict[str, list[str]] | None = None
|
||||
) -> tuple[set[str], set[int]] | None:
|
||||
"""SQL 里指代遮罩列的名字（小写），和结果（width 列）里值取自遮罩列的那几列的位置。解析不了返回 None。
|
||||
|
||||
名字从遮罩列本身起，一层层传下去：选出列的值取自这些名字的（buyer_phone AS 手机号、UPPER(buyer_phone) AS x），它的
|
||||
@@ -653,7 +881,10 @@ def _tainted(sql: str, kind: str | None, masked: set[str], width: int) -> tuple[
|
||||
后面的从末尾倒着数。认不出的——两个 * 之间的项、派生来源里没起别名的表达式（外层的 * 会把它原样带出来，列名
|
||||
又认不准）——整份快照的列都算（同解析不了）。
|
||||
FROM、JOIN、LATERAL 里的表值函数（json_each、unnest、JSON_TABLE、LATERAL VALUES）参数里写着这些名字的，它出的列
|
||||
都算（_function_columns）：遮罩列是 JSON 时，json_each 拆开是最自然的写法。出哪几列说不清的整份快照都算。"""
|
||||
都算（_function_columns）：遮罩列是 JSON 时，json_each 拆开是最自然的写法。出哪几列说不清的整份快照都算。
|
||||
* 展开的列换了名字（写了列名表的派生来源、UNION 后一支的 *）：* 带出这些名字的值时，它展开落到的那几列照位置算
|
||||
（_star_places，展开数据源里的表的按 tables 数位置，数不清的整段算）；标量子查询里的 * 按 _value_columns 算进外层的
|
||||
表达式。"""
|
||||
from sqlglot import exp
|
||||
|
||||
from app.data.sql_semantics import _parse, _read
|
||||
@@ -676,29 +907,37 @@ def _tainted(sql: str, kind: str | None, masked: set[str], width: int) -> tuple[
|
||||
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):
|
||||
if name and not _is_star(item) and name not in names and _reads(names, _value_columns(item)):
|
||||
names.add(name)
|
||||
for top, listed in sources:
|
||||
names |= _source_names(top, listed, names)[0]
|
||||
names |= _source_names(top, listed, names, tables)[0]
|
||||
for source in functions:
|
||||
if names & _value_columns(source):
|
||||
if _reads(names, _value_columns(source)):
|
||||
columns = _function_columns(source)
|
||||
unreadable = unreadable or columns is None
|
||||
names |= columns or set()
|
||||
if len(names) == before:
|
||||
break
|
||||
everything = unreadable or any(_source_names(top, listed, names)[1] for top, listed in sources)
|
||||
everything = unreadable or any(_source_names(top, listed, names, tables)[1] for top, listed in sources)
|
||||
positions: set[int] = set()
|
||||
for top, role, _listed in blocks:
|
||||
if role != "result":
|
||||
continue
|
||||
for branch in _branches(top):
|
||||
branches = [leaf for leaf in _leaves(top) if isinstance(leaf, exp.Select)]
|
||||
for k, select_ in enumerate(branches):
|
||||
branch = list(select_.expressions)
|
||||
for item, place in zip(branch, _places(branch, width)):
|
||||
if names & _value_columns(item):
|
||||
if not _is_star(item) and _reads(names, _value_columns(item)):
|
||||
if place is None or place >= width:
|
||||
everything = True
|
||||
else:
|
||||
positions.add(place)
|
||||
# UNION 后一支的 * 按位置落进第一支起的列名里；第一支、单独一条 SELECT 的 * 照原名带出，按名字认得出
|
||||
spread = _star_places(select_, width, names, tables) if k else set()
|
||||
if spread is None or any(place >= width for place in spread):
|
||||
everything = True
|
||||
else:
|
||||
positions |= spread
|
||||
return names, set(range(width)) if everything else positions
|
||||
except Exception: # noqa: BLE001 - 分析本身的缺陷不能让回答出错：按解析不了处理（宁可多遮）
|
||||
logger.warning("质询的 SQL 遮罩分析失败", exc_info=True)
|
||||
@@ -709,8 +948,10 @@ def _mentions(sql: str, names: Iterable[str]) -> bool:
|
||||
return any(re.search(rf"(?<![0-9A-Za-z_]){re.escape(n)}(?![0-9A-Za-z_])", sql, re.I) for n in names if n)
|
||||
|
||||
|
||||
def _derived_hidden(snapshot: dict[str, Any], hidden: set[str], kind: str | None) -> tuple[set[str], set[str]]:
|
||||
"""(这份快照要遮的列, SQL 里指代遮罩列的名字)，都是小写。
|
||||
def _derived_hidden(snapshot: dict[str, Any], hidden: set[str], kind: str | None,
|
||||
tables: dict[str, list[str]] | None = None) -> tuple[set[str], set[str]]:
|
||||
"""(这份快照要遮的列, SQL 里指代遮罩列的名字)，都是小写。tables 是查询当时冻结的表结构（_schema_columns），* 展开
|
||||
数据源里的表时按它数位置；没给的按快照记着的表结构快照（schema_artifact）认，和工具边界遮的时候是同一份。
|
||||
|
||||
证据面板只按列名遮（hidden）；agent 的查询起了别名、套了函数，结果列名就对不上。本仓库的查询惯例是起中文别名
|
||||
（region AS 区域），所以按 SQL 认出改了名字的遮罩列（_tainted）。SQL 解析不了、里面又写着遮罩列的，整份快照的
|
||||
@@ -720,7 +961,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, len(columns))
|
||||
found = _tainted(sql, kind, base, len(columns), _snapshot_columns(snapshot) if tables is None else tables)
|
||||
if found is None:
|
||||
return (base | {c.lower() for c in columns}, base) if _mentions(sql, base) else (base, base)
|
||||
names, positions = found
|
||||
@@ -729,12 +970,13 @@ def _derived_hidden(snapshot: dict[str, Any], hidden: set[str], kind: str | None
|
||||
|
||||
def derived_columns(snapshot: dict[str, Any], hidden: set[str], kind: str | None) -> set[str] | None:
|
||||
"""按 SQL 认出值取自遮罩列（hidden，小写）的结果列（小写，列名直接对上的也算）。SQL 解析不了、合并查询的快照（不看
|
||||
SQL）为 None：认不准，宽严由调用方定（_derived_hidden 整份遮；双盲复算照常派盲算，由盲算自己的查询去认）。"""
|
||||
SQL）为 None：认不准，宽严由调用方定（_derived_hidden 整份遮；双盲复算照常派盲算，由盲算自己的查询去认）。表结构同
|
||||
_derived_hidden，按快照记着的表结构快照认。"""
|
||||
sql = snapshot.get("sql")
|
||||
if not hidden or not isinstance(sql, str) or not sql.strip() or snapshot.get("source") == MERGE_SOURCE:
|
||||
return None
|
||||
columns = [str(c) for c in snapshot.get("columns") or []]
|
||||
found = _tainted(sql, kind, set(hidden), len(columns))
|
||||
found = _tainted(sql, kind, set(hidden), len(columns), _snapshot_columns(snapshot))
|
||||
if found is None:
|
||||
return None
|
||||
names, positions = found
|
||||
@@ -753,6 +995,23 @@ def _snapshot_labels(snapshot: dict[str, Any], kind: str | None, tables: dict[st
|
||||
return _label_columns(sql, kind, [str(c) for c in snapshot.get("columns") or []], tables)
|
||||
|
||||
|
||||
def _snapshot_columns(snapshot: dict[str, Any]) -> dict[str, list[str]] | None:
|
||||
"""快照记着的表结构快照（schema_artifact，查询当时冻结的，和工具边界遮的时候用的是同一份）里按次序的列名。"""
|
||||
artifact = snapshot.get("schema_artifact")
|
||||
return _frozen_columns(artifact) if isinstance(artifact, str) else None
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=64)
|
||||
def _frozen_columns(artifact: str) -> dict[str, list[str]] | None:
|
||||
"""表结构快照工件里每张表按次序的列名（_schema_columns）。同 _frozen_tables：按内容寻址，记下来；取不到的为 None。"""
|
||||
from app.core import artifact_store
|
||||
|
||||
try:
|
||||
return _schema_columns(artifact_store.load(artifact))
|
||||
except Exception: # noqa: BLE001 - 取不到就当不知道表结构（* 展开表的整段遮，宁可多遮）
|
||||
return None
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=64)
|
||||
def _frozen_tables(artifact: str) -> dict[str, set[str]] | None:
|
||||
"""表结构快照工件里每张表的列名（_schema_tables）。工件按内容寻址、不会变，记下来；取不到、被改过的为 None。"""
|
||||
@@ -1489,16 +1748,39 @@ class QueryMask:
|
||||
secrets: _Secrets
|
||||
|
||||
|
||||
@dataclass
|
||||
class SeenValues:
|
||||
"""一次核查里认出的遮罩值，只在内存里：given 是调用方交进来的（审批质询、反方审稿收的报告快照里的值），found 是
|
||||
agent 自己的查询在工具边界遮掉的（tools/datasource 每次查询并进来，add）。两份分开记：合起来（pooled）时 agent 查到的
|
||||
在前，到了上限 SECRETS_MAX 挤掉的是报告里的值——agent 查到的就在它眼前这几份结果里，后面的查询、它写的回答和理由
|
||||
最可能再碰到。以前先并报告的值，收满了上限，agent 查到的一个也并不进去。"""
|
||||
|
||||
given: _Secrets = field(default_factory=_Secrets)
|
||||
found: _Secrets = field(default_factory=_Secrets)
|
||||
|
||||
def add(self, secrets: _Secrets) -> None:
|
||||
"""并进一次查询时遮掉的值（mask_query_result 交回的 secrets）。"""
|
||||
self.found.merge(secrets)
|
||||
|
||||
def pooled(self) -> _Secrets:
|
||||
"""合起来的另一份（改它不影响这两份）：agent 查到的在前，调用方交进来的在后；调用方免遮的文字（问题、被审的那句
|
||||
结论原文里写着的值）对两边都算。上万个值合一遍要几十毫秒，调用方放进线程。"""
|
||||
out = self.found.copy()
|
||||
out.merge(self.given)
|
||||
return out
|
||||
|
||||
|
||||
def mask_query_result(payload: dict[str, Any], masked: Iterable[str], kind: str | None,
|
||||
known: _Secrets | None = None, *, schema: Any = None) -> QueryMask:
|
||||
"""一份查询结果（tools/datasource 的查询工具交回的样子：columns / rows / column_types / sql）在交给核查 agent、落进
|
||||
快照之前先遮。masked 是数据源设了遮罩的列，kind 是数据源的类型（SQL 按它的方言解析）。known 是此前已经认出的遮罩值：
|
||||
这个 agent 前面的查询遮掉的，加上调用方交进来的（审批质询、反方审稿收的报告快照里的值）。schema 是数据源的表结构
|
||||
（schema_cache，和随快照冻结的是同一份），认分档标签时用（_label_columns 的 tables）。
|
||||
（schema_cache，和随快照冻结的是同一份）：认分档标签（_label_columns 的 tables）、数 * 展开数据源里的表落在第几列
|
||||
（_derived_hidden 的 tables）时用。
|
||||
|
||||
- 整列遮：列名直接对上遮罩列的，加上按 SQL 认出值取自遮罩列的（_derived_hidden：别名、函数、子查询、WITH、
|
||||
SELECT *、表值函数；COUNT、只在条件里用到、分档当分组键的不算）；SQL 解析不了、又写着遮罩列的，整份结果的列都遮。
|
||||
这几列的格子换成「已遮罩」，类型记成文本
|
||||
SELECT *（展开的列换了名字的也算）、表值函数；COUNT、只在条件里用到、分档当分组键的不算）；SQL 解析不了、又写着
|
||||
遮罩列的，整份结果的列都遮。这几列的格子换成「已遮罩」，类型记成文本
|
||||
- 别的列按值换：整格就是遮罩值的换掉，文字格里夹着的那一截换掉（换了写法的号码也认，_Secrets.loose_spans）。遮罩值
|
||||
是这份结果里的加上 known：先查号码、再查写着号码的备注、没设遮罩的联系方式表（别的库里的也一样）时，后一份同样
|
||||
遮好。门槛同按值遮（SECRET_DIGITS、SECRET_MIN）；只装着分档标签的列里的标签不算遮罩值，这份结果自己写的标签也不
|
||||
@@ -1514,7 +1796,7 @@ def mask_query_result(payload: dict[str, Any], masked: Iterable[str], kind: str
|
||||
if not (base or given) or not columns or not _query_shaped(payload):
|
||||
return QueryMask(payload, [], _Secrets())
|
||||
try:
|
||||
hidden = _derived_hidden(payload, base, kind)[0] if base else set()
|
||||
hidden = _derived_hidden(payload, base, kind, _schema_columns(schema))[0] if base else set()
|
||||
own = _Secrets()
|
||||
labels = _snapshot_labels(payload, kind, _schema_tables(schema))
|
||||
if any(c.lower() in hidden for c in columns):
|
||||
|
||||
@@ -393,11 +393,13 @@ def _masked_rationale(outcome: VerifyOutcome, secrets: Any, *, given: Iterable[s
|
||||
换的办法同审批质询的回答（_scrubbed_source）：原文里的值先换掉，再照原来的目录重新组装理由文档；引用到的格子按
|
||||
遮好的快照渲染（文字格里夹着的遮罩值换掉那一截）。理由里认不出遮罩值的原样交回。"""
|
||||
from app.engine.evidence import compose_doc
|
||||
from app.engine.masking import _masked_view, _scrubbed_source
|
||||
from app.engine.masking import _masked_view, _scrubbed_source, _Secrets
|
||||
|
||||
doc, result = outcome.rationale, outcome.result
|
||||
pool = secrets.copy()
|
||||
# agent 查询时遮掉的值先收（masked_values 里已经排在报告的值前面）：报告的值收满了上限 SECRETS_MAX，也不丢 agent 查到的
|
||||
pool = _Secrets()
|
||||
pool.merge(outcome.masked_values)
|
||||
pool.merge(secrets)
|
||||
pool.protect(given)
|
||||
if result is None or not isinstance(doc, dict) or not pool.strong:
|
||||
return outcome
|
||||
|
||||
@@ -7,12 +7,14 @@
|
||||
（tools/datasource.py）。这一路没有审批关卡，不能指望人在半路拦下一条 DELETE。
|
||||
|
||||
**先遮再交。** 建查询工具时还打开遮罩开关（tools/datasource 的 mask_results）：数据源设了遮罩时，取自遮罩列的格子
|
||||
（按 SQL 认，起别名、套函数、经子查询或 WITH 转一道、SELECT *、表值函数拆开的都算）在交给模型、落进查询快照之前就
|
||||
换成「已遮罩」，别的格子按值换：这份结果里的遮罩值、前面的查询遮掉的、调用方交进来的（masked_values：审批质询、
|
||||
反方审稿收的报告快照里的值），文字格里夹着的换那一截（engine/masking.py 的 mask_query_result）；查询失败时报错原话
|
||||
里的值也换。模型看不到遮罩值，就不会把它写进回答、理由、SQL；以前是事后按值认，每一轮复核都找到新的漏法。快照里的
|
||||
masked_at_query 记下查询时就遮了的列，格子引用指到这几列、指到按值换掉的格子的都不采信。遮掉的值只在内存里交回
|
||||
调用方（VerifyOutcome.masked_values，审批质询、反方审稿按值遮的第二道防线用），不进快照、不交给模型、不进 as_dict。
|
||||
（按 SQL 认，起别名、套函数、经子查询或 WITH 转一道、SELECT *（展开的列换了名字的也算）、表值函数拆开的都算）在交给
|
||||
模型、落进查询快照之前就换成「已遮罩」，别的格子按值换：这份结果里的遮罩值、前面的查询遮掉的、调用方交进来的
|
||||
（masked_values：审批质询、反方审稿收的报告快照里的值），文字格里夹着的换那一截（engine/masking.py 的
|
||||
mask_query_result）；查询失败时报错原话里的值也换。模型看不到遮罩值，就不会把它写进回答、理由、SQL；以前是事后
|
||||
按值认，每一轮复核都找到新的漏法。快照里的 masked_at_query 记下查询时就遮了的列，格子引用指到这几列、指到按值换掉
|
||||
的格子的都不采信。遮掉的值只在内存里交回调用方（VerifyOutcome.masked_values，审批质询、反方审稿按值遮的第二道防线
|
||||
用），不进快照、不交给模型、不进 as_dict。调用方交进来的和查询时遮掉的分开记（masking.SeenValues），合起来时 agent
|
||||
查到的在前：合起来到了上限 SECRETS_MAX，挤掉的是报告里的值。
|
||||
|
||||
**数从快照取。** 每次成功的查询是一份查询快照（工具自己落的 put_json），按这个 agent 自己的顺序编 Q1、Q2…。
|
||||
最后交结果时，调用方结构里的数值字段一律改成格子引用 {alias, row, column}，由代码按引用从快照里取值：模型自己
|
||||
@@ -67,7 +69,7 @@ from app.engine import spend
|
||||
from app.engine.evidence import build_catalog, cell_eid, compose_doc, query_entry_fields, schema_ledger_entry
|
||||
from app.engine.expressions import CellError, column_kind, locate_cell, numeric_text
|
||||
from app.engine.judge import FALLBACK_MAX_TOKENS, MAX_TOKENS, _json_in, _tokens, _usd
|
||||
from app.engine.masking import _masked_word, _Secrets
|
||||
from app.engine.masking import SeenValues, _masked_word, _Secrets
|
||||
from app.engine.nodes.llm import _DIGEST_BUDGET, _digest, _split_structured, _typed, explain_model_error
|
||||
from app.engine.state import message_text
|
||||
from app.engine.toolcalls import ToolTimeout, limit_fields, limit_of, run_bounded
|
||||
@@ -160,8 +162,9 @@ class VerifyOutcome:
|
||||
output_tokens, total_tokens, cost_usd, calls, unpriced_calls}。limit：触顶原因（LIMITS 之一），没触顶为 None。
|
||||
answered：模型是不是自己停下来的（没被上限掐断，也没因为输出到了上限被截断）。被截断的一律没有结果，error 是
|
||||
TRUNCATED_ERROR。
|
||||
masked_values：查询时在工具边界遮掉的值（engine/masking 的 _Secrets），连同调用方交进来的。只在内存里交给调用方按值遮
|
||||
（审批质询的回答、反方审稿的理由），不进 as_dict、不往外交。
|
||||
masked_values：查询时在工具边界遮掉的值（engine/masking 的 _Secrets），连同调用方交进来的：agent 查到的在前，合起来
|
||||
到了上限 SECRETS_MAX 挤掉的是调用方的（SeenValues.pooled）。只在内存里交给调用方按值遮（审批质询的回答、反方审稿的
|
||||
理由），不进 as_dict、不往外交。
|
||||
"""
|
||||
|
||||
result: dict[str, Any] | None
|
||||
@@ -670,9 +673,9 @@ class _State:
|
||||
queries: list[dict[str, Any]] = field(default_factory=list)
|
||||
ledger: list[dict[str, Any]] = field(default_factory=list)
|
||||
model: str | None = None
|
||||
#: 查询工具在工具边界遮掉的值（经 ToolContext.extra 并进来，只在内存里），起头是调用方交进来的。后面的查询按它们
|
||||
#: 按值遮
|
||||
seen: _Secrets = field(default_factory=_Secrets)
|
||||
#: 认出的遮罩值（只在内存里）：调用方交进来的，加上查询工具在工具边界遮掉的（经 ToolContext.extra 并进来）。后面的
|
||||
#: 查询按它们按值遮，最后合起来交回调用方（run_verify_agent），agent 查到的在前
|
||||
seen: SeenValues = field(default_factory=SeenValues)
|
||||
|
||||
def outcome(self, **fields: Any) -> VerifyOutcome:
|
||||
usage = self.usage
|
||||
@@ -682,8 +685,7 @@ class _State:
|
||||
for q in self.queries]
|
||||
base: dict[str, Any] = {"result": None, "cells": [], "queries": shown, "rationale": None, "usage": usage,
|
||||
"notes": self.notes, "model": self.model,
|
||||
"duration_ms": int((time.monotonic() - self.started) * 1000),
|
||||
"masked_values": self.seen}
|
||||
"duration_ms": int((time.monotonic() - self.started) * 1000)}
|
||||
return VerifyOutcome(**{**base, **fields})
|
||||
|
||||
|
||||
@@ -720,20 +722,25 @@ async def run_verify_agent(
|
||||
data_versions；不在运行里传 Versions.CURRENT；缺省 UNSET 时上传源一律报错，不拿当前版本顶上）。
|
||||
- masked_values：调用方已经认出的遮罩值（engine/masking 的 _Secrets，审批质询、反方审稿收的报告快照里的值）。
|
||||
查询工具从第一条查询起就按它们按值遮：agent 查到的备注、别的表里写着同一个号码时，模型看不到。随查询时遮掉的值
|
||||
一起交回（VerifyOutcome.masked_values）。
|
||||
一起交回（VerifyOutcome.masked_values）。不改它。
|
||||
"""
|
||||
state = _State(started=time.monotonic())
|
||||
state = _State(started=time.monotonic(), seen=SeenValues(
|
||||
given=masked_values if masked_values is not None else _Secrets()))
|
||||
try:
|
||||
if masked_values is not None:
|
||||
# 上万个值并一遍要几十毫秒，放进线程，别挡住同一进程里别的运行
|
||||
await asyncio.to_thread(state.seen.merge, masked_values)
|
||||
return await _verify(state, spec=spec, system=system, task=task, sources=sources,
|
||||
output_schema=output_schema, max_queries=max_queries, timeout_s=timeout_s,
|
||||
max_cost_usd=max_cost_usd, daily_max_usd=daily_max_usd, max_steps=max_steps, emit=emit,
|
||||
run_id=run_id, node_id=node_id, data_versions=data_versions)
|
||||
out = await _verify(state, spec=spec, system=system, task=task, sources=sources,
|
||||
output_schema=output_schema, max_queries=max_queries, timeout_s=timeout_s,
|
||||
max_cost_usd=max_cost_usd, daily_max_usd=daily_max_usd, max_steps=max_steps, emit=emit,
|
||||
run_id=run_id, node_id=node_id, data_versions=data_versions)
|
||||
except Exception as e: # noqa: BLE001 - 永远不抛：库写不进去、建工具出错这类意外也落在返回值里
|
||||
logger.warning("核查出错", exc_info=True)
|
||||
return state.outcome(error=f"核查出错（{describe_exception(e)}）")
|
||||
out = state.outcome(error=f"核查出错（{describe_exception(e)}）")
|
||||
try:
|
||||
# agent 查到的在前、调用方交进来的在后：上万个值合一遍要几十毫秒，放进线程，别挡住同一进程里别的运行
|
||||
pooled = await asyncio.to_thread(state.seen.pooled)
|
||||
except Exception: # noqa: BLE001 - 永远不抛：合不起来就只交调用方交进来的（agent 的查询结果在工具边界已经遮过）
|
||||
logger.warning("核查遮掉的值合并出错", exc_info=True)
|
||||
pooled = state.seen.given
|
||||
return replace(out, masked_values=pooled)
|
||||
|
||||
|
||||
async def _verify(state: _State, *, spec: ModelSpec, system: str, task: str, sources: Iterable[str],
|
||||
|
||||
@@ -53,8 +53,9 @@ RUN_VERSIONS_KEY = "data_versions"
|
||||
DATA_UNAVAILABLE = "查询失败（数据不可用）："
|
||||
#: 运行中途补固定一个源时发的那条日志事件的 code（_pin_late）
|
||||
LATE_PIN_CODE = "data_version_pinned"
|
||||
#: ToolContext.extra 里放「查询时遮掉的值」的键（engine/masking 的 _Secrets）：核查 agent（engine/verify_agent.py）放进去，
|
||||
#: 打开了遮罩开关的查询工具把每次遮掉的值并进去。只在内存里给调用方按值遮用，不进快照、不交给模型
|
||||
#: ToolContext.extra 里放「认出的遮罩值」的键（engine/masking 的 SeenValues：调用方交进来的，加上查询时遮掉的）：核查
|
||||
#: agent（engine/verify_agent.py）放进去，打开了遮罩开关的查询工具按它们按值遮，再把每次遮掉的值并进去。只在内存里给
|
||||
#: 调用方按值遮用，不进快照、不交给模型
|
||||
MASKED_VALUES_KEY = "masked_values"
|
||||
|
||||
#: 每次查询之后基于数据目录的 SQL 检查最多等这么久。正常一条 SQL 几毫秒；等不到就放弃这一次检查，
|
||||
@@ -498,9 +499,11 @@ def _make_query_tool(source: Any, ctx: ToolContext, *, fixed: bool = False,
|
||||
return f"SQL 被拒绝：{FORCED_READONLY}"
|
||||
limits = QueryLimits(timeout_seconds=seconds, **({"max_rows": int(limit)} if limit else {}))
|
||||
mask = masked_columns(source.options)
|
||||
# 此前认出的遮罩值：交给线程的是一份拷贝，线程里读的时候这边不会再往里并
|
||||
# 此前认出的遮罩值合起来的另一份（agent 前面的查询遮掉的在前，调用方交进来的在后，见 masking.SeenValues）：线程里
|
||||
# 读的时候这边不会再往里并。上万个值合一遍要几十毫秒，放进线程
|
||||
sink = ctx.extra.get(MASKED_VALUES_KEY) if mask_results else None
|
||||
known = sink.copy() if sink is not None and sink.strong else None
|
||||
known = await asyncio.to_thread(sink.pooled) if sink is not None else None
|
||||
known = known if known is not None and known.strong else None
|
||||
try:
|
||||
result = await run_query(source, sql, limits=limits)
|
||||
except SqlRejected as e:
|
||||
@@ -538,8 +541,8 @@ def _make_query_tool(source: Any, ctx: ToolContext, *, fixed: bool = False,
|
||||
if mask or done.payload is not payload:
|
||||
payload, masked_now = done.payload, done.columns
|
||||
if sink is not None:
|
||||
# 上万个值并一遍要几十毫秒（免遮的文字要重新过一遍），也放进线程
|
||||
await asyncio.to_thread(sink.merge, done.secrets)
|
||||
# 并进 agent 查到的那一份（不和调用方交进来的混在一起，上限收满了也不丢），也放进线程
|
||||
await asyncio.to_thread(sink.add, done.secrets)
|
||||
# 对照数据目录检查这条 SQL（一对多关联后重复计算、存量跨期求和……）。结果交给模型（Agent 据此改写），
|
||||
# 也随查询快照存下：证据面板看得到，口径卡据此给指标标「存疑」、出具降档。没查出问题不加这个键，
|
||||
# 快照和以前一字不差
|
||||
|
||||
@@ -1405,3 +1405,26 @@ async def test_the_node_refuses_the_same_bad_settings_at_run_time():
|
||||
with pytest.raises(NodeError) as caught:
|
||||
await challenge_mod.run_challenge({"nodes": {}}, ctx)
|
||||
assert str(caught.value) == message
|
||||
|
||||
|
||||
def test_values_the_agent_found_are_masked_in_the_rationale_when_the_report_fills_the_cap(monkeypatch):
|
||||
"""复核第 2 轮：报告快照里的遮罩值收满了上限 SECRETS_MAX，反方 agent 查询时遮掉的号码照样认得出（agent 查到的先收）。
|
||||
以前先收报告的值、满了就停，agent 查到的号码写进理由原样存下。"""
|
||||
from app.engine import masking
|
||||
from app.engine.evidence import compose_doc
|
||||
from app.engine.verify_agent import VerifyOutcome
|
||||
|
||||
monkeypatch.setattr(masking, "SECRETS_MAX", 40)
|
||||
secrets = masking._Secrets()
|
||||
secrets.add({"columns": ["p"], "rows": [[str(13800000000 + i)] for i in range(40)]}, {"p"})
|
||||
assert len(secrets.weak) == masking.SECRETS_MAX
|
||||
found = masking._Secrets()
|
||||
found.add({"columns": ["p"], "rows": [["19914742785"]]}, {"p"})
|
||||
# 和 run_verify_agent 交回的一样：agent 查到的在前，调用方交进去的在后
|
||||
seen = masking.SeenValues(given=secrets)
|
||||
seen.add(found)
|
||||
doc = compose_doc("号码 199-1474-2785 的客户在华南。", {}, node_id="chal", run_id="r1")
|
||||
outcome = VerifyOutcome(result={"verdict": "stands", "rationale": doc["markdown"]}, cells=[], queries=[],
|
||||
rationale=doc, usage={}, masked_values=seen.pooled())
|
||||
masked = challenge_mod._masked_rationale(outcome, secrets, given=[])
|
||||
assert masked.result["rationale"] == "号码 已遮罩 的客户在华南。"
|
||||
|
||||
@@ -933,6 +933,49 @@ async def test_the_inquiry_agent_sees_and_stores_its_queries_masked(monkeypatch,
|
||||
assert content == artifact_store.load(query["artifact"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql, rows", [
|
||||
pytest.param("SELECT '合计' AS 区域, SUM(amount) AS 金额 FROM sales UNION ALL SELECT * FROM "
|
||||
"(SELECT buyer_phone, amount FROM sales ORDER BY amount DESC LIMIT 2)",
|
||||
[["已遮罩", 2800], ["已遮罩", 1200], ["已遮罩", 900]], id="UNION 后一支的 *"),
|
||||
pytest.param("WITH t(a, b, c, d, e) AS (SELECT * FROM sales) SELECT c AS 区域, d AS 金额 FROM t "
|
||||
"ORDER BY d DESC LIMIT 2", [["已遮罩", 1200], ["已遮罩", 900]], id="CTE 列名表套 *"),
|
||||
])
|
||||
async def test_a_star_renaming_the_masked_column_never_shows(monkeypatch, tmp_path, client, sql, rows):
|
||||
"""复核第 2 轮的 blocker：* 把手机号带进了别的列名（UNION 后一支的 *，结果的列名取第一支的「区域」；CTE 的列名表把
|
||||
SELECT * 展开的列换了名字）。以前交给模型的工具结果里就有原值，回答「金额最大的一笔来自 13900000001」照常交出，事件、
|
||||
质询工件、通用工件接口都给原值。现在 * 展开落到的那几列查询时就遮掉：只展开一个子查询的按位置只遮号码那一列；展开
|
||||
数据源里的表的按查询当时的表结构数出号码排在第几，也只遮那一列。"""
|
||||
name, source_id = await _source(tmp_path, options={"mask_columns": "buyer_phone"}, script=SALES)
|
||||
plays = {"最大": Play([sql], "金额最大的一笔来自 [[v:Q1.r1.区域]]。[[see:Q1]]")}
|
||||
Models(monkeypatch, name, writer="华东销售额 [[v:Q1.r0.销售额]]。[[see:Q1]]", plays=plays)
|
||||
seen: list[str] = []
|
||||
answering = MockChatModel._decide
|
||||
|
||||
def spy(model, messages):
|
||||
seen.extend(str(m.content) for m in messages if isinstance(m, ToolMessage))
|
||||
return answering(model, messages)
|
||||
|
||||
monkeypatch.setattr(MockChatModel, "_decide", spy)
|
||||
try:
|
||||
paused = await pause(name, graph=fetching(name, QUERY))
|
||||
[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"]["ok"], item["answer"]["text"]) == (True, "金额最大的一笔来自 已遮罩。")
|
||||
assert seen and all("1390000000" not in s for s in seen), "交给模型的查询结果里就是遮好的"
|
||||
[query] = item["queries"]
|
||||
for url in (f"/api/runs/{paused.id}/inquiries/artifacts/{query['artifact']}",
|
||||
f"/api/artifacts/{query['artifact']}"):
|
||||
got = await client.get(url)
|
||||
assert got.status_code == 200 and got.json()["content"]["rows"] == rows and "1390000000" not in got.text
|
||||
[event] = await events(paused.id, "approval.inquiry")
|
||||
assert "1390000000" not in json.dumps(event.data, ensure_ascii=False)
|
||||
finally:
|
||||
await engines.invalidate(source_id)
|
||||
|
||||
|
||||
async def test_a_masked_column_under_another_name_is_masked_everywhere(monkeypatch, masked_shop, client):
|
||||
"""遮罩列起了中文别名（本仓库查询的惯例）：回答里引用到的格子、事件、回答文档、质询工件里都不露出原值。
|
||||
报告自己的查询没选手机号，原值哪里都认不出来，只能按 SQL 认出「手机号」就是 buyer_phone。"""
|
||||
@@ -1552,6 +1595,40 @@ async def test_a_number_from_the_agents_other_query_is_masked_when_the_number_ca
|
||||
await engines.invalidate(source_id)
|
||||
|
||||
|
||||
async def test_a_number_the_agent_found_is_masked_even_when_report_values_fill_the_cap(monkeypatch, tmp_path, client):
|
||||
"""复核第 2 轮：报告快照里的遮罩值收满了上限（这里把 SECRETS_MAX 调成报告正好收满的大小）。agent 先查出华东的号码
|
||||
（查询时遮掉），再查回访备注：备注里的号码照样在交给模型之前换掉，回答、事件、质询工件里都没有。以前开始回答时先并
|
||||
报告的值、收满就停，agent 查到的号码并不进去：备注原样交给模型、原样落库，回答引用这一格时第二道防线也认不出，原值
|
||||
写进了运行记录。"""
|
||||
from app.engine import masking
|
||||
|
||||
name, source_id = await _source(tmp_path, options={"mask_columns": "buyer_phone"}, script=CALLBACKS)
|
||||
report_sql = "SELECT buyer_phone, amount AS 金额 FROM sales WHERE region <> '华东' ORDER BY id"
|
||||
plays = {"回访": Play([PHONES_EAST, NOTES], "华东的回访备注是 [[v:Q2.r0.备注]]。[[see:Q2]]")}
|
||||
Models(monkeypatch, name, writer="华南一笔金额 [[v:Q1.r0.金额]]。[[see:Q1]]", plays=plays)
|
||||
try:
|
||||
paused = await pause(name, graph=fetching(name, report_sql))
|
||||
[approval] = await pending(paused.id)
|
||||
report = [artifact_store.load(e.data["query_artifact"]) for e in await events(paused.id, "tool.end")
|
||||
if e.data.get("query_artifact")]
|
||||
full = masking._Secrets()
|
||||
full.add(report[0], {"buyer_phone"})
|
||||
monkeypatch.setattr(masking, "SECRETS_MAX", len(full.weak))
|
||||
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"]["ok"], item["answer"]["text"]) == (True, "华东的回访备注是 客户已遮罩来电改址。")
|
||||
phones, notes = item["queries"]
|
||||
assert artifact_store.load(notes["artifact"])["rows"] == [["华东", "客户已遮罩来电改址"], ["华南", "例行回访"]]
|
||||
[event] = await events(paused.id, "approval.inquiry")
|
||||
assert "1390000000" not in json.dumps(event.data, ensure_ascii=False)
|
||||
got = await client.get(f"/api/runs/{paused.id}/inquiries/artifacts/{notes['artifact']}")
|
||||
assert got.status_code == 200 and "1390000000" not in got.text
|
||||
finally:
|
||||
await engines.invalidate(source_id)
|
||||
|
||||
|
||||
async def test_a_number_from_the_agents_other_query_is_masked_when_the_number_came_last(monkeypatch, tmp_path, client):
|
||||
"""先查备注、后查号码：查备注时还不知道里面写着遮罩值，原样交给模型、原样存下；回答引用这一格，渲染出来才露，整份
|
||||
不交。回答时按内存里的值换掉的位置记进事件（不含值），事后取质询工件照着再换一遍：质询工件里那一截是「已遮罩」，
|
||||
|
||||
@@ -0,0 +1,251 @@
|
||||
"""按 SQL 认取自遮罩列的结果列（engine/masking.py 的 _derived_hidden）：* 把遮罩列带进了别的列名（复核第 2 轮的 blocker）。
|
||||
|
||||
* 展开的列按位置落进结果：WITH t(a, b, …) AS (SELECT * FROM sales) 的列名表、FROM 里子查询的列名表会给它们换上别的
|
||||
名字，UNION 后一支的 * 落进第一支起的列名里，标量子查询里的 * 则把值交给外层的表达式。以前 * 本身不算取自哪一列，这些
|
||||
写法的原值原样交给核查 agent、落进快照，审批质询按值遮的第二道防线也认不出。现在 * 带出的值取自遮罩列时，它展开落到的
|
||||
那几列都算：只展开一个来源、那里的每一列都数得清的（子查询按选出项，数据源里的表按查询当时的表结构），按位置只算取自
|
||||
遮罩列的那几列；数不清的（没有表结构、列数对不上）整段都算（多遮）。不改名的 *（单独一条 SELECT 的 *、UNION 第一支的 *）
|
||||
照旧按名字认；* 只带出没碰遮罩列的值、EXISTS 里的 * 都不算。
|
||||
|
||||
查询都在内存里的 SQLite 上真跑一遍，看遮完的结果里还有没有原值。数据都是虚构的：区域、手机号、姓名、金额。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sqlite3
|
||||
|
||||
import pytest
|
||||
|
||||
from app.db import models # noqa: F401 - 建表（conftest）之前先登记模型：表结构快照要存进工件库
|
||||
from app.engine import masking
|
||||
|
||||
MASKED_WORD = "已遮罩"
|
||||
#: 设了遮罩的列
|
||||
MASKED = ["buyer_phone", "buyer_name"]
|
||||
#: 遮罩列里的值：遮完的结果里不能出现
|
||||
SECRETS = ("1390000000", "张伟", "李娜", "王芳")
|
||||
SCRIPT = ("CREATE TABLE sales (id INTEGER, region TEXT, buyer_phone TEXT, buyer_name TEXT, amount INTEGER, note TEXT);"
|
||||
"INSERT INTO sales VALUES (1, '华东', '13900000001', '张伟', 1200, '改址'),"
|
||||
" (2, '华南', '13900000002', '李娜', 300, '回访'), (3, '华北', '13900000003', '王芳', 900, '');"
|
||||
# 只有一列号码的表：标量子查询、UNION 后一支的 * 直接展开它
|
||||
"CREATE TABLE phones (buyer_phone TEXT);"
|
||||
"INSERT INTO phones VALUES ('13900000009');")
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def db():
|
||||
conn = sqlite3.connect(":memory:")
|
||||
conn.executescript(SCRIPT)
|
||||
yield conn
|
||||
conn.close()
|
||||
|
||||
|
||||
def masked_run(db, sql: str) -> masking.QueryMask:
|
||||
"""真跑一遍 SQL，按工具边界的办法遮（mask_query_result）。"""
|
||||
cursor = db.execute(sql)
|
||||
columns = [d[0] for d in cursor.description]
|
||||
rows = [list(r) for r in cursor.fetchall()]
|
||||
payload = {"columns": columns, "rows": rows, "row_count": len(rows), "truncated": False, "sql": sql}
|
||||
return masking.mask_query_result(payload, MASKED, "sqlite")
|
||||
|
||||
|
||||
def _leaks(done: masking.QueryMask) -> bool:
|
||||
text = json.dumps(done.payload["rows"], ensure_ascii=False)
|
||||
return any(s in text for s in SECRETS)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql, hidden, shown", [
|
||||
pytest.param("WITH t(a, b, c, d, e, f) AS (SELECT * FROM sales) SELECT * FROM t", {"c", "d"}, set(),
|
||||
id="CTE 列名表套 * 展开表"),
|
||||
pytest.param("WITH t(a, b, c, d, e, f) AS (SELECT * FROM sales) SELECT c AS 号码, e AS 金额 FROM t "
|
||||
"ORDER BY e DESC LIMIT 1", {"号码"}, set(), id="CTE 列名表套 * 再改名"),
|
||||
pytest.param("SELECT 0 AS a, '' AS b, '' AS c, '' AS d, 0 AS e, '' AS f UNION ALL SELECT * FROM sales",
|
||||
{"c", "d"}, set(), id="UNION 后一支的 * 展开表"),
|
||||
pytest.param("SELECT region FROM sales UNION ALL SELECT * FROM phones", {"region"}, set(),
|
||||
id="UNION 后一支的 * 展开只有号码的表"),
|
||||
pytest.param("WITH a AS (SELECT buyer_phone FROM sales), b(x) AS (SELECT * FROM a) SELECT x FROM b", {"x"},
|
||||
set(), id="列名表套 * 展开另一段 CTE"),
|
||||
pytest.param("SELECT r FROM (SELECT region AS r FROM sales UNION ALL "
|
||||
"SELECT * FROM (SELECT buyer_phone FROM sales))", {"r"}, set(), id="子查询里 UNION 后一支的 *"),
|
||||
pytest.param("SELECT * FROM (SELECT region FROM sales UNION ALL SELECT * FROM (SELECT buyer_phone FROM sales))",
|
||||
{"region"}, set(), id="外层 * 取子查询里 UNION 后一支的 *"),
|
||||
pytest.param("WITH t(k) AS (SELECT * FROM (SELECT buyer_name FROM sales)) SELECT k AS 名 FROM t", {"名"}, set(),
|
||||
id="列名表套 * 展开子查询"),
|
||||
pytest.param("SELECT region FROM sales UNION ALL SELECT s.* FROM (SELECT buyer_name FROM sales) s", {"region"},
|
||||
set(), id="UNION 后一支的 t.*"),
|
||||
pytest.param("SELECT region FROM sales UNION SELECT * FROM (SELECT buyer_phone FROM sales)", {"region"}, set(),
|
||||
id="UNION 去重"),
|
||||
# 只展开一个派生来源、每一列都数得清：按位置只遮取自遮罩列的那一列，金额照常
|
||||
pytest.param("SELECT '合计' AS 区域, SUM(amount) AS 金额 FROM sales UNION ALL SELECT * FROM "
|
||||
"(SELECT buyer_phone, amount FROM sales ORDER BY amount DESC LIMIT 2)", {"区域"}, {"金额"},
|
||||
id="合计行 UNION 后一支的 * 按位置"),
|
||||
pytest.param("SELECT region AS 区域, amount AS 金额 FROM sales UNION ALL SELECT * FROM "
|
||||
"(SELECT * FROM (SELECT buyer_name, amount FROM sales))", {"区域"}, {"金额"},
|
||||
id="套了两层的 * 按位置"),
|
||||
pytest.param("WITH t(r, x) AS (SELECT * FROM (SELECT buyer_phone, amount FROM sales)) SELECT r AS 号码, x AS 金额 "
|
||||
"FROM t", {"号码"}, {"金额"}, id="列名表套 * 按位置"),
|
||||
pytest.param("SELECT (SELECT * FROM (SELECT buyer_phone FROM sales) LIMIT 1) AS x", {"x"}, set(),
|
||||
id="标量子查询里的 * 展开子查询"),
|
||||
pytest.param("SELECT region, (SELECT * FROM (SELECT s2.buyer_name FROM sales s2 WHERE s2.id = s.id)) AS 名 "
|
||||
"FROM sales s", {"名"}, {"region"}, id="相关子查询里的 *"),
|
||||
pytest.param("SELECT (SELECT * FROM phones LIMIT 1) AS x", {"x"}, set(), id="标量子查询里的 * 展开表"),
|
||||
])
|
||||
def test_a_star_that_renames_a_masked_column_is_masked(db, sql, hidden, shown):
|
||||
"""* 把遮罩列的值带进了别的列名：这几列整列遮，原值不交出去。按位置认得出的，别的列照常。"""
|
||||
done = masked_run(db, sql)
|
||||
columns = {c.lower() for c in done.columns}
|
||||
assert hidden <= columns, done.columns
|
||||
assert not shown & columns, done.columns
|
||||
assert not _leaks(done), done.payload["rows"]
|
||||
assert done.secrets.strong, "遮掉的值交回调用方（后面的查询、回答里按值遮）"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql", [
|
||||
"SELECT * FROM (SELECT region, amount FROM sales) t",
|
||||
"WITH t AS (SELECT * FROM sales) SELECT region, SUM(amount) AS 合计 FROM t GROUP BY region",
|
||||
"SELECT region AS 区域 FROM sales UNION ALL SELECT * FROM (SELECT region FROM sales)",
|
||||
"WITH t(a, b) AS (SELECT * FROM (SELECT region, amount FROM sales)) SELECT a, SUM(b) AS 合计 FROM t GROUP BY a",
|
||||
"SELECT region, EXISTS(SELECT * FROM sales s2 WHERE s2.id = s.id AND s2.amount > 500) AS 大单 FROM sales s",
|
||||
"SELECT '合计' AS 区域, SUM(amount) AS 金额 FROM sales UNION ALL SELECT * FROM "
|
||||
"(SELECT region, SUM(amount) FROM sales GROUP BY region)",
|
||||
])
|
||||
def test_a_star_that_brings_no_masked_values_stays_visible(db, sql):
|
||||
"""* 只带出没碰遮罩列的值、EXISTS 里的 *（只说有没有）：一列都不遮，结果原样。"""
|
||||
done = masked_run(db, sql)
|
||||
assert done.columns == [] and MASKED_WORD not in json.dumps(done.payload, ensure_ascii=False)
|
||||
|
||||
|
||||
def test_a_plain_star_still_masks_by_name(db):
|
||||
"""不改名的 *（单独一条 SELECT）照旧只遮名字对上遮罩列的那几列。"""
|
||||
done = masked_run(db, "SELECT * FROM sales")
|
||||
assert done.columns == ["buyer_phone", "buyer_name"] and not _leaks(done)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql, columns, hidden", [
|
||||
pytest.param("SELECT x FROM (SELECT * FROM (SELECT buyer_phone FROM sales) s) AS q(x)", ["x"], {"x"},
|
||||
id="子查询列名表套 *"),
|
||||
pytest.param("SELECT l.x FROM sales s CROSS JOIN LATERAL (SELECT * FROM (SELECT s.buyer_phone) z) AS l(x)", ["x"],
|
||||
{"x"}, id="LATERAL 列名表套 *"),
|
||||
pytest.param("SELECT ARRAY(SELECT * FROM (SELECT buyer_phone FROM sales) s) AS arr", ["arr"], {"arr"},
|
||||
id="ARRAY 子查询里的 *"),
|
||||
])
|
||||
def test_a_star_renaming_in_postgres_is_masked(sql, columns, hidden):
|
||||
snapshot = {"source": "shop", "sql": sql, "columns": columns, "rows": []}
|
||||
assert hidden <= masking._derived_hidden(snapshot, set(MASKED), "postgres")[0]
|
||||
|
||||
|
||||
def test_a_metric_summed_over_a_star_renamed_masked_column_is_not_sent_to_blind_check():
|
||||
"""双盲复算按同一套认法挡下取自遮罩列的指标（derived_columns）：遮罩的金额经 CTE 列名表套 * 改名之后汇总，照样认得出，
|
||||
区域照常。"""
|
||||
sql = "WITH q(r, x) AS (SELECT * FROM (SELECT region, amount FROM sales)) SELECT r AS 区域, SUM(x) AS gmv FROM q " \
|
||||
"GROUP BY r"
|
||||
snapshot = {"source": "shop", "sql": sql, "columns": ["区域", "gmv"], "rows": []}
|
||||
assert masking.derived_columns(snapshot, {"amount"}, "sqlite") == {"gmv"}
|
||||
|
||||
|
||||
def test_a_recursive_cte_with_a_star_is_read_without_looping(db):
|
||||
"""递归的 CTE 里 * 展开它自己：照样认得出、不会绕不出来。"""
|
||||
sql = ("WITH RECURSIVE r(n, p) AS (SELECT 1, buyer_phone FROM sales WHERE id = 1 UNION ALL "
|
||||
"SELECT n + 1, p FROM r WHERE n < 2) SELECT * FROM (SELECT * FROM r) t")
|
||||
done = masked_run(db, sql)
|
||||
assert "p" in {c.lower() for c in done.columns} and not _leaks(done)
|
||||
loop = "WITH RECURSIVE r AS (SELECT buyer_phone AS p FROM sales UNION ALL SELECT * FROM r) SELECT p AS 号 FROM r"
|
||||
snapshot = {"source": "shop", "sql": loop, "columns": ["号"], "rows": []}
|
||||
assert "号" in masking._derived_hidden(snapshot, set(MASKED), "sqlite")[0]
|
||||
|
||||
|
||||
def test_a_star_with_a_replace_is_not_counted_by_position():
|
||||
"""* 带着 REPLACE / EXCLUDE 这类改动（DuckDB 一类的写法）：展开成的列不再是来源那几列的原样，不按位置数，整段都算。
|
||||
这里 REPLACE 把手机号换进了「amount」那一列。"""
|
||||
sql = ("SELECT region, amount FROM sales UNION ALL SELECT * REPLACE (buyer_phone AS amount) FROM "
|
||||
"(SELECT buyer_phone, amount FROM sales) t")
|
||||
snapshot = {"source": "shop", "sql": sql, "columns": ["region", "amount"], "rows": []}
|
||||
assert {"region", "amount"} <= masking._derived_hidden(snapshot, set(MASKED), "duckdb")[0]
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# * 展开数据源里的表：按查询当时的表结构数位置
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
#: 两个月的订单表列一样、次序一样（上传表格里的月度工作表常这样）；三月那张的号码和区域换了次序
|
||||
MONTHS = ("CREATE TABLE jan (id INTEGER, region TEXT, buyer_phone TEXT, amount INTEGER);"
|
||||
"CREATE TABLE feb (id INTEGER, region TEXT, buyer_phone TEXT, amount INTEGER);"
|
||||
"CREATE TABLE mar (id INTEGER, buyer_phone TEXT, region TEXT, amount INTEGER);"
|
||||
"INSERT INTO jan VALUES (1, '华东', '13900000001', 1200), (2, '华南', '13900000002', 300);"
|
||||
"INSERT INTO feb VALUES (3, '华东', '13900000003', 900);"
|
||||
"INSERT INTO mar VALUES (4, '13900000004', '华北', 400);")
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def months():
|
||||
conn = sqlite3.connect(":memory:")
|
||||
conn.executescript(MONTHS)
|
||||
yield conn
|
||||
conn.close()
|
||||
|
||||
|
||||
def schema_of(conn) -> dict:
|
||||
"""数据源的表结构（schema_cache 的样子）：每张表的列按表里的次序。"""
|
||||
names = [name for (name,) in conn.execute("SELECT name FROM sqlite_master WHERE type = 'table' ORDER BY name")]
|
||||
return {"tables": {name: {"columns": [{"name": row[1]} for row in conn.execute(f"PRAGMA table_info({name})")]}
|
||||
for name in names}}
|
||||
|
||||
|
||||
def run_months(conn, sql: str, schema: dict | None) -> masking.QueryMask:
|
||||
cursor = conn.execute(sql)
|
||||
payload = {"columns": [d[0] for d in cursor.description], "rows": [list(r) for r in cursor.fetchall()], "sql": sql}
|
||||
return masking.mask_query_result(payload, ["buyer_phone"], "sqlite", schema=schema)
|
||||
|
||||
|
||||
def test_a_union_of_tables_with_the_same_layout_is_masked_by_name_with_the_schema(months):
|
||||
"""UNION 的几支都是 SELECT * FROM 表、表的列一样：按表结构数得出号码落在哪一列，只遮那一列，汇总照常。没有表结构时
|
||||
说不清号码排在第几，* 展开落到的那一段整段遮（多遮）。"""
|
||||
schema = schema_of(months)
|
||||
total = "SELECT region, SUM(amount) AS 合计 FROM (SELECT * FROM jan UNION ALL SELECT * FROM feb) t GROUP BY region"
|
||||
done = run_months(months, total, schema)
|
||||
assert done.columns == [] and done.payload["rows"] == [["华东", 2100], ["华南", 300]]
|
||||
assert {"region", "合计"} <= set(run_months(months, total, None).columns)
|
||||
both = run_months(months, "SELECT * FROM jan UNION ALL SELECT * FROM feb", schema)
|
||||
assert both.columns == ["buyer_phone"] and not _leaks(both)
|
||||
|
||||
|
||||
def test_a_union_of_tables_in_another_order_masks_where_the_values_land(months):
|
||||
"""三月那张表号码排在第二列：UNION 之后它落进第一支的「region」，按表结构认得出，这一列遮；金额照常。"""
|
||||
schema = schema_of(months)
|
||||
both = run_months(months, "SELECT * FROM jan UNION ALL SELECT * FROM mar", schema)
|
||||
assert {"region", "buyer_phone"} <= set(both.columns) and "amount" not in both.columns and not _leaks(both)
|
||||
total = "SELECT region AS 区域, SUM(amount) AS 合计 FROM (SELECT * FROM jan UNION ALL SELECT * FROM mar) t " \
|
||||
"GROUP BY region"
|
||||
done = run_months(months, total, schema)
|
||||
assert done.columns == ["区域"] and not _leaks(done)
|
||||
|
||||
|
||||
def test_a_column_list_over_a_table_is_masked_by_position_with_the_schema(months):
|
||||
"""列名表套 * 展开数据源里的表：按表结构只遮号码换了名字的那一列。"""
|
||||
done = run_months(months, "WITH t(a, b, c, d) AS (SELECT * FROM jan) SELECT * FROM t", schema_of(months))
|
||||
assert done.columns == ["c"] and [row[3] for row in done.payload["rows"]] == [1200, 300]
|
||||
|
||||
|
||||
def test_a_schema_that_does_not_match_the_table_masks_the_whole_stretch(months):
|
||||
"""表结构里没有这张表，或者列数和实际展开的对不上（表结构过期了）：照旧整段遮。"""
|
||||
stale = schema_of(months)
|
||||
stale["tables"]["feb"]["columns"] = stale["tables"]["feb"]["columns"][:3]
|
||||
for schema in (stale, {"tables": {"jan": schema_of(months)["tables"]["jan"]}}):
|
||||
both = run_months(months, "SELECT * FROM jan UNION ALL SELECT * FROM feb", schema)
|
||||
assert set(both.columns) == {"id", "region", "buyer_phone", "amount"} and not _leaks(both)
|
||||
|
||||
|
||||
async def test_a_snapshot_is_read_with_its_frozen_schema(months):
|
||||
"""事后按快照认（双盲复算挡指标、审批质询收值）用快照记着的表结构快照（schema_artifact），和查询当时工具边界的认法
|
||||
一样：几张同样的月度表 UNION 起来汇总的金额不算取自遮罩列。"""
|
||||
from app.core import artifact_store
|
||||
|
||||
sql = "SELECT region, SUM(amount) AS 合计 FROM (SELECT * FROM jan UNION ALL SELECT * FROM feb) t GROUP BY region"
|
||||
schema = await artifact_store.put_json({"source": "shop", **schema_of(months)}, kind="schema_snapshot",
|
||||
run_id="r-star", node_id="n1")
|
||||
snapshot = {"source": "shop", "sql": sql, "columns": ["region", "合计"], "rows": [], "schema_artifact": schema}
|
||||
assert masking.derived_columns(snapshot, {"buyer_phone"}, "sqlite") == set()
|
||||
assert masking._derived_hidden(snapshot, {"buyer_phone"}, "sqlite")[0] == {"buyer_phone"}
|
||||
assert masking.derived_columns({**snapshot, "schema_artifact": "missing"}, {"buyer_phone"}, "sqlite") == {
|
||||
"region", "合计"}
|
||||
@@ -107,3 +107,55 @@ def test_a_masked_column_only_in_the_condition_of_an_aggregated_case_is_like_a_f
|
||||
])
|
||||
def test_a_case_returning_the_masked_value_or_not_aggregated_is_still_derived(sql):
|
||||
assert "total" in hidden_of(sql, ["total"])
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 函数参数里的 JSON 箭头（复核第 2 轮）
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind, sql, columns, shown", [
|
||||
pytest.param("sqlite", "SELECT j.value AS v FROM sales, json_each(profile->'$.tels') AS j", ["v"], ["v"],
|
||||
id="SQLite json_each 参数带箭头"),
|
||||
pytest.param("postgres", "SELECT e AS 号 FROM sales, json_array_elements_text(profile->'tels') AS e", ["号"],
|
||||
["号"], id="Postgres json_array_elements_text"),
|
||||
pytest.param("postgres", "SELECT e.value AS 号 FROM sales, jsonb_array_elements(profile->'tels') AS e", ["号"],
|
||||
["号"], id="Postgres jsonb_array_elements"),
|
||||
pytest.param("postgres", "SELECT t.tel FROM sales, jsonb_to_recordset(profile->'contacts') AS t(tel text)",
|
||||
["tel"], ["tel"], id="Postgres jsonb_to_recordset 列名表"),
|
||||
pytest.param("postgres", "SELECT k, v FROM sales, jsonb_each_text(profile->'contact') AS x(k, v)", ["k", "v"],
|
||||
["k", "v"], id="Postgres jsonb_each_text 列名表"),
|
||||
pytest.param("postgres", "SELECT jsonb_array_elements_text(profile->'tels') AS 号 FROM sales", ["号"], ["号"],
|
||||
id="Postgres 选出项里拆开"),
|
||||
pytest.param("postgres", "SELECT jsonb_array_elements_text(profile -> 'contact' -> 'tels') AS 号 FROM sales",
|
||||
["号"], ["号"], id="Postgres 连着两个箭头"),
|
||||
pytest.param("mysql", "SELECT JSON_UNQUOTE(profile->'$.tel') AS 电话 FROM sales", ["电话"], ["电话"],
|
||||
id="MySQL JSON_UNQUOTE"),
|
||||
pytest.param("mysql", "SELECT CONCAT(profile->'$.tel', '') AS 电话 FROM sales", ["电话"], ["电话"],
|
||||
id="MySQL CONCAT"),
|
||||
])
|
||||
def test_a_json_arrow_inside_function_arguments_is_read_as_the_masked_column(kind, sql, columns, shown):
|
||||
"""函数参数里没写表名的 col->'path' 被解析器当成以 col 为参数的 lambda（不是一列）：以前认不出遮罩列，拆开、取出来的
|
||||
值原样交给模型、落库。遮罩列是 JSON 时这几种是最自然的写法。写了表名、用 ->> 的本来就认得出。"""
|
||||
assert {c.lower() for c in shown} <= hidden_of(sql, columns, kind)
|
||||
|
||||
|
||||
def test_a_json_arrow_over_an_unmasked_column_stays_visible():
|
||||
"""箭头取的不是遮罩列：一列都不遮。真的 lambda 的参数照列名收，只有和遮罩列同名时才多遮。"""
|
||||
assert hidden_of("SELECT j.value AS 标签 FROM sales, json_each(tags->'$.list') AS j", ["标签"]) == MASKED
|
||||
assert hidden_of("SELECT transform(scores, x -> x + 1) AS 分 FROM sales", ["分"], "postgres") == MASKED
|
||||
|
||||
|
||||
def test_a_json_arrow_inside_json_each_is_masked_when_really_run():
|
||||
"""SQLite 里真跑一遍：json_each(profile->'$.tels') 拆出来的号码在交给模型之前就遮掉。"""
|
||||
import sqlite3
|
||||
|
||||
db = sqlite3.connect(":memory:")
|
||||
db.executescript("CREATE TABLE sales (id INTEGER, profile TEXT);"
|
||||
"INSERT INTO sales VALUES (1, '{\"tels\": [\"13900000001\", \"13900000002\"]}');")
|
||||
sql = "SELECT j.value AS v FROM sales, json_each(profile->'$.tels') AS j"
|
||||
cursor = db.execute(sql)
|
||||
payload = {"columns": [d[0] for d in cursor.description], "rows": [list(r) for r in cursor.fetchall()], "sql": sql}
|
||||
db.close()
|
||||
done = masking.mask_query_result(payload, ["profile"], "sqlite")
|
||||
assert done.columns == ["v"] and done.payload["rows"] == [["已遮罩"], ["已遮罩"]]
|
||||
|
||||
@@ -58,10 +58,10 @@ async def clean_spend():
|
||||
yield
|
||||
|
||||
|
||||
async def _source(tmp_path, options: dict | None) -> tuple[str, int]:
|
||||
async def _source(tmp_path, options: dict | None, script: str = SALES) -> tuple[str, int]:
|
||||
path = tmp_path / f"{uuid.uuid4().hex[:8]}.db"
|
||||
db = sqlite3.connect(path)
|
||||
db.executescript(SALES)
|
||||
db.executescript(script)
|
||||
db.commit()
|
||||
db.close()
|
||||
name = f"shop_{uuid.uuid4().hex[:8]}"
|
||||
@@ -385,3 +385,108 @@ def test_labels_written_in_this_query_are_not_masked_by_values_from_earlier_ones
|
||||
payload = {"sql": "SELECT CASE WHEN amount > 1000 THEN '大单' ELSE '小单' END AS 档, COUNT(*) AS 笔数 FROM sales "
|
||||
"GROUP BY 1", "columns": ["档", "笔数"], "rows": [["大单", 1], ["小单", 3]]}
|
||||
assert masking.mask_query_result(payload, ["amount"], "sqlite", known).payload["rows"] == [["大单", 1], ["小单", 3]]
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# * 把遮罩列带进别的列名、函数参数里的 JSON 箭头（复核第 2 轮）
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql, masked", [
|
||||
pytest.param("SELECT region, amount FROM sales WHERE id < 0 UNION ALL "
|
||||
"SELECT * FROM (SELECT buyer_phone, amount FROM sales ORDER BY id)", ["region"],
|
||||
id="UNION 后一支的 *"),
|
||||
pytest.param("WITH t(k, v) AS (SELECT * FROM (SELECT buyer_name, amount FROM sales ORDER BY id)) "
|
||||
"SELECT k AS 姓名, v AS amount FROM t", ["姓名"], id="CTE 列名表套 *"),
|
||||
pytest.param("SELECT (SELECT * FROM (SELECT s2.buyer_phone FROM sales s2 WHERE s2.id = s.id)) AS 号码, s.amount "
|
||||
"FROM sales s ORDER BY s.id", ["号码"], id="标量子查询里的 *"),
|
||||
])
|
||||
async def test_a_star_renaming_a_masked_column_is_masked_before_the_model_sees_it(monkeypatch, masked_shop, sql,
|
||||
masked):
|
||||
"""* 把遮罩列的值带进了别的列名（UNION 后一支、写了列名表的 CTE、标量子查询）：交给模型的、落库的、通用接口给的都是
|
||||
「已遮罩」，金额照常。以前 * 本身不算取自哪一列，原值原样交给模型、落进快照。"""
|
||||
name = masked_shop
|
||||
script = Script(monkeypatch, [q(name, sql, "c1")],
|
||||
{"value": {"alias": "Q1", "row": 0, "column": "amount"}, "rationale": "见 [[see:Q1]]。"})
|
||||
out, _ = await ask(name)
|
||||
assert out.error is None and out.result["value"] == 1200
|
||||
seen = script.tool_results()["c1"]
|
||||
assert MASKED in seen and _clean(seen), seen
|
||||
snapshot = artifact_store.load(out.queries[0]["artifact"])
|
||||
assert snapshot["masked_at_query"] == masked and set(_cells(snapshot, masked[0])) == {MASKED}
|
||||
assert _cells(snapshot, "amount") == [1200, 300, 900, 400]
|
||||
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
|
||||
got = await client.get(f"/api/artifacts/{out.queries[0]['artifact']}")
|
||||
assert got.status_code == 200 and got.json()["content"] == snapshot and _clean(got.text)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql, masked", [
|
||||
pytest.param("SELECT c.id AS 编号, j.value AS 值 FROM customers c, json_each(profile->'$') AS j "
|
||||
"ORDER BY c.id, j.key", ["值"], id="json_each 参数带箭头"),
|
||||
pytest.param("SELECT id AS 编号, upper(profile->'$.name') AS 姓名 FROM customers ORDER BY id", ["姓名"],
|
||||
id="函数参数里的箭头"),
|
||||
])
|
||||
async def test_a_json_arrow_over_the_masked_column_is_masked_before_the_model_sees_it(monkeypatch, json_shop, sql,
|
||||
masked):
|
||||
"""遮罩的是 JSON 列：函数参数里没写表名的 profile->'…' 被解析器当成 lambda，以前认不出，拆开、取出来的号码和姓名原样
|
||||
交给模型、落库。"""
|
||||
name = json_shop
|
||||
script = Script(monkeypatch, [q(name, sql, "c1")],
|
||||
{"value": {"alias": "Q1", "row": 0, "column": "编号"}, "rationale": "x"})
|
||||
out, _ = await ask(name)
|
||||
assert out.result["value"] == 1 and _clean(script.tool_results()["c1"])
|
||||
snapshot = artifact_store.load(out.queries[0]["artifact"])
|
||||
assert snapshot["masked_at_query"] == masked and set(_cells(snapshot, masked[0])) == {MASKED}
|
||||
assert _clean(json.dumps(snapshot, ensure_ascii=False))
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 遮罩值的上限：agent 自己查到的值先收（复核第 2 轮）
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_values_the_agent_finds_survive_caller_values_that_fill_the_cap(monkeypatch, masked_shop):
|
||||
"""调用方交进来的值（报告快照里的）收满了上限 SECRETS_MAX：agent 自己查到的号码照样收，后面查到的备注按它遮，交回
|
||||
调用方的 masked_values 里也认得出。以前先并调用方的值、满了就停，agent 查到的一个也并不进去：后面的备注原样交给
|
||||
模型、原样落库，回答时按值遮也认不出。"""
|
||||
monkeypatch.setattr(masking, "SECRETS_MAX", 40)
|
||||
name = masked_shop
|
||||
given = masking._Secrets()
|
||||
given.add({"columns": ["buyer_phone"], "rows": [[str(13800000000 + i)] for i in range(40)]}, {"buyer_phone"})
|
||||
assert len(given.weak) == masking.SECRETS_MAX and not given.hit("13900000001")
|
||||
steps = [q(name, "SELECT buyer_phone, amount AS total FROM sales ORDER BY id", "c1"),
|
||||
q(name, "SELECT id, note FROM sales ORDER BY id", "c2")]
|
||||
script = Script(monkeypatch, steps, {"value": {"alias": "Q1", "row": 0, "column": "total"}, "rationale": "x"})
|
||||
out, _ = await ask(name, masked_values=given)
|
||||
assert out.result["value"] == 1200
|
||||
assert _clean(script.tool_results()["c2"]), script.tool_results()["c2"]
|
||||
notes = artifact_store.load(out.queries[1]["artifact"])
|
||||
assert _cells(notes, "note") == [f"客户{MASKED}来电改址", "例行回访", f"回访：联系 {MASKED}", ""]
|
||||
assert out.masked_values.hit("13900000001") and out.masked_values.hit("13900000003")
|
||||
assert len(out.masked_values.weak) <= masking.SECRETS_MAX, "合起来照样不超过上限：挤掉的是调用方的值"
|
||||
|
||||
|
||||
#: 两个月的订单表，列一样、次序一样（上传表格里的月度工作表常这样）
|
||||
MONTHS = ("CREATE TABLE jan (id INTEGER PRIMARY KEY, region TEXT, buyer_phone TEXT, amount INTEGER);"
|
||||
"CREATE TABLE feb (id INTEGER PRIMARY KEY, region TEXT, buyer_phone TEXT, amount INTEGER);"
|
||||
"INSERT INTO jan VALUES (1, '华东', '13900000001', 1200), (2, '华南', '13900000002', 300);"
|
||||
"INSERT INTO feb VALUES (3, '华东', '13900000003', 900);")
|
||||
|
||||
|
||||
async def test_a_union_of_monthly_tables_is_counted_with_the_sources_schema(monkeypatch, tmp_path):
|
||||
"""几张列一样的月度表用 SELECT * 连起来再汇总：按数据源的表结构数得出号码落在第几列，汇总的金额照常交给模型、照常
|
||||
取值；选出全部列的，只遮号码那一列。* 展开数据源里的表又不按表结构数位置的话，展开落到的那一段整段都遮，金额也拿不到。"""
|
||||
name, source_id = await _source(tmp_path, {"mask_columns": "buyer_phone"}, script=MONTHS)
|
||||
try:
|
||||
steps = [q(name, "SELECT region, SUM(amount) AS total FROM (SELECT * FROM jan UNION ALL SELECT * FROM feb) t "
|
||||
"GROUP BY region ORDER BY total DESC", "c1"),
|
||||
q(name, "SELECT * FROM jan UNION ALL SELECT * FROM feb ORDER BY id", "c2")]
|
||||
script = Script(monkeypatch, steps, {"value": {"alias": "Q1", "row": 0, "column": "total"}, "rationale": "x"})
|
||||
out, _ = await ask(name)
|
||||
assert out.error is None and out.result["value"] == 2100
|
||||
assert all(_clean(text) for text in script.tool_results().values())
|
||||
totals, rows = (artifact_store.load(query["artifact"]) for query in out.queries)
|
||||
assert totals["masked_at_query"] == [] and totals["rows"] == [["华东", 2100], ["华南", 300]]
|
||||
assert rows["masked_at_query"] == ["buyer_phone"] and _cells(rows, "amount") == [1200, 300, 900]
|
||||
finally:
|
||||
await engines.invalidate(source_id)
|
||||
|
||||
Reference in New Issue
Block a user