From b0fdf56eede6d3a3725136cd7943eb169289ecdb Mon Sep 17 00:00:00 2001 From: ThearchyHelios Date: Tue, 6 Oct 2026 21:50:48 +0800 Subject: [PATCH] =?UTF-8?q?K3=20=E5=A4=8D=E6=A0=B8=E4=BF=AE=E6=AD=A3?= =?UTF-8?q?=EF=BC=88=E7=AC=AC=203=20=E8=BD=AE=EF=BC=89=EF=BC=9A=E5=86=99?= =?UTF-8?q?=E4=BA=86=E5=BA=93=E5=90=8D=E7=9A=84=E8=A1=A8=E6=8C=89=E3=80=8C?= =?UTF-8?q?=E5=BA=93=E5=90=8D.=E8=A1=A8=E5=90=8D=E3=80=8D=E6=9F=A5?= =?UTF-8?q?=E8=A1=A8=E7=BB=93=E6=9E=84=EF=BC=9BUNION=20=E7=AC=AC=E4=B8=80?= =?UTF-8?q?=E6=94=AF=E6=8C=89=E8=A1=A8=E7=BB=93=E6=9E=84=E6=95=B0=E5=87=BA?= =?UTF-8?q?=E7=9A=84=E5=88=97=E5=90=8D=E8=A6=81=E5=92=8C=E5=90=8E=E9=9D=A2?= =?UTF-8?q?=E6=AF=8F=E4=B8=80=E6=94=AF=E4=B8=80=E6=A0=B7=E5=AE=BD=E6=89=8D?= =?UTF-8?q?=E7=94=A8=EF=BC=9B=E6=95=B4=E8=A1=8C=E5=BC=95=E7=94=A8=E3=80=81?= =?UTF-8?q?=E5=87=BD=E6=95=B0=E5=8F=82=E6=95=B0=E9=87=8C=E7=9A=84=20*?= =?UTF-8?q?=E3=80=81=E8=A1=A8=E7=9A=84=E5=88=97=E5=90=8D=E8=A1=A8=E3=80=81?= =?UTF-8?q?TABLE=20x=E3=80=81PIVOT=20=E8=BF=99=E7=B1=BB=E8=AE=A4=E5=BE=97?= =?UTF-8?q?=E5=87=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 写了库名的表（复核 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 探针、端到端 在临时库上重跑，都不再漏。 --- backend/app/engine/inquiry.py | 20 +- backend/app/engine/masking.py | 284 +++++++++++++++++++-- backend/tests/test_masking_star_renames.py | 127 +++++++++ backend/tests/test_masking_whole_rows.py | 244 ++++++++++++++++++ backend/tests/test_verify_agent_masking.py | 133 ++++++++++ 5 files changed, 776 insertions(+), 32 deletions(-) create mode 100644 backend/tests/test_masking_whole_rows.py diff --git a/backend/app/engine/inquiry.py b/backend/app/engine/inquiry.py index 1ca86ad7..696ad70d 100644 --- a/backend/app/engine/inquiry.py +++ b/backend/app/engine/inquiry.py @@ -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 分组再按档数个数也看得出（同按阈值查出某几行）。只有只装着标签的列不收： diff --git a/backend/app/engine/masking.py b/backend/app/engine/masking.py index b5775c11..58f3bac0 100644 --- a/backend/app/engine/masking.py +++ b/backend/app/engine/masking.py @@ -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"(? 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）；只装着分档标签的列里的标签不算遮罩值，这份结果自己写的标签也不 diff --git a/backend/tests/test_masking_star_renames.py b/backend/tests/test_masking_star_renames.py index 36c93230..aebab95a 100644 --- a/backend/tests/test_masking_star_renames.py +++ b/backend/tests/test_masking_star_renames.py @@ -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] diff --git a/backend/tests/test_masking_whole_rows.py b/backend/tests/test_masking_whole_rows.py new file mode 100644 index 00000000..b3fc11a3 --- /dev/null +++ b/backend/tests/test_masking_whole_rows.py @@ -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() diff --git a/backend/tests/test_verify_agent_masking.py b/backend/tests/test_verify_agent_masking.py index ea6f2ab3..01a4968a 100644 --- a/backend/tests/test_verify_agent_masking.py +++ b/backend/tests/test_verify_agent_masking.py @@ -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)