diff --git a/backend/app/tools/datasource.py b/backend/app/tools/datasource.py index 575fc847..88aa8b03 100644 --- a/backend/app/tools/datasource.py +++ b/backend/app/tools/datasource.py @@ -35,7 +35,7 @@ from app.data import catalog as data_catalog from app.data import introspect, table_versions from app.data.sqlcheck import SqlChecker from app.data.engine import SnapshotTampered, masked_columns, query_timeout, run_query -from app.data.guard import QueryLimits, SqlRejected, is_write +from app.data.guard import QueryLimits, SqlRejected, check, is_write from app.data.tabular import UNSHAPED_NOTE from app.db.models import DataSource from app.tools.registry import ToolContext, Versions @@ -57,6 +57,10 @@ LATE_PIN_CODE = "data_version_pinned" #: 每次查询之后基于数据目录的 SQL 检查最多等这么久。正常一条 SQL 几毫秒；等不到就放弃这一次检查， #: 查询结果照常交回——检查是附加的，不能拖慢查询 SQL_CHECK_TIMEOUT_S = 2.0 +#: force_readonly 的工具拒掉写语句时交回的原因（接在「SQL 被拒绝：」后面）。不说「数据源是只读的」：源本身 +#: 可能是可写的，是这次调用只许读；也不提「另建一个可写的数据源」——核查用不着写 +FORCED_READONLY = ("这次查询只允许读取数据，这条语句会修改数据，不被允许。" + "只能执行不改数据的 SELECT / WITH / SHOW / DESCRIBE / EXPLAIN") class _QueryArgs(BaseModel): @@ -165,12 +169,13 @@ _FROZEN_SCHEMA_KEYS = ("schema", "synced_at", "truncated", "total", "import_mode CATALOG_SNAPSHOT_KEY = "catalog" -def _query_description(source: Any) -> str: +def _query_description(source: Any, *, readonly: bool | None = None) -> str: + """readonly 不给就按源本身；force_readonly 的工具传 True（源可写，这次调用也只许读）。""" tables = introspect.table_names(source) head = f"在数据源「{source.name}」上执行 SQL 查询。" if source.description: head += f"{source.description}。" - head += "只读。" if source.readonly else "可写（写操作需要人工审批）。" + head += "只读。" if (source.readonly if readonly is None else readonly) else "可写（写操作需要人工审批）。" if tables: listed = "、".join(tables[:25]) more = f" 等 {len(tables)} 张表" if len(tables) > 25 else "" @@ -194,13 +199,16 @@ def _query_description(source: Any) -> str: async def build_datasource_tools( - names: list[str], ctx: ToolContext, session: AsyncSession + names: list[str], ctx: ToolContext, session: AsyncSession, *, force_readonly: bool = False ) -> list[StructuredTool]: """按工具名反查数据源并构造工具。build_tools 的第四个来源。 上传源先按 ctx.data_versions 解析成绑定快照的视图（见 _resolve），工具描述、表结构、 查询都按那个版本来。固定的版本找不到时建一个同名的替身（_unavailable_tool）， 每次调用都说明原因，不拿当前版本顶上。 + + force_readonly：查询工具一律按只读守卫检查，可写源上的写语句也拒绝（见 _make_query_tool）。 + 核查（engine/verify_agent.py）建工具时传它：那一路没有审批关卡，绝不能写。 """ wanted: dict[str, list[str]] = {} for name in names: @@ -230,7 +238,7 @@ async def build_datasource_tools( entries = await _catalog_of(row) for prefix in prefixes: if prefix == QUERY_PREFIX: - tools.append(_make_query_tool(source, ctx, fixed=fixed, catalog=entries)) + tools.append(_make_query_tool(source, ctx, fixed=fixed, catalog=entries, force_readonly=force_readonly)) else: tools.append(_make_schema_tool(source, fixed=fixed, catalog=entries)) return tools @@ -441,11 +449,15 @@ def _unavailable_tool(row: DataSource, prefix: str, reason: str) -> StructuredTo def _make_query_tool(source: Any, ctx: ToolContext, *, fixed: bool = False, - catalog: Mapping[str, data_catalog.CatalogEntry] | None = None) -> StructuredTool: + catalog: Mapping[str, data_catalog.CatalogEntry] | None = None, + force_readonly: bool = False) -> StructuredTool: """source 是手工源的 DataSource，或上传源绑定快照的 SourceView（字段同名）。 fixed：上传源的版本是这次运行固定的（找不到时的说法不同）。 catalog：建工具时读到的数据目录，随表结构快照冻结（_store_schema）。 + force_readonly：不管源可不可写，一律按只读守卫检查，写语句交回 FORCED_READONLY、不执行。只靠守卫这一道： + 连接池按源缓存、不分只读与否（engine.cache_key），可写源借不到连接层只读的连接——别为此另起一个只读池， + 缓存键相同，会把这个源平时用的可写连接也换成只读的。 """ # 数据源在 options.query_timeout_s 里配了就用它，否则取缺省。数据库按它停下语句， # 引擎按它告诉界面上限是多少（metadata["timeout_s"]） @@ -463,6 +475,15 @@ def _make_query_tool(source: Any, ctx: ToolContext, *, fixed: bool = False, async def _run(sql: str, limit: int | None = None) -> str: if gone := _vanished(source, fixed): return f"{DATA_UNAVAILABLE}{gone}" + if force_readonly and not source.readonly: + # 只读源由 run_query 自己按只读检查（原因照旧说「数据源是只读的」）；可写源在这里先按只读判一遍。 + # 先过一遍和 run_query 相同的检查：多条语句、结构变更这类问题照原来的说法交回 + try: + statement = check(sql, readonly=False, source_name=source.name, dialect=(source.kind or "").lower()) + except SqlRejected as e: + return f"SQL 被拒绝：{e}" + if is_write(statement, dialect=(source.kind or "").lower()): + return f"SQL 被拒绝：{FORCED_READONLY}" limits = QueryLimits(timeout_seconds=seconds, **({"max_rows": int(limit)} if limit else {})) try: result = await run_query(source, sql, limits=limits) @@ -520,12 +541,14 @@ def _make_query_tool(source: Any, ctx: ToolContext, *, fixed: bool = False, return StructuredTool( name=f"{QUERY_PREFIX}{source.name}", - description=_query_description(source), + description=_query_description(source, readonly=True if force_readonly else None), args_schema=_QueryArgs, coroutine=_run, func=None, - # 同一个工具，危不危险看这一次的 SQL。审批关卡经 registry.call_is_dangerous 问它 - metadata={"dangerous_if": lambda args: is_dangerous_call(source, str(args.get("sql") or "")), + # 同一个工具，危不危险看这一次的 SQL。审批关卡经 registry.call_is_dangerous 问它。 + # force_readonly 的工具写不了，没有要人批的调用 + metadata={"dangerous_if": (lambda args: False) if force_readonly + else (lambda args: is_dangerous_call(source, str(args.get("sql") or ""))), "timeout_s": seconds}, ) diff --git a/backend/tests/test_datasource_force_readonly.py b/backend/tests/test_datasource_force_readonly.py new file mode 100644 index 00000000..856a5cb3 --- /dev/null +++ b/backend/tests/test_datasource_force_readonly.py @@ -0,0 +1,106 @@ +"""数据源查询工具的 force_readonly：核查（双盲复算、反方审稿、审批质询）只许读数据。 + +核查的 agent 是另起的一路，没有审批关卡兜底；它要查的数据源却可能是可写的。所以建工具时传 force_readonly=True： +不管源本身可不可写，一律按只读守卫检查——可写源上的写语句也拒绝，原因交回模型（它据此改写），数据原样不动； +查询照常。工具描述跟着写「只读」，审批判定也不再把写语句当成要人批的危险调用（反正执行不了）。 +不传的时候一切照旧：可写源上的写语句仍是要审批的危险调用。 + +夹具是一个现造的可写 SQLite 文件，表和数字都是虚构的。 +""" +from __future__ import annotations + +import json +import sqlite3 +import uuid + +import pytest + +from app.data import introspect +from app.data.engine import engines +from app.db.base import SessionLocal +from app.db.models import DataSource +from app.tools.datasource import build_datasource_tools +from app.tools.registry import ToolContext, Versions + + +def rows_in(path: str) -> list[tuple]: + db = sqlite3.connect(path) + try: + return db.execute("SELECT id, amount FROM orders ORDER BY id").fetchall() + finally: + db.close() + + +@pytest.fixture +async def writable(tmp_path): + """一个可写的手工源：orders 两行，金额 120 和 80。""" + path = tmp_path / f"{uuid.uuid4().hex[:8]}.db" + db = sqlite3.connect(path) + db.executescript("CREATE TABLE orders (id INTEGER PRIMARY KEY, amount INTEGER);" + "INSERT INTO orders VALUES (1, 120), (2, 80);") + db.commit() + db.close() + name = f"shop_{uuid.uuid4().hex[:8]}" + async with SessionLocal() as session: + row = DataSource(name=name, kind="sqlite", database=str(path), readonly=False, enabled=True) + session.add(row) + await session.flush() + row.schema_cache = await introspect.introspect(row) + await session.commit() + source_id = row.id + yield name, str(path) + await engines.invalidate(source_id) + + +async def tools(name: str, **kwargs) -> dict: + async with SessionLocal() as session: + built = await build_datasource_tools([f"db_query__{name}", f"db_schema__{name}"], + ToolContext(data_versions=Versions.CURRENT), session, **kwargs) + return {t.name: t for t in built} + + +@pytest.mark.parametrize("sql", [ + "DELETE FROM orders", + "UPDATE orders SET amount = 0", + "INSERT INTO orders VALUES (3, 1)", + "WITH k AS (SELECT 1) DELETE FROM orders WHERE id = 1", +]) +async def test_writes_are_refused_on_a_writable_source(writable, sql): + name, path = writable + query = (await tools(name, force_readonly=True))[f"db_query__{name}"] + got = await query.ainvoke({"sql": sql}) + assert got.startswith("SQL 被拒绝：") and "只允许读取数据" in got, got + assert rows_in(path) == [(1, 120), (2, 80)], "数据原样不动" + + +async def test_structural_changes_are_refused_with_the_usual_reason(writable): + name, path = writable + query = (await tools(name, force_readonly=True))[f"db_query__{name}"] + got = await query.ainvoke({"sql": "DROP TABLE orders"}) + assert got.startswith("SQL 被拒绝：") and "结构变更" in got + assert rows_in(path) == [(1, 120), (2, 80)] + + +async def test_queries_still_work_and_land_as_snapshots(writable): + name, _ = writable + built = await tools(name, force_readonly=True) + got = json.loads(await built[f"db_query__{name}"].ainvoke({"sql": "SELECT SUM(amount) AS total FROM orders"})) + assert got["rows"] == [[200]] and got["artifact"], "查询照常，结果照常落成查询快照" + assert "orders" in await built[f"db_schema__{name}"].ainvoke({}) + + +async def test_the_tool_says_read_only_and_needs_no_approval(writable): + name, _ = writable + forced = (await tools(name, force_readonly=True))[f"db_query__{name}"] + assert "只读。" in forced.description and "可写" not in forced.description + assert forced.metadata["dangerous_if"]({"sql": "DELETE FROM orders"}) is False, "写语句反正执行不了，不用找人批" + + +async def test_without_the_flag_nothing_changes(writable): + name, path = writable + normal = (await tools(name))[f"db_query__{name}"] + assert "可写（写操作需要人工审批）" in normal.description + assert normal.metadata["dangerous_if"]({"sql": "DELETE FROM orders"}) is True + got = json.loads(await normal.ainvoke({"sql": "UPDATE orders SET amount = 90 WHERE id = 2"})) + assert got["rows"] == [[1]], "不开 force_readonly 时可写源照常能写（审批由节点把关）" + assert rows_in(path) == [(1, 120), (2, 90)]