diff --git a/backend/app/engine/inquiry.py b/backend/app/engine/inquiry.py index 894f6a41..f9796a0e 100644 --- a/backend/app/engine/inquiry.py +++ b/backend/app/engine/inquiry.py @@ -71,9 +71,16 @@ 一律算取自遮罩列。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 的整份遮，整行引用这类不写列名的写法认不出 + 遮）。Postgres 先按列认：表结构里查得到和表同名的列（visits 表的 visits 列）才当列，查不到时按整行算（多遮）；LATERAL + 里只看排在它前面的来源，它自己的结果列不算。PIVOT、UNPIVOT、MATCH_RECOGNIZE 读到遮罩列的，整份快照的列都遮（多遮）。 + SQL 解析不了的只看文字：写着遮罩列、*（COUNT(*) 除外）、TABLE x、下面那几类函数和视图的整份遮，整行引用这类不写列名的 + 写法认不出 +- 认不全的写法整份遮（多遮，同一份里没碰遮罩列的编号、金额也遮）：(表达式).* 展开读到遮罩列的（复合类型的列 + (p.profile).*、函数交回的记录 (jsonb_each_text(pj)).*、子查询里的 ROW (q.r).*，展开成几列、各叫什么 SQL 里看不出）； + 表名、SQL 写在字符串里交给函数执行、整表导出的（query_to_xml、table_to_xml、cursor_to_xml、schema_to_xml、 + database_to_xml 这一族，dblink、crosstab、ts_stat，Oracle 的 DBMS_XMLGEN、XMLTYPE，DuckDB 的 query，SQL Server 的 + OPENQUERY）、按页按文件读原始数据的、统计信息和会话记录里存着值的（pg_stats、sqlite_stat4、pg_stat_statements）， + 按名字认（masking._INDIRECT），自己建的同类函数认不出；分析出了意外、一长串 CTE 超出工作量上限（ANALYSIS_STEPS）的 - 按遮罩列分档的标签（CASE、IIF 的结果全是写死的常量）：标签文字是 agent 写在 SQL 里的，不当遮罩值，SQL、回答正文里 照原样；标签列当分组键（每一档一行）照常显示，分档写在子查询、WITH 里、外层（或中间一层）按它分组的也算 （masking._grouped_labels），不分组的逐行标签列整列遮（格子里看不出某一笔落在哪一档），可 agent diff --git a/backend/app/engine/masking.py b/backend/app/engine/masking.py index e7caa989..2e5283dc 100644 --- a/backend/app/engine/masking.py +++ b/backend/app/engine/masking.py @@ -18,14 +18,17 @@ 查不到）。整行和函数参数里的 *（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）。认不出落在哪一列的（PIVOT、UNPIVOT 这类也算）、SQL 解析不了又写着遮罩列（或者 *、 - TABLE x）的，整份快照的列都算。双盲复算按同一套认法在派盲算之前挡下取自遮罩列的指标（derived_columns） + 不算（_value_columns、_band_only）。认不全的写法整份快照的列都算，不再逐个形状补规则：认不出落在哪一列的（PIVOT、 + UNPIVOT 这类也算），(表达式).* 展开读到遮罩列的（复合类型的列、函数交回的记录，展开成几列、叫什么说不清），表名、SQL + 写在字符串里交给函数、统计信息里存着值的（_INDIRECT：query_to_xml 这一族、dblink、pg_stats……），SQL 解析不了又写着 + 遮罩列（或者 *、TABLE x、_INDIRECT 那几类）的，分析出了意外、超出工作量上限（ANALYSIS_STEPS）的。一次分析里派生来源 + 对应的那块查询、一项的值取自哪些列记下来（_analysis），一长串 CTE 不再卡住。双盲复算按同一套认法在派盲算之前挡下 + 取自遮罩列的指标（derived_columns） - 只装着分档标签的列（_label_columns）：标签是 agent 写在 SQL 里的字，不当遮罩值；* 展开数据源里的表时按冻结的表结构 认它带不带同名的列（_schema_tables），不知道表结构的按带得出算（多遮）。分档写在派生来源里、外层按它分组的（一档 一行，_grouped_labels）同当分组键的 CASE 一样照常显示（_shown_labels） - 哪些值（_Secrets）：遮罩列里出现过的值，按值认有门槛（SECRET_MIN、SECRET_DIGITS），占位的写法（PLACEHOLDERS：未知、 - 暂无……）不收；换的时候先按词、按开头几个字筛出 - 正文里真有的那几个值，不现编上万个分支的正则；正文、文字格、SQL 的字符串和注释里换了写法的长号码（199-1474-2785、 + 暂无……）不收；换的时候先按词、按开头几个字筛出正文里真有的那几个值，不现编上万个分支的正则；正文、文字格、SQL 的字符串和注释里换了写法的长号码（199-1474-2785、 +86 …、全角数字）按纯数字写法在集合里查（loose_spans）；已经摆在读者眼前的值免遮，之后并进来的照样免（protected） - 快照遮好的样子（_masked_view），回答时才认得出、事后要照位置重放的地方（view_edits、replay_edits） - 一次核查里认出的遮罩值（SeenValues）：调用方交进来的和 agent 查询时遮掉的分开记，合起来时 agent 查到的在前 @@ -40,6 +43,7 @@ import functools import json import logging import re +import threading from collections.abc import Callable, Iterable, Iterator from dataclasses import dataclass, field from decimal import Decimal @@ -89,6 +93,49 @@ def _masked_word() -> str: #: _value_columns 里「值可能取自数据源里一张表的任何一列」：套在里面的子查询里的 * 展开了表，SQL 里看不出这张表有哪几列。 #: 按它认（_reads）一律算取自遮罩列（宁可多遮）；真实的列名不会是这个 _ANY_COLUMN = "*" +#: 一次遮罩分析（_tainted 认遮罩列）最多走这么多步：展开一个派生来源、认一项的值取自哪些列各算一步，记下的结果再用不算。 +#: 一般的查询几百到几千步，两百段的普通 CTE 几万步；一长串 CTE 一段套一段地整行引用、展开 *（2 万字的上限里写得下几百段） +#: 以前要几十秒，超出就不再分析，整份遮（工具边界每次查询都要分析，不能为一条病态的 SQL 卡住） +ANALYSIS_STEPS = 400_000 + + +class _Overworked(Exception): + """遮罩分析超出了工作量上限（ANALYSIS_STEPS）。""" + + +class _Analysis: + """一次遮罩分析里记下的中间结果：派生来源对应的那块查询（_derived_block）、一项的值取自哪些列（_value_columns）， + 按语法树节点记（分析期间树不再改写，节点都活着）；走过的步数（超出 ANALYSIS_STEPS 抛 _Overworked）。""" + + def __init__(self) -> None: + self.steps = 0 + self.blocks: dict[int, Any] = {} + self.values: dict[int, set[str]] = {} + + def step(self) -> None: + self.steps += 1 + if self.steps > ANALYSIS_STEPS: + raise _Overworked() + + +#: 当前线程上正在做的那次分析（_analysis）。按线程分开：工具边界的遮罩在线程里做，几个核查 agent 同时查询时互不相干 +_CURRENT = threading.local() + + +@contextlib.contextmanager +def _analysis() -> Iterator[_Analysis]: + """在一次分析（一棵解析、改写完的语法树）里记下中间结果、数步数；出来就丢掉。里面再开一次的另起一份，出来还原： + 记下的结果按节点的 id 认，不同的树不能混用。""" + outer = getattr(_CURRENT, "state", None) + state = _CURRENT.state = _Analysis() + try: + yield state + finally: + _CURRENT.state = outer + + +def _current() -> _Analysis | None: + return getattr(_CURRENT, "state", None) def _value_columns(node: Any) -> set[str]: @@ -104,12 +151,21 @@ def _value_columns(node: Any) -> set[str]: col->'path'（JSON 箭头）被解析器当成以 col 为参数的 lambda，参数名照列名收（真的 lambda 的参数多算一列无害）。""" from sqlglot import exp + state = _current() + if state is not None and id(node) in state.values: + return state.values[id(node)] out: set[str] = set() stack = [node] # 已经展开过的派生来源：递归的 CTE 里 * 展开它自己时不再绕回去 expanded = {id(node)} while stack: item = stack.pop() + if state is not None: + state.step() + if item is not node and id(item) in state.values: + # 这一项前面认过：直接并进来（它自己展开过的不必再走一遍） + out |= state.values[id(item)] + continue 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: @@ -145,6 +201,8 @@ def _value_columns(node: Any) -> set[str]: stack.extend(fresh) else: stack.extend(item.iter_expressions()) + if state is not None: + state.values[id(node)] = out return out @@ -693,6 +751,18 @@ def _from_sources(select: Any) -> list[tuple[str, Any]]: def _derived_block(source: Any) -> tuple[Any, list[str] | None] | None: """FROM、JOIN 里的一个来源 → (它的那块查询, 写的列名表)：FROM 里的子查询、LATERAL、指向 WITH 里一段的表名。数据源里 的表、VALUES、函数这类为 None（出哪些列、装着什么，SQL 里看不出）。""" + state = _current() + if state is None: + return _derived_block_of(source) + if id(source) not in state.blocks: + state.step() + state.blocks[id(source)] = _derived_block_of(source) + return state.blocks[id(source)] + + +def _derived_block_of(source: Any) -> tuple[Any, list[str] | None] | None: + """_derived_block 真去找的那一步（不看记下的结果）：指向 WITH 的名字要由里往外数看得见的 CTE（_visible_ctes），一长串 + CTE 时每次都数一遍很慢，所以分析里记下来。""" from sqlglot import exp from app.data.sql_semantics import _visible_ctes @@ -1053,20 +1123,26 @@ def _star_places(select: Any, width: int | None, names: set[str], tables: dict[s _ROW_DIALECTS = frozenset({"postgres", "duckdb"}) -def _normalized(root: Any, read: str | None, tables: dict[str, list[str]] | None) -> None: +def _normalized(root: Any, read: str | None, tables: dict[str, list[str]] | None) -> list[set[str]]: """几种写法先就地改写成认得出的样子，再认遮罩列、标签列（_tainted、_label_columns）。只改解析出来的语法树，原文不动； 改写前后结果有哪几列、每一列的值取自哪里都一样： - TABLE x（Postgres、MySQL）：解析器认不出，CTE、子查询里的读成了一个叫 TABLE 的列、表。改成 SELECT * FROM x；带库名、 ORDER BY 的解析不了，按解析不了处理（_implicit） - - 选出项套了括号的 *（(s.*)）、(s).*：Postgres 照样展开成好几列，改成 s.*，不然按一列数位置。括号里不是来源的 - （(t.profile).* 展开一个复合类型的列）改成展开一个对不上任何来源的 *：说不清展开成几列，整段遮 + - 选出项套了括号的 *（(s.*)）、(s).*（s 是这一层 FROM 里的来源，不是哪个来源的列）：Postgres 照样展开成好几列，改成 + s.*，不然按一列数位置 - 带列名表的表（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。""" + - 别的 (表达式).*（复合类型的列 (p.profile).*、函数交回的记录 (jsonb_each_text(pj)).*、子查询里的 ROW (q.r).*）：展开 + 成几列、各叫什么（复合类型的字段名、函数的出列名）SQL 里看不出，改成一个对不上任何来源的 *（_star_sources 为 + None），只用来数别的项的位置。表达式的值取自哪些列（_value_columns，整行引用已经改写过）记下来交回：读到遮罩列的， + 整份快照都算（_tainted）——以前只改成对不上来源的 *，最外层、第一支、派生来源里没写列名表的 * 按「照原名带出」认， + 字段名和遮罩列对不上，号码原样交了出去 + tables 是查询当时冻结的表结构（_schema_columns，按次序的列名），认整行引用时看来源出不出某一列；没有为 None。 + 交回每个 (表达式).* 的值取自哪些列。""" from sqlglot import exp for cte in list(root.find_all(exp.CTE)): @@ -1083,7 +1159,7 @@ def _normalized(root: Any, read: str | None, tables: dict[str, list[str]] | None 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: + if (star := _bare_star(item, select, tables)) is not None: item.replace(star) for select in selects: for _name, source in _from_sources(select): @@ -1104,6 +1180,18 @@ def _normalized(root: Any, read: str | None, tables: dict[str, list[str]] | None # 直接是选出项的起回原名：结果的列名照旧（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) + opaque: list[set[str]] = [] + for select in list(root.find_all(exp.Select)): + for item in list(select.expressions): + if isinstance(item, exp.Dot) and isinstance(item.expression, exp.Star): + # 先认值取自哪些列（还挂在树上，* 和整行引用按作用域认），再换掉 + opaque.append(set(_value_columns(item.this))) + item.replace(exp.Column(this=exp.Star(), table=exp.to_identifier(_OPAQUE_STAR))) + return opaque + + +#: (表达式).* 改写成的 * 写的「表名」：对不上任何来源（_star_sources 为 None），真实的表名、别名不会是它 +_OPAQUE_STAR = "*" def _table_keyword(node: Any) -> bool: @@ -1114,8 +1202,9 @@ def _table_keyword(node: Any) -> bool: 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）。""" +def _bare_star(item: Any, select: Any, tables: dict[str, list[str]] | None) -> Any: + """select 的选出项是套了括号的 *（(s.*)、((*))），或者 (s).*、s 是这一层 FROM 里一个来源的名字、又不是哪个来源说得准 + 出的一列（Postgres 先按列认，_has_column）时，改写成的 *、s.*；别的为 None（_normalized，别的 (表达式).* 另外处理）。""" from sqlglot import exp node = item @@ -1128,10 +1217,14 @@ def _bare_star(item: Any) -> Any: 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("*")) + if not (isinstance(inner, exp.Column) and not inner.table and isinstance(inner.this, exp.Identifier)): + return None + name = inner.name.lower() + sources = _from_sources(select) + [((lateral.alias_or_name or "").lower(), lateral) + for lateral in select.args.get("laterals") or []] + if not any(alias == name for alias, _ in sources) or any(_has_column(s, name, tables) for _, s in sources): + return None + return exp.Column(this=exp.Star(), table=inner.this.copy()) def _carries_value(column: Any) -> bool: @@ -1149,13 +1242,13 @@ def _carries_value(column: Any) -> bool: def _row_source(column: Any, tables: dict[str, list[str]] | None) -> str | None: """Postgres 里不带表名的名字 column 指整行时，它指的来源的名字（别名，没起别名的是表名）；是列、不是任何来源的为 None。 - Postgres 先按列认：由里往外看得见的每一层（_outer_selects）的来源里，有一个出得了这一列（_has_column）就是列；都不是、 - 又有一个来源叫这个名字，才是整行（取最里面的那个）。说不清出不出这一列的（表结构里查不到的表、* 展开的列）按出不了算， - 当整行（宁可多遮）。""" + Postgres 先按列认：由里往外看得见的每一层（_outer_selects）的来源里（LATERAL 里只看排在它前面的，_visible_sources）， + 有一个出得了这一列（_has_column）就是列；都不是、又有一个来源叫这个名字，才是整行（取最里面的那个）。说不清出不出这一列 + 的（表结构里查不到的表、* 展开的列）按出不了算，当整行（宁可多遮）。""" name = column.name.lower() found: str | None = None for select in _outer_selects(column): - for alias, source in _from_sources(select): + for alias, source in _visible_sources(select, column): if _has_column(source, name, tables): return None if found is None and alias == name: @@ -1163,6 +1256,27 @@ def _row_source(column: Any, tables: dict[str, list[str]] | None) -> str | None: return found +def _visible_sources(select: Any, node: Any) -> list[tuple[str, Any]]: + """select 的 FROM、JOIN 里，node（select 里面的一个节点）看得见的来源（_from_sources 的样子）。node 在 LATERAL 里的，只看 + 得见排在它前面的来源：看不见 LATERAL 自己（它的结果列不是 node 能取的列），也看不见后面的。以前把 LATERAL 自己的结果列 + （SELECT s …、1 AS s）当成看得见的列，里面的整行引用 s 就被认成了列，号码原样交了出去。""" + from sqlglot import exp + + sources = _from_sources(select) + child = node + while child.parent is not None and child.parent is not select: + child = child.parent + if child.parent is not select: + return sources + if isinstance(child, exp.From): + return [] + if isinstance(child, exp.Join): + joins = list(select.args.get("joins") or []) + at = next((i for i, join in enumerate(joins) if join is child), None) + return sources[:1 + at] if at is not None else sources + return sources + + def _has_column(source: Any, name: str, tables: dict[str, list[str]] | None) -> bool: """FROM 里的这个来源说得准出 name 这一列：派生来源看列名表、第一支里起了名字的选出项（列名表改了名的不算），数据源里 的表看表结构（_table_columns），表值函数这类看它出的列名（_function_columns）。* 展开的列说不准，不算。""" @@ -1293,9 +1407,9 @@ def _tainted(sql: str, kind: str | None, masked: set[str], width: int, tables: d JSON_OBJECT(*)）、整行引用（row_to_json(s)）按 _value_columns 算进外层的表达式。表的列名表、TABLE x、套了括号的 * 先改写成上面这些认得出的样子（_normalized）。 PIVOT、UNPIVOT、MATCH_RECOGNIZE 读到这些名字的：值转进了它们自己起名的列（UNPIVOT 的值列、PIVOT 按 IN 里的值起名的列、 - MEASURES），出的是哪几列不好认，整份快照的列都算。""" - from sqlglot import exp - + MEASURES），出的是哪几列不好认，整份快照的列都算。(表达式).* 读到这些名字的、_INDIRECT 那几类写法同样整份算。 + 解析之后的分析出了意外、超出工作量上限（ANALYSIS_STEPS）的也整份算——不退回「解析不了」只看文字：没写遮罩列、 + 没有 * 的整行引用（SELECT s）那样就放过去了。""" from app.data.sql_semantics import _parse, _read try: @@ -1305,54 +1419,69 @@ 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) - sources = [(top, listed) for top, role, listed in blocks if role == "source"] - functions = [source for select_ in selects for source in _function_sources(select_)] - unreadable = False - for _ in range(len(selects) + len(sources) + len(functions) + 1): - before = len(names) - for select_ in selects: - for item in select_.expressions: - name = item.alias_or_name.lower() - if name and not _is_star(item) and name not in names and _reads(names, _value_columns(item)): - names.add(name) - for top, listed in sources: - names |= _source_names(top, listed, names, tables)[0] - for source in functions: - if _reads(names, _value_columns(source)): - columns = _function_columns(source) - unreadable = unreadable or columns is None - names |= columns or set() - if len(names) == before: - break - everything = unreadable or any(_source_names(top, listed, names, 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": - continue - branches = [leaf for leaf in _leaves(top) if isinstance(leaf, exp.Select)] - for k, select_ in enumerate(branches): - branch = list(select_.expressions) - for item, place in zip(branch, _places(branch, width)): - if not _is_star(item) and _reads(names, _value_columns(item)): - if place is None or place >= width: - everything = True - else: - positions.add(place) - # UNION 后一支的 * 按位置落进第一支起的列名里；第一支、单独一条 SELECT 的 * 照原名带出，按名字认得出 - spread = _star_places(select_, width, names, tables) if k else set() - if spread is None or any(place >= width for place in spread): - everything = True - else: - positions |= spread - return names, set(range(width)) if everything else positions, everything - except Exception: # noqa: BLE001 - 分析本身的缺陷不能让回答出错：按解析不了处理（宁可多遮） - logger.warning("质询的 SQL 遮罩分析失败", exc_info=True) - return None + opaque = _normalized(root, _read(kind), tables) + with _analysis(): + return _tainted_in(root, sql, masked, width, tables, opaque) + except _Overworked: + # 一长串 CTE 一段套一段地整行引用、展开 *：分析不过来，认不准，整份遮 + logger.warning("SQL 遮罩分析超出工作量上限（%d 步），整份快照按遮罩处理", ANALYSIS_STEPS) + except Exception: # noqa: BLE001 - 分析本身的缺陷不能让回答出错，也不能把原值放过去：整份遮（宁可多遮） + logger.warning("SQL 遮罩分析失败，整份快照按遮罩处理", exc_info=True) + return set(masked), set(range(width)), True + + +def _tainted_in(root: Any, sql: str, masked: set[str], width: int, tables: dict[str, list[str]] | None, + opaque: list[set[str]]) -> tuple[set[str], set[int], bool]: + """_tainted 解析、改写（_normalized）之后真去认的那一步。opaque 是 (表达式).* 的值取自哪些列（_normalized 交回的）。""" + from sqlglot import exp + + names = set(masked) + selects = list(root.find_all(exp.Select)) + blocks = _blocks(root) + sources = [(top, listed) for top, role, listed in blocks if role == "source"] + functions = [source for select_ in selects for source in _function_sources(select_)] + unreadable = False + for _ in range(len(selects) + len(sources) + len(functions) + 1): + before = len(names) + for select_ in selects: + for item in select_.expressions: + name = item.alias_or_name.lower() + if name and not _is_star(item) and name not in names and _reads(names, _value_columns(item)): + names.add(name) + for top, listed in sources: + names |= _source_names(top, listed, names, tables)[0] + for source in functions: + if _reads(names, _value_columns(source)): + columns = _function_columns(source) + unreadable = unreadable or columns is None + names |= columns or set() + if len(names) == before: + break + # 认不出的整份遮：表值函数出哪几列说不清、派生来源里有认不出名字的列、PIVOT 这类读到遮罩列，(表达式).* 读到遮罩列 + # （展开成几列、叫什么都说不清），SQL 写在字符串里交给函数执行、整表导出、统计信息里的值（_INDIRECT） + 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)) \ + or any(_reads(names, cols) for cols in opaque) or bool(_INDIRECT.search(sql)) + positions: set[int] = set() + for top, role, _listed in blocks: + if role != "result": + continue + branches = [leaf for leaf in _leaves(top) if isinstance(leaf, exp.Select)] + for k, select_ in enumerate(branches): + branch = list(select_.expressions) + for item, place in zip(branch, _places(branch, width)): + if not _is_star(item) and _reads(names, _value_columns(item)): + if place is None or place >= width: + everything = True + else: + positions.add(place) + # UNION 后一支的 * 按位置落进第一支起的列名里；第一支、单独一条 SELECT 的 * 照原名带出，按名字认得出 + spread = _star_places(select_, width, names, tables) if k else set() + if spread is None or any(place >= width for place in spread): + everything = True + else: + positions |= spread + return names, set(range(width)) if everything else positions, everything def _mentions(sql: str, names: Iterable[str]) -> bool: @@ -1366,10 +1495,27 @@ _COUNT_STAR = re.compile(r"\bcount\s*\(\s*(?:all\s+|distinct\s+)?\*\s*\)", re.I) _IMPLICIT = re.compile(r"(?:\bselect|\bdistinct|\ball|[,(.])\s*\*|\btable\s+[\w\"`]", re.I) +#: 不点列名也读得出表里内容的写法，SQL 里认不出读的是哪一列：表名、SQL 写在字符串里交给函数执行或整表导出（Postgres 的 +#: query_to_xml、table_to_xml 这一族、ts_stat、tablefunc 的 crosstab / connectby、dblink，Oracle 的 DBMS_XMLGEN、XMLTYPE、 +#: XQuery，DuckDB 的 query / query_table，SQL Server 的 OPENQUERY、OPENROWSET），按页、按文件读出的原始数据，统计信息、 +#: 会话记录里存着的值（pg_stats 的常见值、sqlite_stat4 的样本、pg_stat_statements 记下的 SQL 原文）。出现就整份遮 +#: （_tainted、_implicit），写在字符串、注释里的也算（多遮）。函数要跟着左括号才算（列名碰巧叫 query 的不算）。 +#: information_schema、pg_catalog 只有表结构，不在里面 +_INDIRECT = re.compile(r"(? bool: - """解析不了的 SQL 里有没有不写列名也会带出遮罩列的写法（*、TABLE x）：它们能把遮罩列带进别的列名（列名表、UNION 的 - 后一支），同写着遮罩列一样整份遮（_derived_hidden）。只看文字，字符串、注释里的也算（多遮）。""" - return bool(_IMPLICIT.search(_COUNT_STAR.sub("", sql))) + """解析不了的 SQL 里有没有不写列名也会带出遮罩列的写法（*、TABLE x，_INDIRECT 那几类）：它们能把遮罩列带进别的列名 + （列名表、UNION 的后一支）、整表交出来，同写着遮罩列一样整份遮（_derived_hidden）。只看文字，字符串、注释里的也算 + （多遮）。""" + return bool(_IMPLICIT.search(_COUNT_STAR.sub("", sql)) or _INDIRECT.search(sql)) def _derived_hidden(snapshot: dict[str, Any], hidden: set[str], kind: str | None, diff --git a/backend/tests/test_masking_opaque_sql.py b/backend/tests/test_masking_opaque_sql.py new file mode 100644 index 00000000..5e58f16a --- /dev/null +++ b/backend/tests/test_masking_opaque_sql.py @@ -0,0 +1,282 @@ +"""认不全的 SQL 写法整份遮（engine/masking.py 的 _tainted），不再逐个形状补规则（K3 第 3 轮复验留下的几处）： + +- (表达式).* 展开：复合类型的列 (p.profile).*、函数交回的记录 (jsonb_each_text(pj)).*、(jsonb_populate_record(…)).*、 + 子查询里的匿名 ROW (q.r).*。展开成几列、各叫什么 SQL 里看不出，表达式读到遮罩列的，整份结果都遮——最外层、子查询、 + CTE、LATERAL 里都一样。以前改写成一个对不上来源的 *，最外层、没写列名表的派生来源里按「照原名带出」认，字段名 + （tel、key/value、f1/f2）和遮罩列对不上，号码原样交给模型、落库。(s).*、s 是 FROM 里的来源时照旧按 s.* 认 +- 表名、SQL 写在字符串里交给函数执行、整表导出的（query_to_xml、table_to_xml、cursor_to_xml、schema_to_xml、 + database_to_xml、dblink，Oracle 的 DBMS_XMLGEN），统计信息里存着的值（pg_stats）：整份遮，解析不了的也一样 +- LATERAL 看不见它自己：里面的整行引用 s 不因为 LATERAL 自己的结果列叫 s 就被当成列 +- 分析不过来、分析出了意外：整份遮。派生来源对应的那块查询、一项的值取自哪些列在一次分析里记下来，一长串 CTE 一段套 + 一段地整行引用（250 段，2 万字的上限里写得下）以前要二十几秒，现在零点几秒；超出工作量上限（ANALYSIS_STEPS）的整份遮 + +这里没有真的 Postgres：结果按它交回的样子造（复合类型、记录展开成几列）。数据都是虚构的：手机号、地址、区域。 +""" +from __future__ import annotations + +import json +import time +import uuid + +import pytest +from httpx import ASGITransport, AsyncClient + +from app.core import artifact_store +from app.data.engine import engines +from app.db.base import SessionLocal +from app.db.models import DataSource +from app.engine import masking +from app.main import app +from tests.test_verify_agent import Script, q +from tests.test_verify_agent_masking import ask, clean_spend # noqa: F401 - clean_spend 是自动用上的夹具 + +MASKED_WORD = "已遮罩" +#: people.profile 是复合类型（tel、addr），people2.pj 是 JSON（tel、addr），sales 是订单：设了遮罩的是 profile、pj、buyer_phone +SCHEMA = {"tables": { + "people": {"qualified": "people", "schema": None, "columns": [{"name": c} for c in ["id", "profile", "address"]]}, + "people2": {"qualified": "people2", "schema": None, "columns": [{"name": c} for c in ["id", "pj", "city"]]}, + "sales": {"qualified": "sales", "schema": None, + "columns": [{"name": c} for c in ["id", "region", "buyer_phone", "amount"]]}, +}} +MASKED = ["profile", "pj", "buyer_phone"] +PHONE = "13900000001" + + +def _leaks(done: masking.QueryMask) -> bool: + return "1390000000" in json.dumps(done.payload, ensure_ascii=False) + + +def masked(sql: str, columns: list[str], rows: list[list], kind: str = "postgresql") -> masking.QueryMask: + payload = {"columns": columns, "rows": rows, "sql": sql} + return masking.mask_query_result(payload, MASKED, kind, schema=SCHEMA) + + +# -------------------------------------------------------------------------- +# (表达式).* 展开 +# -------------------------------------------------------------------------- + + +#: 三种展开（复合类型的列、函数交回的记录、子查询里的匿名 ROW）× 四个位置（最外层、子查询、CTE、LATERAL）： +#: (SQL, 结果列, 一行) +EXPANSIONS = [ + pytest.param("SELECT p.id AS n, (p.profile).* FROM people p", ["n", "tel", "addr"], [1, PHONE, "某路 1 号"], + id="复合类型-最外层"), + pytest.param("SELECT x.* FROM (SELECT p.id AS n, (p.profile).* FROM people p) x", ["n", "tel", "addr"], + [1, PHONE, "某路 1 号"], id="复合类型-子查询"), + pytest.param("WITH x AS (SELECT p.id AS n, (p.profile).* FROM people p) SELECT * FROM x", ["n", "tel", "addr"], + [1, PHONE, "某路 1 号"], id="复合类型-CTE"), + pytest.param("SELECT x.* FROM people p, LATERAL (SELECT p.id AS n, (p.profile).*) x", ["n", "tel", "addr"], + [1, PHONE, "某路 1 号"], id="复合类型-LATERAL"), + pytest.param("SELECT (jsonb_each_text(pj)).* FROM people2", ["key", "value"], ["tel", PHONE], + id="记录函数-最外层"), + pytest.param("SELECT kv.value AS tel FROM (SELECT (jsonb_each_text(pj)).* FROM people2) kv WHERE kv.key = 'tel'", + ["tel"], [PHONE], id="记录函数-子查询"), + pytest.param("WITH kv AS (SELECT (jsonb_each_text(pj)).* FROM people2) SELECT value AS 号码 FROM kv " + "WHERE key = 'tel'", ["号码"], [PHONE], id="记录函数-CTE"), + pytest.param("SELECT x.value FROM people2 p, LATERAL (SELECT (jsonb_each_text(p.pj)).*) x", ["value"], [PHONE], + id="记录函数-LATERAL"), + pytest.param("SELECT (q.r).* FROM (SELECT ROW(buyer_phone, region) AS r FROM sales) q", ["f1", "f2"], + [PHONE, "华东"], id="匿名 ROW-最外层"), + pytest.param("SELECT y.* FROM (SELECT (q.r).* FROM (SELECT ROW(buyer_phone, region) AS r FROM sales) q) y", + ["f1", "f2"], [PHONE, "华东"], id="匿名 ROW-子查询"), + pytest.param("WITH q AS (SELECT ROW(buyer_phone, region) AS r FROM sales) SELECT (q.r).* FROM q", ["f1", "f2"], + [PHONE, "华东"], id="匿名 ROW-CTE"), + pytest.param("SELECT x.* FROM (SELECT ROW(buyer_phone, region) AS r FROM sales) q, LATERAL (SELECT (q.r).*) x", + ["f1", "f2"], [PHONE, "华东"], id="匿名 ROW-LATERAL"), + pytest.param("SELECT p.id AS n, (jsonb_populate_record(NULL::contact, p.pj)).* FROM people2 p", + ["n", "tel", "addr"], [1, PHONE, "某路 1 号"], id="按类型展开 JSON"), +] + + +@pytest.mark.parametrize("sql, columns, row", EXPANSIONS) +def test_an_expression_star_reading_a_masked_column_masks_the_whole_result(sql, columns, row): + """(表达式).* 读到遮罩列：展开成几列、叫什么说不清，整份结果的每一列都遮（同一份里的编号也遮，多遮）。""" + done = masked(sql, columns, [list(row)]) + assert done.columns == columns + assert done.payload["rows"] == [[MASKED_WORD] * len(columns)] and not _leaks(done) + + +@pytest.mark.parametrize("sql, columns, row", [ + pytest.param("SELECT p.id AS n, (p.address).* FROM people p", ["n", "street", "city"], [1, "某路", "某市"], + id="没设遮罩的复合类型"), + pytest.param("SELECT (jsonb_each_text(to_jsonb(c))).* FROM (SELECT id, region FROM sales) c", ["key", "value"], + ["region", "华东"], id="记录里没有遮罩列"), +]) +def test_an_expression_star_that_reads_no_masked_column_stays_visible(sql, columns, row): + done = masked(sql, columns, [list(row)]) + assert done.columns == [] and done.payload["rows"] == [list(row)] + + +def test_a_parenthesised_source_star_is_still_read_as_its_columns(): + """(s).*、s 是 FROM 里的来源：照旧当 s.* 认，只遮号码那一列，区域、金额照常。""" + done = masked("SELECT (s).* FROM sales s", ["id", "region", "buyer_phone", "amount"], [[1, "华东", PHONE, 1200]]) + assert done.columns == ["buyer_phone"] and done.payload["rows"] == [[1, "华东", MASKED_WORD, 1200]] + + +#: 造出来的 Postgres 数据源：表结构同上，查询交给换掉的 run_query（不连库） +@pytest.fixture +async def composite_shop(): + made: list[int] = [] + + async def make() -> str: + name = f"pg_{uuid.uuid4().hex[:8]}" + async with SessionLocal() as session: + row = DataSource(name=name, kind="postgresql", host="127.0.0.1", port=5432, username="probe", + database="shop", readonly=True, enabled=True, + options={"mask_columns": ", ".join(MASKED)}, schema_cache=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("sql, columns, rows", [ + ("SELECT p.id AS n, (p.profile).* FROM people p", ["n", "tel", "addr"], [[1, PHONE, "某路 1 号"]]), + ("SELECT x.* FROM (SELECT p.id AS n, (p.profile).* FROM people p) x", ["n", "tel", "addr"], + [[1, PHONE, "某路 1 号"]]), + ("SELECT p.id AS n, (jsonb_populate_record(NULL::contact, p.pj)).* FROM people2 p", ["n", "tel", "addr"], + [[1, PHONE, "某路 1 号"]]), +], ids=["复合类型", "子查询里的复合类型", "按类型展开 JSON"]) +async def test_an_expression_star_never_reaches_the_model_or_the_store(monkeypatch, composite_shop, sql, columns, + rows): + """端到端（核查 agent 的真实查询工具）：交给模型的、落库的、通用工件接口给的都只有「已遮罩」，masked_at_query 是 + 全部的列。以前三条的 masked_at_query 都是空的，号码原样交给模型、落库。""" + from app.data.engine import QueryResult + from app.tools import datasource + + async def driver(source, statement, **kwargs): + assert statement == sql + return QueryResult(columns=list(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 composite_shop() + 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_WORD in seen and "1390000000" not in seen, seen + snapshot = artifact_store.load(out.queries[0]["artifact"]) + assert snapshot["masked_at_query"] == columns and "1390000000" not in 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 "1390000000" not in got.text + + +# -------------------------------------------------------------------------- +# LATERAL 看不见它自己 +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("sql, columns, row", [ + pytest.param("SELECT y.* FROM sales t2, LATERAL (SELECT s FROM sales s LIMIT 1) y", ["s"], + [f"(1,华东,{PHONE},1200)"], id="选出整行"), + pytest.param("SELECT y.j FROM sales t2 CROSS JOIN LATERAL (SELECT 1 AS s, row_to_json(s) AS j FROM sales s " + "LIMIT 1) y", ["j"], [json.dumps({"id": 1, "buyer_phone": PHONE})], id="结果列也叫 s"), + pytest.param("SELECT y.j FROM sales t2, LATERAL (SELECT s.id AS s, row_to_json(s) AS j FROM sales s LIMIT 1) y", + ["j"], [json.dumps({"id": 1, "buyer_phone": PHONE})], id="编号起名叫 s"), +]) +def test_a_whole_row_inside_a_lateral_is_not_taken_for_its_own_column(sql, columns, row): + """LATERAL 里的 s 是它自己 FROM 里的整行：LATERAL 的结果列碰巧叫 s，也不把它当成列（Postgres 里 LATERAL 看不见 + 自己）。以前当成了列，整行原样交了出去。""" + done = masked(sql, columns, [list(row)]) + assert done.columns == columns and not _leaks(done) + + +def test_a_lateral_still_sees_the_sources_before_it(): + """排在 LATERAL 前面的来源照样看得见：前面的表有一列叫 s 时，s 是那一列（没设遮罩），不是整行。""" + schema = {"tables": {**SCHEMA["tables"], "codes": {"qualified": "codes", "schema": None, + "columns": [{"name": "id"}, {"name": "s"}]}}} + payload = {"columns": ["v"], "rows": [["A1"]], + "sql": "SELECT y.v FROM codes c, LATERAL (SELECT s AS v FROM sales x WHERE x.id = c.id) y"} + done = masking.mask_query_result(payload, MASKED, "postgresql", schema=schema) + assert done.columns == [] and done.payload["rows"] == [["A1"]] + + +# -------------------------------------------------------------------------- +# 表名、SQL 写在字符串里交给函数，统计信息里的值 +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("sql", [ + "SELECT query_to_xml('select * from sales', true, false, '') AS x", + "SELECT table_to_xml('sales', true, false, '') AS x", + "SELECT query_to_xml_and_xmlschema('select buyer_phone from sales', true, false, '') AS x", + "SELECT cursor_to_xml('c1', 10, false, true, '') AS x", + "SELECT schema_to_xml('public', true, false, '') AS x", + "SELECT database_to_xml(true, false, '') AS x", + "SELECT x FROM dblink('dbname=shop', 'select buyer_phone from sales') AS t(x text)", + "SELECT DBMS_XMLGEN.GETXML('select * from sales') AS x FROM dual", + "SELECT array_to_string(most_common_vals, ',') AS x FROM pg_stats WHERE tablename = 'sales'", +]) +def test_tables_and_sql_handed_to_functions_as_text_mask_the_whole_result(sql): + """表名、SQL 写在字符串里交给函数执行、整表导出，统计信息里存着的值：读的是哪一列认不出，整份遮。以前 masked=[]， + 号码原样。""" + done = masked(sql, ["x"], [[f"{PHONE}"]], kind="oracle" + if "DBMS" in sql else "postgresql") + assert done.columns == ["x"] and not _leaks(done) + + +def test_a_table_handed_to_a_function_masks_the_whole_result_even_when_the_sql_cannot_be_read(): + sql = "SELECT table_to_xml('sales', true, false, '') AS x WHERE (" + hidden, _ = masking._derived_hidden({"sql": sql, "columns": ["x"], "rows": []}, {"buyer_phone"}, "postgresql") + assert "x" in hidden + + +# -------------------------------------------------------------------------- +# 分析不过来、分析出了意外：整份遮；一长串 CTE 不再卡住 +# -------------------------------------------------------------------------- + + +def _chain(k: int, tail: str) -> str: + """一长串 CTE，一段套一段地整行引用、展开上一段（复验时 250 段要二十几秒）。""" + parts = ["c0 AS (SELECT * FROM sales)"] + parts += [f"c{i} AS (SELECT row_to_json(c{i - 1}) AS j, c{i - 1}.* FROM c{i - 1})" for i in range(1, k)] + return "WITH " + ", ".join(parts) + f" {tail.format(last=f'c{k - 1}')}" + + +def _plain_chain(k: int) -> str: + parts = ["c0 AS (SELECT id, region, amount FROM sales)"] + parts += [f"c{i} AS (SELECT id, region, amount + {i} AS amount FROM c{i - 1})" for i in range(1, k)] + return "WITH " + ", ".join(parts) + f" SELECT region, SUM(amount) AS total FROM c{k - 1} GROUP BY region" + + +def test_a_long_chain_of_whole_rows_is_read_quickly_and_completely(): + """250 段一段套一段地整行引用：零点几秒认完（以前二十几秒），不靠工作量上限兜底——只数个数的结果照常不遮。""" + columns = masking._schema_columns(SCHEMA) + sql = _chain(250, "SELECT COUNT(*) AS n FROM {last}") + assert len(sql) < 20_000 + started = time.perf_counter() + hidden, _ = masking._derived_hidden({"sql": sql, "columns": ["n"], "rows": []}, {"buyer_phone"}, "postgresql", + columns) + assert time.perf_counter() - started < 5 + assert "n" not in hidden + hidden, _ = masking._derived_hidden({"sql": _chain(250, "SELECT j FROM {last}"), "columns": ["j"], "rows": []}, + {"buyer_phone"}, "postgresql", columns) + assert "j" in hidden + + +def test_an_analysis_over_the_work_limit_masks_the_whole_result(monkeypatch): + """超出工作量上限：不再分析，整份遮（结果本来没碰遮罩列也一样）。没超的照常。""" + sql = _plain_chain(40) + snapshot = {"sql": sql, "columns": ["region", "total"], "rows": []} + assert masking._derived_hidden(snapshot, {"buyer_phone"}, "postgresql")[0] == {"buyer_phone"} + monkeypatch.setattr(masking, "ANALYSIS_STEPS", 50) + assert masking._derived_hidden(snapshot, {"buyer_phone"}, "postgresql")[0] == {"buyer_phone", "region", "total"} + payload = {"sql": sql, "columns": ["region", "total"], "rows": [["华东", 1200]]} + done = masking.mask_query_result(payload, ["buyer_phone"], "postgresql") + assert done.columns == ["region", "total"] and done.payload["rows"] == [[MASKED_WORD, MASKED_WORD]] + + +def test_a_failing_analysis_masks_the_whole_result_instead_of_reading_the_text(monkeypatch): + """解析之后的分析出了意外：整份遮。以前按「解析不了」只看文字，没写遮罩列、没有 * 的整行引用（SELECT s）就放过去了。""" + def broken(*args, **kwargs): + raise RuntimeError("分析里的意外") + + monkeypatch.setattr(masking, "_tainted_in", broken) + snapshot = {"sql": "SELECT s AS r FROM sales s", "columns": ["r"], "rows": []} + assert "r" in masking._derived_hidden(snapshot, {"buyer_phone"}, "postgresql")[0]