diff --git a/backend/app/api/datasources.py b/backend/app/api/datasources.py new file mode 100644 index 00000000..0fd56e73 --- /dev/null +++ b/backend/app/api/datasources.py @@ -0,0 +1,206 @@ +"""数据源的增删改查、连接测试、结构探查。 + +密码的处理规则和 Provider 完全一致(app/api/settings.py):入库加密、 +出站只给掩码、PATCH 时不传表示不动、传空串表示清空。 +""" +from __future__ import annotations + +from datetime import datetime, timezone +from typing import Any + +from fastapi import APIRouter, Depends, HTTPException +from pydantic import BaseModel, Field +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from app.core.crypto import encrypt, mask +from app.data import introspect as introspect_mod +from app.data.engine import SUPPORTED_KINDS, engines, test_connection +from app.db.base import get_session +from app.db.models import DataSource +from app.tools.datasource import tool_names + +router = APIRouter(prefix="/api/datasources", tags=["datasources"]) + +# 名字会成为工具名的一部分(db_query__),而模型的 function name 只认 +# ASCII。所以这里收紧成 slug,中文说明放 description 字段。 +_NAME_PATTERN = r"^[a-z][a-z0-9_]{0,40}$" + + +class DataSourceIn(BaseModel): + name: str = Field(pattern=_NAME_PATTERN, + description="英文标识,会成为工具名的一部分,如 sales") + kind: str + host: str | None = None + port: int | None = None + database: str | None = None + username: str | None = None + password: str | None = None + options: dict[str, Any] = Field(default_factory=dict) + readonly: bool = True + description: str = "" + enabled: bool = True + + +class DataSourcePatch(BaseModel): + kind: str | None = None + host: str | None = None + port: int | None = None + database: str | None = None + username: str | None = None + password: str | None = None # 空串=清空;不传=不动 + options: dict[str, Any] | None = None + readonly: bool | None = None + description: str | None = None + enabled: bool | None = None + + +class DataSourceOut(BaseModel): + id: str + name: str + kind: str + host: str | None + port: int | None + database: str | None + username: str | None + options: dict[str, Any] + readonly: bool + description: str + enabled: bool + password_masked: str = "" + has_password: bool = False + table_count: int = 0 + schema_synced_at: str | None = None + tools: list[str] = Field(default_factory=list) + + +def _to_out(row: DataSource) -> DataSourceOut: + tables = (row.schema_cache or {}).get("tables") or {} + return DataSourceOut( + id=row.id, name=row.name, kind=row.kind, host=row.host, port=row.port, + database=row.database, username=row.username, options=row.options or {}, + readonly=row.readonly, description=row.description, enabled=row.enabled, + password_masked=mask(row.password), + has_password=bool(row.password), + table_count=len(tables), + schema_synced_at=row.schema_synced_at.isoformat() if row.schema_synced_at else None, + tools=tool_names(row), + ) + + +async def _get_or_404(session: AsyncSession, source_id: str) -> DataSource: + row = await session.get(DataSource, source_id) + if not row: + raise HTTPException(404, "数据源不存在") + return row + + +@router.get("", response_model=list[DataSourceOut]) +async def list_sources(session: AsyncSession = Depends(get_session)) -> list[DataSourceOut]: + rows = (await session.execute(select(DataSource).order_by(DataSource.name))).scalars() + return [_to_out(r) for r in rows] + + +@router.get("/kinds") +async def list_kinds() -> dict[str, Any]: + """支持哪些数据库,以及各自需要填什么——前端表单据此渲染。""" + return { + "kinds": [ + {"value": "mysql", "label": "MySQL / MariaDB", "default_port": 3306, + "needs": ["host", "database", "username", "password"], + "hint": "options 可填 charset(默认 utf8mb4)"}, + {"value": "postgres", "label": "PostgreSQL", "default_port": 5432, + "needs": ["host", "database", "username", "password"], "hint": ""}, + {"value": "oracle", "label": "Oracle", "default_port": 1521, + "needs": ["host", "username", "password"], + "hint": "库名填 service_name,或在 options 里给 service_name / sid。" + "驱动走 thin 模式,不需要装 Instant Client"}, + {"value": "sqlite", "label": "SQLite(文件)", "default_port": None, + "needs": ["database"], "hint": "库名填数据库文件的绝对路径"}, + ], + "supported": list(SUPPORTED_KINDS), + } + + +@router.post("", response_model=DataSourceOut, status_code=201) +async def create_source( + payload: DataSourceIn, session: AsyncSession = Depends(get_session) +) -> DataSourceOut: + if payload.kind not in SUPPORTED_KINDS: + raise HTTPException(400, f"不支持的类型 {payload.kind},支持:{', '.join(SUPPORTED_KINDS)}") + exists = (await session.execute( + select(DataSource).where(DataSource.name == payload.name) + )).scalar_one_or_none() + if exists: + raise HTTPException(409, f"已经有叫 {payload.name} 的数据源了") + + row = DataSource( + **payload.model_dump(exclude={"password"}), + password=encrypt(payload.password), + ) + session.add(row) + await session.commit() + await session.refresh(row) + return _to_out(row) + + +@router.patch("/{source_id}", response_model=DataSourceOut) +async def update_source( + source_id: str, payload: DataSourcePatch, session: AsyncSession = Depends(get_session) +) -> DataSourceOut: + row = await _get_or_404(session, source_id) + data = payload.model_dump(exclude_unset=True) + if "password" in data: + # 空串表示清空,非空表示换新的 + data["password"] = encrypt(data["password"]) if data["password"] else None + for key, value in data.items(): + setattr(row, key, value) + await session.commit() + await session.refresh(row) + # 连接参数可能变了,旧连接池不能再用——否则改完密码还在用旧连接,很难排查 + await engines.invalidate(source_id) + return _to_out(row) + + +@router.delete("/{source_id}", status_code=204) +async def delete_source(source_id: str, session: AsyncSession = Depends(get_session)) -> None: + row = await _get_or_404(session, source_id) + await session.delete(row) + await session.commit() + await engines.invalidate(source_id) + + +@router.post("/{source_id}/test") +async def test_source(source_id: str, session: AsyncSession = Depends(get_session)) -> dict[str, Any]: + """测连接。失败时把驱动的原始报错带出来——这类问题九成靠错误信息定位。""" + row = await _get_or_404(session, source_id) + return await test_connection(row) + + +@router.post("/{source_id}/introspect", response_model=DataSourceOut) +async def introspect_source( + source_id: str, schema: str | None = None, session: AsyncSession = Depends(get_session) +) -> DataSourceOut: + """探查结构并缓存。显式动作,不做后台轮询——生产库不该被实验工具定时扫。""" + row = await _get_or_404(session, source_id) + try: + row.schema_cache = await introspect_mod.introspect(row, schema=schema) + except Exception as e: # noqa: BLE001 + raise HTTPException(400, f"探查失败:{type(e).__name__}: {e}") from e + row.schema_synced_at = datetime.now(timezone.utc) + await session.commit() + await session.refresh(row) + return _to_out(row) + + +@router.get("/{source_id}/schema") +async def get_schema(source_id: str, table: str | None = None, + session: AsyncSession = Depends(get_session)) -> dict[str, Any]: + row = await _get_or_404(session, source_id) + if table: + return {"table": table, "detail": introspect_mod.describe_table(row, table)} + return { + "tables": introspect_mod.table_names(row), + "summary": introspect_mod.summary(row), + "synced_at": row.schema_synced_at.isoformat() if row.schema_synced_at else None, + } diff --git a/backend/app/data/__init__.py b/backend/app/data/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/backend/app/data/engine.py b/backend/app/data/engine.py new file mode 100644 index 00000000..a3f82b7f --- /dev/null +++ b/backend/app/data/engine.py @@ -0,0 +1,241 @@ +"""数据源连接层。各家数据库的方言差异全部收敛在这个文件里。 + +上层(工具、探查、API)只认 DataSource 这个模型,不需要知道 Oracle 要 service_name、 +MySQL 要 charset、asyncpg 不认 sslmode 这些事。哪天要加一种库,改这里一处。 +""" +from __future__ import annotations + +import asyncio +import json +from dataclasses import dataclass +from typing import Any + +from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine +from sqlalchemy import text + +from app.core.crypto import decrypt +from app.data.guard import QueryLimits, SqlRejected, check + +# 各方言的 SQLAlchemy 驱动。都选 asyncio 原生驱动,没有一个需要装数据库客户端: +# oracledb 走 thin 模式(纯 Python 协议实现),省掉了 Oracle Instant Client 这个 +# 历来最劝退的一步。 +_DRIVERS = { + "mysql": "mysql+aiomysql", + "mariadb": "mysql+aiomysql", + "postgres": "postgresql+asyncpg", + "postgresql": "postgresql+asyncpg", + "oracle": "oracle+oracledb", + "sqlite": "sqlite+aiosqlite", +} + +_DEFAULT_PORTS = { + "mysql": 3306, "mariadb": 3306, + "postgres": 5432, "postgresql": 5432, + "oracle": 1521, +} + +SUPPORTED_KINDS = tuple(sorted(set(_DRIVERS))) + + +@dataclass +class QueryResult: + columns: list[str] + rows: list[list[Any]] + row_count: int + truncated: bool + elapsed_ms: int + sql: str + + def to_payload(self) -> dict[str, Any]: + return { + "columns": self.columns, + "rows": self.rows, + "row_count": self.row_count, + "truncated": self.truncated, + "elapsed_ms": self.elapsed_ms, + "sql": self.sql, + } + + +def build_url(source: Any, *, reveal: bool = False) -> str: + """拼连接串。reveal=False 时把密码换成星号——日志和错误信息里用这个。 + + 密码永远不该出现在日志、事件流、或者返回给前端的任何地方。默认遮掉, + 真要连库时才显式 reveal=True。 + """ + kind = (source.kind or "").lower() + driver = _DRIVERS.get(kind) + if not driver: + raise ValueError(f"不支持的数据库类型:{source.kind}(支持 {', '.join(SUPPORTED_KINDS)})") + + options: dict[str, Any] = dict(source.options or {}) + + if kind == "sqlite": + # SQLite 没有主机和账号,database 就是文件路径 + return f"{driver}:///{source.database or ':memory:'}" + + password = decrypt(source.password) or "" if reveal else "***" + from urllib.parse import quote_plus + + user = quote_plus(source.username or "") + pw = quote_plus(password) if reveal else password + auth = f"{user}:{pw}@" if user else "" + port = source.port or _DEFAULT_PORTS.get(kind) + host = f"{source.host or 'localhost'}:{port}" if port else (source.host or "localhost") + + if kind == "oracle": + # Oracle 用 service_name 或 sid,不是"数据库名"。这是最常见的配错点, + # 所以两种都认,并且在都没填时给出明确提示而不是让驱动报晦涩的 ORA 错误。 + service = options.pop("service_name", None) or source.database + sid = options.pop("sid", None) + if not service and not sid: + raise ValueError("Oracle 需要在 options 里给 service_name 或 sid(也可以填在库名里)") + tail = f"?service_name={service}" if service else f"?sid={sid}" + query = "&".join(f"{k}={v}" for k, v in options.items()) + return f"{driver}://{auth}{host}/{tail}{('&' + query) if query else ''}" + + query = "&".join(f"{k}={v}" for k, v in options.items()) + return f"{driver}://{auth}{host}/{source.database or ''}{('?' + query) if query else ''}" + + +class EngineCache: + """按数据源 id 缓存 engine。连接池的建立不便宜,不该每次查询都重来一遍。 + + 配置变了要显式 invalidate——改了密码却还在用旧连接,排查起来很费神。 + """ + + def __init__(self) -> None: + self._engines: dict[str, AsyncEngine] = {} + self._lock = asyncio.Lock() + + async def get(self, source: Any) -> AsyncEngine: + key = str(source.id) + cached = self._engines.get(key) + if cached is not None: + return cached + async with self._lock: + if key in self._engines: + return self._engines[key] + engine = create_async_engine( + build_url(source, reveal=True), + pool_size=3, + max_overflow=2, + pool_pre_ping=True, # 长时间空闲后连接会被数据库掐掉,先探活再用 + pool_recycle=1800, + ) + self._engines[key] = engine + return engine + + async def invalidate(self, source_id: str) -> None: + engine = self._engines.pop(str(source_id), None) + if engine is not None: + await engine.dispose() + + async def close(self) -> None: + for engine in list(self._engines.values()): + await engine.dispose() + self._engines.clear() + + +engines = EngineCache() + + +async def test_connection(source: Any) -> dict[str, Any]: + """测连接。失败时把驱动的原始错误带出来——这类问题九成靠错误信息定位。""" + import time + + started = time.perf_counter() + try: + engine = await engines.get(source) + async with engine.connect() as conn: + probe = "SELECT 1 FROM DUAL" if (source.kind or "").lower() == "oracle" else "SELECT 1" + await conn.execute(text(probe)) + return { + "ok": True, + "elapsed_ms": int((time.perf_counter() - started) * 1000), + "url": build_url(source), # 遮掉密码的版本 + } + except Exception as e: # noqa: BLE001 + return { + "ok": False, + "error": f"{type(e).__name__}: {e}", + "elapsed_ms": int((time.perf_counter() - started) * 1000), + "url": build_url(source), + } + + +async def run_query( + source: Any, sql: str, *, limits: QueryLimits | None = None +) -> QueryResult: + """执行一条查询。守卫先过,再谈执行。 + + 注意执行的是 guard.check 的返回值而不是原始输入——守卫会把尾分号之类 + 规范掉,执行原文等于绕过了守卫。 + """ + import time + + limits = limits or QueryLimits() + statement = check(sql, readonly=bool(source.readonly), source_name=source.name) + + engine = await engines.get(source) + started = time.perf_counter() + + async def _exec() -> tuple[list[str], list[list[Any]], bool]: + async with engine.connect() as conn: + cursor = await conn.stream(text(statement)) + columns = list(cursor.keys()) + rows: list[list[Any]] = [] + truncated = False + size = 0 + async for row in cursor: + # 逐行累加,边收边判上限:一次性 fetchall 一个亿级表就晚了 + values = [_jsonable(v) for v in row] + rows.append(values) + size += len(json.dumps(values, ensure_ascii=False, default=str)) + if len(rows) >= limits.max_rows or size >= limits.max_bytes: + truncated = True + break + return columns, rows, truncated + + try: + columns, rows, truncated = await asyncio.wait_for( + _exec(), timeout=limits.timeout_seconds + ) + except asyncio.TimeoutError as e: + raise SqlRejected( + f"查询超过 {limits.timeout_seconds}s 被中断。加上 WHERE 条件或 LIMIT 缩小范围。" + ) from e + + return QueryResult( + columns=columns, + rows=rows, + row_count=len(rows), + truncated=truncated, + elapsed_ms=int((time.perf_counter() - started) * 1000), + sql=statement, + ) + + +def _jsonable(value: Any) -> Any: + """把驱动返回的类型转成能进 JSON 的形态。 + + Decimal、date、datetime、bytes 这些直接 json.dumps 会炸,而它们在真实业务 + 数据里遍地都是——金额是 Decimal、时间是 datetime,不处理等于不可用。 + """ + from datetime import date, datetime, time as dtime + from decimal import Decimal + + if value is None or isinstance(value, (str, int, float, bool)): + return value + if isinstance(value, Decimal): + # 转 float 会丢精度,而金额恰恰不能丢。字符串保真,下游要算再自己转。 + return str(value) + if isinstance(value, (datetime, date, dtime)): + return value.isoformat() + if isinstance(value, (bytes, bytearray, memoryview)): + raw = bytes(value) + try: + return raw.decode() + except UnicodeDecodeError: + return f"" + return str(value) diff --git a/backend/app/data/guard.py b/backend/app/data/guard.py new file mode 100644 index 00000000..4b6022d1 --- /dev/null +++ b/backend/app/data/guard.py @@ -0,0 +1,146 @@ +"""SQL 安全守卫。 + +这是数据源接入里唯一不能出错的地方:SQL 是 Copilot 或 agent 生成的,不是人写的。 +模型会写出什么,事前无法穷举——所以防线不能建立在"模型应该不会这么干"上。 + +三层,从外到内: + +1. **只读判定**:只读数据源只放行确定无副作用的语句。用白名单而不是黑名单—— + 黑名单永远漏,`TRUNCATE`、`MERGE`、`CALL`、`GRANT`、各家方言的私货写不完。 +2. **多语句拦截**:`SELECT 1; DROP TABLE users` 这种拼接必须在到达驱动之前就断掉。 + DBAPI 层的 `execute` 多数驱动只跑第一条,但"多数"不是"全部",MySQL 开了 + `CLIENT_MULTI_STATEMENTS` 就会全跑。不赌驱动的默认值。 +3. **结果集上限**:行数 + 字节 + 超时。一条 `SELECT * FROM 亿级表` 不该把后端打爆, + 也不该把几百 MB 塞进模型的上下文。 + +刻意不引入 sqlparse 之类的重型依赖:完整解析 SQL 方言是个无底洞,而这里需要的 +判断很窄——"第一个关键字是什么"和"有没有第二条语句"。窄问题用窄解法, +把复杂度留在能看懂的范围内。 +""" +from __future__ import annotations + +import re +from dataclasses import dataclass + +# 只读数据源放行的首关键字。只有确定不产生副作用的才进这个名单。 +# 注意 SHOW/DESCRIBE/EXPLAIN 要留着:agent 探查表结构靠它们。 +_READONLY_VERBS = frozenset({ + "select", "with", "show", "describe", "desc", "explain", "pragma", "values", +}) + +# 即便在可写数据源上也一律拒绝的语句。这类操作的破坏是不可逆的, +# 不该由一个实验工具里的模型来发起——要做请人去数据库客户端里做。 +_ALWAYS_DENIED = frozenset({ + "drop", "truncate", "grant", "revoke", "alter", "create", "shutdown", + "attach", "detach", +}) + +_LINE_COMMENT = re.compile(r"--[^\n]*") +_BLOCK_COMMENT = re.compile(r"/\*.*?\*/", re.S) + + +class SqlRejected(ValueError): + """SQL 没通过守卫。消息直接给模型看,所以要说清楚为什么被拒、该怎么改。""" + + +@dataclass(frozen=True) +class QueryLimits: + max_rows: int = 1000 + max_bytes: int = 1_000_000 + timeout_seconds: int = 30 + + +def strip_comments(sql: str) -> str: + """去掉注释再判断。否则 `/*x*/ DROP TABLE t` 的首关键字会被看成注释内容。""" + return _BLOCK_COMMENT.sub(" ", _LINE_COMMENT.sub(" ", sql)) + + +def split_statements(sql: str) -> list[str]: + """按分号切语句,但要认得字符串字面量里的分号。 + + `SELECT ';'` 只是一条语句。简单 split(";") 会把它切成两半,然后把后半截 + 当成"第二条语句"拒掉——那是误杀,比漏杀更让人无从下手。 + """ + out: list[str] = [] + buf: list[str] = [] + quote: str | None = None + i = 0 + text = strip_comments(sql) + while i < len(text): + ch = text[i] + if quote: + buf.append(ch) + if ch == quote: + # SQL 里连续两个引号是转义('' 表示一个单引号) + if i + 1 < len(text) and text[i + 1] == quote: + buf.append(text[i + 1]) + i += 2 + continue + quote = None + elif ch in ("'", '"', "`"): + quote = ch + buf.append(ch) + elif ch == ";": + stmt = "".join(buf).strip() + if stmt: + out.append(stmt) + buf = [] + else: + buf.append(ch) + i += 1 + tail = "".join(buf).strip() + if tail: + out.append(tail) + return out + + +def first_verb(sql: str) -> str: + """取语句的首关键字(小写)。括号开头的 `(SELECT ...)` 也要认出来。""" + text = strip_comments(sql).lstrip().lstrip("(").lstrip() + match = re.match(r"[A-Za-z_]+", text) + return match.group(0).lower() if match else "" + + +def check(sql: str, *, readonly: bool, source_name: str = "") -> str: + """校验并返回规范化后的单条 SQL。不通过就抛 SqlRejected。 + + 返回值是去掉尾分号的那一条语句——调用方应该执行这个返回值, + 而不是原始输入,否则守卫等于没做。 + """ + if not sql or not sql.strip(): + raise SqlRejected("SQL 是空的") + + statements = split_statements(sql) + if not statements: + raise SqlRejected("SQL 里没有可执行的语句") + if len(statements) > 1: + raise SqlRejected( + f"一次只能执行一条语句,这里有 {len(statements)} 条。" + "拆成多次调用——分号拼接的语句不会被执行。" + ) + + stmt = statements[0] + verb = first_verb(stmt) + if not verb: + raise SqlRejected("看不出这条 SQL 要做什么(没有识别到关键字)") + + if verb in _ALWAYS_DENIED: + raise SqlRejected( + f"{verb.upper()} 属于结构变更或权限操作,任何数据源上都不允许。" + "这类操作请在数据库客户端里人工执行。" + ) + + if readonly and verb not in _READONLY_VERBS: + where = f"「{source_name}」" if source_name else "这个数据源" + raise SqlRejected( + f"{where}是只读的,{verb.upper()} 不被允许。" + "只能执行 SELECT / WITH / SHOW / DESCRIBE / EXPLAIN。" + "确实需要写入的话,请在设置里另建一个关闭了只读的数据源。" + ) + + return stmt + + +def is_write(sql: str) -> bool: + """是不是写操作。用来决定要不要走人工审批。""" + return first_verb(sql) not in _READONLY_VERBS diff --git a/backend/app/data/introspect.py b/backend/app/data/introspect.py new file mode 100644 index 00000000..3c14045f --- /dev/null +++ b/backend/app/data/introspect.py @@ -0,0 +1,135 @@ +"""探查数据库结构。 + +产出两种粒度,对应两种消费者: + +- `summary()`:表名 + 一句话,几百 token。常驻 Copilot 的 system prompt, + 让它知道"有哪些库、每个库大概有什么"。 +- `describe_table()`:某张表的完整字段。按需取,Copilot 决定要查某张表时才拉。 + +分两种粒度是因为 token 预算:一个中等业务库几百张表、每表几十个字段, +全量塞进 prompt 轻松几万 token,既贵又会把真正的需求淹掉。 + +用 SQLAlchemy 的 Inspector 而不是手写各家的 information_schema 查询—— +方言差异它已经处理过了,我们不必再趟一遍。 +""" +from __future__ import annotations + +from datetime import datetime, timezone +from typing import Any + +from sqlalchemy import inspect as sa_inspect + +from app.data.engine import engines + +# 一次探查最多带回多少张表。生产库动辄上千张,全拉回来既慢又没人看得完。 +_MAX_TABLES = 200 +# summary 里每张表列几个代表字段。够 Copilot 判断"这表是不是我要的"即可。 +_PREVIEW_COLUMNS = 6 + + +async def introspect(source: Any, *, schema: str | None = None) -> dict[str, Any]: + """探查整个库的结构,结果直接存进 DataSource.schema_cache。""" + engine = await engines.get(source) + + def _collect(sync_conn: Any) -> dict[str, Any]: + inspector = sa_inspect(sync_conn) + target_schema = schema or None + names = inspector.get_table_names(schema=target_schema) + views = inspector.get_view_names(schema=target_schema) + truncated = len(names) + len(views) > _MAX_TABLES + picked = (names + views)[:_MAX_TABLES] + + tables: dict[str, Any] = {} + for name in picked: + try: + columns = inspector.get_columns(name, schema=target_schema) + except Exception: # noqa: BLE001 - 单表失败不该让整次探查失败 + continue + try: + pk = (inspector.get_pk_constraint(name, schema=target_schema) or {}).get( + "constrained_columns" + ) or [] + except Exception: # noqa: BLE001 + pk = [] + try: + comment = (inspector.get_table_comment(name, schema=target_schema) or {}).get("text") + except Exception: # noqa: BLE001 - 不是所有方言都支持表注释 + comment = None + + tables[name] = { + "columns": [ + { + "name": c["name"], + "type": str(c.get("type") or ""), + "nullable": bool(c.get("nullable", True)), + "comment": c.get("comment") or None, + } + for c in columns + ], + "primary_key": pk, + "comment": comment, + "is_view": name in views, + } + return {"tables": tables, "truncated": truncated, "total": len(names) + len(views)} + + async with engine.connect() as conn: + payload = await conn.run_sync(_collect) + + payload["synced_at"] = datetime.now(timezone.utc).isoformat() + payload["schema"] = schema + return payload + + +def summary(source: Any, *, max_tables: int = 40) -> str: + """给 Copilot 的紧凑摘要。没探查过就明说,别让它对着空气编表名。""" + cache = source.schema_cache or {} + tables = cache.get("tables") or {} + if not tables: + return f"- {source.name}({source.kind}):{source.description or '未填说明'}|尚未探查结构" + + lines = [ + f"- {source.name}({source.kind},{'只读' if source.readonly else '可写'}):" + f"{source.description or '未填说明'}" + ] + for name, meta in list(tables.items())[:max_tables]: + cols = meta.get("columns") or [] + preview = "、".join(c["name"] for c in cols[:_PREVIEW_COLUMNS]) + more = f" 等 {len(cols)} 字段" if len(cols) > _PREVIEW_COLUMNS else "" + note = f"({meta['comment']})" if meta.get("comment") else "" + lines.append(f" · {name}{note}: {preview}{more}") + if len(tables) > max_tables: + lines.append(f" · …另有 {len(tables) - max_tables} 张表,用 db_schema 工具查看") + return "\n".join(lines) + + +def describe_table(source: Any, table: str) -> str: + """单表的完整字段说明。agent 调 db_schema 工具时返回这个。""" + cache = source.schema_cache or {} + tables = cache.get("tables") or {} + meta = tables.get(table) + if not meta: + available = "、".join(list(tables)[:30]) or "(还没探查过结构)" + return f"数据源「{source.name}」里没有表 {table}。现有的表:{available}" + + head = f"表 {table}" + if meta.get("comment"): + head += f"({meta['comment']})" + if meta.get("is_view"): + head += " [视图]" + lines = [head] + pk = set(meta.get("primary_key") or []) + for col in meta.get("columns") or []: + marks = [] + if col["name"] in pk: + marks.append("主键") + if not col.get("nullable", True): + marks.append("非空") + if col.get("comment"): + marks.append(col["comment"]) + suffix = f" {'、'.join(marks)}" if marks else "" + lines.append(f" {col['name']} {col['type']}{suffix}") + return "\n".join(lines) + + +def table_names(source: Any) -> list[str]: + return list((source.schema_cache or {}).get("tables") or {}) diff --git a/backend/app/db/models.py b/backend/app/db/models.py index 20893b68..4d980cf9 100644 --- a/backend/app/db/models.py +++ b/backend/app/db/models.py @@ -29,6 +29,36 @@ class Provider(Base, TimestampMixin): extra: Mapped[dict[str, Any]] = mapped_column(default=dict) +class DataSource(Base, TimestampMixin): + """一个数据库接入点。password 落库前经 app.core.crypto 加密,和 Provider 同一套。 + + readonly 默认为真,是刻意的安全基线:这里的 SQL 由模型生成,不是人写的。 + 要写库就另建一个显式关掉只读的源,并且写操作会走人工审批——把"能写"这件事 + 变成一个需要两次明确动作的决定,而不是默认状态。 + """ + + __tablename__ = "data_sources" + + id: Mapped[str] = mapped_column(String(32), primary_key=True, default=new_id) + name: Mapped[str] = mapped_column(String(100), unique=True) + kind: Mapped[str] = mapped_column(String(32)) # mysql | postgres | oracle | sqlite + host: Mapped[str | None] = mapped_column(String(255), default=None) + port: Mapped[int | None] = mapped_column(Integer, default=None) + database: Mapped[str | None] = mapped_column(String(255), default=None) + username: Mapped[str | None] = mapped_column(String(128), default=None) + password: Mapped[str | None] = mapped_column(Text, default=None) # 密文 + # 驱动差异塞这里:Oracle 的 service_name/sid、MySQL 的 charset、SSL 参数… + # 不给每种数据库加一列,否则表会长成各家方言的并集 + options: Mapped[dict[str, Any]] = mapped_column(default=dict) + readonly: Mapped[bool] = mapped_column(default=True) + # 给 Copilot 看的一句话:"销售库,含订单/客户/产品"。它据此判断该查哪个源 + description: Mapped[str] = mapped_column(Text, default="") + # 探查到的表结构。缓存下来,否则每次建图都要连一次生产库 + schema_cache: Mapped[dict[str, Any]] = mapped_column(default=dict) + schema_synced_at: Mapped[datetime | None] = mapped_column(default=None) + enabled: Mapped[bool] = mapped_column(default=True) + + class Setting(Base, TimestampMixin): """键值形式的用户设置,前端 Settings 页直接读写。""" diff --git a/backend/app/main.py b/backend/app/main.py index 9f0b91f4..04d08b7e 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -8,7 +8,7 @@ from fastapi import FastAPI, Request from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse -from app.api import artifacts, copilot, governance, knowledge, runs, sandbox, settings as settings_api, tools, workflows +from app.api import datasources, artifacts, copilot, governance, knowledge, runs, sandbox, settings as settings_api, tools, workflows from app.core.config import settings from app.db.base import init_db from app.engine.runner import run_manager @@ -72,6 +72,7 @@ for router in ( copilot.router, artifacts.router, governance.router, + datasources.router, ): app.include_router(router) diff --git a/backend/app/tools/datasource.py b/backend/app/tools/datasource.py new file mode 100644 index 00000000..e138ac39 --- /dev/null +++ b/backend/app/tools/datasource.py @@ -0,0 +1,155 @@ +"""把数据源变成 agent 能调的工具。 + +每个启用的数据源生成一对工具,而不是"一个通用工具带 datasource 参数": + +- 工具名自带语义(`db_query__sales`),模型不会把查销售库的 SQL 发给财务库 +- 工具 description 里直接塞该库的表清单,模型一眼看到有什么可查, + 省掉"先问有哪些表"那一轮 +- 权限边界跟着工具走:只读源生成的工具就是只读的,dangerous 标记也各自独立 + +代价是工具数量随数据源线性增长。真到了几十个数据源的规模,再考虑按需装载。 +""" +from __future__ import annotations + +import json +from typing import Any + +from langchain_core.tools import StructuredTool +from pydantic import BaseModel, Field +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from app.data import introspect +from app.data.engine import run_query +from app.data.guard import QueryLimits, SqlRejected, is_write +from app.db.models import DataSource +from app.tools.registry import ToolContext + +QUERY_PREFIX = "db_query__" +SCHEMA_PREFIX = "db_schema__" + + +class _QueryArgs(BaseModel): + sql: str = Field(description="要执行的 SQL,一次一条") + limit: int | None = Field(default=None, description="最多返回多少行,默认 1000") + + +class _SchemaArgs(BaseModel): + table: str | None = Field(default=None, description="表名;留空则列出所有表") + + +def tool_names(source: DataSource) -> list[str]: + return [f"{QUERY_PREFIX}{source.name}", f"{SCHEMA_PREFIX}{source.name}"] + + +def _query_description(source: DataSource) -> str: + tables = introspect.table_names(source) + head = f"在数据源「{source.name}」上执行 SQL 查询。" + if source.description: + head += f"{source.description}。" + head += "只读。" if source.readonly else "可写(写操作需要人工审批)。" + if tables: + listed = "、".join(tables[:25]) + more = f" 等 {len(tables)} 张表" if len(tables) > 25 else "" + head += f" 可用表:{listed}{more}。" + else: + head += " 结构尚未探查,先用同名的 db_schema 工具看看有哪些表。" + head += " 字段不确定时先查结构,不要猜字段名。" + return head + + +async def build_datasource_tools( + names: list[str], ctx: ToolContext, session: AsyncSession +) -> list[StructuredTool]: + """按工具名反查数据源并构造工具。build_tools 的第四个来源。""" + wanted: dict[str, list[str]] = {} + for name in names: + for prefix in (QUERY_PREFIX, SCHEMA_PREFIX): + if name.startswith(prefix): + wanted.setdefault(name[len(prefix):], []).append(prefix) + + if not wanted: + return [] + + rows = list(( + await session.execute( + select(DataSource).where( + DataSource.name.in_(list(wanted)), DataSource.enabled.is_(True) + ) + ) + ).scalars()) + + tools: list[StructuredTool] = [] + for row in rows: + for prefix in wanted.get(row.name, []): + if prefix == QUERY_PREFIX: + tools.append(_make_query_tool(row, ctx)) + else: + tools.append(_make_schema_tool(row)) + return tools + + +def _make_query_tool(source: DataSource, ctx: ToolContext) -> StructuredTool: + async def _run(sql: str, limit: int | None = None) -> str: + limits = QueryLimits(max_rows=int(limit)) if limit else QueryLimits() + try: + result = await run_query(source, sql, limits=limits) + except SqlRejected as e: + # 被守卫拒掉要把原因原样交回模型——它据此改写 SQL 重试, + # 给一句笼统的"失败了"只会让它瞎猜 + return f"SQL 被拒绝:{e}" + except Exception as e: # noqa: BLE001 + return f"查询失败:{type(e).__name__}: {e}" + + payload = result.to_payload() + payload["source"] = source.name + # 查询快照进工件库:出具体系的数字回指要能下钻到"这个数是哪条 SQL 查出来的" + try: + from app.core.artifact_store import put_json + + payload["artifact"] = await put_json( + payload, kind="query_snapshot", + run_id=ctx.run_id or "", node_id=ctx.node_id or "", + ) + except Exception: # noqa: BLE001 - 存不下不影响查询本身 + pass + return json.dumps(payload, ensure_ascii=False, default=str) + + return StructuredTool( + name=f"{QUERY_PREFIX}{source.name}", + description=_query_description(source), + args_schema=_QueryArgs, + coroutine=_run, + func=None, + ) + + +def _make_schema_tool(source: DataSource) -> StructuredTool: + async def _run(table: str | None = None) -> str: + if table: + return introspect.describe_table(source, table) + tables = introspect.table_names(source) + if not tables: + return ( + f"数据源「{source.name}」还没有探查过结构。" + "请在设置页里对它执行一次结构探查。" + ) + return f"数据源「{source.name}」共 {len(tables)} 张表:\n" + "\n".join( + f" {n}" for n in tables + ) + + return StructuredTool( + name=f"{SCHEMA_PREFIX}{source.name}", + description=( + f"查看数据源「{source.name}」的表结构。" + "传 table 看某张表的字段;不传则列出所有表。写 SQL 前先用它确认字段名。" + ), + args_schema=_SchemaArgs, + coroutine=_run, + func=None, + ) + + +def is_dangerous_call(source: DataSource, sql: str) -> bool: + """这次调用要不要走人工审批:非只读源上的写操作要。""" + return not source.readonly and is_write(sql) diff --git a/backend/app/tools/registry.py b/backend/app/tools/registry.py index e23582ea..7d506bf2 100644 --- a/backend/app/tools/registry.py +++ b/backend/app/tools/registry.py @@ -109,10 +109,18 @@ def build_tool(spec: ToolSpec, ctx: ToolContext) -> BaseTool: async def build_tools( names: list[str], ctx: ToolContext, *, session: Any = None ) -> list[BaseTool]: - """按名字列表组装工具。支持三种来源:内置、MCP(mcp:/)、自定义。""" + """按名字列表组装工具。 + + 四种来源:内置(静态注册)、MCP(mcp:/)、数据源 + (db_query__<源名> / db_schema__<源名>)、自定义(数据库里的 CustomTool)。 + 后两种都是运行时从库里读出来动态构造的,不在静态注册表里。 + """ + from app.tools.datasource import QUERY_PREFIX, SCHEMA_PREFIX + specs = all_specs() tools: list[BaseTool] = [] mcp_names: list[str] = [] + datasource_names: list[str] = [] custom_names: list[str] = [] for name in names: @@ -120,9 +128,16 @@ async def build_tools( mcp_names.append(name) elif name in specs: tools.append(build_tool(specs[name], ctx)) + elif name.startswith((QUERY_PREFIX, SCHEMA_PREFIX)): + datasource_names.append(name) else: custom_names.append(name) + if datasource_names and session is not None: + from app.tools.datasource import build_datasource_tools + + tools.extend(await build_datasource_tools(datasource_names, ctx, session)) + if mcp_names: from app.tools.mcp_manager import mcp_manager diff --git a/backend/pyproject.toml b/backend/pyproject.toml index fc5b21a5..d06f16ef 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -33,6 +33,9 @@ dependencies = [ microvm = ["microsandbox>=0.6"] # 测试:conda 环境默认不装,跑测试前 pip install 'agentlab-backend[dev]' dev = ["pytest>=8.0", "pytest-asyncio>=0.24"] +# 数据库接入。三个驱动都是 asyncio 原生且免装数据库客户端—— +# oracledb 走 thin 模式,省掉了 Oracle Instant Client 这个历来最劝退的一步。 +db = ["aiomysql>=0.2", "oracledb>=2.0", "asyncpg>=0.30"] [build-system] requires = ["hatchling"] diff --git a/backend/tests/test_datasource.py b/backend/tests/test_datasource.py new file mode 100644 index 00000000..c1eb236f --- /dev/null +++ b/backend/tests/test_datasource.py @@ -0,0 +1,151 @@ +"""数据源连接层与工具层。 + +用真实的 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 + + +@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_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")) diff --git a/backend/tests/test_sql_guard.py b/backend/tests/test_sql_guard.py new file mode 100644 index 00000000..1cd5233c --- /dev/null +++ b/backend/tests/test_sql_guard.py @@ -0,0 +1,106 @@ +"""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 diff --git a/environment.yml b/environment.yml index b89b8ae8..80608c6a 100644 --- a/environment.yml +++ b/environment.yml @@ -15,4 +15,4 @@ dependencies: # [microvm] 带上独立内核的沙箱后端(约 35MB wheel,首次运行还要下载 # 约 50MB 运行时,外加一次 OCI 镜像拉取)。不需要就把方括号去掉, # 沙箱会退回 Seatbelt —— 代价是内存限额在 macOS 上限不住。 - - -e ./backend[microvm,dev] + - -e ./backend[microvm,dev,db]