99 lines
3.5 KiB
Python
99 lines
3.5 KiB
Python
"""执行已通过安全校验的只读 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,
|
|
parameters: dict[str, object] | 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:
|
|
statement = text(validated_sql.sql)
|
|
execution = (
|
|
session.execute(statement, parameters)
|
|
if parameters
|
|
else session.execute(statement)
|
|
)
|
|
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,
|
|
)
|