mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 01:11:59 +00:00
审批质询遮罩收口（第 3 轮复核）：标签列按作用域认取自哪个来源，UNION 另一支、子查询里的遮罩列流进标签列照收；取模、整除、整数相除按算出来的值认；CAST 的类型参数不当边界换
复核留下的一条 major（第 2 轮只修了一部分）和两条次要问题： - 分档标签列：以前按名字认，从「所有名字都只装着标签」起往下剔。遮罩列原样选出（SELECT vip_level）时名字就是它自己， 永远剔不掉；结果的 UNION 按位置那一遍只往里加标签，另一支取的是遮罩列也不剔；派生来源里的 UNION 根本不按位置看。 于是标签在前、遮罩列在后（不起别名、别名不同、外层再套子查询、CTE 里的标签再 UNION），遮罩列经按原名选出的子查询、 以标签的别名传出，或者两个子查询 JOIN 起来同名，等级 1 撞上 THEN 1 都不收，WHERE vip_level = 1 在查询清单、事件、 质询工件里照原样封存（对基线是回退）。现在 _label_columns 按作用域认每一项取自哪个来源：结果的这一列在 UNION 的 每一支里都要是结果全为写死常量的 CASE / IIF，或者取自子查询、CTE、LATERAL 里只装着标签的一列（派生来源里的 UNION 也按位置看每一支，写了列名表的按列名表）；取自数据源里的表的一律不算，哪怕和标签同名。没写表名的列，派生来源出 这一列的就是它（SQL 执行过，别的来源就没有同名的列，不然报列名有歧义），按这一列 USING、NATURAL 合成的除外；递归的 CTE、UNION 里落在某一支 * 里的列认不准，照收（多遮） - 取模、整除：_arithmetic、_worked_out 补上 %、MOD、DIV、//（往零取整，同 SQL）。顺带 / 照数据源的算法算：SQLite、 Postgres 里两个写成整数的数相除只留整数部分（sqlglot 记在 typed 上），以前一律按小数算，age >= 25 AND age < 53 / 2 在 SQLite 里夹住了 25 却认不出；写成小数的（53.0 / 2）、MySQL、DuckDB 照小数算 - CAST 的类型参数（第 2 轮引入，只多遮）：_halves 只收往上走得到这一边的字面量（_lifts），CAST(30000 AS DECIMAL(10, 2)) 的 10、2 不再当成算出边界的数换掉，SQL 检查结果的片段里别处的 10、2 也不跟着换；_wraps 里 CAST 只算被转换的那个值， Oracle 的 CAST … DEFAULT 0 ON CONVERSION ERROR 里的 0 照原样 - 模块说明的已知局限照实更新：标签列认得出、认不准的几种写法；认得出的算式补上取模、整除和整数相除，认不出的补上 位运算、POWER、CAST 改了值的 测试：单元补标签在前、遮罩列在后的八种写法（不起别名、别名不同、子查询里的 UNION、写了列名表的 CTE、CTE 里的标签再 UNION、经按原名选出的子查询传出、两个子查询 JOIN 同名、RIGHT JOIN USING 同名列），按作用域仍认得出的五种标签（外层连上 数据源里的表不写表名、按别的列 USING、两层子查询写了表名、UNION 两支、派生来源的 UNION 经 * 带出）；取模、整除、整数 相除的六种写法和照小数算的三种；CAST 的类型参数和 DEFAULT。端到端补 UNION 标签在前的两种写法、取模和整数相除算出来的 边界。每处实现都逐个去掉验证过测试变红（十一个变异全红）
This commit is contained in:
ThearchyHelios
+248
-74
@@ -43,9 +43,11 @@
|
||||
汇总里套着结果全是写死常量的 CASE（SUM(CASE WHEN amount <= 20000 THEN 1 ELSE 0 END)）只是按条件数个数，同 COUNT
|
||||
- 按遮罩列分档的标签（CASE、IIF 的结果全是写死的常量）：标签文字是 agent 写在 SQL 里的，不当遮罩值，SQL、回答正文里
|
||||
照原样；标签列当分组键（每一档一行）照常显示，不分组的逐行标签列整列遮（格子里看不出某一笔落在哪一档），可 agent
|
||||
在正文里写出「这一笔是大单」遮不住，按 id 分组再按档数个数也看得出（同按阈值查出某几行）。只有只装着标签的列不收
|
||||
（选出项本身就是这样的 CASE / IIF，或者经别名、子查询原样传下去的），同一个名字在哪一处装着别的就照收；标签恰好和
|
||||
别的列里的遮罩值相同时，正文里的档名、SQL 里写成字符串的标签也换（多遮）
|
||||
在正文里写出「这一笔是大单」遮不住，按 id 分组再按档数个数也看得出（同按阈值查出某几行）。只有只装着标签的列不收：
|
||||
结果的这一列在 UNION 的每一支里都是这样的 CASE / IIF，或者按作用域认得出取自子查询、CTE 里只装着标签的一列；有一支
|
||||
取自数据源里的表（遮罩列原样选出、经子查询按原名传出，哪怕起了标签的名字）就照收。认不准的也照收（多遮），比如
|
||||
按这一列 USING、NATURAL 合成的同名列，递归的 CTE，UNION 里落在某一支 * 里的列，写死的字符串、和列一起算的值。
|
||||
标签恰好和别的列里的遮罩值相同时，正文里的档名、SQL 里写成字符串的标签也换（多遮）
|
||||
- 「夹住某一个值」（_pinned）：范围比较在整条 SQL 里配对，只有 CASE 的不同分支、SELECT 里不筛行的不同项（选出列、
|
||||
GROUP BY、ORDER BY）、UNION 的不同支各算各的，两头误差都有：把找值拆进这几处（CASE WHEN age > 25 THEN … ELSE CASE
|
||||
WHEN age >= 25 …，两个选出列 SUM(age >= 25)、SUM(age > 25) 再自己相减）认不出，边界照原样留着，正文里复述也不遮；
|
||||
@@ -58,9 +60,11 @@
|
||||
隔开的一律按夹住算，分档很多时会成批多遮
|
||||
- 找某一个值的写法按语法树认常见的几种：等值、IN、简单 CASE、把它夹住的范围，被比较的一边套了 CAST、函数、汇总
|
||||
（CAST(age AS INTEGER) = 25、HAVING MAX(age) BETWEEN 25 AND 25），值换个写法（25.0、025、2.5e1）、写成几个数的
|
||||
算式（13900000000 + 1、BETWEEN 20 + 5 AND 30 - 5，算式里的数都换）。更绕的写法认不出：值藏进和列一起算的算式
|
||||
（age - 25 = 0、buyer_phone - 1 = 13900000000）、拼接（'1390' || '0000001'）、函数拼出来的值（REPLACE、SUBSTR），
|
||||
短的值照原样封存，长号码只在整个写出来时按值换
|
||||
算式（加减乘除、取模、整除：13900000000 + 1、BETWEEN 20 + 5 AND 30 - 5、125 % 100、50 DIV 2，算式里的数都换；
|
||||
两个整数相除照数据源的算法，SQLite、Postgres 里 53 / 2 是 26，MySQL 里是 26.5）。更绕的写法认不出：值藏进和列
|
||||
一起算的算式（age - 25 = 0、buyer_phone - 1 = 13900000000）、位运算（24 | 1）、拼接（'1390' || '0000001'）、
|
||||
函数拼出来的值（REPLACE、SUBSTR、POWER）、CAST 改了值的（CAST(25.6 AS INTEGER) 按 25.6 认，CAST(53 AS REAL) / 2
|
||||
按整数相除算），短的值照原样封存，长号码只在整个写出来时按值换
|
||||
- 质询工件里 SQL 检查结果的片段是检查器按语法树重新写出来的，按整条 SQL 里整个换掉了的字面量照值换：同一个值在
|
||||
片段里别处出现也换（多遮，算式里换掉的 0、1 这类也是）；截断处半截的值去掉（停在字符串、带引号的名字、注释里的，
|
||||
从开头的引号、注释起整个去掉）。说明里码值表列出的码值和别的字一样按值换，恰好是遮罩值的码值也换（多遮）
|
||||
@@ -751,12 +755,12 @@ def _label_columns(sql: Any, kind: str | None, columns: list[str]) -> dict[str,
|
||||
"""结果里只装着分档标签的列：列名（小写）→ 标签在格子里的写法。标签（CASE、IIF 写死的结果，分档的名字、档位编号）
|
||||
是 agent 自己写在 SQL 里的字，不是数据源里的值：这几列派生自遮罩列也不收进遮罩值（_Secrets.add），不然 SQL 里的
|
||||
标签、回答正文里的档名都会被换成「已遮罩」。
|
||||
只装着标签：这一列的选出项本身是结果全为写死常量的 CASE / IIF，或者经别名、子查询、CTE 原样传下去的这样一列
|
||||
（_label_texts）；同一个名字在哪一处装着别的（vip_level AS 等级、UNION 另一支取的是遮罩列），就不算——改了名字的
|
||||
遮罩列的 1、2、3 撞上同一条 SQL 里按条件数个数的 THEN 1 ELSE 0，照样是遮罩值。没起别名的标签列按位置认（每一支的
|
||||
那一项都是标签）。解析不了的一个都不算（标签照旧当遮罩值，宁可多遮）。"""
|
||||
from sqlglot import exp
|
||||
|
||||
只装着标签：结果的这一列在 UNION 这类的每一支里，都是结果全为写死常量的 CASE / IIF，或者取自派生来源（子查询、
|
||||
CTE、LATERAL）里只装着标签的一列（_item_label）。一项取自哪个来源按作用域认，不按名字猜：按名字认时，遮罩列原样
|
||||
选出（SELECT vip_level）的名字就是它自己，永远算只装着标签，UNION 另一支、子查询里的遮罩列流进标签列也认不出，
|
||||
等级 1 撞上 THEN 1 就不收。取自数据源里的表的（哪怕和标签同名）、有一支装着别的、认不准的（说不清取自哪个来源、
|
||||
递归的 CTE、几支的 * 展开在哪说不清）都不算，照收（宁可多遮）；同名的几列有一列不算就都不算。解析不了的一个都
|
||||
不算。"""
|
||||
from app.data.sql_semantics import _parse, _read
|
||||
|
||||
try:
|
||||
@@ -766,47 +770,53 @@ def _label_columns(sql: Any, kind: str | None, columns: list[str]) -> dict[str,
|
||||
if root is None:
|
||||
return {}
|
||||
try:
|
||||
items = [item for select in root.find_all(exp.Select) for item in select.expressions]
|
||||
# 先当每个名字都只装着标签，再去掉装着别的，直到不再变：名字一层层往下传，先认哪一层说不准
|
||||
known: dict[str, set[str]] = {name: set() for item in items if (name := _output_name(item))}
|
||||
for _ in range(len(items) + 2):
|
||||
found: dict[str, set[str]] = {}
|
||||
spoiled: set[str] = set()
|
||||
for item in items:
|
||||
if not (name := _output_name(item)):
|
||||
continue
|
||||
texts = _label_texts(item, known)
|
||||
if texts is None:
|
||||
spoiled.add(name)
|
||||
else:
|
||||
found.setdefault(name, set()).update(texts)
|
||||
found = {name: texts for name, texts in found.items() if name not in spoiled}
|
||||
if found == known:
|
||||
break
|
||||
known = found
|
||||
out = {c.lower(): known[c.lower()] for c in columns if c.lower() in known}
|
||||
width = len(columns)
|
||||
for top, role, _listed in _blocks(root):
|
||||
if role != "result":
|
||||
continue
|
||||
branches = _branches(top)
|
||||
seen: dict[int, list[set[str] | None]] = {}
|
||||
for branch in branches:
|
||||
for item, place in zip(branch, _places(branch, width)):
|
||||
if place is not None and place < width:
|
||||
seen.setdefault(place, []).append(_label_texts(item, known))
|
||||
for place, texts in seen.items():
|
||||
if len(texts) == len(branches) and all(t is not None for t in texts):
|
||||
out.setdefault(columns[place].lower(), set()).update(*texts)
|
||||
return out
|
||||
result = next((top for top, role, _listed in _blocks(root) if role == "result"), None)
|
||||
if result is None:
|
||||
return {}
|
||||
memo: _LabelMemo = {}
|
||||
found: dict[str, list[set[str] | None]] = {}
|
||||
for place, column in enumerate(columns):
|
||||
name = column.lower()
|
||||
found.setdefault(name, []).append(_result_label(result, place, name, len(columns), memo))
|
||||
return {name: texts for name, parts in found.items() if (texts := _all_labels(parts)) is not None}
|
||||
except Exception: # noqa: BLE001 - 分析本身的缺陷不能让回答出错：一个都不算（宁可多遮）
|
||||
logger.warning("质询的 SQL 标签列分析失败", exc_info=True)
|
||||
return {}
|
||||
|
||||
|
||||
def _label_texts(item: Any, known: dict[str, set[str]]) -> set[str] | None:
|
||||
"""一个选出项只装着标签时，标签在格子里的写法：结果全为写死常量的 CASE / IIF 取它的各个结果（数按两种写法：25、
|
||||
25.0），单独一列、名字在 known 里（经别名、子查询传下来的标签列）取那一列的；不是为 None。"""
|
||||
#: 标签列分析里「这个来源不出这一列」（_block_label、_source_label）：和 None（出这一列，但装着别的、认不准）分开
|
||||
_ABSENT: Any = object()
|
||||
#: 标签列分析记下的结果：(id(那块查询), 列名) → 标签的写法、None 或 _ABSENT
|
||||
_LabelMemo = dict[tuple[int, str], Any]
|
||||
|
||||
|
||||
def _result_label(top: Any, place: int, name: str, width: int, memo: _LabelMemo) -> set[str] | None:
|
||||
"""结果（width 列）的第 place 列（列名 name）只装着标签时标签的写法，不然为 None：UNION 这类的每一支，那一项都要
|
||||
只装着标签。落在 * 里的（认不出是哪一项）按列名经 * 往下找，只有一支时才找（几支的 * 各展开成什么说不清）。"""
|
||||
from sqlglot import exp
|
||||
|
||||
leaves = _leaves(top)
|
||||
if not all(isinstance(leaf, exp.Select) for leaf in leaves):
|
||||
return None
|
||||
out: set[str] = set()
|
||||
for select in leaves:
|
||||
items = list(select.expressions)
|
||||
item = next((i for i, p in zip(items, _places(items, width)) if p == place), None)
|
||||
if item is not None:
|
||||
texts = _item_label(item, select, memo)
|
||||
elif len(leaves) == 1:
|
||||
texts = _star_label(select, name, memo)
|
||||
else:
|
||||
return None
|
||||
if texts is None or texts is _ABSENT:
|
||||
return None
|
||||
out |= texts
|
||||
return out
|
||||
|
||||
|
||||
def _item_label(item: Any, select: Any, memo: _LabelMemo) -> set[str] | None:
|
||||
"""select 的一个选出项只装着标签时，标签在格子里的写法：结果全为写死常量的 CASE / IIF 取它的各个结果（数按两种
|
||||
写法：25、25.0）；单独一列取它在 FROM 里取自的那一列的（_source_label）；别的为 None。"""
|
||||
from sqlglot import exp
|
||||
|
||||
node = item
|
||||
@@ -824,22 +834,165 @@ def _label_texts(item: Any, known: dict[str, set[str]]) -> set[str] | None:
|
||||
elif value is not None:
|
||||
out.add(_plain_text(value))
|
||||
return out
|
||||
if isinstance(node, exp.Column) and not isinstance(node.this, exp.Star) and node.name.lower() in known:
|
||||
return known[node.name.lower()]
|
||||
if isinstance(node, exp.Column) and not isinstance(node.this, exp.Star) and node.name:
|
||||
texts = _source_label(select, node.table.lower() or None, node.name.lower(), memo)
|
||||
return None if texts is _ABSENT else texts
|
||||
return None
|
||||
|
||||
|
||||
def _star_label(select: Any, name: str, memo: _LabelMemo) -> Any:
|
||||
"""select 经 *（t.* 只看 t）带出的 name 这一列只装着标签时标签的写法；没有 *、* 带不出这一列为 _ABSENT，带出来的
|
||||
装着别的、认不准为 None。"""
|
||||
from sqlglot import exp
|
||||
|
||||
out: set[str] = set()
|
||||
seen = False
|
||||
for item in select.expressions:
|
||||
if not _is_star(item):
|
||||
continue
|
||||
qualifier = (item.table.lower() or None) if isinstance(item, exp.Column) else None
|
||||
texts = _source_label(select, qualifier, name, memo)
|
||||
if texts is None:
|
||||
return None
|
||||
if texts is not _ABSENT:
|
||||
out |= texts
|
||||
seen = True
|
||||
return out if seen else _ABSENT
|
||||
|
||||
|
||||
def _source_label(select: Any, qualifier: str | None, name: str, memo: _LabelMemo) -> Any:
|
||||
"""select 的 FROM、JOIN 里读到的那一列（qualifier 是写的表名或别名，没写为 None）只装着标签时标签的写法；没有哪个
|
||||
来源出这一列为 _ABSENT，出这一列但装着别的、认不准为 None。
|
||||
- 写了表名的只看那个来源（对不上的是外层的表，认不准）；
|
||||
- 没写的看所有来源：派生来源出这一列的就是它——SQL 执行过，别的来源就没有同名的列，不然会报列名有歧义。按这一列
|
||||
连起来的（USING 里写着它、NATURAL）例外：几个来源的这一列合成一列，外连接时取的可能是另一边的值；
|
||||
- 数据源里的表一律不算（遮罩列本身、别的列原样选出来，哪怕和标签同名），VALUES、函数这类来源也认不准。"""
|
||||
sources = _from_sources(select)
|
||||
if qualifier is not None:
|
||||
sources = [(alias, source) for alias, source in sources if alias == qualifier]
|
||||
if len(sources) != 1:
|
||||
return None
|
||||
found: list[set[str]] = []
|
||||
opaque = False
|
||||
for _alias, source in sources:
|
||||
derived = _derived_block(source)
|
||||
if derived is None:
|
||||
opaque = True
|
||||
continue
|
||||
texts = _block_label(derived[0], derived[1], name, memo)
|
||||
if texts is None:
|
||||
return None
|
||||
if texts is not _ABSENT:
|
||||
found.append(texts)
|
||||
if not found:
|
||||
return None if opaque else _ABSENT
|
||||
if opaque and _merging(select, name):
|
||||
return None
|
||||
return set().union(*found)
|
||||
|
||||
|
||||
def _block_label(block: Any, listed: list[str] | None, name: str, memo: _LabelMemo) -> Any:
|
||||
"""派生来源的那块查询（listed 是写的列名表）出的 name 这一列只装着标签时标签的写法；不出这一列为 _ABSENT，出这一列
|
||||
但装着别的、认不准为 None。一条 SELECT 按选出项的名字找，没有的经 * 往下找；UNION 这类按位置看每一支，列名取写的
|
||||
列名表，没写的取第一支的。"""
|
||||
key = (id(block), name)
|
||||
if key not in memo:
|
||||
# 先记成 None：递归的 CTE 绕回自己时按认不准算
|
||||
memo[key] = None
|
||||
memo[key] = _block_label_of(block, listed, name, memo)
|
||||
return memo[key]
|
||||
|
||||
|
||||
def _block_label_of(block: Any, listed: list[str] | None, name: str, memo: _LabelMemo) -> Any:
|
||||
"""_block_label 真去找的那一步（不看记下的结果）。"""
|
||||
from sqlglot import exp
|
||||
|
||||
leaves = _leaves(block)
|
||||
if not all(isinstance(leaf, exp.Select) for leaf in leaves):
|
||||
return None
|
||||
if listed is None and len(leaves) == 1:
|
||||
named = [item for item in leaves[0].expressions if _output_name(item) == name]
|
||||
if not named:
|
||||
return _star_label(leaves[0], name, memo)
|
||||
return _all_labels([_item_label(item, leaves[0], memo) for item in named])
|
||||
width: int | None = None
|
||||
if listed is not None:
|
||||
if name not in listed:
|
||||
return _ABSENT
|
||||
place, width = (listed.index(name) if listed.count(name) == 1 else None), len(listed)
|
||||
else:
|
||||
first = list(leaves[0].expressions)
|
||||
places = [p for item, p in zip(first, _places(first, None)) if _output_name(item) == name]
|
||||
if not places:
|
||||
# 第一支没明写这一列：有 * 时说不清它在不在、在哪个位置
|
||||
return None if any(_is_star(item) for item in first) else _ABSENT
|
||||
place = places[0] if len(places) == 1 else None
|
||||
if place is None:
|
||||
return None
|
||||
texts: list[set[str] | None] = []
|
||||
for select in leaves:
|
||||
items = list(select.expressions)
|
||||
item = next((i for i, p in zip(items, _places(items, width)) if p == place), None)
|
||||
texts.append(_item_label(item, select, memo) if item is not None else None)
|
||||
return _all_labels(texts)
|
||||
|
||||
|
||||
def _all_labels(texts: list[set[str] | None]) -> set[str] | None:
|
||||
"""几项都只装着标签时，标签的写法合起来；有一项不是为 None。"""
|
||||
return set().union(*texts) if all(t is not None for t in texts) else None
|
||||
|
||||
|
||||
def _from_sources(select: Any) -> list[tuple[str, Any]]:
|
||||
"""一条 SELECT 的 FROM、JOIN 里的来源：(名字, 来源节点)，名字是别名，没起别名的是表名（小写）。"""
|
||||
from_ = select.args.get("from_")
|
||||
nodes = [from_.this] if from_ is not None and from_.this is not None else []
|
||||
nodes += [join.this for join in select.args.get("joins") or [] if join.this is not None]
|
||||
return [((node.alias_or_name or "").lower(), node) for node in nodes]
|
||||
|
||||
|
||||
def _derived_block(source: Any) -> tuple[Any, list[str] | None] | None:
|
||||
"""FROM、JOIN 里的一个来源 → (它的那块查询, 写的列名表)：FROM 里的子查询、LATERAL、指向 WITH 里一段的表名。数据源里
|
||||
的表、VALUES、函数这类为 None（出哪些列、装着什么，SQL 里看不出）。"""
|
||||
from sqlglot import exp
|
||||
|
||||
from app.data.sql_semantics import _visible_ctes
|
||||
|
||||
node, alias = source, None
|
||||
if isinstance(node, exp.Lateral):
|
||||
node, alias = node.this, node.args.get("alias")
|
||||
if isinstance(node, exp.Table):
|
||||
if node.args.get("db") or node.args.get("catalog") or not isinstance(node.this, exp.Identifier):
|
||||
return None
|
||||
cte = _visible_ctes(node).get(node.name.lower())
|
||||
if cte is None:
|
||||
return None
|
||||
node, alias = cte.this, cte.args.get("alias")
|
||||
while isinstance(node, exp.Subquery):
|
||||
alias = alias or node.args.get("alias")
|
||||
node = node.this
|
||||
if not isinstance(node, (exp.Select, exp.SetOperation)):
|
||||
return None
|
||||
return node, ([c.name.lower() for c in alias.columns] if alias is not None and alias.columns else None)
|
||||
|
||||
|
||||
def _merging(select: Any, name: str) -> bool:
|
||||
"""FROM 里按同名的 name 这一列连起来（USING 里写着它，或者 NATURAL）：几个来源的这一列合成一列。"""
|
||||
return any(str(join.args.get("method") or "").upper() == "NATURAL"
|
||||
or any(str(getattr(u, "name", "")).lower() == name for u in join.args.get("using") or [])
|
||||
for join in select.args.get("joins") or [])
|
||||
|
||||
|
||||
def _is_star(item: Any) -> bool:
|
||||
from sqlglot import exp
|
||||
|
||||
return isinstance(item, exp.Star) or (isinstance(item, exp.Column) and isinstance(item.this, exp.Star))
|
||||
|
||||
|
||||
def _branches(root: Any) -> list[list[Any]]:
|
||||
"""一块查询选出的列：一条 SELECT 一支，UNION 这类的每一支各一支（结果的第 i 列来自每一支的第 i 项）。"""
|
||||
def _leaves(root: Any) -> list[Any]:
|
||||
"""一块查询的每一支：UNION 这类连起来的按先后摊开，套在外面的括号去掉。一般是 SELECT，也可能是 VALUES 这类。"""
|
||||
from sqlglot import exp
|
||||
|
||||
out: list[list[Any]] = []
|
||||
out: list[Any] = []
|
||||
stack = [root]
|
||||
while stack:
|
||||
node = stack.pop()
|
||||
@@ -847,11 +1000,18 @@ def _branches(root: Any) -> list[list[Any]]:
|
||||
stack.append(node.this)
|
||||
elif isinstance(node, exp.SetOperation):
|
||||
stack.extend((node.expression, node.this))
|
||||
elif isinstance(node, exp.Select):
|
||||
out.append(list(node.expressions))
|
||||
else:
|
||||
out.append(node)
|
||||
return out
|
||||
|
||||
|
||||
def _branches(root: Any) -> list[list[Any]]:
|
||||
"""一块查询选出的列：一条 SELECT 一支，UNION 这类的每一支各一支（结果的第 i 列来自每一支的第 i 项）。"""
|
||||
from sqlglot import exp
|
||||
|
||||
return [list(node.expressions) for node in _leaves(root) if isinstance(node, exp.Select)]
|
||||
|
||||
|
||||
def _output_name(item: Any) -> str | None:
|
||||
"""一个选出项在结果里的列名（小写）：起了别名的是别名，直接选一列的是列名。别的表达式（函数、运算、CAST、括号、
|
||||
子查询）的列名各数据库起法不同（SQLite 照写法原样记「SUBSTR(buyer_phone, 4)」，Postgres 记函数名或 ?column?），
|
||||
@@ -1472,10 +1632,11 @@ def _use_of(node: Any, values: dict[int, Any], lifts: dict[int, Any]) -> tuple[s
|
||||
|
||||
|
||||
def _wraps(parent: Any, child: Any) -> bool:
|
||||
"""parent 只是 child 的外壳，不改变它是哪个值：括号、负号、CAST、字符集前缀（_utf8mb4'x'）、COLLATE。"""
|
||||
"""parent 只是 child 的外壳，不改变它是哪个值：括号、负号、CAST（只算被转换的那个值，类型参数、DEFAULT 后面的值
|
||||
不算）、字符集前缀（_utf8mb4'x'）、COLLATE。"""
|
||||
from sqlglot import exp
|
||||
|
||||
return isinstance(parent, (exp.Paren, exp.Neg, exp.Cast)) or (
|
||||
return (isinstance(parent, (exp.Paren, exp.Neg, exp.Cast)) and child is parent.this) or (
|
||||
isinstance(parent, exp.Introducer) and child is parent.expression) or (
|
||||
isinstance(parent, exp.Collate) and child is parent.this)
|
||||
|
||||
@@ -1497,10 +1658,11 @@ def _exact_ops() -> dict[Any, str]:
|
||||
|
||||
|
||||
def _arithmetic() -> tuple[Any, ...]:
|
||||
"""加减乘除：几个写死的数算出来的整个算式按算出来的那一个值认（13900000000 + 1 就是 13900000001）。"""
|
||||
"""加减乘除、取模（%、MOD）、整除（MySQL 的 DIV、DuckDB 的 //）：几个写死的数算出来的整个算式按算出来的那一个值认
|
||||
（13900000000 + 1 就是 13900000001，125 % 100 就是 25）。"""
|
||||
from sqlglot import exp
|
||||
|
||||
return exp.Add, exp.Sub, exp.Mul, exp.Div
|
||||
return exp.Add, exp.Sub, exp.Mul, exp.Div, exp.Mod, exp.IntDiv
|
||||
|
||||
|
||||
#: 数的写法里指数最大多少（1e999999 这样的写法不按数认，免得展开成上百万位）
|
||||
@@ -1557,19 +1719,27 @@ def _constants(root: Any) -> dict[int, Any]:
|
||||
|
||||
|
||||
def _worked_out(node: Any, a: Any, b: Any) -> Decimal | None:
|
||||
"""一步加减乘除的结果；不是两个数、除以零、结果太大为 None。"""
|
||||
"""一步加减乘除、取模、整除的结果；不是两个数、除以零、结果太大为 None。取模、整除照 SQL 往零取整（-7 % 2 是 -1，
|
||||
-7 DIV 2 是 -3，同 Decimal 的 % 和 //）。/ 在 SQLite、Postgres 这类方言里两个整数相除只留整数部分（解析时记在
|
||||
typed 上：53 / 2 是 26），写成小数的（53.0 / 2）、别的方言照小数算。以前一律按小数算，SQLite 里 age >= 25 AND
|
||||
age < 53 / 2 夹住了 25 却认不出。"""
|
||||
from sqlglot import exp
|
||||
|
||||
if not (isinstance(a, Decimal) and isinstance(b, Decimal)):
|
||||
return None
|
||||
try:
|
||||
if isinstance(node, exp.Div):
|
||||
value = a / b if b else None
|
||||
if isinstance(node, (exp.Div, exp.IntDiv, exp.Mod)):
|
||||
if not b:
|
||||
return None
|
||||
# 写成整数的数（和它们加减乘出来的）指数正好是 0；53.0、1e1 这样写的在 SQL 里不是整数
|
||||
whole = isinstance(node, exp.IntDiv) or (
|
||||
node.args.get("typed") and a.as_tuple().exponent == 0 and b.as_tuple().exponent == 0)
|
||||
value = a % b if isinstance(node, exp.Mod) else a // b if whole else a / b
|
||||
else:
|
||||
value = a + b if isinstance(node, exp.Add) else a - b if isinstance(node, exp.Sub) else a * b
|
||||
except ArithmeticError:
|
||||
return None
|
||||
return value if value is not None and value.is_finite() and abs(value.adjusted()) <= _EXPONENT_MAX else None
|
||||
return value if value.is_finite() and abs(value.adjusted()) <= _EXPONENT_MAX else None
|
||||
|
||||
|
||||
def _lifts(root: Any, values: dict[int, Any]) -> dict[int, Any]:
|
||||
@@ -1588,16 +1758,19 @@ def _lifts(root: Any, values: dict[int, Any]) -> dict[int, Any]:
|
||||
|
||||
|
||||
#: 一个比较的一半：(算出边界的那几个字面量, 另一边的键, 写死的值, 「键 运算符 值」的运算符)。边界是几个数算出来的
|
||||
#: （13900000002 - 1），算式里的字面量都记下：夹住了某一个值，留着哪一个都可能算得出它，都换
|
||||
#: （13900000002 - 1），算式里的字面量都记下：夹住了某一个值，留着哪一个都可能算得出它，都换。不参与算值的字面量
|
||||
#: （CAST(30000 AS DECIMAL(10, 2)) 的类型参数 10、2）不记，见 _halves
|
||||
_Half = tuple[tuple[Any, ...], str, Any, str]
|
||||
#: 一个比较拆成的几半，加上几半怎么连（False：且，BETWEEN；True：或，IN）、是不是等值一类（不等、IN）
|
||||
_Leaf = tuple[list[_Half | None], bool, bool]
|
||||
|
||||
|
||||
def _halves(node: Any, values: dict[int, Any]) -> list[_Half | None]:
|
||||
"""一个比较拆成的几半（values 是 _constants 算好的写死的值）：BETWEEN 是两半的「且」（不低于下界、不高于上界），
|
||||
> >= < <= 是一半，= <> 这类是一半，IN 是一串「=」的「或」。边界不是写死的值（另一列、子查询）或者两边都写死了的，
|
||||
这一半记 None：它的真假跟哪一列的取值都对不上，当成不知道真假。"""
|
||||
def _halves(node: Any, values: dict[int, Any], lifts: dict[int, Any]) -> list[_Half | None]:
|
||||
"""一个比较拆成的几半（values 是 _constants 算好的写死的值，lifts 见 _lifts）：BETWEEN 是两半的「且」（不低于下界、
|
||||
不高于上界），> >= < <= 是一半，= <> 这类是一半，IN 是一串「=」的「或」。边界不是写死的值（另一列、子查询）或者
|
||||
两边都写死了的，这一半记 None：它的真假跟哪一列的取值都对不上，当成不知道真假。算出边界的字面量是往上走得到这一边
|
||||
的那几个（_lifts）：CAST(30000 AS DECIMAL(10, 2)) 里只有 30000，类型参数 10、2 不算出值，换了只会让封存的 SQL 难读，
|
||||
SQL 检查结果的片段里别处的 10、2 也会跟着换。"""
|
||||
from sqlglot import exp
|
||||
|
||||
if isinstance(node, exp.Between):
|
||||
@@ -1615,8 +1788,9 @@ def _halves(node: Any, values: dict[int, Any]) -> list[_Half | None]:
|
||||
out: list[_Half | None] = []
|
||||
for operand, other, op in sides:
|
||||
value = values.get(id(operand)) if operand is not None else None
|
||||
out.append((tuple(operand.find_all(*_literal_types())), _key_of(other), value, op)
|
||||
if value is not None else None)
|
||||
literals = () if value is None else tuple(
|
||||
literal for literal in operand.find_all(*_literal_types()) if lifts.get(id(literal)) is operand)
|
||||
out.append((literals, _key_of(other), value, op) if value is not None else None)
|
||||
return out
|
||||
|
||||
|
||||
@@ -1725,9 +1899,9 @@ def _marks(value: Any, op: str) -> list[tuple[Any, int]]:
|
||||
return [(value, 0 if op in (">=", "<") else 1), *steps]
|
||||
|
||||
|
||||
def _pinned(root: Any, values: dict[int, Any]) -> dict[int, tuple[str, ...]]:
|
||||
def _pinned(root: Any, values: dict[int, Any], lifts: dict[int, Any]) -> dict[int, tuple[str, ...]]:
|
||||
"""把某一个值夹住了的边界字面量：id(字面量节点) → 夹住的值（_point_text 的写法）。values 是 _constants 算好的
|
||||
写死的值。
|
||||
写死的值，lifts 见 _lifts。
|
||||
|
||||
一个值 v 两边都被隔开（_marks），条件才可能只认 v 这一个值。范围比较在整条 SQL 里配对，只有分档的几种写法各算
|
||||
各的（_apart：CASE 的不同分支、SELECT 里不筛行的不同项、UNION 的不同支），它们共用的边界是分档的界，不是在找哪
|
||||
@@ -1754,7 +1928,7 @@ def _pinned(root: Any, values: dict[int, Any]) -> dict[int, tuple[str, ...]]:
|
||||
nearby: dict[tuple[int, str], set[Decimal]] = {}
|
||||
for node in root.find_all(exp.Between, *_range_ops(), *exact_types):
|
||||
is_exact = isinstance(node, exact_types)
|
||||
parts = _halves(node, values)
|
||||
parts = _halves(node, values, lifts)
|
||||
if is_exact and all(h is None for h in parts):
|
||||
continue # 连接条件（两边都是列）这类：和夹住哪个值无关
|
||||
top = _condition_root(node)
|
||||
@@ -2006,7 +2180,7 @@ def _literal_uses(sql: str, kind: str | None) -> dict[int, _Use] | None:
|
||||
try:
|
||||
values = _constants(root)
|
||||
lifts = _lifts(root, values)
|
||||
pinned = _pinned(root, values)
|
||||
pinned = _pinned(root, values, lifts)
|
||||
for node in root.walk():
|
||||
start = (node.meta or {}).get("start") if isinstance(node, _literal_types()) else None
|
||||
if not isinstance(start, int):
|
||||
|
||||
@@ -14,12 +14,12 @@
|
||||
别名的、和 * 写在一起没起别名的表达式也算）的格子渲染成「已遮罩」，回答正文、查询清单、事件、质询工件里都不出现
|
||||
遮罩列的值；换着写法找某一个客户（简单 CASE、IN (VALUES …)、两头相同的范围、行值、选出列里写着同一个号码，两半
|
||||
拆到 ON 和 WHERE、WHERE 和 HAVING、INTERSECT 两支、CTE 里，边界、等值写成算式，两头不是整数，范围再排除一头，
|
||||
号码前面补零……）也一样，年龄这类短的值换着写法找（套了 CAST、汇总值两头相同、补零……）SQL 里也遮；遮罩值不误伤
|
||||
阈值、天数、LIMIT（长的数也一样）、提问人自己写的数、报告口径卡里的指标，分档共用的边界、和日期、天数修饰、标签里
|
||||
写着的同一个数也不牵连阈值，分档的文字标签不当遮罩值（当分组键的档位列照常显示），改了名字的遮罩列撞上标签照样遮；
|
||||
SQL 里照原样留着的阈值只在回答正文里不遮，格子照旧按值遮，回答和质询工件对得上；质询工件按值遮只出现在报告快照里
|
||||
的遮罩值；一格文字里夹着遮罩值时回答整份不交，质询工件里只换掉那一截；质询工件里 SQL 检查结果的片段、说明照整条
|
||||
SQL 换（短的数也换），给模型的那句改法不交出去
|
||||
号码前面补零……）也一样，年龄这类短的值换着写法找（套了 CAST、汇总值两头相同、补零、取模和整数相除算出来的边界……）
|
||||
SQL 里也遮；遮罩值不误伤阈值、天数、LIMIT（长的数也一样）、提问人自己写的数、报告口径卡里的指标，分档共用的边界、
|
||||
和日期、天数修饰、标签里写着的同一个数也不牵连阈值，分档的文字标签不当遮罩值（当分组键的档位列照常显示），改了
|
||||
名字的遮罩列、UNION 另一支取的遮罩列撞上标签照样遮；SQL 里照原样留着的阈值只在回答正文里不遮，格子照旧按值遮，
|
||||
回答和质询工件对得上；质询工件按值遮只出现在报告快照里的遮罩值；一格文字里夹着遮罩值时回答整份不交，质询工件里
|
||||
只换掉那一截；质询工件里 SQL 检查结果的片段、说明照整条 SQL 换（短的数也换），给模型的那句改法不交出去
|
||||
- 接线：设置里的上限、运行固定的数据版本交到 agent 手里（上传表格重传之后查的仍是运行固定的那一版）；同名重建的
|
||||
数据源、没固定过版本的上传表格不查
|
||||
- 触顶、失败照样记一条事件，说法给提问人看得懂、说下一步；轮数用完不算查询次数到顶，查询失败、引用解析不了在末尾注明
|
||||
@@ -1051,11 +1051,19 @@ async def test_a_threshold_matching_numbers_left_alone_elsewhere_is_sealed_as_wr
|
||||
pytest.param("SELECT COUNT(*) AS 人数 FROM sales WHERE age BETWEEN 20 + 5 AND 30 - 5",
|
||||
"SELECT COUNT(*) AS 人数 FROM sales WHERE age BETWEEN 已遮罩 + 已遮罩 AND 已遮罩 - 已遮罩",
|
||||
id="算出来的边界"),
|
||||
# 第 3 轮复核：取模、整数相除（SQLite 里只留整数部分）算出来的边界
|
||||
pytest.param("SELECT COUNT(*) AS 人数 FROM sales WHERE age >= 25 AND age <= 125 % 100",
|
||||
"SELECT COUNT(*) AS 人数 FROM sales WHERE age >= 已遮罩 AND age <= 已遮罩 % 已遮罩",
|
||||
id="取模算出来的边界"),
|
||||
pytest.param("SELECT COUNT(*) AS 人数 FROM sales WHERE age >= 25 AND age < 53 / 2",
|
||||
"SELECT COUNT(*) AS 人数 FROM sales WHERE age >= 已遮罩 AND age < 已遮罩 / 已遮罩",
|
||||
id="整数相除算出来的边界"),
|
||||
])
|
||||
async def test_a_short_masked_value_looked_up_in_other_shapes_is_sealed_masked(monkeypatch, tmp_path, client, sql,
|
||||
sealed):
|
||||
"""年龄设了遮罩，agent 换着写法找某一个年龄（套了 CAST、汇总值两头相同、补零、两头不是整数、范围再排除一头、
|
||||
算出来的边界）：第 2 轮复核时都照原样封存。现在查询清单、事件、质询工件里的 SQL 都换掉这个年龄，三处一字不差。"""
|
||||
算出来的边界，取模、整数相除算出来的边界）：第 2、3 轮复核时都照原样封存。现在查询清单、事件、质询工件里的 SQL
|
||||
都换掉这个年龄，三处一字不差。"""
|
||||
name, source_id = await _source(tmp_path, options={"mask_columns": "age"}, script=AGED)
|
||||
graph = chain(node("start", "input"),
|
||||
node("fetch", "tool", tool=f"db_query__{name}",
|
||||
@@ -1124,6 +1132,42 @@ async def test_a_renamed_masked_column_keeps_values_that_match_labels_written_el
|
||||
await engines.invalidate(source_id)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mixed", [
|
||||
pytest.param("SELECT CASE WHEN amount > 1000 THEN 1 ELSE 0 END AS 等级 FROM sales UNION ALL "
|
||||
"SELECT vip_level FROM sales", id="标签在前、遮罩列在后"),
|
||||
pytest.param("SELECT CASE WHEN amount > 1000 THEN 1 ELSE 0 END AS 等级 FROM sales UNION ALL "
|
||||
"SELECT b.vip_level AS 等级 FROM (SELECT vip_level FROM sales) b", id="遮罩列经子查询按原名传出"),
|
||||
])
|
||||
async def test_a_masked_column_after_a_label_in_a_union_keeps_its_values(monkeypatch, tmp_path, client, mixed):
|
||||
"""第 3 轮复核：UNION 的第一支是按条件数大单的标签（THEN 1 ELSE 0），另一支取的是等级（遮罩列原样选出，或者经按
|
||||
原名选出的子查询、以标签的别名传出）。以前这一列被当成只装着标签，等级 1 撞上标签 1 不收，之后 WHERE vip_level = 1
|
||||
在查询清单、事件、质询工件里都照原样，= 2 却遮了。现在这一列在一支里取的是遮罩列，就照收，两条都遮。"""
|
||||
name, source_id = await _source(tmp_path, options={"mask_columns": "vip_level"}, script=LEVELS)
|
||||
first, second = (f"SELECT COUNT(*) AS 人数 FROM sales WHERE vip_level = {level}" for level in (1, 2))
|
||||
plays = {"各等级": Play([mixed, first, second],
|
||||
"问到的两个等级各有 [[v:Q2.r0.人数]] 人、[[v:Q3.r0.人数]] 人。[[see:Q2]]")}
|
||||
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="各等级的客户有几位？",
|
||||
approval_id=approval.id)
|
||||
assert r.status_code == 200, r.text
|
||||
item = r.json()["inquiry"]
|
||||
assert (item["answer"]["ok"], item["answer"]["text"]) == (True, "问到的两个等级各有 2 人、1 人。")
|
||||
sealed = [mixed, *["SELECT COUNT(*) AS 人数 FROM sales WHERE vip_level = 已遮罩"] * 2]
|
||||
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
|
||||
for query, sql in zip(item["queries"], sealed):
|
||||
got = await client.get(f"/api/runs/{paused.id}/inquiries/artifacts/{query['artifact']}")
|
||||
assert got.status_code == 200 and got.json()["content"]["sql"] == sql, "质询工件的 SQL 和事件里的一字不差"
|
||||
got = await client.get(f"/api/runs/{paused.id}/inquiries/artifacts/{item['queries'][0]['artifact']}")
|
||||
assert {cell for [cell] in got.json()["content"]["rows"]} == {"已遮罩"} and got.json()["masked"] == ["等级"]
|
||||
finally:
|
||||
await engines.invalidate(source_id)
|
||||
|
||||
|
||||
STAR_TAIL = "SELECT *, SUBSTR(buyer_phone, 4) FROM sales ORDER BY amount DESC LIMIT 1"
|
||||
|
||||
|
||||
|
||||
@@ -9,8 +9,9 @@
|
||||
- SQL 里数字照原样留着的只有范围比较的边界（没把某一个值夹住的）和 LIMIT / OFFSET 的行数；在找某一个值的（等值、IN、
|
||||
LIKE、简单 CASE、IS [NOT] DISTINCT FROM，套了括号、写成行值、VALUES、ARRAY 的也算，还有 BETWEEN v AND v、
|
||||
>= v AND <= v 这类夹住某一个值的范围）和遮罩列比较时短的数也换，被比较的一边套了 CAST、汇总也算；同一个数换个
|
||||
写法（25.0、025、2.5e1）、写成几个数的算式（25 + 0、13900000000 + 1）照样认，算式里的数都换；选出列、函数参数、
|
||||
算式里够长的遮罩值也换；字符串（范围边界也一样）、注释、名字按值换；SQL 解析不了时认不准哪个是阈值，一律不免
|
||||
写法（25.0、025、2.5e1）、写成几个数的算式（25 + 0、13900000000 + 1、125 % 100、50 DIV 2，整数相除照数据源的算法）
|
||||
照样认，算式里的数都换（CAST 的类型参数不算）；选出列、函数参数、算式里够长的遮罩值也换；字符串（范围边界也一样）、
|
||||
注释、名字按值换；SQL 解析不了时认不准哪个是阈值，一律不免
|
||||
- 「夹住某一个值」在整条 SQL 里认，只有分档的几种写法各算各的：CASE 的不同分支、不同的选出列（连同 GROUP BY、
|
||||
ORDER BY）、UNION 的不同支，共用的边界是阈值；同一行集上的条件（ON 和 WHERE、WHERE 和 HAVING、INTERSECT 两支、
|
||||
包在 CASE / 函数 / 算式里的一半、子查询、CTE）照样认得出在找哪一个值。同一个条件里按合起来的真假认（旁边的别的
|
||||
@@ -22,8 +23,9 @@
|
||||
同一个短的数，没设遮罩的列上的等值，都不牵连阈值；不留的照等值的办法换，别的列上的短阈值、LIMIT 撞上别处换掉的
|
||||
短的数照原样（已知局限）
|
||||
- 按遮罩列分档的标签（CASE、IIF 的结果都是写死的常量）是 agent 自己写在 SQL 里的，不当遮罩值：SQL、回答正文里照原样；
|
||||
当分组键的（每一档一行）不算派生值，格子照常显示，逐行的标签列照旧整列遮。只有只装着标签的列不收（经别名、子查询
|
||||
传下去、没起别名的也算），改了名字的遮罩列撞上同一条 SQL 里的 THEN 1 ELSE 0 照收
|
||||
当分组键的（每一档一行）不算派生值，格子照常显示，逐行的标签列照旧整列遮。只有只装着标签的列不收（按作用域认取自
|
||||
哪个来源：经别名、子查询、CTE 传下去、没起别名的也算），改了名字的遮罩列撞上同一条 SQL 里的 THEN 1 ELSE 0 照收，
|
||||
UNION 另一支、子查询里的遮罩列流进标签列也照收
|
||||
- 快照按列遮，另按值遮：别的列里写着遮罩值的格子也换成「已遮罩」；列名里写着遮罩值的，列名也换；质询工件里，
|
||||
没设遮罩的列里一格文字夹着的遮罩值按子串换（回答一路照旧整份不交），SQL 检查结果的片段、说明照整条 SQL 换（截断的
|
||||
片段从没闭合的引号、注释起去掉，说明里的码值表和别的字一样按值换）
|
||||
@@ -278,6 +280,75 @@ def test_labels_passed_on_under_their_own_name_are_still_not_masked_values(sql,
|
||||
assert plan.secrets.scrub("小单 3 笔、大单 2 笔") == "小单 3 笔、大单 2 笔"
|
||||
|
||||
|
||||
#: 按条件数大单写的 1、0：和会员等级（1、2、3）撞在一起
|
||||
BIG_ONE = "CASE WHEN amount > 1000 THEN 1 ELSE 0 END"
|
||||
#: 标签那一支的 1、0 在前，等级 1、2、3 在后
|
||||
LEVEL_ROWS = [[1], [0], [0], [1], [1], [2], [3], [1]]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql, columns, rows", [
|
||||
# UNION 标签在前、遮罩列在后：遮罩列没起别名，或者起了别的名字
|
||||
(f"SELECT {BIG_ONE} AS 等级 FROM sales UNION ALL SELECT vip_level FROM sales", ["等级"], LEVEL_ROWS),
|
||||
(f"SELECT {BIG_ONE} AS 等级 FROM sales UNION ALL SELECT vip_level AS lv FROM sales", ["等级"], LEVEL_ROWS),
|
||||
# 子查询里的 UNION，外层按名字取；写了列名表的 CTE 里的 UNION
|
||||
(f"SELECT 等级, COUNT(*) AS 人数 FROM (SELECT {BIG_ONE} AS 等级 FROM sales UNION ALL SELECT vip_level FROM sales) t "
|
||||
"GROUP BY 等级", ["等级", "人数"], [[0, 2], [1, 4], [2, 1], [3, 1]]),
|
||||
(f"WITH t(lv) AS (SELECT {BIG_ONE} FROM sales UNION ALL SELECT vip_level FROM sales) SELECT lv AS 等级 FROM t",
|
||||
["等级"], LEVEL_ROWS),
|
||||
# 标签写在 CTE 里，结果的 UNION 另一支取遮罩列
|
||||
(f"WITH b AS (SELECT {BIG_ONE} AS lvl FROM sales) SELECT lvl FROM b UNION ALL SELECT vip_level FROM sales",
|
||||
["lvl"], LEVEL_ROWS),
|
||||
# 遮罩列经按原名选出的子查询、以标签的别名传出
|
||||
(f"SELECT {BIG_ONE} AS 档 FROM sales UNION ALL SELECT b.vip_level AS 档 FROM (SELECT vip_level FROM sales) b",
|
||||
["档"], LEVEL_ROWS),
|
||||
# 不用 UNION：两个子查询连起来，一边是标签，一边是改了名字的遮罩列，名字相同
|
||||
(f"SELECT y.档 AS 档 FROM (SELECT id, {BIG_ONE} AS 档 FROM sales) x "
|
||||
"JOIN (SELECT id, vip_level AS 档 FROM (SELECT id, vip_level FROM sales) z) y ON x.id = y.id", ["档"],
|
||||
[[1], [2], [3], [1]]),
|
||||
# USING 连起来的同名列：右连接时，左边没对上的行取的是右边那张表里的值
|
||||
(f"WITH b AS (SELECT id, {BIG_ONE} AS vip_level FROM sales) SELECT vip_level AS 等级 FROM b "
|
||||
"RIGHT JOIN members USING (vip_level)", ["等级"], [[1], [2], [3]]),
|
||||
])
|
||||
def test_a_column_that_also_carries_the_masked_column_is_not_a_label_column(sql, columns, rows):
|
||||
"""第 3 轮复核：只装着标签的列以前按名字认，从「所有名字都只装着标签」起往下剔，遮罩列原样选出（SELECT vip_level）
|
||||
时名字就是它自己，永远剔不掉；结果的 UNION 按位置那一遍只往里加标签，另一支取的是遮罩列也不剔，派生来源里的
|
||||
UNION 根本不按位置看。于是标签在前、遮罩列在后时，等级 1 撞上 THEN 1 就不收，之后 WHERE vip_level = 1 照原样
|
||||
封存。现在按作用域认每一项取自哪个来源，结果的这一列在每一支里都只装着标签才算，取自数据源里的表的一律不算。"""
|
||||
plan = plan_of({"source": "shop", "sql": sql, "columns": columns, "rows": rows}, {"vip_level"})
|
||||
assert inquiry_mod._label_columns(sql, "sqlite", columns) == {}
|
||||
assert {"1", "2", "3"} <= plan.secrets.weak
|
||||
for level in (1, 2, 3):
|
||||
assert inquiry_mod._scrub_sql(f"SELECT COUNT(*) AS 人数 FROM sales WHERE vip_level = {level}", plan.secrets,
|
||||
plan.tainted["a1"], "sqlite") == (
|
||||
"SELECT COUNT(*) AS 人数 FROM sales WHERE vip_level = 已遮罩")
|
||||
|
||||
|
||||
_EXTRA = "CASE WHEN amount > 60000 THEN '特大单' END"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql, columns, texts", [
|
||||
# CTE 里分档，外层连上数据源里的表、按档位汇总：档位没写表名也认得出取自 CTE（SQL 能执行，别的表就没有同名的列，
|
||||
# 不然会报列名有歧义）
|
||||
(f"WITH banded AS (SELECT id, {BANDS} AS 档位 FROM sales) SELECT 档位, SUM(s.visits) AS 访问 FROM banded b "
|
||||
"JOIN sales s ON s.id = b.id GROUP BY 档位", ["档位", "访问"], {"小单", "大单"}),
|
||||
# 按别的列连起来（USING (id)）：档位没合成一列，照样取自 CTE
|
||||
(f"WITH banded AS (SELECT id, {BANDS} AS 档位 FROM sales) SELECT 档位, SUM(o.amount) AS 金额 FROM banded "
|
||||
"JOIN orders o USING (id) GROUP BY 档位", ["档位", "金额"], {"小单", "大单"}),
|
||||
# 写了表名、经两层子查询传下来
|
||||
(f"SELECT t.档位 FROM (SELECT b.档位 FROM (SELECT {BANDS} AS 档位 FROM sales) b) t", ["档位"], {"小单", "大单"}),
|
||||
# UNION 两支都是标签（第二支没写 ELSE、没起别名）：两支的标签都算
|
||||
(f"SELECT {BANDS} AS 档位 FROM sales UNION ALL SELECT {_EXTRA} FROM sales", ["档位"], {"小单", "大单", "特大单"}),
|
||||
# 派生来源里的 UNION 两支都是标签，外层经 * 带出来
|
||||
(f"SELECT * FROM (SELECT {BANDS} AS 档位 FROM sales UNION ALL SELECT {_EXTRA} FROM sales) t", ["档位"],
|
||||
{"小单", "大单", "特大单"}),
|
||||
])
|
||||
def test_label_columns_are_traced_to_where_they_are_written(sql, columns, texts):
|
||||
"""改成按作用域认之后，写在别处、原样传下来的标签照样认得出：外层连上数据源里的表没写表名、按别的列 USING、写了
|
||||
表名经两层子查询；UNION 两支都是标签的，每一支的标签都算（以前派生来源里的 UNION 按名字认，只算起了这个名字的
|
||||
第一支）。"""
|
||||
assert inquiry_mod._label_columns(sql, "sqlite", columns) == {"档位": texts}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 按值认：门槛、边界
|
||||
# --------------------------------------------------------------------------
|
||||
@@ -505,6 +576,55 @@ def test_a_threshold_worked_out_from_several_numbers_is_kept_like_one_number(amo
|
||||
assert inquiry_mod._shown_literals([(swapped, "sqlite", {"amount"})], amounts) == set()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind, sql, sealed", [
|
||||
# 取模、整除算出来的边界和等值：125 % 100、MOD(125, 100) 是 25，50 DIV 2、51 // 2 也是
|
||||
("sqlite", "SELECT COUNT(*) FROM sales WHERE age >= 25 AND age <= 125 % 100",
|
||||
"SELECT COUNT(*) FROM sales WHERE age >= 已遮罩 AND age <= 已遮罩 % 已遮罩"),
|
||||
("postgresql", "SELECT COUNT(*) FROM sales WHERE age = MOD(125, 100)",
|
||||
"SELECT COUNT(*) FROM sales WHERE age = MOD(已遮罩, 已遮罩)"),
|
||||
("mysql", "SELECT COUNT(*) FROM sales WHERE age >= 25 AND age <= 50 DIV 2",
|
||||
"SELECT COUNT(*) FROM sales WHERE age >= 已遮罩 AND age <= 已遮罩 DIV 已遮罩"),
|
||||
("duckdb", "SELECT COUNT(*) FROM sales WHERE age >= 25 AND age <= 51 // 2",
|
||||
"SELECT COUNT(*) FROM sales WHERE age >= 已遮罩 AND age <= 已遮罩 // 已遮罩"),
|
||||
# SQLite、Postgres 里两个整数相除只留整数部分：53 / 2 是 26，age < 26 和 age >= 25 夹住 25
|
||||
("sqlite", "SELECT COUNT(*) FROM sales WHERE age >= 25 AND age < 53 / 2",
|
||||
"SELECT COUNT(*) FROM sales WHERE age >= 已遮罩 AND age < 已遮罩 / 已遮罩"),
|
||||
("postgresql", "SELECT COUNT(*) FROM sales WHERE buyer_phone >= 13900000001 AND buyer_phone < 27800000005 / 2",
|
||||
"SELECT COUNT(*) FROM sales WHERE buyer_phone >= 已遮罩 AND buyer_phone < 已遮罩 / 已遮罩"),
|
||||
])
|
||||
def test_a_value_worked_out_with_remainders_or_whole_division_is_read_as_that_value(phones, kind, sql, sealed):
|
||||
"""第 3 轮复核：取模（%、MOD）、整除（MySQL 的 DIV、DuckDB 的 //）写成的算式以前不按算出来的值认，age >= 25 AND
|
||||
age <= 125 % 100 照原样封存；整数相除在 SQLite、Postgres 里只留整数部分，以前按小数算（53 / 2 算成 26.5），
|
||||
夹住的值也认不出。现在都按数据源算出来的那一个值认，算式里的数都换，正文里也不免遮。"""
|
||||
assert inquiry_mod._scrub_sql(sql, phones, {"buyer_phone", "age"}, kind) == sealed
|
||||
assert not inquiry_mod._shown_literals([(sql, kind, {"buyer_phone", "age"})], phones) & {"25", "13900000001"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind, divided", [("mysql", "53 / 2"), ("duckdb", "53 / 2"), ("sqlite", "53.0 / 2")])
|
||||
def test_division_that_keeps_the_fraction_is_read_as_written(phones, kind, divided):
|
||||
"""MySQL、DuckDB 的 / 照小数算，SQLite 里写成小数的也是：53 / 2、53.0 / 2 是 26.5，age >= 25 AND age < 26.5 里
|
||||
25、26 都成立，没夹住，照原样。"""
|
||||
sql = f"SELECT COUNT(*) FROM sales WHERE age >= 25 AND age < {divided}"
|
||||
assert inquiry_mod._scrub_sql(sql, phones, {"buyer_phone", "age"}, kind) == sql
|
||||
|
||||
|
||||
def test_type_parameters_of_a_cast_bound_are_left_alone(amounts):
|
||||
"""第 3 轮复核：夹住某一个值的边界写成 CAST(v AS DECIMAL(10, 2)) 时，类型参数 10、2 以前也当成算出边界的数换掉，
|
||||
封存的 SQL 成了 DECIMAL(已遮罩, 已遮罩)，SQL 检查结果的片段里别处的 10、2 也跟着换。只换算出边界的那个数；Oracle
|
||||
的 CAST … DEFAULT 0 ON CONVERSION ERROR 里的 0 也不算出值，照原样。"""
|
||||
sql = ("SELECT COUNT(*) FROM sales WHERE amount >= CAST(30000 AS DECIMAL(10, 2)) "
|
||||
"AND amount <= CAST(30000 AS DECIMAL(10, 2))")
|
||||
assert inquiry_mod._scrub_sql(sql, amounts, {"amount"}, "sqlite") == sql.replace("30000", MASKED)
|
||||
assert inquiry_mod._literal_swaps(sql, amounts, {"amount"}, "sqlite") == {"30000"}
|
||||
checks = [{"code": "stock_summed", "level": "warning", "table": "sales", "message": "存量跨日期求和",
|
||||
"sql_excerpt": "ROUND(SUM(amount), 2) >= 10"}]
|
||||
[check] = inquiry_mod._masked_checks(checks, sql, amounts, {"amount"}, "sqlite")
|
||||
assert check["sql_excerpt"] == "ROUND(SUM(amount), 2) >= 10"
|
||||
fallback = ("SELECT COUNT(*) FROM sales WHERE amount >= CAST('30000' AS NUMBER DEFAULT 0 ON CONVERSION ERROR) "
|
||||
"AND amount <= CAST('30000' AS NUMBER DEFAULT 0 ON CONVERSION ERROR)")
|
||||
assert inquiry_mod._scrub_sql(fallback, amounts, {"amount"}, "oracle") == fallback.replace("30000", MASKED)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql, sealed", [
|
||||
# 遮罩列套了 CAST、汇总：被比较的一边取的还是它的值，短的数照样换
|
||||
("SELECT COUNT(*) FROM sales WHERE CAST(age AS INTEGER) = 25",
|
||||
|
||||
Reference in New Issue
Block a user