mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 00:11:59 +00:00
K3 复核修正（第 3 轮）：写了库名的表按「库名.表名」查表结构；UNION 第一支按表结构数出的列名要和后面每一支一样宽才用；整行引用、函数参数里的 *、表的列名表、TABLE x、PIVOT 这类认得出
- 写了库名的表（复核 major 1）：表结构以前按表名查，查询写的是另一个库里的同名表（archive.sales，列序不同）时照探查
的那张数位置，号码原样交给模型、落库。现在表结构另按「库名.表名」记一份（库名取探查时记下的 schema、qualified 的
前缀，masking._schema_keys），写了库名的引用按它查（_table_key），查不到的按数不清整段遮；探查时没设 schema、查询又
写了库名的也查不到（多遮）。数位置（_star_columns、_branch_names）、认标签列（_table_lacks）都照这个查法
- 过期的表结构（复核 major 2，对 099473c 的回退）：UNION 第一支是 * 时，按表结构数出的列名以前不核对列数就拿来对后一支
的位置，表结构过期、列数已经变了也照对，遮罩值记到错的列名上。现在后面每一支数得出的列数都和它一样宽才用
（_source_names、_width），对不上、数不清的退回整段遮
- 整行和别的 * 写法（复核 major 3）：写在函数参数、表达式里的 *、s.*（to_json(s.*)、concat_ws(',', s.*)、ROW(s.*)、
Oracle 的 JSON_OBJECT(*)）按它展开的来源认（_value_columns、_scoped_star_inputs，相关子查询、FROM 里的子查询照样
找得到外层的来源）。几种写法先改写语法树再认（_normalized，原文不动）：Postgres 的整行引用（row_to_json(s)、
json_agg(t)、SELECT s、s::text）按 Postgres 先认列的规矩，不是列的改成 ROW(s.*)；带列名表的表（FROM sales AS
s(i, r, p)，引用 CTE 时写的也算）改成派生来源；解析器读错的 TABLE x 改成 SELECT * FROM x；选出项套了括号的 *
（(s.*)、(s).*）按 * 数位置。PIVOT、UNPIVOT、MATCH_RECOGNIZE 读到遮罩列的整份遮。SQL 解析不了、里面有 *（COUNT(*)
除外）或 TABLE x 的（UNION ALL TABLE x、带库名的 TABLE），同写着遮罩列一样整份遮（_implicit）
- Postgres 的列名表只写前几列时，不按表结构补全：别的方言里列名表比列少只能是表结构过期了，补全就丢了列数这道核对；
这时改了名的几列整段遮（多遮），照原名的按名字认
- inquiry 模块说明的已知局限补上：写了库名的表、过期表结构的列数核对、Postgres 列名表只写前几列时整段遮、探查的不是
默认库时没写库名的同名表可能数错、整行引用按表结构先认列（查不到时按整行多遮）、解析不了的 SQL 里的整行引用认不出
测试：新增 test_masking_whole_rows（整行、函数参数里的 *、列名表、TABLE x、Oracle 的 JSON_OBJECT(*) / UNPIVOT / PIVOT /
MATCH_RECOGNIZE、解析不了的 SQL、不碰遮罩列的照常）；test_masking_star_renames 补写了库名的表（SQLite ATTACH 真跑）、
过期的表结构；test_verify_agent_masking 补核查 agent 的真实查询工具：过期的表结构，Postgres / Oracle 形状的结果（换掉
驱动）交给模型、落库、通用工件接口都遮好。每处实现都做过去掉之后测试变红的验证；复核给的 Postgres 18 探针、端到端
在临时库上重跑，都不再漏。
This commit is contained in:
ThearchyHelios
@@ -58,12 +58,20 @@
|
||||
遮罩列的，按出列名认它出的列（常见函数的出列名、写的列名表、别名、函数名，_function_columns），同名的别的列一并遮
|
||||
（多遮）；不认得的函数又没写列名表的，整份快照都遮。函数参数里没写表名的 col->'path'（JSON 箭头，解析器当成以 col
|
||||
为参数的 lambda）按 col 认，真的 lambda（transform(arr, x -> …)）的参数和遮罩列同名时也算（多遮）
|
||||
- * 把遮罩列带进别的列名（WITH t(a, b, …) AS (SELECT * …) 这类列名表、UNION 后一支的 *、标量子查询里的 *）：只展开一个
|
||||
来源、那里每一列都数得清的，按位置只遮取自遮罩列的那几列——子查询按它的选出项数，数据源里的表按查询当时冻结的表结构
|
||||
数（几张列一样的月度表 SELECT * 连起来汇总，金额照常）。表结构里查不到这张表、数出来的列数和实际展开的对不上，或者
|
||||
* 连着几个来源的，* 展开落到的那一段整段遮，同一段里的金额这类别的列一并遮（多遮）；表结构过期、列数没变却换了次序
|
||||
或改了名的，可能数错（同按表结构认标签列）。标量子查询里的 * 展开数据源里的表一律算取自遮罩列。EXCEPT 后一支只用来
|
||||
去掉行，也照 UNION 的后一支算（多遮）
|
||||
- * 把遮罩列带进别的列名（WITH t(a, b, …) AS (SELECT * …) 这类列名表、表的列名表 FROM sales AS s(i, r, p)、UNION 后一支
|
||||
的 *、TABLE x、标量子查询里的 *）：只展开一个来源、那里每一列都数得清的，按位置只遮取自遮罩列的那几列——子查询按它的
|
||||
选出项数，数据源里的表按查询当时冻结的表结构数（几张列一样的月度表 SELECT * 连起来汇总，金额照常）。表结构里查不到
|
||||
这张表（写了库名的按「库名.表名」查：另一个库里的同名表 archive.sales 不按探查的 sales 数，探查时没设 schema、查询
|
||||
又写了库名的也一样）、数出来的列数和实际展开的对不上（UNION 第一支的 * 按表结构数出的列名，后面每一支都要一样宽才用；
|
||||
Postgres 的列名表只写前几列、后面的照原名时也对不上），或者 * 连着几个来源的，* 展开落到的那一段整段遮，同一段里的
|
||||
金额这类别的列一并遮（多遮）。表结构过期、列数没变却换了次序或改了名的，可能数错（同按表结构认标签列）；没写库名的表
|
||||
按探查的那张数，探查的不是默认的库、默认的库里又有一张列序不同的同名表时也会数错。标量子查询里的 * 展开数据源里的表
|
||||
一律算取自遮罩列。EXCEPT 后一支只用来去掉行，也照 UNION 的后一支算（多遮）
|
||||
- 整行和函数参数里的 *（Postgres 的 row_to_json(s)、json_agg(t)、SELECT s、s::text、to_json(s.*)、concat_ws(',', s.*)，
|
||||
Oracle 的 JSON_OBJECT(*)）按 * 展开的来源算：取自子查询、CTE 的看它选出的列，取自数据源里的表的一律算取自遮罩列（整列
|
||||
遮）。Postgres 先按列认：表结构里查得到和表同名的列（visits 表的 visits 列）才当列，查不到时按整行算（多遮）。PIVOT、
|
||||
UNPIVOT、MATCH_RECOGNIZE 读到遮罩列的，整份快照的列都遮（多遮）。SQL 解析不了的只看文字：写着遮罩列、*（COUNT(*) 除外）、
|
||||
TABLE x 的整份遮，整行引用这类不写列名的写法认不出
|
||||
- 按遮罩列分档的标签（CASE、IIF 的结果全是写死的常量）：标签文字是 agent 写在 SQL 里的，不当遮罩值，SQL、回答正文里
|
||||
照原样；标签列当分组键（每一档一行）照常显示，不分组的逐行标签列整列遮（格子里看不出某一笔落在哪一档），可 agent
|
||||
在正文里写出「这一笔是大单」遮不住，按 id 分组再按档数个数也看得出（同按阈值查出某几行）。只有只装着标签的列不收：
|
||||
|
||||
+258
-26
@@ -13,10 +13,12 @@
|
||||
- 哪些列（_derived_hidden）：从遮罩列本身起，按 SQL 一层层认出值取自遮罩列的结果列（别名、函数、子查询、WITH、UNION 的
|
||||
另一支、和 * 写在同一支里没起别名的表达式、FROM 里拆开遮罩列的表值函数、函数参数里的 JSON 箭头 col->'path'）；* 把
|
||||
遮罩列带进别的列名的（写了列名表的派生来源、UNION 后一支的 *、标量子查询里的 *）按它展开落到的位置算（_star_places，
|
||||
展开数据源里的表的按查询当时冻结的表结构数位置，_schema_columns）。
|
||||
展开数据源里的表的按查询当时冻结的表结构数位置，_schema_columns；写了库名的表按「库名.表名」查，另一个库里的同名表
|
||||
查不到）。整行和函数参数里的 *（Postgres 的 row_to_json(s)、SELECT s、to_json(s.*)，Oracle 的 JSON_OBJECT(*)）按 *
|
||||
展开的来源算；表的列名表（FROM sales AS s(i, r, p)）、TABLE x、套了括号的 * 先改写成认得出的样子（_normalized）。
|
||||
COUNT 只数个数、窗口的分区键只决定次序、按遮罩列分档又当分组键或套在汇总里数个数的、汇总里的 CASE 只在条件里用到的
|
||||
不算（_value_columns、_band_only）。认不出落在哪一列的、SQL 解析不了又写着遮罩列的，整份快照的列都算。双盲复算按同一套
|
||||
认法在派盲算之前挡下取自遮罩列的指标（derived_columns）
|
||||
不算（_value_columns、_band_only）。认不出落在哪一列的（PIVOT、UNPIVOT 这类也算）、SQL 解析不了又写着遮罩列（或者 *、
|
||||
TABLE x）的，整份快照的列都算。双盲复算按同一套认法在派盲算之前挡下取自遮罩列的指标（derived_columns）
|
||||
- 只装着分档标签的列（_label_columns）：标签是 agent 写在 SQL 里的字，不当遮罩值；* 展开数据源里的表时按冻结的表结构
|
||||
认它带不带同名的列（_schema_tables）
|
||||
- 哪些值（_Secrets）：遮罩列里出现过的值，按值认有门槛（SECRET_MIN、SECRET_DIGITS），换的时候先按词、按开头几个字筛出
|
||||
@@ -35,7 +37,7 @@ import functools
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from collections.abc import Callable, Iterable
|
||||
from collections.abc import Callable, Iterable, Iterator
|
||||
from dataclasses import dataclass, field
|
||||
from decimal import Decimal
|
||||
from typing import Any
|
||||
@@ -87,7 +89,9 @@ def _value_columns(node: Any) -> set[str]:
|
||||
同写在 WHERE 里（SUM(CASE WHEN buyer_phone IS NOT NULL THEN amount ELSE 0 END) 的值全部来自金额）；不汇总的逐行
|
||||
CASE 照旧连条件一起算（每一行落在哪一支本身就是这一行的信息）。
|
||||
套在里面的子查询选的是 *（标量子查询、ARRAY(SELECT * …)）：值取自 * 展开的来源（_star_inputs），派生来源看它的各个
|
||||
选出项，表值函数看参数，展开了数据源里的表的记成 _ANY_COLUMN；EXISTS 里的 * 只说有没有，不算。函数参数里的
|
||||
选出项，表值函数看参数，展开了数据源里的表的记成 _ANY_COLUMN；EXISTS 里的 * 只说有没有，不算。写在函数参数、表达式
|
||||
里的 *（to_json(s.*)、concat_ws(',', s.*)、ROW(s.*)、Oracle 的 JSON_OBJECT(*)，Postgres 的整行引用改写成的 ROW(s.*)，
|
||||
见 _normalized）同样按它展开的来源算（_scoped_star_inputs）；COUNT(*)、count(s.*) 只数个数，不算。函数参数里的
|
||||
col->'path'（JSON 箭头）被解析器当成以 col 为参数的 lambda，参数名照列名收（真的 lambda 的参数多算一列无害）。"""
|
||||
from sqlglot import exp
|
||||
|
||||
@@ -97,8 +101,16 @@ def _value_columns(node: Any) -> set[str]:
|
||||
expanded = {id(node)}
|
||||
while stack:
|
||||
item = stack.pop()
|
||||
if isinstance(item, exp.Column):
|
||||
if not isinstance(item.this, exp.Star) and item.name:
|
||||
if isinstance(item, exp.Star) or (isinstance(item, exp.Column) and isinstance(item.this, exp.Star)):
|
||||
inputs = _scoped_star_inputs(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)
|
||||
elif isinstance(item, exp.Column):
|
||||
if item.name:
|
||||
out.add(item.name.lower())
|
||||
elif isinstance(item, exp.Count) or (isinstance(item, (exp.Case, exp.If)) and _band_only(item)):
|
||||
continue
|
||||
@@ -217,6 +229,8 @@ def _label_columns(sql: Any, kind: str | None, columns: list[str], tables: dict[
|
||||
if root is None:
|
||||
return {}
|
||||
try:
|
||||
# 和认遮罩列时一样先改写（套了括号的 * 按 * 数位置）。这里的表结构只有列名的集合，认整行引用时不看它（多遮）
|
||||
_normalized(root, _read(kind), None)
|
||||
result = next((top for top, role, _listed in _blocks(root) if role == "result"), None)
|
||||
if result is None:
|
||||
return {}
|
||||
@@ -243,9 +257,23 @@ class _LabelMemo(dict): # type: ignore[type-arg]
|
||||
tables: dict[str, set[str]] | None = None
|
||||
|
||||
|
||||
def _schema_keys(key: Any, meta: dict[str, Any]) -> set[str]:
|
||||
"""表结构里的一张表按哪些名字查（小写，查法见 _table_key）：没写库名的引用按登记的名字、不带库名的表名查；写了库名的
|
||||
（archive.sales）按「库名.表名」查，库名取表结构记下的 schema 和 qualified 的前缀。表结构探查的是一个库（schema），
|
||||
查询写了别的库名的同名表是另一张表（列的次序可能不一样），查不到，按说不清处理；探查时没设 schema 的，表结构里没有
|
||||
库名，写了库名的一律查不到（宁可多遮）。"""
|
||||
name = str(key).lower()
|
||||
qualified = str(meta.get("qualified") or key).lower()
|
||||
bare = qualified.rsplit(".", 1)[-1]
|
||||
keys = {name, name.rsplit(".", 1)[-1], bare, qualified}
|
||||
if meta.get("schema"):
|
||||
keys.add(f"{str(meta['schema']).lower()}.{bare}")
|
||||
return keys
|
||||
|
||||
|
||||
def _schema_tables(schema: Any) -> dict[str, set[str]] | None:
|
||||
"""冻结的表结构（数据源的 schema_cache、表结构快照的内容：{"tables": {表名: {"columns": [{"name": …}]}}}）→
|
||||
{表名（小写）: 列名（小写）}，按登记的名字和不带库名的表名各记一份（几个库里同名的表合起来算，宁可多遮）。表结构
|
||||
{表名（小写）: 列名（小写）}，按 _schema_keys 的几个名字各记一份（几个库里同名的表合起来算，宁可多遮）。表结构
|
||||
里没列出的表不记（截断了的表结构也一样），查不到就当说不清；整份拿不到的为 None。"""
|
||||
tables = schema.get("tables") if isinstance(schema, dict) else None
|
||||
if not isinstance(tables, dict):
|
||||
@@ -256,15 +284,14 @@ def _schema_tables(schema: Any) -> dict[str, set[str]] | None:
|
||||
if not isinstance(columns, list):
|
||||
continue
|
||||
names = {str(c["name"]).lower() for c in columns if isinstance(c, dict) and c.get("name") is not None}
|
||||
for name in {str(key).lower(), str(key).lower().rsplit(".", 1)[-1],
|
||||
str(meta.get("qualified") or key).lower().rsplit(".", 1)[-1]}:
|
||||
for name in _schema_keys(key, meta):
|
||||
out.setdefault(name, set()).update(names)
|
||||
return out
|
||||
|
||||
|
||||
def _schema_columns(schema: Any) -> dict[str, list[str]] | None:
|
||||
"""冻结的表结构 → {表名（小写）: 按表里次序的列名（小写）}，名字的记法同 _schema_tables；几个库里同名、列又不一样的表
|
||||
不记（说不清 * 展开的是哪一张）。* 展开数据源里的表时按它数位置（_star_columns）；整份拿不到的为 None。"""
|
||||
不按表名记（说不清 * 展开的是哪一张）。* 展开数据源里的表时按它数位置（_star_columns）；整份拿不到的为 None。"""
|
||||
tables = schema.get("tables") if isinstance(schema, dict) else None
|
||||
if not isinstance(tables, dict):
|
||||
return None
|
||||
@@ -275,20 +302,26 @@ def _schema_columns(schema: Any) -> dict[str, list[str]] | 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]}:
|
||||
for name in _schema_keys(key, meta):
|
||||
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_key(source: Any) -> str:
|
||||
"""FROM 里一张数据源里的表按什么名字查表结构（_schema_keys）：没写库名的是表名，写了的连着库名（archive.sales，写了
|
||||
三段的连着 catalog），小写。"""
|
||||
return ".".join(part for part in (source.catalog, source.db, source.name) if part).lower()
|
||||
|
||||
|
||||
def _table_lacks(source: Any, name: str, memo: _LabelMemo) -> bool:
|
||||
"""FROM 里的这个来源是数据源里的一张表，冻结的表结构里查得到它、又没有 name 这一列：* 展开它带不出 name。"""
|
||||
"""FROM 里的这个来源是数据源里的一张表，冻结的表结构里查得到它、又没有 name 这一列：* 展开它带不出 name。写了库名的
|
||||
按「库名.表名」查（_table_key）。"""
|
||||
from sqlglot import exp
|
||||
|
||||
if not memo.tables or not isinstance(source, exp.Table) or not isinstance(source.this, exp.Identifier):
|
||||
return False
|
||||
columns = memo.tables.get(source.name.lower())
|
||||
columns = memo.tables.get(_table_key(source))
|
||||
return columns is not None and name not in columns
|
||||
|
||||
|
||||
@@ -669,6 +702,32 @@ def _star_inputs(star: Any, select: Any) -> list[Any] | None:
|
||||
return out
|
||||
|
||||
|
||||
def _outer_selects(node: Any) -> Iterator[Any]:
|
||||
"""由里往外，node 所在的位置看得见 FROM 里的来源的那几层 SELECT。CTE、FROM / JOIN 里的子查询（不是 LATERAL）看不见
|
||||
同一层 FROM 里的别的来源：WITH、FROM 属于的那一层跳过，再外层的照样看得见（相关子查询）。"""
|
||||
from sqlglot import exp
|
||||
|
||||
skip = None
|
||||
node = node.parent
|
||||
while node is not None:
|
||||
if isinstance(node, exp.Select) and node is not skip:
|
||||
yield node
|
||||
if isinstance(node, exp.CTE) or (isinstance(node, exp.Subquery)
|
||||
and isinstance(node.parent, (exp.From, exp.Join))):
|
||||
skip = node.parent.parent
|
||||
node = node.parent
|
||||
|
||||
|
||||
def _scoped_star_inputs(star: Any) -> list[Any] | None:
|
||||
"""写在表达式里的 *（函数参数里的 s.*、JSON_OBJECT(*)）带出的值取自哪里（_star_inputs）：由里往外找 FROM 里有这个来源的
|
||||
那一层 SELECT（_outer_selects，相关子查询里的 s.* 指外层的 s），没写表名的 * 展开最近一层的全部来源。找不到为 None
|
||||
（说不清，按取自遮罩列算）。"""
|
||||
for select in _outer_selects(star):
|
||||
if _star_sources(star, select) is not None:
|
||||
return _star_inputs(star, select)
|
||||
return None
|
||||
|
||||
|
||||
def _star_feeds(star: Any, select: Any, names: set[str]) -> bool:
|
||||
"""select 里的这个 * 带出的值里有没有取自 names 的（展开了数据源里的表的一律算，见 _star_inputs）。"""
|
||||
inputs = _star_inputs(star, select)
|
||||
@@ -676,12 +735,13 @@ def _star_feeds(star: Any, select: Any, names: set[str]) -> bool:
|
||||
|
||||
|
||||
def _table_columns(source: Any, tables: dict[str, list[str]] | None) -> list[str] | None:
|
||||
"""FROM 里的这个来源是数据源里的一张表、冻结的表结构里查得到它时，它按次序的列名（_schema_columns）；别的为 None。"""
|
||||
"""FROM 里的这个来源是数据源里的一张表、冻结的表结构里查得到它时，它按次序的列名（_schema_columns）；别的为 None。
|
||||
写了库名的按「库名.表名」查（_table_key）：另一个库里的同名表查不到，不拿探查的那张表的列序去数。"""
|
||||
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())
|
||||
return tables.get(_table_key(source))
|
||||
|
||||
|
||||
def _single_source(star: Any, select: Any, depth: int) -> Any:
|
||||
@@ -764,6 +824,12 @@ def _branch_names(select: Any, tables: dict[str, list[str]] | None, depth: int =
|
||||
return out
|
||||
|
||||
|
||||
def _width(select: Any, tables: dict[str, list[str]] | None) -> int | None:
|
||||
"""一支（一条 SELECT）结果的列数，按 _branch_names 数；* 说不清展开成几列的为 None。"""
|
||||
names = _branch_names(select, tables)
|
||||
return len(names) if names is not None else None
|
||||
|
||||
|
||||
def _star_places(select: Any, width: int | None, names: set[str], tables: dict[str, list[str]] | None = None
|
||||
) -> set[int] | None:
|
||||
"""select 这一支里的 * 把取自 names 的值带到了结果（width 列）的哪几列。* 展开的列落在第一个 * 起、最后一个 * 后面
|
||||
@@ -788,6 +854,145 @@ def _star_places(select: Any, width: int | None, names: set[str], tables: dict[s
|
||||
return set(range(first, end))
|
||||
|
||||
|
||||
#: 有整行引用的方言（sqlglot 的方言名）：不带表名的名字是 FROM 里一个来源时，可能指它的整行（Postgres 的
|
||||
#: row_to_json(s)、SELECT s，DuckDB 同）。SQLite、MySQL、Oracle 里这样写只能是列
|
||||
_ROW_DIALECTS = frozenset({"postgres", "duckdb"})
|
||||
|
||||
|
||||
def _normalized(root: Any, read: str | None, tables: dict[str, list[str]] | None) -> None:
|
||||
"""几种写法先就地改写成认得出的样子，再认遮罩列、标签列（_tainted、_label_columns）。只改解析出来的语法树，原文不动；
|
||||
改写前后结果有哪几列、每一列的值取自哪里都一样：
|
||||
- TABLE x（Postgres、MySQL）：解析器认不出，CTE、子查询里的读成了一个叫 TABLE 的列、表。改成 SELECT * FROM x；带库名、
|
||||
ORDER BY 的解析不了，按解析不了处理（_implicit）
|
||||
- 选出项套了括号的 *（(s.*)）、(s).*：Postgres 照样展开成好几列，改成 s.*，不然按一列数位置。括号里不是来源的
|
||||
（(t.profile).* 展开一个复合类型的列）改成展开一个对不上任何来源的 *：说不清展开成几列，整段遮
|
||||
- 带列名表的表（Postgres 的 FROM sales AS s(i, r, p)，引用 CTE 时再写列名表的也算）：改成 (SELECT * FROM sales) AS
|
||||
s(i, r, p)，按派生来源的列名表认改了名的列。Postgres 的列名表可以比列少（只改前几列的名字，后面的照原名），这时
|
||||
* 展开的列数和列名表对不上，列名表里的几列整段算（_star_places），后面照原名的按名字认。不按表结构补全列名表：
|
||||
别的方言里列名表比列少只能是表结构过期了，补全就丢了这一道核对
|
||||
- 整行引用（read 是 _ROW_DIALECTS 之一，_row_source）：row_to_json(s)、json_agg(t)、SELECT s、s::text 里的 s 改成
|
||||
ROW(s.*)，按 * 展开的来源认（_value_columns）
|
||||
tables 是查询当时冻结的表结构（_schema_columns，按次序的列名），认整行引用时看来源出不出某一列；没有为 None。"""
|
||||
from sqlglot import exp
|
||||
|
||||
for cte in list(root.find_all(exp.CTE)):
|
||||
body = cte.this
|
||||
if isinstance(body, exp.Alias) and _table_keyword(body.this) and isinstance(body.args.get("alias"),
|
||||
exp.Identifier):
|
||||
cte.set("this", exp.select("*").from_(exp.Table(this=body.args["alias"].copy())))
|
||||
for sub in list(root.find_all(exp.Subquery)):
|
||||
body = sub.this
|
||||
alias = body.args.get("alias") if isinstance(body, exp.Table) else None
|
||||
if alias is not None and _table_keyword(body.this) and not alias.columns \
|
||||
and isinstance(alias.this, exp.Identifier):
|
||||
sub.set("this", exp.select("*").from_(exp.Table(this=alias.this.copy())))
|
||||
selects = list(root.find_all(exp.Select))
|
||||
for select in selects:
|
||||
for item in list(select.expressions):
|
||||
if (star := _bare_star(item)) is not None:
|
||||
item.replace(star)
|
||||
for select in selects:
|
||||
for _name, source in _from_sources(select):
|
||||
alias = source.args.get("alias")
|
||||
if isinstance(source, exp.Table) and isinstance(source.this, exp.Identifier) and alias is not None \
|
||||
and alias.columns:
|
||||
inner = source.copy()
|
||||
inner.set("alias", None)
|
||||
source.replace(exp.Subquery(this=exp.select("*").from_(inner), alias=alias.copy()))
|
||||
if read in _ROW_DIALECTS:
|
||||
for column in list(root.find_all(exp.Column)):
|
||||
if column.table or not isinstance(column.this, exp.Identifier) or not _carries_value(column):
|
||||
continue
|
||||
source = _row_source(column, tables)
|
||||
if source is None:
|
||||
continue
|
||||
row = exp.Anonymous(this="ROW", expressions=[exp.Column(this=exp.Star(), table=exp.to_identifier(source))])
|
||||
# 直接是选出项的起回原名：结果的列名照旧（SELECT s 的列叫 s）
|
||||
bare = isinstance(column.parent, exp.Select) and column.arg_key == "expressions"
|
||||
column.replace(exp.alias_(row, column.this.copy()) if bare else row)
|
||||
|
||||
|
||||
def _table_keyword(node: Any) -> bool:
|
||||
"""解析器把 TABLE x 读成的那个名字：不带引号、不带表名的 TABLE。"""
|
||||
from sqlglot import exp
|
||||
|
||||
ident = node.this if isinstance(node, exp.Column) and not node.table else node
|
||||
return isinstance(ident, exp.Identifier) and not ident.quoted and ident.name.upper() == "TABLE"
|
||||
|
||||
|
||||
def _bare_star(item: Any) -> Any:
|
||||
"""选出项是套了括号的 *（(s.*)、((*))）或 (s).* 时，改写成的 *、s.*；别的为 None（_normalized）。"""
|
||||
from sqlglot import exp
|
||||
|
||||
node = item
|
||||
while isinstance(node, exp.Paren):
|
||||
node = node.this
|
||||
if node is not item and _is_star(node):
|
||||
return node.copy()
|
||||
if not (isinstance(item, exp.Dot) and isinstance(item.expression, exp.Star)):
|
||||
return None
|
||||
inner = item.this
|
||||
while isinstance(inner, exp.Paren):
|
||||
inner = inner.this
|
||||
if isinstance(inner, exp.Column) and not inner.table and isinstance(inner.this, exp.Identifier):
|
||||
return exp.Column(this=exp.Star(), table=inner.this.copy())
|
||||
# 展开的不是 FROM 里的来源（复合类型的列、函数交回的记录）：记成一个对不上任何来源的 *（_star_sources 为 None）
|
||||
return exp.Column(this=exp.Star(), table=exp.to_identifier("*"))
|
||||
|
||||
|
||||
def _carries_value(column: Any) -> bool:
|
||||
"""列引用写在选出项里、FROM 里（函数参数、LATERAL）：值会流进结果。WHERE、GROUP BY、ORDER BY、JOIN 的 ON / USING
|
||||
里的不算（_normalized 只改写会带出值的整行引用）。"""
|
||||
from sqlglot import exp
|
||||
|
||||
node = column
|
||||
while node.parent is not None and not isinstance(node.parent, exp.Select):
|
||||
if isinstance(node.parent, exp.Join) and node.arg_key in ("on", "using"):
|
||||
return False
|
||||
node = node.parent
|
||||
return node.parent is not None and node.arg_key in ("expressions", "from", "from_", "joins", "laterals")
|
||||
|
||||
|
||||
def _row_source(column: Any, tables: dict[str, list[str]] | None) -> str | None:
|
||||
"""Postgres 里不带表名的名字 column 指整行时，它指的来源的名字（别名，没起别名的是表名）；是列、不是任何来源的为 None。
|
||||
Postgres 先按列认：由里往外看得见的每一层（_outer_selects）的来源里，有一个出得了这一列（_has_column）就是列；都不是、
|
||||
又有一个来源叫这个名字，才是整行（取最里面的那个）。说不清出不出这一列的（表结构里查不到的表、* 展开的列）按出不了算，
|
||||
当整行（宁可多遮）。"""
|
||||
name = column.name.lower()
|
||||
found: str | None = None
|
||||
for select in _outer_selects(column):
|
||||
for alias, source in _from_sources(select):
|
||||
if _has_column(source, name, tables):
|
||||
return None
|
||||
if found is None and alias == name:
|
||||
found = source.alias_or_name
|
||||
return found
|
||||
|
||||
|
||||
def _has_column(source: Any, name: str, tables: dict[str, list[str]] | None) -> bool:
|
||||
"""FROM 里的这个来源说得准出 name 这一列：派生来源看列名表、第一支里起了名字的选出项（列名表改了名的不算），数据源里
|
||||
的表看表结构（_table_columns），表值函数这类看它出的列名（_function_columns）。* 展开的列说不准，不算。"""
|
||||
from sqlglot import exp
|
||||
|
||||
derived = _derived_block(source)
|
||||
if derived is not None:
|
||||
block, listed = derived
|
||||
first = _leaves(block)[0]
|
||||
items = list(first.expressions) if isinstance(first, exp.Select) else []
|
||||
if listed is not None:
|
||||
if name in listed:
|
||||
return True
|
||||
if any(_is_star(item) for item in items):
|
||||
return False
|
||||
items = items[len(listed):]
|
||||
return any(_output_name(item) == name for item in items if not _is_star(item))
|
||||
if isinstance(source, exp.Table) and isinstance(source.this, exp.Identifier):
|
||||
columns = _table_columns(source, tables)
|
||||
return columns is not None and name in columns
|
||||
columns = _function_columns(source)
|
||||
return columns is not None and name in columns
|
||||
|
||||
|
||||
def _blocks(root: Any) -> list[tuple[Any, str, list[str] | None]]:
|
||||
"""SQL 里的每一块查询（一条 SELECT，或者 UNION 这类连起来的一串）和它的结果流到哪：
|
||||
- "result"：最外层，结果就是快照的列；
|
||||
@@ -825,7 +1030,8 @@ def _source_names(top: Any, listed: list[str] | None, names: set[str],
|
||||
写了列名表的按位置对上列名表；没写的，单独一条 SELECT 取各项自己的列名，UNION 这类按位置对上第一支的列名。
|
||||
* 带出遮罩列的值、展开的列又换了名字的（写了列名表、UNION 后一支的 *），它展开落到的那几列照位置对上（_star_places）；
|
||||
单独一条 SELECT、第一支的 * 不改名，遮罩列照原名传出去，按名字就认得出。第一支的 * 展开成哪几列数得清的（_branch_names），
|
||||
后一支落进去的位置按它对上列名。tables 是查询当时冻结的表结构（_schema_columns），* 展开数据源里的表时按它数位置。"""
|
||||
后面每一支数出的列数又都和它一样宽时，后一支落进去的位置按它对上列名；对不上的（表结构过期了）不按它对。tables 是
|
||||
查询当时冻结的表结构（_schema_columns），* 展开数据源里的表时按它数位置。"""
|
||||
from sqlglot import exp
|
||||
|
||||
selects = [leaf for leaf in _leaves(top) if isinstance(leaf, exp.Select)]
|
||||
@@ -858,6 +1064,11 @@ def _source_names(top: Any, listed: list[str] | None, names: set[str],
|
||||
take(_output_name(item))
|
||||
if len(selects) > 1:
|
||||
header = _branch_names(selects[0], tables)
|
||||
# 第一支的 * 按表结构数出来的列名：后面每一支数得出的列数都和它一样宽才按它对位置。表结构过期了（这张表后来多了、
|
||||
# 少了列），按它数出的列名和实际的不一样宽，对上的名字是错的；有一支对不上、数不清的，退回不按表结构数（* 展开的
|
||||
# 位置认不出，整段遮）
|
||||
if header is not None and any(_width(select, tables) != len(header) for select in selects[1:]):
|
||||
header = None
|
||||
if header is not None:
|
||||
width, heads = len(header), dict(enumerate(header))
|
||||
else:
|
||||
@@ -883,8 +1094,11 @@ def _tainted(sql: str, kind: str | None, masked: set[str], width: int, tables: d
|
||||
FROM、JOIN、LATERAL 里的表值函数（json_each、unnest、JSON_TABLE、LATERAL VALUES）参数里写着这些名字的，它出的列
|
||||
都算（_function_columns）：遮罩列是 JSON 时，json_each 拆开是最自然的写法。出哪几列说不清的整份快照都算。
|
||||
* 展开的列换了名字（写了列名表的派生来源、UNION 后一支的 *）：* 带出这些名字的值时，它展开落到的那几列照位置算
|
||||
（_star_places，展开数据源里的表的按 tables 数位置，数不清的整段算）；标量子查询里的 * 按 _value_columns 算进外层的
|
||||
表达式。"""
|
||||
（_star_places，展开数据源里的表的按 tables 数位置，数不清的整段算）；标量子查询里的 *、函数参数里的 *（to_json(s.*)、
|
||||
JSON_OBJECT(*)）、整行引用（row_to_json(s)）按 _value_columns 算进外层的表达式。表的列名表、TABLE x、套了括号的 *
|
||||
先改写成上面这些认得出的样子（_normalized）。
|
||||
PIVOT、UNPIVOT、MATCH_RECOGNIZE 读到这些名字的：值转进了它们自己起名的列（UNPIVOT 的值列、PIVOT 按 IN 里的值起名的列、
|
||||
MEASURES），出的是哪几列不好认，整份快照的列都算。"""
|
||||
from sqlglot import exp
|
||||
|
||||
from app.data.sql_semantics import _parse, _read
|
||||
@@ -896,6 +1110,7 @@ def _tainted(sql: str, kind: str | None, masked: set[str], width: int, tables: d
|
||||
if root is None:
|
||||
return None
|
||||
try:
|
||||
_normalized(root, _read(kind), tables)
|
||||
names = set(masked)
|
||||
selects = list(root.find_all(exp.Select))
|
||||
blocks = _blocks(root)
|
||||
@@ -918,7 +1133,8 @@ def _tainted(sql: str, kind: str | None, masked: set[str], width: int, tables: d
|
||||
names |= columns or set()
|
||||
if len(names) == before:
|
||||
break
|
||||
everything = unreadable or any(_source_names(top, listed, names, tables)[1] for top, listed in sources)
|
||||
everything = unreadable or any(_source_names(top, listed, names, tables)[1] for top, listed in sources) \
|
||||
or any(_reads(names, _value_columns(node)) for node in root.find_all(exp.Pivot, exp.MatchRecognize))
|
||||
positions: set[int] = set()
|
||||
for top, role, _listed in blocks:
|
||||
if role != "result":
|
||||
@@ -948,14 +1164,28 @@ 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)
|
||||
|
||||
|
||||
#: COUNT(*)：只数个数，不把列带出来（_implicit 先去掉它）
|
||||
_COUNT_STAR = re.compile(r"\bcount\s*\(\s*(?:all\s+|distinct\s+)?\*\s*\)", re.I)
|
||||
#: SQL 文字里不写列名也会带出遮罩列的写法：选出项的 *、t.*、函数参数里的 *，TABLE x 语句。乘号前面是值（列名、数、右
|
||||
#: 括号），对不上
|
||||
_IMPLICIT = re.compile(r"(?:\bselect|\bdistinct|\ball|[,(.])\s*\*|\btable\s+[\w\"`]", re.I)
|
||||
|
||||
|
||||
def _implicit(sql: str) -> bool:
|
||||
"""解析不了的 SQL 里有没有不写列名也会带出遮罩列的写法（*、TABLE x）：它们能把遮罩列带进别的列名（列名表、UNION 的
|
||||
后一支），同写着遮罩列一样整份遮（_derived_hidden）。只看文字，字符串、注释里的也算（多遮）。"""
|
||||
return bool(_IMPLICIT.search(_COUNT_STAR.sub("", 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 解析不了、里面又写着遮罩列的，整份快照的
|
||||
列都遮掉；合并查询的快照不看 SQL，它记着自己的来历（_hidden_columns 已按 _merge_masks 遮好）。"""
|
||||
（region AS 区域），所以按 SQL 认出改了名字的遮罩列（_tainted）。SQL 解析不了、里面又写着遮罩列的（*、TABLE x 不写
|
||||
列名也会把遮罩列带进别的列名，同样算，_implicit），整份快照的列都遮掉；合并查询的快照不看 SQL，它记着自己的来历
|
||||
（_hidden_columns 已按 _merge_masks 遮好）。"""
|
||||
base = set(hidden)
|
||||
sql = snapshot.get("sql")
|
||||
if not base or not isinstance(sql, str) or not sql.strip() or snapshot.get("source") == MERGE_SOURCE:
|
||||
@@ -963,7 +1193,8 @@ def _derived_hidden(snapshot: dict[str, Any], hidden: set[str], kind: str | None
|
||||
columns = [str(c) for c in snapshot.get("columns") or []]
|
||||
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)
|
||||
hides = _mentions(sql, base) or _implicit(sql)
|
||||
return (base | {c.lower() for c in columns}, base) if hides else (base, base)
|
||||
names, positions = found
|
||||
return names | {columns[i].lower() for i in positions if i < len(columns)}, names
|
||||
|
||||
@@ -1779,8 +2010,9 @@ def mask_query_result(payload: dict[str, Any], masked: Iterable[str], kind: str
|
||||
（_derived_hidden 的 tables）时用。
|
||||
|
||||
- 整列遮：列名直接对上遮罩列的，加上按 SQL 认出值取自遮罩列的（_derived_hidden：别名、函数、子查询、WITH、
|
||||
SELECT *（展开的列换了名字的也算）、表值函数；COUNT、只在条件里用到、分档当分组键的不算）；SQL 解析不了、又写着
|
||||
遮罩列的，整份结果的列都遮。这几列的格子换成「已遮罩」，类型记成文本
|
||||
SELECT *（展开的列换了名字的也算）、表值函数、整行引用和函数参数里的 *、列名表、TABLE x；COUNT、只在条件里用到、
|
||||
分档当分组键的不算）；SQL 解析不了、又写着遮罩列（或者 *、TABLE x）的，整份结果的列都遮。这几列的格子换成
|
||||
「已遮罩」，类型记成文本
|
||||
- 别的列按值换：整格就是遮罩值的换掉，文字格里夹着的那一截换掉（换了写法的号码也认，_Secrets.loose_spans）。遮罩值
|
||||
是这份结果里的加上 known：先查号码、再查写着号码的备注、没设遮罩的联系方式表（别的库里的也一样）时，后一份同样
|
||||
遮好。门槛同按值遮（SECRET_DIGITS、SECRET_MIN）；只装着分档标签的列里的标签不算遮罩值，这份结果自己写的标签也不
|
||||
|
||||
@@ -249,3 +249,130 @@ async def test_a_snapshot_is_read_with_its_frozen_schema(months):
|
||||
assert masking._derived_hidden(snapshot, {"buyer_phone"}, "sqlite")[0] == {"buyer_phone"}
|
||||
assert masking.derived_columns({**snapshot, "schema_artifact": "missing"}, {"buyer_phone"}, "sqlite") == {
|
||||
"region", "合计"}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 写了库名的表、过期的表结构（复核第 3 轮）
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
#: 同名的两张表：主库的 sales 和另一个库（archive）里的 sales，号码、区域换了次序。表结构探查的是主库
|
||||
TWO_SALES = ("CREATE TABLE sales (id INTEGER, region TEXT, buyer_phone TEXT, buyer_name TEXT, amount INTEGER, "
|
||||
"note TEXT);"
|
||||
"INSERT INTO sales VALUES (1, '华东', '13900000001', '张伟', 1200, '改址');"
|
||||
"CREATE TABLE archive.sales (id INTEGER, buyer_phone TEXT, region TEXT, buyer_name TEXT, amount INTEGER, "
|
||||
"note TEXT);"
|
||||
"INSERT INTO archive.sales VALUES (7, '13900000007', '华北', '王芳', 500, '回访');")
|
||||
SALES_LAYOUT = ["id", "region", "buyer_phone", "buyer_name", "amount", "note"]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def two_sales():
|
||||
conn = sqlite3.connect(":memory:")
|
||||
conn.execute("ATTACH DATABASE ':memory:' AS archive")
|
||||
conn.executescript(TWO_SALES)
|
||||
yield conn
|
||||
conn.close()
|
||||
|
||||
|
||||
def probed(schema: str | None = None) -> dict:
|
||||
"""探查主库得到的表结构（schema_cache 的样子）：没设 schema 的探查当前库，qualified 就是表名。"""
|
||||
meta = {"qualified": f"{schema}.sales" if schema else "sales", "schema": schema,
|
||||
"columns": [{"name": c} for c in SALES_LAYOUT]}
|
||||
return {"tables": {"sales": meta}, "schema": schema}
|
||||
|
||||
|
||||
def run_two(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, MASKED, "sqlite", schema=schema)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql", [
|
||||
pytest.param("WITH t(a, b, c, d, e, f) AS (SELECT * FROM archive.sales) SELECT b AS j, a AS n FROM t",
|
||||
id="列名表套 * 展开另一个库里的同名表"),
|
||||
pytest.param("SELECT 0 AS a, '' AS b, '' AS c, '' AS d, 0 AS e, '' AS f WHERE 0 UNION ALL "
|
||||
"SELECT * FROM archive.sales", id="UNION 后一支的 * 展开另一个库里的同名表"),
|
||||
pytest.param("WITH t AS (SELECT * FROM sales UNION ALL SELECT * FROM archive.sales) SELECT region AS 区域 FROM t",
|
||||
id="派生来源里 UNION 后一支展开另一个库里的同名表"),
|
||||
])
|
||||
def test_a_table_in_another_schema_is_not_counted_by_the_probed_one(two_sales, sql):
|
||||
"""表结构探查的是主库：查询写了别的库名的同名表（archive.sales，号码排在第二列），不能按主库那张表的列序数位置，
|
||||
按数不清处理，* 展开落到的那一段整段遮。以前按表名查表结构，号码落进的那一列被当成区域，原样交给模型、落库。"""
|
||||
done = run_two(two_sales, sql, probed())
|
||||
assert not _leaks(done), done.payload["rows"]
|
||||
|
||||
|
||||
def test_a_table_named_with_its_probed_schema_is_still_counted(two_sales):
|
||||
"""库名和表结构记下的对得上（探查时设了 schema、qualified 带着它）：照样按表结构数位置，只遮号码那一列。没写库名的
|
||||
照旧按表名查。库名对不上表结构记下的（这里探查时没设 schema，查询写了 main.）按数不清处理，整段遮（多遮）。"""
|
||||
precise = "WITH t(a, b, c, d, e, f) AS (SELECT * FROM {}) SELECT c AS j, a AS n FROM t"
|
||||
assert run_two(two_sales, precise.format("main.sales"), probed("main")).columns == ["j"]
|
||||
assert run_two(two_sales, precise.format("sales"), probed()).columns == ["j"]
|
||||
assert set(run_two(two_sales, precise.format("main.sales"), probed()).columns) == {"j", "n"}
|
||||
|
||||
|
||||
def test_schema_columns_are_also_keyed_by_the_probed_schema():
|
||||
"""表结构按「库名.表名」另记一份（库名取 schema、qualified 的前缀）；写了别的库名的查不到。"""
|
||||
tables = masking._schema_columns(probed("public"))
|
||||
assert tables["sales"] == tables["public.sales"] == SALES_LAYOUT and "archive.sales" not in tables
|
||||
assert "public.sales" not in masking._schema_columns(probed())
|
||||
assert "archive.sales" not in (masking._schema_tables(probed("public")) or {})
|
||||
|
||||
|
||||
#: 一月的表结构是早先探查的（那时还没有 extra 这一列），表里现在多了一列，排在第二；二月那张号码排在第二列
|
||||
STALE = ("CREATE TABLE jan (id INTEGER, extra TEXT, region TEXT, buyer_phone TEXT, amount INTEGER);"
|
||||
"CREATE TABLE feb (id INTEGER, buyer_phone TEXT, region TEXT, amount INTEGER, extra TEXT);"
|
||||
"INSERT INTO jan VALUES (1, '加急', '华东', '13900000001', 1200);"
|
||||
"INSERT INTO feb VALUES (2, '13900000002', '华南', 300, '普通');")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def stale():
|
||||
conn = sqlite3.connect(":memory:")
|
||||
conn.executescript(STALE)
|
||||
yield conn
|
||||
conn.close()
|
||||
|
||||
|
||||
def layout(**tables: list[str]) -> dict:
|
||||
return {"tables": {name: {"qualified": name, "columns": [{"name": c} for c in columns]}
|
||||
for name, columns in tables.items()}}
|
||||
|
||||
|
||||
#: 过期的表结构：一月那张少了 extra（列数和表里实际的不一样）
|
||||
OLD_JAN = layout(jan=["id", "region", "buyer_phone", "amount"], feb=["id", "buyer_phone", "region", "amount", "extra"])
|
||||
#: 准的表结构
|
||||
NOW = layout(jan=["id", "extra", "region", "buyer_phone", "amount"],
|
||||
feb=["id", "buyer_phone", "region", "amount", "extra"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql", [
|
||||
pytest.param("WITH t AS (SELECT * FROM jan UNION ALL SELECT id, buyer_phone, region, amount, extra FROM feb) "
|
||||
"SELECT extra AS 备注 FROM t", id="后一支明写各列"),
|
||||
pytest.param("WITH t AS (SELECT * FROM jan UNION ALL SELECT * FROM feb) SELECT extra AS 备注 FROM t",
|
||||
id="后一支也是 *"),
|
||||
])
|
||||
def test_a_first_branch_counted_by_a_stale_schema_is_not_trusted(stale, sql):
|
||||
"""UNION 第一支是 *、按表结构数出的列名（表结构过期了，少一列）和后一支的列数对不上：不按它对列名，* 所在的位置按认不出
|
||||
算，整份遮。以前照过期的列名表对，二月的号码落进的「extra」没认出来，原样交回。"""
|
||||
cursor = stale.execute(sql)
|
||||
payload = {"columns": [d[0] for d in cursor.description], "rows": [list(r) for r in cursor.fetchall()], "sql": sql}
|
||||
for schema in (OLD_JAN, None, NOW):
|
||||
done = masking.mask_query_result(payload, ["buyer_phone"], "sqlite", schema=schema)
|
||||
assert done.columns == ["备注"] and not _leaks(done), (schema, done.payload["rows"])
|
||||
|
||||
|
||||
def test_a_first_branch_counted_by_an_accurate_schema_still_places_the_values(stale):
|
||||
"""表结构准、每一支数得出的列数都对得上：照样按第一支的列名对上位置，只遮号码落进的那一列，金额照常。"""
|
||||
sql = ("WITH t AS (SELECT * FROM jan UNION ALL SELECT id, buyer_phone, region, buyer_phone, amount FROM feb) "
|
||||
"SELECT region AS 区域, amount AS 金额 FROM t")
|
||||
cursor = stale.execute(sql)
|
||||
payload = {"columns": [d[0] for d in cursor.description], "rows": [list(r) for r in cursor.fetchall()], "sql": sql}
|
||||
assert masking.mask_query_result(payload, ["buyer_phone"], "sqlite", schema=NOW).columns == []
|
||||
moved = sql.replace("SELECT id, buyer_phone, region,", "SELECT id, region, buyer_phone,")
|
||||
cursor = stale.execute(moved)
|
||||
payload = {"columns": [d[0] for d in cursor.description], "rows": [list(r) for r in cursor.fetchall()],
|
||||
"sql": moved}
|
||||
done = masking.mask_query_result(payload, ["buyer_phone"], "sqlite", schema=NOW)
|
||||
assert done.columns == ["区域"] and [row[1] for row in done.payload["rows"]] == [1200, 300]
|
||||
|
||||
@@ -0,0 +1,244 @@
|
||||
"""按 SQL 认取自遮罩列的结果列（engine/masking.py 的 _derived_hidden）：整行、函数参数里的 *、列名表把遮罩列带进别的列名
|
||||
（复核第 3 轮）。
|
||||
|
||||
Postgres 里不带表名、名字是 FROM 里一个来源的列引用指它的整行（row_to_json(s)、json_agg(t)、SELECT s、s::text），函数
|
||||
参数里的 s.*（to_json(s.*)、concat_ws(',', s.*)、ROW(s.*)）也把整行交给函数；数据源里的表可以带列名表（FROM sales AS
|
||||
s(i, r, p)，只改前几列的名字），TABLE sales 等于 SELECT * FROM sales；Oracle 的 JSON_OBJECT(*) 把整行拼成 JSON，UNPIVOT、
|
||||
PIVOT、MATCH_RECOGNIZE 把遮罩列的值转进它们自己起名的列。这些写法以前都认不出：遮罩列的值落进了别的列名，原样交给核查
|
||||
agent、落进快照。现在整行、函数参数里的 * 按 * 展开的来源认（数据源里的表一律算带出了遮罩列），列名表按位置对上（表结构
|
||||
查得到这张表的只遮号码改了名字的那一列），TABLE x 当 SELECT * FROM x，PIVOT 这类读到遮罩列的整份遮。
|
||||
|
||||
这里没有真的 Postgres、Oracle：结果按它们交回的样子造（行转成 JSON、记录转成文字），只看遮完的结果里还有没有原值。数据
|
||||
都是虚构的：区域、手机号、姓名、金额。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sqlite3
|
||||
|
||||
import pytest
|
||||
|
||||
from app.engine import masking
|
||||
|
||||
MASKED_WORD = "已遮罩"
|
||||
MASKED = ["buyer_phone", "buyer_name"]
|
||||
SECRETS = ("1390000000", "张伟", "李娜")
|
||||
LAYOUT = ["id", "region", "buyer_phone", "buyer_name", "amount", "note"]
|
||||
#: 探查得到的表结构（没设 schema，探查的是当前库）
|
||||
SCHEMA = {"tables": {"sales": {"qualified": "sales", "schema": None, "columns": [{"name": c} for c in LAYOUT]}}}
|
||||
ROWS = [[1, "华东", "13900000001", "张伟", 1200, "改址"], [2, "华南", "13900000002", "李娜", 300, "回访"]]
|
||||
|
||||
|
||||
def as_json(row: list) -> str:
|
||||
return json.dumps(dict(zip(LAYOUT, row)), ensure_ascii=False)
|
||||
|
||||
|
||||
def as_record(row: list) -> str:
|
||||
return "(" + ",".join(str(v) for v in row) + ")"
|
||||
|
||||
|
||||
def hidden_of(sql: str, columns: list[str], kind: str = "postgres", schema: dict | None = None) -> set[str]:
|
||||
snapshot = {"source": "shop", "sql": sql, "columns": columns, "rows": []}
|
||||
return masking._derived_hidden(snapshot, set(MASKED), kind, masking._schema_columns(schema))[0]
|
||||
|
||||
|
||||
def masked(sql: str, columns: list[str], rows: list[list], kind: str = "postgres",
|
||||
schema: dict | None = SCHEMA) -> masking.QueryMask:
|
||||
payload = {"columns": columns, "rows": rows, "sql": sql}
|
||||
return masking.mask_query_result(payload, MASKED, kind, schema=schema)
|
||||
|
||||
|
||||
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, cells", [
|
||||
pytest.param("SELECT row_to_json(s) AS j, s.id AS n FROM sales s", [as_json(r) for r in ROWS], id="row_to_json(s)"),
|
||||
pytest.param("SELECT to_jsonb(s) AS j, s.id AS n FROM sales s", [as_json(r) for r in ROWS], id="to_jsonb(s)"),
|
||||
pytest.param("SELECT s AS j, s.id AS n FROM sales s", [as_record(r) for r in ROWS], id="SELECT s"),
|
||||
pytest.param("SELECT s::text AS j, s.id AS n FROM sales s", [as_record(r) for r in ROWS], id="s::text"),
|
||||
pytest.param("SELECT to_json(s.*) AS j, s.id AS n FROM sales s", [as_json(r) for r in ROWS], id="to_json(s.*)"),
|
||||
pytest.param("SELECT concat_ws(',', s.*) AS j, s.id AS n FROM sales s",
|
||||
[",".join(str(v) for v in r) for r in ROWS], id="concat_ws(',', s.*)"),
|
||||
pytest.param("SELECT ROW(s.*) AS j, s.id AS n FROM sales s", [as_record(r) for r in ROWS], id="ROW(s.*)"),
|
||||
pytest.param("SELECT (s).buyer_phone AS j, s.id AS n FROM sales s", [r[2] for r in ROWS], id="(s).buyer_phone"),
|
||||
pytest.param("SELECT (s.*).buyer_name AS j, s.id AS n FROM sales s", [r[3] for r in ROWS], id="(s.*).buyer_name"),
|
||||
pytest.param("SELECT (SELECT row_to_json(s)) AS j, s.id AS n FROM sales s", [as_json(r) for r in ROWS],
|
||||
id="相关子查询里的整行"),
|
||||
pytest.param("SELECT kv.value AS j, s.id AS n FROM sales s, LATERAL jsonb_each_text(to_jsonb(s)) AS kv "
|
||||
"WHERE kv.key = 'buyer_phone'", [r[2] for r in ROWS], id="LATERAL 拆开整行"),
|
||||
])
|
||||
def test_a_whole_row_carrying_masked_columns_is_masked(sql, cells):
|
||||
"""整行、函数参数里的 s.* 把遮罩列的值交给了别的列：这一列遮，编号照常。"""
|
||||
done = masked(sql, ["j", "n"], [[cell, row[0]] for cell, row in zip(cells, ROWS)])
|
||||
assert done.columns == ["j"] and not leaks(done), done.payload["rows"]
|
||||
assert [row[1] for row in done.payload["rows"]] == [1, 2]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql", [
|
||||
"SELECT json_agg(t) AS j FROM (SELECT id, buyer_phone FROM sales) t",
|
||||
"SELECT row_to_json(t) AS j FROM (SELECT buyer_name AS 名 FROM sales) t",
|
||||
"WITH t AS (SELECT id, buyer_phone AS p FROM sales) SELECT json_agg(t) AS j FROM t",
|
||||
])
|
||||
def test_a_whole_row_of_a_derived_source_is_read_through_its_columns(sql):
|
||||
"""整行取自子查询、CTE：按它选出的列认，里面有遮罩列（改了名字的也算）就算。"""
|
||||
assert "j" in hidden_of(sql, ["j"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql", [
|
||||
"SELECT (SELECT j FROM (SELECT row_to_json(s) AS j) z) AS j FROM sales s",
|
||||
"SELECT (SELECT j FROM (SELECT json_agg(t) AS j) z) AS j FROM (SELECT id, buyer_phone FROM sales) t",
|
||||
"SELECT (WITH w AS (SELECT row_to_json(s) AS j) SELECT j FROM w) AS j FROM sales s",
|
||||
"SELECT (SELECT j FROM (SELECT to_json(s.*) AS j) z) AS j FROM sales s",
|
||||
])
|
||||
def test_a_whole_row_of_an_outer_query_is_found_from_a_nested_derived_table(sql):
|
||||
"""FROM 里的子查询、CTE 看不见同一层 FROM 里的别的来源，却看得见再外层的查询：写在里面的整行引用、s.* 指外层的来源，
|
||||
照样认得出。"""
|
||||
assert "j" in hidden_of(sql, ["j"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql", [
|
||||
"SELECT json_agg(t) AS j FROM (SELECT id, region, amount FROM sales) t",
|
||||
"SELECT count(s.*) AS j FROM sales s",
|
||||
"SELECT COUNT(*) AS j FROM sales s",
|
||||
"SELECT region AS j, SUM(amount) AS total FROM sales s GROUP BY region",
|
||||
"SELECT EXISTS(SELECT * FROM sales s2 WHERE s2.id = s.id) AS j FROM sales s",
|
||||
])
|
||||
def test_rows_and_stars_that_bring_no_masked_values_stay_visible(sql):
|
||||
"""整行取自没有遮罩列的子查询、count(s.*) 只数个数：不算。"""
|
||||
assert hidden_of(sql, ["j", "total"]) == set(MASKED)
|
||||
|
||||
|
||||
def test_a_column_named_like_its_table_is_a_column_where_the_schema_says_so():
|
||||
"""Postgres 先按列认：表里有一列和表同名（visits 表的 visits 列），SELECT visits 取的是这一列，不是整行。表结构里查得到
|
||||
这一列就不按整行算；SQLite、MySQL 没有整行引用，不看表结构也不按整行算。"""
|
||||
schema = {"tables": {"visits": {"columns": [{"name": c} for c in ("id", "visits", "buyer_phone")]}}}
|
||||
sql = "SELECT visits AS j FROM visits"
|
||||
assert hidden_of(sql, ["j"], schema=schema) == set(MASKED)
|
||||
assert hidden_of(sql, ["j"], kind="sqlite") == set(MASKED)
|
||||
assert "j" in hidden_of(sql, ["j"]), "表结构里查不到这一列：可能是整行，按整行算（多遮）"
|
||||
|
||||
|
||||
def test_a_star_inside_a_function_reads_every_source_it_expands():
|
||||
"""Oracle 的 JSON_OBJECT(*) 把整行拼成 JSON：同 to_json(s.*)。"""
|
||||
done = masked("SELECT JSON_OBJECT(*) AS j FROM sales", ["j"], [[as_json(r)] for r in ROWS], kind="oracle")
|
||||
assert done.columns == ["j"] and not leaks(done)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 列名表、TABLE x
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_a_table_with_a_column_list_masks_the_renamed_masked_column():
|
||||
"""数据源里的表带列名表：号码改叫 p。列名表写全了的，按表结构对上位置，只遮 p 和姓名改叫的 q，编号照常；Postgres 的
|
||||
列名表只写前几列（后面的照原名）时，列名表和表的列数对不上，改了名的几列整段遮（多遮），照原名的姓名按名字遮。"""
|
||||
full = "SELECT p AS j, i AS n FROM sales AS s(i, r, p, q, a, o)"
|
||||
done = masked(full, ["j", "n"], [[r[2], r[0]] for r in ROWS])
|
||||
assert done.columns == ["j"] and not leaks(done)
|
||||
assert set(masked(full, ["j", "n"], [[r[2], r[0]] for r in ROWS], schema=None).columns) == {"j", "n"}, \
|
||||
"表结构里查不到这张表：说不清哪几列改了名，都遮"
|
||||
short = masked("SELECT p AS j, i AS n FROM sales AS s(i, r, p)", ["j", "n"], [[r[2], r[0]] for r in ROWS])
|
||||
assert set(short.columns) == {"j", "n"} and not leaks(short)
|
||||
star = masked("SELECT * FROM sales AS s(i, r, p)", ["i", "r", "p", *LAYOUT[3:]], [list(r) for r in ROWS])
|
||||
assert star.columns == ["i", "r", "p", "buyer_name"] and not leaks(star)
|
||||
|
||||
|
||||
def test_a_column_list_on_a_cte_reference_renames_it_too():
|
||||
"""引用 CTE 时再写列名表（FROM t AS x(p, a)）：同样按位置对上。"""
|
||||
sql = "WITH t AS (SELECT buyer_phone, amount FROM sales) SELECT p AS j, a AS total FROM t AS x(p, a)"
|
||||
assert hidden_of(sql, ["j", "total"]) >= {"j"}
|
||||
assert "total" not in hidden_of(sql, ["j", "total"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql, columns, shown", [
|
||||
pytest.param("WITH t(a, b, c, d, e, f) AS (TABLE sales) SELECT c AS j, a AS n FROM t", ["j", "n"], {"n"},
|
||||
id="CTE 里的 TABLE"),
|
||||
pytest.param("SELECT * FROM (TABLE sales) AS t(a, b, c, d, e, f)", ["a", "b", "c", "d", "e", "f"], {"a", "b", "e"},
|
||||
id="子查询里的 TABLE"),
|
||||
pytest.param("WITH t(a, b, c) AS (TABLE sales) SELECT c AS j FROM t", ["j"], set(), id="列名表只写前几列"),
|
||||
])
|
||||
def test_a_table_statement_is_read_as_select_star(sql, columns, shown):
|
||||
"""TABLE sales（Postgres、MySQL）就是 SELECT * FROM sales：列名表照样按位置对上，号码改叫 c 的那一列遮。"""
|
||||
hidden = hidden_of(sql, columns, schema=SCHEMA)
|
||||
assert {"j", "c"} & hidden and not shown & hidden, hidden
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql", [
|
||||
"SELECT 0 AS a, '' AS b, '' AS c, '' AS d, 0 AS e, '' AS f WHERE false UNION ALL TABLE sales",
|
||||
"WITH t(a, b, c) AS (TABLE archive.sales) SELECT c AS j FROM t",
|
||||
"WITH t(a, b, c) AS (TABLE sales ORDER BY id LIMIT 1) SELECT c AS j FROM t",
|
||||
"WITH t(a, b, c, d, e, f) AS (SELECT * FROM sales) SELECT c AS j FROM t WHERE (",
|
||||
])
|
||||
def test_sql_that_cannot_be_read_but_has_a_star_or_a_table_statement_hides_everything(sql):
|
||||
"""解析器认不出的写法（UNION 后一支的 TABLE、带库名或 ORDER BY 的 TABLE）：SQL 里没写遮罩列，可 TABLE 语句、* 会把它
|
||||
带进别的列名，整份快照的列都算，同解析不了又写着遮罩列的。"""
|
||||
columns = ["a", "b", "c", "j"]
|
||||
assert hidden_of(sql, columns) >= set(columns)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql", [
|
||||
"SELECT COUNT(*) AS j FROM sales WHERE (",
|
||||
"SELECT amount * 2 AS j FROM sales WHERE (",
|
||||
])
|
||||
def test_sql_that_cannot_be_read_with_only_a_count_or_a_product_shows_the_rest(sql):
|
||||
"""只有 COUNT(*)、乘号：不会把遮罩列带出来，照旧只遮遮罩列本身。"""
|
||||
assert hidden_of(sql, ["j"]) == set(MASKED)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# UNPIVOT、PIVOT、MATCH_RECOGNIZE
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql, columns", [
|
||||
pytest.param("SELECT id, val FROM sales UNPIVOT (val FOR col IN (buyer_phone, region))", ["id", "val"],
|
||||
id="UNPIVOT"),
|
||||
pytest.param("SELECT * FROM (SELECT region, buyer_phone FROM sales) PIVOT (MAX(buyer_phone) FOR region IN "
|
||||
"('华东' AS e, '华南' AS s))", ["e", "s"], id="PIVOT"),
|
||||
pytest.param("SELECT * FROM sales MATCH_RECOGNIZE (PARTITION BY region ORDER BY id MEASURES buyer_phone AS x "
|
||||
"ONE ROW PER MATCH PATTERN (a) DEFINE a AS amount > 0)", ["region", "x"], id="MATCH_RECOGNIZE"),
|
||||
])
|
||||
def test_pivots_reading_masked_columns_hide_everything(sql, columns):
|
||||
"""UNPIVOT 的值列、PIVOT 按 IN 里的值起的列、MATCH_RECOGNIZE 的 MEASURES 装着遮罩列的值：出的是哪几列不好认，整份遮。"""
|
||||
assert hidden_of(sql, columns, kind="oracle") >= {c.lower() for c in columns}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql, columns", [
|
||||
("SELECT id, val FROM sales UNPIVOT (val FOR col IN (region, note))", ["id", "val"]),
|
||||
("SELECT * FROM (SELECT region, buyer_phone FROM sales) PIVOT (COUNT(buyer_phone) FOR region IN "
|
||||
"('华东' AS e, '华南' AS s))", ["e", "s"]),
|
||||
])
|
||||
def test_pivots_that_bring_no_masked_values_stay_visible(sql, columns):
|
||||
assert hidden_of(sql, columns, kind="oracle") == set(MASKED)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 选出项套了括号的 *
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_a_star_in_parentheses_is_counted_as_a_star():
|
||||
"""Postgres 里选出项 (s.*)、(s).* 照样展开成好几列：后面没起别名的表达式按结果末尾倒着数，不按一列数位置。"""
|
||||
columns = [*LAYOUT, "upper"]
|
||||
for item in ("(s.*)", "(s).*"):
|
||||
sql = f"SELECT {item}, UPPER(s.buyer_name) FROM sales s"
|
||||
assert hidden_of(sql, columns) == {*MASKED, "upper"}, item
|
||||
|
||||
|
||||
def test_sqlite_runs_of_the_same_shapes_are_unchanged():
|
||||
"""SQLite 没有整行引用、列名表：同样的写法照旧（表和列同名时按列认，不多遮）。"""
|
||||
conn = sqlite3.connect(":memory:")
|
||||
conn.executescript("CREATE TABLE region (id INTEGER, region TEXT, buyer_phone TEXT);"
|
||||
"INSERT INTO region VALUES (1, '华东', '13900000001');")
|
||||
sql = "SELECT region AS 区域, COUNT(*) AS n FROM region GROUP BY region"
|
||||
cursor = conn.execute(sql)
|
||||
payload = {"columns": [d[0] for d in cursor.description], "rows": [list(r) for r in cursor.fetchall()], "sql": sql}
|
||||
assert masking.mask_query_result(payload, MASKED, "sqlite").payload is payload
|
||||
conn.close()
|
||||
@@ -490,3 +490,136 @@ async def test_a_union_of_monthly_tables_is_counted_with_the_sources_schema(monk
|
||||
assert rows["masked_at_query"] == ["buyer_phone"] and _cells(rows, "amount") == [1200, 300, 900]
|
||||
finally:
|
||||
await engines.invalidate(source_id)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 过期的表结构、写了库名的表、整行与列名表（复核第 3 轮）
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
#: 一月那张表后来多了一列 extra（排在第二），数据源的表结构是那之前探查的；二月那张号码排在第二列
|
||||
REBUILT = ("CREATE TABLE jan (id INTEGER PRIMARY KEY, extra TEXT, region TEXT, buyer_phone TEXT, amount INTEGER);"
|
||||
"CREATE TABLE feb (id INTEGER PRIMARY KEY, buyer_phone TEXT, region TEXT, amount INTEGER, extra TEXT);"
|
||||
"INSERT INTO jan VALUES (1, '加急', '华东', '13900000001', 1200);"
|
||||
"INSERT INTO feb VALUES (2, '13900000002', '华南', 300, '普通');")
|
||||
|
||||
|
||||
async def test_a_stale_schema_does_not_place_masked_values_by_its_old_layout(monkeypatch, tmp_path):
|
||||
"""数据源的表结构过期了（一月那张少了 extra）：UNION 第一支的 * 按它数出的列名和后一支的列数对不上，不按它对位置，整份
|
||||
遮。以前照过期的列名对，二月的号码落进的「extra」认不出，原样交给模型、落库。"""
|
||||
name, source_id = await _source(tmp_path, {"mask_columns": "buyer_phone"}, script=REBUILT)
|
||||
async with SessionLocal() as session:
|
||||
row = await session.get(DataSource, source_id)
|
||||
cache = dict(row.schema_cache)
|
||||
jan = dict(cache["tables"]["jan"])
|
||||
jan["columns"] = [c for c in jan["columns"] if c["name"] != "extra"]
|
||||
row.schema_cache = {**cache, "tables": {**cache["tables"], "jan": jan}}
|
||||
await session.commit()
|
||||
try:
|
||||
sql = ("WITH t AS (SELECT * FROM jan UNION ALL SELECT id, buyer_phone, region, amount, extra FROM feb) "
|
||||
"SELECT extra AS 备注, id AS n FROM t ORDER BY id")
|
||||
script = Script(monkeypatch, [q(name, sql, "c1")], {"value": {"alias": "Q1", "row": 0, "column": "n"},
|
||||
"rationale": "x"})
|
||||
out, _ = await ask(name)
|
||||
seen = script.tool_results()["c1"]
|
||||
assert MASKED in seen and _clean(seen), seen
|
||||
snapshot = artifact_store.load(out.queries[0]["artifact"])
|
||||
assert "备注" in snapshot["masked_at_query"] and _clean(json.dumps(snapshot, ensure_ascii=False))
|
||||
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 _clean(got.text)
|
||||
finally:
|
||||
await engines.invalidate(source_id)
|
||||
|
||||
|
||||
#: 造出来的 Postgres、Oracle 数据源的表结构：探查的是当前库（没设 schema），sales 一张表
|
||||
PG_LAYOUT = ["id", "region", "buyer_phone", "buyer_name", "amount", "note"]
|
||||
PG_SCHEMA = {"tables": {"sales": {"qualified": "sales", "schema": None,
|
||||
"columns": [{"name": c, "type": "TEXT"} for c in PG_LAYOUT]}}, "schema": None}
|
||||
PG_ROWS = [[1, "华东", "13900000001", "张伟", 1200, "改址"], [2, "华南", "13900000002", "李娜", 300, "回访"]]
|
||||
|
||||
|
||||
def _json_row(row: list) -> str:
|
||||
return json.dumps(dict(zip(PG_LAYOUT, row)), ensure_ascii=False)
|
||||
|
||||
|
||||
#: (数据源类型, SQL, 数据库交回的行, 查询时就整列遮掉的列, 编号那一格取到的值)。结果列按 SQL 里起的别名（j、n）。行照
|
||||
#: 驱动交回的样子造：JSON 是字符串，记录转成文字；archive.sales 是另一个库里的同名表，号码排在第二列。表的列名表只写了
|
||||
#: 前几列（Postgres 后面的列照原名）、另一个库里的同名表：说不清号码改叫了哪个名字，那一段整段遮，编号也遮
|
||||
DRIVER_CASES = [
|
||||
("postgresql", "SELECT json_agg(t) AS j, 1 AS n FROM (SELECT id, buyer_phone FROM sales) t",
|
||||
[[json.dumps([{"id": r[0], "buyer_phone": r[2]} for r in PG_ROWS]), 1]], ["j"], 1),
|
||||
("postgresql", "SELECT row_to_json(s) AS j, s.id AS n FROM sales s ORDER BY s.id",
|
||||
[[_json_row(r), r[0]] for r in PG_ROWS], ["j"], 1),
|
||||
("postgresql", "SELECT to_json(s.*) AS j, s.id AS n FROM sales s ORDER BY s.id",
|
||||
[[_json_row(r), r[0]] for r in PG_ROWS], ["j"], 1),
|
||||
("postgresql", "SELECT s AS j, s.id AS n FROM sales s ORDER BY s.id",
|
||||
[["(" + ",".join(map(str, r)) + ")", r[0]] for r in PG_ROWS], ["j"], 1),
|
||||
("postgresql", "SELECT p AS j, i AS n FROM sales AS s(i, r, p) ORDER BY i", [[r[2], r[0]] for r in PG_ROWS],
|
||||
["j", "n"], None),
|
||||
("postgresql", "WITH t(a, b, c, d, e, f) AS (TABLE sales) SELECT c AS j, a AS n FROM t ORDER BY a",
|
||||
[[r[2], r[0]] for r in PG_ROWS], ["j"], 1),
|
||||
("postgresql", "WITH t(a, b, c, d, e, f) AS (SELECT * FROM archive.sales) SELECT b AS j, a AS n FROM t",
|
||||
[["13900000007", 7]], ["j", "n"], None),
|
||||
("postgresql", "WITH t(a, b, c, d, e, f) AS (SELECT * FROM sales) SELECT c AS j, a AS n FROM t ORDER BY a",
|
||||
[[r[2], r[0]] for r in PG_ROWS], ["j"], 1),
|
||||
("oracle", "SELECT JSON_OBJECT(*) AS j, id AS n FROM sales ORDER BY id", [[_json_row(r), r[0]] for r in PG_ROWS],
|
||||
["j"], 1),
|
||||
("oracle", "SELECT id AS n, val AS j FROM sales UNPIVOT (val FOR col IN (buyer_phone, region)) ORDER BY id",
|
||||
[[1, "13900000001"], [1, "华东"], [2, "13900000002"], [2, "华南"]], ["n", "j"], None),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def driver_shop():
|
||||
"""造一个 Postgres / Oracle 数据源（查询交给下面换掉的 run_query，不连库）：买家手机号、姓名设了遮罩。返回建源的函数。"""
|
||||
made: list[int] = []
|
||||
|
||||
async def make(kind: str) -> str:
|
||||
name = f"pg_{uuid.uuid4().hex[:8]}"
|
||||
async with SessionLocal() as session:
|
||||
row = DataSource(name=name, kind=kind, host="127.0.0.1", port=5432, username="probe", database="shop",
|
||||
readonly=True, enabled=True, options={"mask_columns": "buyer_phone, buyer_name"},
|
||||
schema_cache=PG_SCHEMA)
|
||||
session.add(row)
|
||||
await session.commit()
|
||||
made.append(row.id)
|
||||
return name
|
||||
|
||||
yield make
|
||||
for source_id in made:
|
||||
await engines.invalidate(source_id)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind, sql, rows, masked, value", DRIVER_CASES,
|
||||
ids=["json_agg(t)", "row_to_json(s)", "to_json(s.*)", "SELECT s", "表的列名表", "CTE 里的 TABLE",
|
||||
"另一个库里的同名表", "本库的表按位置数", "JSON_OBJECT(*)", "UNPIVOT"])
|
||||
async def test_whole_rows_and_renamed_columns_from_other_databases_are_masked(monkeypatch, driver_shop, kind, sql,
|
||||
rows, masked, value):
|
||||
"""Postgres 的整行（json_agg(t)、row_to_json(s)、to_json(s.*)、SELECT s）、表的列名表、CTE 里的 TABLE sales、另一个库里的
|
||||
同名表，Oracle 的 JSON_OBJECT(*)、UNPIVOT：交给模型的、落库的、通用接口给的都是遮好的，编号照常（认不出位置的整份遮）。
|
||||
以前这些写法 masked_at_query 是空的，号码、姓名原样交给模型、落库。"""
|
||||
from app.data.engine import QueryResult
|
||||
from app.tools import datasource
|
||||
|
||||
async def driver(source, statement, **kwargs):
|
||||
assert statement == sql
|
||||
columns = ["n", "j"] if statement.startswith("SELECT id AS n") else ["j", "n"]
|
||||
return QueryResult(columns=columns, rows=[list(r) for r in rows], row_count=len(rows), truncated=False,
|
||||
elapsed_ms=1, sql=statement)
|
||||
|
||||
monkeypatch.setattr(datasource, "run_query", driver)
|
||||
name = await driver_shop(kind)
|
||||
script = Script(monkeypatch, [q(name, sql, "c1")], {"value": {"alias": "Q1", "row": 0, "column": "n"},
|
||||
"rationale": "x"})
|
||||
out, events = await ask(name)
|
||||
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 _clean(json.dumps(snapshot, ensure_ascii=False))
|
||||
assert out.result["value"] == value
|
||||
[end] = [d for k, d in events if k == EventType.TOOL_END]
|
||||
assert _clean(end["preview"])
|
||||
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 _clean(got.text)
|
||||
|
||||
Reference in New Issue
Block a user