mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 23:01:59 +00:00
ThearchyHelios
71561a0133
从 run 554a0f92 查出来的。那次问「userID=2 是谁」，9 次工具调用花掉 6 次、
前 29 秒没碰到一点真数据：
+ 2.1s db_schema → "还没有探查过结构，请在设置页里执行一次"
+ 4.9s db_schema ×7 个猜的表名 → "没有 user / users / member / …"
+ 7.7s SELECT 1 → 不信任前面的结果，先探一下连通性
+12.8s information_schema.tables ┐
+18.3s information_schema.tables ├ 自己把 schema 重建一遍
+21.5s information_schema.columns ┘
+29.2s 终于开始查真数据
根因：shop 的探查**失败过**（Lost connection to MySQL server during query），
但失败被存成了一个"成功但没表"的缓存，schema_synced_at 照常盖戳：
{"tables": {}, "total": 0,
"error": "OperationalError: (2013, 'Lost connection...')",
"available_schemas": ["prod_shop", "shop_dev", "toolsdb", …10 个]}
于是三处同时说假话：界面显示"结构未探查"、工具对 agent 说"还没有探查过"、
而那 10 个候选 schema 就躺在缓存里没人交出去。**用户明明点过探查，是超时了。**
**分清"还没探查"和"探查失败了"。** introspect 的结果加 failed 标记；
失败照样把缓存存下来（候选是下一步的线索），但不盖"已同步"的戳；
界面把红字和候选列出来；探查按钮的提示也分清是超时还是 schema 填错。
**给 agent 一个它做得到的下一步。** 以前那句"请在设置页里执行一次结构探查"
——agent 进不了设置页。现在换成实情 + 候选 schema + "不要猜表名，直接查
information_schema / sqlite_master"。那条自救路线 agent 最后也走通了，
只是自己摸索了三轮才找到，直接告诉它就行。
describe_table 同样：缓存空的时候说"数据源里没有 user"是在误导，真相是我们
什么都不知道。那次一口气问了 7 个表名，收到 7 条一模一样的
"没有 X。现有的对象：（还没探查过结构）"，白花一整步。
顺带把逐表反射换成 SQLAlchemy 2.0 的 get_multi_*，但**别指望它治超时**：
实测 60 张表逐表 121 次往返、批量 122 次——反射缓存本来就把 get_columns 和
get_pk_constraint 合并了，而 MySQL 方言压根没实现 get_multi_*，落到基类还是
每表一次 SHOW CREATE TABLE。换它是因为这是 2.0 的正经接口、少三层
try/except，不是因为它快。shop 那次超时多半是偶发的，所以真正的对策
不在那儿，在"失败要如实记成失败"。
重新探查了一次 shop：153 张表、38 秒、成功。之前 agent 要花 6 次调用、
29 秒拼出来的东西，现在工具描述第一眼就给了。
验证：后端 269 passed（test_datasource.py 新增 4 条，守"失败不能被说成没探查过"
和"给 agent 的下一步它做得到"）；tsc 干净；三个 check 不回归。
337 lines
13 KiB
Python
337 lines
13 KiB
Python
"""数据源连接层与工具层。
|
||
|
||
用真实的 SQLite 库跑——SQLAlchemy 统一了方言，SQLite 走的是和 MySQL/Oracle
|
||
同一套代码路径，所以这里验证的不只是 SQLite。各驱动特有的差异（Oracle 的
|
||
service_name、MySQL 的 charset）在 URL 构造那组用例里覆盖。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import sqlite3
|
||
from types import SimpleNamespace
|
||
|
||
import pytest
|
||
|
||
from app.data.engine import build_url, run_query, test_connection
|
||
from app.data.guard import QueryLimits, SqlRejected
|
||
from app.data.introspect import describe_table, introspect, summary, table_names
|
||
from app.main import app as _fastapi_app
|
||
from httpx import ASGITransport, AsyncClient
|
||
|
||
|
||
@pytest.fixture
|
||
async def client():
|
||
transport = ASGITransport(app=_fastapi_app)
|
||
async with AsyncClient(transport=transport, base_url="http://test") as c:
|
||
yield c
|
||
|
||
|
||
@pytest.fixture
|
||
def shop_db(tmp_path):
|
||
"""一个最小但真实的业务库：地区 / 商品 / 订单。"""
|
||
path = tmp_path / "shop.db"
|
||
db = sqlite3.connect(path)
|
||
db.executescript("""
|
||
CREATE TABLE regions (id INTEGER PRIMARY KEY, name TEXT NOT NULL);
|
||
CREATE TABLE orders (
|
||
id INTEGER PRIMARY KEY, region_id INTEGER NOT NULL,
|
||
qty INTEGER NOT NULL, amount REAL NOT NULL, ordered_at TEXT NOT NULL
|
||
);
|
||
""")
|
||
db.executemany("INSERT INTO regions(id,name) VALUES(?,?)",
|
||
[(1, "华东"), (2, "华北")])
|
||
db.executemany(
|
||
"INSERT INTO orders(id,region_id,qty,amount,ordered_at) VALUES(?,?,?,?,?)",
|
||
[(i, 1 + i % 2, i, i * 10.5, "2026-06-15") for i in range(1, 101)],
|
||
)
|
||
db.commit()
|
||
db.close()
|
||
return str(path)
|
||
|
||
|
||
@pytest.fixture
|
||
def source(shop_db):
|
||
return SimpleNamespace(
|
||
id="test-src", name="sales", kind="sqlite", database=shop_db,
|
||
host=None, port=None, username=None, password=None, options={},
|
||
readonly=True, description="测试销售库", schema_cache={},
|
||
)
|
||
|
||
|
||
async def test_connection_and_introspect(source):
|
||
assert (await test_connection(source))["ok"]
|
||
|
||
source.schema_cache = await introspect(source)
|
||
assert set(table_names(source)) == {"regions", "orders"}
|
||
|
||
orders = source.schema_cache["tables"]["orders"]
|
||
assert [c["name"] for c in orders["columns"]] == [
|
||
"id", "region_id", "qty", "amount", "ordered_at"
|
||
]
|
||
assert orders["primary_key"] == ["id"]
|
||
|
||
|
||
async def test_summary_is_compact_enough_for_a_prompt(source):
|
||
source.schema_cache = await introspect(source)
|
||
text = summary(source)
|
||
assert "orders" in text and "amount" in text
|
||
# 摘要要能塞进 system prompt——超出这个量级就该改成按需探查
|
||
assert len(text) < 2000
|
||
|
||
|
||
async def test_summary_says_so_when_never_introspected(source):
|
||
"""没探查过就明说，别让 Copilot 对着空气编表名。"""
|
||
assert "还没有探查过" in summary(source)
|
||
|
||
|
||
async def test_describe_table_accepts_qualified_name(source):
|
||
"""模型可能传 SCHEMA.TABLE 全名，也可能只传表名，两种都要认。"""
|
||
source.schema_cache = await introspect(source)
|
||
assert "amount" in describe_table(source, "orders")
|
||
assert "amount" in describe_table(source, "main.orders")
|
||
assert "amount" in describe_table(source, "ORDERS") # Oracle 惯用大写
|
||
|
||
|
||
async def test_describe_table_lists_unknown_table_options(source):
|
||
source.schema_cache = await introspect(source)
|
||
out = describe_table(source, "没有这张表")
|
||
assert "现有的对象" in out and "orders" in out
|
||
|
||
|
||
async def test_query_returns_rows(source):
|
||
result = await run_query(source, "SELECT COUNT(*) AS c FROM orders")
|
||
assert result.rows[0][0] == 100
|
||
assert result.columns == ["c"]
|
||
|
||
|
||
async def test_readonly_guard_applies_on_real_connection(source):
|
||
"""守卫不能只在单测里生效——真实连接路径上也必须拦住。"""
|
||
for sql in ("UPDATE orders SET qty=0", "DROP TABLE orders", "SELECT 1; DROP TABLE orders"):
|
||
with pytest.raises(SqlRejected):
|
||
await run_query(source, sql)
|
||
# 数据确实没被动过
|
||
assert (await run_query(source, "SELECT COUNT(*) FROM orders")).rows[0][0] == 100
|
||
|
||
|
||
async def test_result_set_is_capped(source):
|
||
"""一条 SELECT * 不该把内存和模型上下文一起打爆。"""
|
||
result = await run_query(source, "SELECT * FROM orders", limits=QueryLimits(max_rows=10))
|
||
assert result.row_count == 10
|
||
assert result.truncated is True
|
||
|
||
|
||
async def test_rows_are_json_serializable(source):
|
||
"""Decimal / date / bytes 在真实业务数据里遍地都是，不处理等于不可用。"""
|
||
result = await run_query(source, "SELECT amount, ordered_at FROM orders LIMIT 3")
|
||
json.dumps(result.to_payload()) # 不抛就算过
|
||
|
||
|
||
# --------------------------------------------------------------------------
|
||
# URL 构造：各驱动的方言差异
|
||
# --------------------------------------------------------------------------
|
||
|
||
def _src(kind, **kw):
|
||
base = dict(id="x", name="n", readonly=True, host="db", port=None,
|
||
database=None, username="u", password="p", options={})
|
||
base.update(kw)
|
||
return SimpleNamespace(kind=kind, **base)
|
||
|
||
|
||
@pytest.mark.parametrize("kind,expect", [
|
||
("mysql", "mysql+aiomysql://"),
|
||
("postgres", "postgresql+asyncpg://"),
|
||
("oracle", "oracle+oracledb://"),
|
||
])
|
||
def test_url_uses_async_drivers(kind, expect):
|
||
opts = {"service_name": "ORCL"} if kind == "oracle" else {}
|
||
assert build_url(_src(kind, database="d", options=opts)).startswith(expect)
|
||
|
||
|
||
def test_password_is_masked_by_default():
|
||
"""密码不该出现在日志、事件流或任何返回给前端的地方。"""
|
||
url = build_url(_src("mysql", database="d"))
|
||
assert "***" in url and ":p@" not in url
|
||
|
||
|
||
def test_oracle_requires_service_name_or_sid():
|
||
"""Oracle 最常见的配错点，要给明确提示而不是让驱动报晦涩的 ORA 错误。"""
|
||
with pytest.raises(ValueError, match="service_name"):
|
||
build_url(_src("oracle"))
|
||
|
||
|
||
def test_oracle_accepts_database_as_service_name():
|
||
assert "service_name=ORCLPDB1" in build_url(_src("oracle", database="ORCLPDB1"))
|
||
|
||
|
||
def test_unknown_kind_is_rejected():
|
||
with pytest.raises(ValueError, match="不支持"):
|
||
build_url(_src("mongodb"))
|
||
|
||
|
||
# --------------------------------------------------------------------------
|
||
# schema 发现：各家的 "schema" 根本不是同一个东西
|
||
# MySQL schema == database（SHOW schemas）
|
||
# Oracle schema == user（all_users）
|
||
# PG database > schema（pg_namespace）
|
||
# SQLite 只有 main + attached
|
||
# 所以这组用例守的是"不管哪一家，探不到东西时都要给出可走的下一步"。
|
||
# --------------------------------------------------------------------------
|
||
|
||
async def test_sqlite_degrades_without_information_schema(source):
|
||
"""SQLite 没有 information_schema，统计路径要安静地降级而不是报错。"""
|
||
from app.data.introspect import list_schemas, schema_stats
|
||
|
||
assert await schema_stats(source) == []
|
||
assert isinstance(await list_schemas(source), list)
|
||
|
||
|
||
async def test_qualified_name_omits_schema_when_absent(source):
|
||
"""没有 schema 概念时不要硬拼前缀，否则生成的 SQL 反而错。"""
|
||
source.schema_cache = await introspect(source)
|
||
assert source.schema_cache["tables"]["orders"]["qualified"] == "orders"
|
||
|
||
|
||
async def test_qualified_name_includes_schema_when_given(source):
|
||
"""有 schema 时必须带前缀——Oracle 上少了它就是 ORA-00942。"""
|
||
cache = await introspect(source, schema="main")
|
||
assert cache["tables"]["orders"]["qualified"] == "main.orders"
|
||
|
||
|
||
async def test_bad_schema_returns_candidates_instead_of_raising(source, monkeypatch):
|
||
"""schema 填错是高频事故（大小写、拼写），要给可选项而不是抛驱动异常。"""
|
||
from app.data import introspect as mod
|
||
|
||
async def fake_list(_source):
|
||
return ["ANALYTICS", "STAGING"]
|
||
|
||
monkeypatch.setattr(mod, "list_schemas", fake_list)
|
||
cache = await mod.introspect(source, schema="根本不存在")
|
||
assert cache["tables"] == {}
|
||
assert cache["available_schemas"] == ["ANALYTICS", "STAGING"]
|
||
|
||
|
||
async def test_summary_points_at_candidate_schemas(source, monkeypatch):
|
||
"""摘要要把候选说给 Copilot 听，否则它只会对着空结构编表名。"""
|
||
from app.data import introspect as mod
|
||
|
||
async def fake_list(_source):
|
||
return ["ANALYTICS", "STAGING"]
|
||
|
||
monkeypatch.setattr(mod, "list_schemas", fake_list)
|
||
source.schema_cache = await mod.introspect(source, schema="根本不存在")
|
||
text = summary(source)
|
||
assert "ANALYTICS" in text and "该账号名下没有对象" in text
|
||
|
||
|
||
def test_schema_stats_sql_covers_every_supported_kind():
|
||
"""新增数据库类型时别忘了这里——漏了就退化成黑名单猜测。"""
|
||
from app.data.engine import SUPPORTED_KINDS
|
||
from app.data.introspect import _SCHEMA_STATS_SQL
|
||
|
||
# sqlite 显式短路，其余要么有专用 SQL，要么落到 information_schema 标准路径
|
||
assert "oracle" in _SCHEMA_STATS_SQL
|
||
assert "_standard" in _SCHEMA_STATS_SQL
|
||
for kind in SUPPORTED_KINDS:
|
||
assert kind in _SCHEMA_STATS_SQL or kind in ("sqlite", "mysql", "mariadb",
|
||
"postgres", "postgresql")
|
||
|
||
|
||
# --------------------------------------------------------------------------
|
||
# 探查失败 ≠ 还没探查
|
||
# --------------------------------------------------------------------------
|
||
#
|
||
# run 554a0f92：shop 的探查因为 "Lost connection to MySQL server during query"
|
||
# 失败了，但失败被存成了一个"成功但没表"的缓存，schema_synced_at 照常盖戳。
|
||
# 于是 db_schema 工具对 agent 说"还没有探查过结构，请在设置页里执行一次"——
|
||
# 一句错话，配一个 agent 根本做不到的动作。它只好自己摸：猜 7 个表名、发一条
|
||
# SELECT 1 试连通性、再用三条 information_schema 把 schema 重建一遍。
|
||
# 9 次工具调用花掉 6 次，前 29 秒没碰到一点真数据。
|
||
#
|
||
# 而 available_schemas 那 10 个候选一直躺在缓存里，没人交给它。
|
||
|
||
|
||
class _FailedCache:
|
||
"""探查失败之后 schema_cache 的样子。"""
|
||
|
||
name = "shop"
|
||
kind = "mysql"
|
||
database = "prod_shop"
|
||
description = ""
|
||
readonly = True
|
||
schema_cache = {
|
||
"tables": {}, "total": 0, "failed": True,
|
||
"error": "OperationalError: (2013, 'Lost connection to MySQL server during query')",
|
||
"available_schemas": ["prod_shop", "shop_dev", "toolsdb"],
|
||
}
|
||
|
||
|
||
class _NeverRun(_FailedCache):
|
||
"""从来没探查过的样子。和上面必须说不一样的话。"""
|
||
|
||
schema_cache: dict = {}
|
||
|
||
|
||
def test_why_empty_tells_the_two_cases_apart():
|
||
from app.data.introspect import why_empty
|
||
|
||
assert "失败" in why_empty(_FailedCache.schema_cache)
|
||
assert "Lost connection" in why_empty(_FailedCache.schema_cache)
|
||
assert why_empty({}) == "还没有探查过结构"
|
||
|
||
|
||
def test_schema_tool_tells_the_agent_the_truth():
|
||
"""给 agent 的那句话要是实情，而且要给一个它做得到的下一步。"""
|
||
from app.tools.datasource import _no_schema_note
|
||
|
||
note = _no_schema_note(_FailedCache)
|
||
assert "失败" in note and "Lost connection" in note
|
||
# 候选一直在缓存里，以前没人交出去
|
||
assert "prod_shop" in note and "toolsdb" in note
|
||
# agent 进不了设置页。它做得到的是自己查数据字典——那条路它最后也走通了，
|
||
# 只是自己摸索了三轮
|
||
assert "information_schema" in note
|
||
assert "不要猜表名" in note
|
||
# 不能再说"请到设置页"
|
||
assert "设置页" not in note
|
||
|
||
fresh = _no_schema_note(_NeverRun)
|
||
assert "还没有探查过" in fresh
|
||
assert "失败" not in fresh
|
||
|
||
|
||
def test_describe_table_does_not_claim_the_table_is_missing():
|
||
"""缓存是空的时候，说"里面没有这张表"是在误导——真相是我们什么都不知道。"""
|
||
from app.data.introspect import describe_table
|
||
|
||
out = describe_table(_FailedCache, "user")
|
||
assert "无法判断" in out
|
||
assert "失败" in out
|
||
# 以前是 "数据源「shop」里没有 user。现有的对象：（还没探查过结构）"
|
||
assert "现有的对象" not in out
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_introspect_failure_does_not_get_a_synced_stamp(client, monkeypatch):
|
||
"""失败要把缓存留下（候选是下一步的线索），但不能盖"已同步"的戳。
|
||
|
||
以前盖了，于是界面说同步过、工具说还没探查过，两边都不说真话。
|
||
"""
|
||
import app.api.datasources as mod
|
||
|
||
async def failed(source, schema=None):
|
||
return dict(_FailedCache.schema_cache)
|
||
|
||
created = await client.post("/api/datasources", json={
|
||
"name": "introspect_fail", "kind": "sqlite", "database": ":memory:",
|
||
})
|
||
assert created.status_code == 201, created.text
|
||
sid = created.json()["id"]
|
||
|
||
monkeypatch.setattr(mod.introspect_mod, "introspect", failed)
|
||
out = (await client.post(f"/api/datasources/{sid}/introspect")).json()
|
||
|
||
assert out["schema_synced_at"] is None, "失败不该盖已同步的戳"
|
||
assert "Lost connection" in out["schema_error"]
|
||
assert out["available_schemas"] == ["prod_shop", "shop_dev", "toolsdb"]
|
||
assert out["table_count"] == 0
|