package vip.mate.tool.builtin; import cn.hutool.json.JSONArray; import cn.hutool.json.JSONObject; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; import org.springframework.ai.tool.annotation.Tool; import org.springframework.ai.tool.annotation.ToolParam; import org.springframework.stereotype.Component; import vip.mate.datasource.model.DatasourceEntity; import vip.mate.datasource.service.DatasourceConnectionManager; import vip.mate.datasource.service.DatasourceService; import vip.mate.datasource.service.EChartsOptionBuilder; import vip.mate.datasource.service.SqlValidationService; import java.sql.*; import java.util.ArrayList; import java.util.List; /** * 内置工具:SQL 查询执行 *
* 仅允许 SELECT 语句。自动注入 LIMIT 保护。
* 查询超时 30 秒。结果格式化为 Markdown 表格或 JSON。
* 自动分析数据特征并生成 ECharts 图表配置。
*
* @author MateClaw Team
*/
@Slf4j
@Component
@RequiredArgsConstructor
public class SqlQueryTool {
private final DatasourceService datasourceService;
private final DatasourceConnectionManager connectionManager;
private final SqlValidationService sqlValidationService;
private static final int QUERY_TIMEOUT_SECONDS = 30;
private static final int MAX_ROWS = 500;
private static final int MARKDOWN_TABLE_THRESHOLD = 20;
@Tool(description = """
在外部数据源上执行只读 SQL 查询。
仅允许 SELECT 语句,禁止 INSERT/UPDATE/DELETE/DROP 等写操作。
如果 SQL 没有 LIMIT 子句会自动添加 LIMIT 500。
返回查询结果(Markdown 表格或 JSON 格式)以及行数和执行耗时。
如果数据适合可视化,会自动附带一个 echarts 图表配置代码块,前端会自动渲染为交互式图表。
""")
public String execute_sql(
@ToolParam(description = "目标数据源 ID") Long datasourceId,
@ToolParam(description = "要执行的 SQL 查询(仅允许 SELECT)") String sql) {
try {
// 1. 验证并规范化 SQL
String safeSql = sqlValidationService.validateAndNormalize(sql);
log.info("执行 SQL 查询 [数据源 {}]: {}", datasourceId, safeSql);
// 2. 获取数据源连接
DatasourceEntity entity = datasourceService.getDecrypted(datasourceId);
// 3. 执行查询
long startTime = System.currentTimeMillis();
try (Connection conn = connectionManager.getConnection(entity);
Statement stmt = conn.createStatement()) {
stmt.setQueryTimeout(QUERY_TIMEOUT_SECONDS);
stmt.setMaxRows(MAX_ROWS);
ResultSet rs = stmt.executeQuery(safeSql);
long elapsed = System.currentTimeMillis() - startTime;
return formatResult(rs, safeSql, elapsed);
}
} catch (Exception e) {
log.error("SQL 查询执行失败: {}", e.getMessage(), e);
JSONObject result = new JSONObject();
result.set("error", e.getMessage());
result.set("sql", sql);
return result.toStringPretty();
}
}
private String formatResult(ResultSet rs, String sql, long elapsedMs) throws SQLException {
ResultSetMetaData meta = rs.getMetaData();
int colCount = meta.getColumnCount();
// 收集列名
List> rows = new ArrayList<>();
while (rs.next() && rows.size() < MAX_ROWS) {
List