mirror of
https://gitee.com/mateos/mateclaw.git
synced 2026-09-13 03:13:41 +08:00
210 lines
9.7 KiB
Java
210 lines
9.7 KiB
Java
package vip.mate.tool.builtin;
|
||
|
||
import cn.hutool.json.JSONObject;
|
||
import cn.hutool.json.JSONUtil;
|
||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||
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 java.sql.*;
|
||
import java.util.ArrayList;
|
||
import java.util.LinkedHashMap;
|
||
import java.util.List;
|
||
import java.util.Map;
|
||
import java.util.regex.Pattern;
|
||
|
||
/**
|
||
* 内置工具：数据源发现
|
||
* <p>
|
||
* 提供数据源列表查询、表列表查询、表结构查询三个动作，
|
||
* 供 Agent 在查数场景下发现可用数据源和表结构。
|
||
*
|
||
* @author MateClaw Team
|
||
*/
|
||
@Slf4j
|
||
@Component
|
||
@RequiredArgsConstructor
|
||
public class DatasourceTool {
|
||
|
||
private final DatasourceService datasourceService;
|
||
private final DatasourceConnectionManager connectionManager;
|
||
/**
|
||
* The application ObjectMapper, which serializes every {@code Long} as a JSON
|
||
* string (see {@code JacksonConfig}). Tool output goes through it so 19-digit
|
||
* Snowflake ids reach the model as strings — exactly like the HTTP API — and
|
||
* never as JSON numbers that lose their low digits in a double/JS-number
|
||
* round-trip on the way back into a tool call.
|
||
*/
|
||
private final ObjectMapper objectMapper;
|
||
|
||
/** SQL identifier whitelist: letters, digits, underscore, dot, hyphen only */
|
||
private static final Pattern SAFE_IDENTIFIER = Pattern.compile("^[a-zA-Z0-9_][a-zA-Z0-9_.\\-]{0,127}$");
|
||
|
||
private static String sanitizeIdentifier(String name) {
|
||
if (name == null || !SAFE_IDENTIFIER.matcher(name).matches()) {
|
||
throw new IllegalArgumentException("Invalid SQL identifier: " + name);
|
||
}
|
||
return name;
|
||
}
|
||
|
||
@Tool(description = """
|
||
查询外部数据源的元数据。支持三种动作：
|
||
1. action='list_datasources' — 列出所有可用数据源（无需其他参数）
|
||
2. action='list_tables' — 列出指定数据源中的所有表（需要 datasourceId）
|
||
3. action='describe_table' — 查看指定表的列详情（需要 datasourceId 和 tableName）
|
||
""")
|
||
public String query_datasource(
|
||
@ToolParam(description = "动作：list_datasources / list_tables / describe_table") String action,
|
||
@ToolParam(description = "数据源 ID（list_tables 和 describe_table 时必填）", required = false) Long datasourceId,
|
||
@ToolParam(description = "表名（describe_table 时必填）", required = false) String tableName) {
|
||
|
||
try {
|
||
return switch (action) {
|
||
case "list_datasources" -> listDatasources();
|
||
case "list_tables" -> listTables(datasourceId);
|
||
case "describe_table" -> describeTable(datasourceId, tableName);
|
||
default -> error("未知动作: " + action + "，支持: list_datasources / list_tables / describe_table");
|
||
};
|
||
} catch (Exception e) {
|
||
log.error("数据源查询失败: {}", e.getMessage(), e);
|
||
return error(e.getMessage());
|
||
}
|
||
}
|
||
|
||
private String listDatasources() throws Exception {
|
||
List<DatasourceEntity> list = datasourceService.listEnabled();
|
||
List<Map<String, Object>> rows = new ArrayList<>();
|
||
for (DatasourceEntity ds : list) {
|
||
Map<String, Object> obj = new LinkedHashMap<>();
|
||
// ds.getId() is a Long; the shared ObjectMapper renders it as a JSON
|
||
// string so the model copies an exact id back into list_tables /
|
||
// execute_sql / describe_table calls.
|
||
obj.put("id", ds.getId());
|
||
obj.put("name", ds.getName());
|
||
obj.put("dbType", ds.getDbType());
|
||
obj.put("databaseName", ds.getDatabaseName());
|
||
obj.put("description", ds.getDescription());
|
||
rows.add(obj);
|
||
}
|
||
Map<String, Object> result = new LinkedHashMap<>();
|
||
result.put("datasources", rows);
|
||
result.put("count", rows.size());
|
||
return objectMapper.writerWithDefaultPrettyPrinter().writeValueAsString(result);
|
||
}
|
||
|
||
private String listTables(Long datasourceId) throws SQLException {
|
||
if (datasourceId == null) {
|
||
return error("list_tables 需要 datasourceId 参数");
|
||
}
|
||
DatasourceEntity entity = datasourceService.getDecrypted(datasourceId);
|
||
String dbType = entity.getDbType().toLowerCase();
|
||
|
||
String sql = switch (dbType) {
|
||
case "mysql", "mariadb" -> String.format(
|
||
"SELECT TABLE_NAME, TABLE_COMMENT, TABLE_ROWS FROM information_schema.TABLES WHERE TABLE_SCHEMA = '%s' ORDER BY TABLE_NAME",
|
||
sanitizeIdentifier(entity.getDatabaseName()));
|
||
case "postgresql" -> String.format(
|
||
"SELECT tablename AS table_name, obj_description(c.oid) AS table_comment " +
|
||
"FROM pg_tables t LEFT JOIN pg_class c ON c.relname = t.tablename " +
|
||
"WHERE t.schemaname = '%s' ORDER BY tablename",
|
||
sanitizeIdentifier(entity.getSchemaName() != null ? entity.getSchemaName() : "public"));
|
||
case "clickhouse" -> "SHOW TABLES";
|
||
default -> throw new IllegalArgumentException("不支持的数据库类型: " + dbType);
|
||
};
|
||
|
||
try (Connection conn = connectionManager.getConnection(entity);
|
||
Statement stmt = conn.createStatement()) {
|
||
stmt.setQueryTimeout(15);
|
||
ResultSet rs = stmt.executeQuery(sql);
|
||
return formatResultSet(rs, 200);
|
||
}
|
||
}
|
||
|
||
private String describeTable(Long datasourceId, String tableName) throws SQLException {
|
||
if (datasourceId == null || tableName == null || tableName.isBlank()) {
|
||
return error("describe_table 需要 datasourceId 和 tableName 参数");
|
||
}
|
||
DatasourceEntity entity = datasourceService.getDecrypted(datasourceId);
|
||
String dbType = entity.getDbType().toLowerCase();
|
||
|
||
String sql = switch (dbType) {
|
||
case "mysql", "mariadb" -> String.format(
|
||
"SELECT COLUMN_NAME, COLUMN_TYPE, IS_NULLABLE, COLUMN_KEY, COLUMN_DEFAULT, COLUMN_COMMENT " +
|
||
"FROM information_schema.COLUMNS WHERE TABLE_SCHEMA = '%s' AND TABLE_NAME = '%s' ORDER BY ORDINAL_POSITION",
|
||
sanitizeIdentifier(entity.getDatabaseName()), sanitizeIdentifier(tableName));
|
||
case "postgresql" -> String.format(
|
||
"SELECT c.column_name, c.data_type, c.is_nullable, " +
|
||
"CASE WHEN pk.column_name IS NOT NULL THEN 'PRI' ELSE '' END AS column_key, " +
|
||
"c.column_default, pgd.description AS column_comment " +
|
||
"FROM information_schema.columns c " +
|
||
"LEFT JOIN (SELECT ku.column_name FROM information_schema.table_constraints tc " +
|
||
"JOIN information_schema.key_column_usage ku ON tc.constraint_name = ku.constraint_name " +
|
||
"WHERE tc.table_name = '%s' AND tc.constraint_type = 'PRIMARY KEY') pk ON c.column_name = pk.column_name " +
|
||
"LEFT JOIN pg_catalog.pg_statio_all_tables st ON st.relname = c.table_name " +
|
||
"LEFT JOIN pg_catalog.pg_description pgd ON pgd.objoid = st.relid AND pgd.objsubid = c.ordinal_position " +
|
||
"WHERE c.table_schema = '%s' AND c.table_name = '%s' ORDER BY c.ordinal_position",
|
||
sanitizeIdentifier(tableName),
|
||
sanitizeIdentifier(entity.getSchemaName() != null ? entity.getSchemaName() : "public"),
|
||
sanitizeIdentifier(tableName));
|
||
case "clickhouse" -> String.format("DESCRIBE TABLE %s", sanitizeIdentifier(tableName));
|
||
default -> throw new IllegalArgumentException("不支持的数据库类型: " + dbType);
|
||
};
|
||
|
||
try (Connection conn = connectionManager.getConnection(entity);
|
||
Statement stmt = conn.createStatement()) {
|
||
stmt.setQueryTimeout(15);
|
||
ResultSet rs = stmt.executeQuery(sql);
|
||
return formatResultSet(rs, 500);
|
||
}
|
||
}
|
||
|
||
/**
|
||
* 将 ResultSet 格式化为 Markdown 表格
|
||
*/
|
||
private String formatResultSet(ResultSet rs, int maxRows) throws SQLException {
|
||
ResultSetMetaData meta = rs.getMetaData();
|
||
int colCount = meta.getColumnCount();
|
||
|
||
// 表头
|
||
StringBuilder sb = new StringBuilder();
|
||
List<String> headers = new ArrayList<>();
|
||
for (int i = 1; i <= colCount; i++) {
|
||
headers.add(meta.getColumnLabel(i));
|
||
}
|
||
sb.append("| ").append(String.join(" | ", headers)).append(" |\n");
|
||
sb.append("| ").append("--- | ".repeat(colCount)).append("\n");
|
||
|
||
// 数据行
|
||
int rowCount = 0;
|
||
while (rs.next() && rowCount < maxRows) {
|
||
sb.append("| ");
|
||
for (int i = 1; i <= colCount; i++) {
|
||
String val = rs.getString(i);
|
||
sb.append(val != null ? val.replace("|", "\\|") : "NULL");
|
||
if (i < colCount) sb.append(" | ");
|
||
}
|
||
sb.append(" |\n");
|
||
rowCount++;
|
||
}
|
||
|
||
if (rowCount == 0) {
|
||
return "查询结果为空";
|
||
}
|
||
sb.append("\n共 ").append(rowCount).append(" 条记录");
|
||
if (rowCount >= maxRows) {
|
||
sb.append("（已截断，实际可能更多）");
|
||
}
|
||
return sb.toString();
|
||
}
|
||
|
||
private String error(String message) {
|
||
return JSONUtil.toJsonStr(new JSONObject().set("error", message));
|
||
}
|
||
}
|