feat:新增投顾agent和nl2sqlagent

2026-09-13 16:19:24 +08:00
parent c80c6acac0
commit 163192bf55
122 changed files with 7488 additions and 362 deletions
@@ -0,0 +1,92 @@
+"""执行已通过安全校验的只读 SQL。"""
+from __future__ import annotations
+
+import asyncio
+import time
+
+from sqlalchemy import text
+
+from nl2sql.contracts import DataQueryResult
+from nl2sql.masking import mask_rows
+from nl2sql.rendering import render_markdown
+from nl2sql.runtime import QueryRuntimeRegistry, query_runtime_registry
+from nl2sql.runtime import kill_mysql_query
+from nl2sql.sql_security import ValidatedSql
+
+
+class QueryExecutionError(RuntimeError):
+    """只读 SQL 执行失败。"""
+
+
+async def execute_readonly_sql(
+    session,
+    validated_sql: ValidatedSql,
+    *,
+    query_id: str,
+    trace_id: str,
+    masks: dict[tuple[str, str], str] | None = None,
+    max_rows: int | None = None,
+    timeout_seconds: float | None = None,
+    user_id: int = 0,
+    connection_id: int | None = None,
+    kill_query=None,
+    runtime_registry: QueryRuntimeRegistry | None = query_runtime_registry,
+) -> DataQueryResult:
+    """执行安全 SQL 并转换为统一查询结果，不接收未校验的原始 SQL。"""
+    started_at = time.perf_counter()
+    if runtime_registry is not None:
+        if connection_id is None and hasattr(session, "connection"):
+            try:
+                connection = await session.connection()
+                connection_result = await connection.execute(text("SELECT CONNECTION_ID()"))
+                connection_id = int(connection_result.scalar_one())
+            except Exception:  # noqa: BLE001 连接号仅用于运维，不影响正常查询
+                connection_id = None
+        await runtime_registry.register(
+            query_id=query_id,
+            user_id=user_id,
+            sql=validated_sql.sql,
+            connection_id=connection_id,
+        )
+    try:
+        execution = session.execute(text(validated_sql.sql))
+        result = (
+            await asyncio.wait_for(execution, timeout_seconds)
+            if timeout_seconds is not None
+            else await execution
+        )
+        columns = [str(key) for key in result.keys()]
+        rows = [dict(row) for row in result.mappings().all()]
+        if masks:
+            rows = mask_rows(rows, masks)
+        truncated = max_rows is not None and len(rows) > max_rows
+        if truncated:
+            rows = rows[:max_rows]
+    except asyncio.TimeoutError as exc:
+        killer = kill_query or kill_mysql_query
+        if connection_id is not None:
+            try:
+                await killer(connection_id)
+            except Exception:  # noqa: BLE001 中止失败仍需返回统一超时错误
+                pass
+        if runtime_registry is not None:
+            await runtime_registry.complete(query_id, status="timeout")
+        raise QueryExecutionError("查询超时") from exc
+    except Exception as exc:  # noqa: BLE001 统一收敛数据库异常
+        if runtime_registry is not None:
+            await runtime_registry.complete(query_id, status="failed")
+        raise QueryExecutionError("查询执行失败") from exc
+    elapsed_ms = (time.perf_counter() - started_at) * 1000
+    if runtime_registry is not None:
+        await runtime_registry.complete(query_id)
+    return DataQueryResult(
+        query_id=query_id,
+        trace_id=trace_id,
+        columns=columns,
+        rows=rows,
+        row_count=len(rows),
+        truncated=bool(truncated),
+        markdown=render_markdown(columns, rows),
+        sql=validated_sql.sql,
+        elapsed_ms=elapsed_ms,
+    )