mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 02:21:58 +00:00
ThearchyHelios
39e7a5b014
阶段①。产品原先能取数的只有 http_request 和沙箱里自己写代码——面对"分析上季度
华东区销量下滑"这类需求，用户得自己知道数据在哪张表、字段叫什么，再手写 SQL 塞进
代码节点。最难的部分留给了用户。
## 连接层（app/data/）
方言差异全收敛在 engine.py 一处：Oracle 要 service_name/sid、MySQL 要 charset、
各家默认端口不同。三个驱动都是 asyncio 原生且免装客户端——oracledb 走 thin 模式，
不需要 Instant Client。密码复用 app/core/crypto 的 Fernet，和 Provider 同一套；
build_url 默认把密码替换成星号，日志和错误信息里用这个版本。
introspect.py 产出两种粒度：summary()（表名+代表字段，几百 token，给 Copilot 的
system prompt 用）和 describe_table()（单表完整字段，按需取）。分两种是因为一个
中等业务库全量塞进 prompt 轻松几万 token，既贵又会把真正的需求淹掉。
## 守卫（app/data/guard.py）
这是唯一不能出错的地方：SQL 由模型生成，模型会写出什么事前无法穷举，防线不能
建立在"它应该不会这么干"上。三层：
- 只读判定用**白名单**而非黑名单——黑名单永远漏，MERGE/CALL/GRANT/各家私货写不完
- 多语句在到达驱动前就断掉，不赌驱动默认只跑第一条（MySQL 开了
CLIENT_MULTI_STATEMENTS 就会全跑）
- 行数/字节/超时三重上限，一条 SELECT * 不该把内存和模型上下文一起打爆
刻意不引 sqlparse：完整解析 SQL 方言是无底洞，而这里的判断很窄——"首关键字是
什么"和"有没有第二条语句"。窄问题用窄解法。难点是不能误杀：`SELECT ';'` 只是
一条语句，简单 split(";") 会把它切两半然后拒掉，那比漏杀更让人无从下手。
## 工具（app/tools/datasource.py）
每个数据源动态生成一对 db_query__<源> / db_schema__<源>，而不是一个通用工具带
datasource 参数：工具名自带语义，模型不会把查销售库的 SQL 发给财务库；description
里直接塞表清单，省掉"先问有哪些表"那一轮。接进 registry.build_tools 作为第四来源
（前三个是内置/MCP/自定义）。
查询结果落 put_json(kind="query_snapshot")，把 SQL 原文 + 结果集 + 数据源一起存，
出具体系的数字回指能下钻到"这个数是哪条 SQL 查出来的"。
## 实测
真实 SQLite 业务库（600 行订单，埋了华东 6 月断崖的信号）+ 真实模型：
agent 8 轮自主查询——列表 → 查三张表结构 → 探数据时间范围 → 查地区维度 →
按月汇总 → 按地区+月份下钻，准确定位到华东 6 月从 1187 件跌到 354 件。
只读守卫在 agent 路径上同样生效（UPDATE/DROP/多语句全被拒且数据未变）。
测试 55 项：守卫 30 项（含"字符串里的分号不能误杀"）、连接与探查 15 项、
URL 构造 6 项、Copilot 流式 4 项。
107 lines
3.4 KiB
Python
107 lines
3.4 KiB
Python
"""SQL 守卫。
|
|
|
|
这是数据源接入里唯一不能出错的地方：SQL 由模型生成，不是人写的。
|
|
守卫放松一分，模型某次抽风就可能打到生产库上——所以这组用例只增不减。
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from app.data.guard import SqlRejected, check, first_verb, split_statements
|
|
|
|
|
|
def rejected(sql: str, *, readonly: bool = True) -> str | None:
|
|
try:
|
|
check(sql, readonly=readonly, source_name="测试库")
|
|
return None
|
|
except SqlRejected as e:
|
|
return str(e)
|
|
|
|
|
|
@pytest.mark.parametrize("sql", [
|
|
"SELECT * FROM orders",
|
|
" select 1 ",
|
|
"WITH t AS (SELECT 1) SELECT * FROM t",
|
|
"(SELECT a FROM b)",
|
|
"SHOW TABLES",
|
|
"DESCRIBE orders",
|
|
"EXPLAIN SELECT * FROM orders",
|
|
"/* 注释 */ SELECT 1",
|
|
"-- 行注释\nSELECT 2",
|
|
"SELECT 1;", # 尾分号不算第二条语句
|
|
])
|
|
def test_readonly_allows_queries(sql):
|
|
assert rejected(sql) is None, f"不该拒绝：{sql!r}"
|
|
|
|
|
|
@pytest.mark.parametrize("sql", [
|
|
"UPDATE orders SET x=1",
|
|
"DELETE FROM orders",
|
|
"INSERT INTO orders VALUES(1)",
|
|
"MERGE INTO a USING b ON (1=1)", # 白名单式判定才拦得住的方言私货
|
|
"CALL some_proc()",
|
|
"REPLACE INTO t VALUES(1)",
|
|
])
|
|
def test_readonly_blocks_writes(sql):
|
|
reason = rejected(sql)
|
|
assert reason and "只读" in reason
|
|
|
|
|
|
@pytest.mark.parametrize("sql", [
|
|
"DROP TABLE orders",
|
|
"TRUNCATE TABLE orders",
|
|
"ALTER TABLE t ADD c INT",
|
|
"GRANT ALL ON db.* TO u",
|
|
"CREATE TABLE t(a INT)",
|
|
])
|
|
def test_destructive_blocked_even_on_writable_source(sql):
|
|
"""结构变更和权限操作在可写源上也不放行——不可逆的事不交给模型发起。"""
|
|
reason = rejected(sql, readonly=False)
|
|
assert reason and "任何数据源" in reason
|
|
|
|
|
|
@pytest.mark.parametrize("sql", [
|
|
"SELECT 1; DROP TABLE users",
|
|
"SELECT 1;DROP TABLE users;",
|
|
"SELECT 1 /*x*/; DELETE FROM t",
|
|
])
|
|
def test_multi_statement_blocked(sql):
|
|
"""分号拼接必须在到达驱动之前断掉，不赌驱动默认不跑第二条。"""
|
|
reason = rejected(sql)
|
|
assert reason and "一条语句" in reason
|
|
|
|
|
|
@pytest.mark.parametrize("sql", [
|
|
"SELECT ';' AS x",
|
|
"SELECT 'a;b' FROM t WHERE c = 'd;e'",
|
|
"SELECT '它说：''你好'';' AS s", # 连续两个引号是转义，不是字符串结束
|
|
])
|
|
def test_semicolon_inside_string_is_not_a_second_statement(sql):
|
|
"""误杀比漏杀更让人无从下手：字符串字面量里的分号不算语句分隔。"""
|
|
assert rejected(sql) is None, f"误杀了：{sql!r}"
|
|
|
|
|
|
def test_writable_source_allows_dml():
|
|
for sql in ("UPDATE orders SET x=1", "INSERT INTO t VALUES(1)", "DELETE FROM t WHERE id=1"):
|
|
assert rejected(sql, readonly=False) is None
|
|
|
|
|
|
@pytest.mark.parametrize("sql", ["", " ", "-- 只有注释", "!!!"])
|
|
def test_empty_and_malformed_rejected(sql):
|
|
assert rejected(sql) is not None
|
|
|
|
|
|
def test_check_returns_statement_without_trailing_semicolon():
|
|
"""调用方必须执行返回值而不是原始输入，否则守卫等于没做。"""
|
|
assert check("SELECT 1;", readonly=True) == "SELECT 1"
|
|
|
|
|
|
def test_first_verb_sees_through_comments_and_parens():
|
|
assert first_verb("/* x */ (SELECT 1)") == "select"
|
|
assert first_verb("-- c\nDROP TABLE t") == "drop"
|
|
|
|
|
|
def test_split_keeps_quoted_semicolons_together():
|
|
assert len(split_statements("SELECT ';'")) == 1
|
|
assert len(split_statements("SELECT 1; SELECT 2")) == 2
|