(state: &AppState, args: Value)
| 29 | } |
| 30 | |
| 31 | pub async fn execute_query(state: &AppState, args: Value) -> Result<ToolResult, String> { |
| 32 | let connection_id = args["connection_id"] |
| 33 | .as_i64() |
| 34 | .ok_or("Missing connection_id")?; |
| 35 | let database = args["database"].as_str().map(|s| s.to_string()); |
| 36 | let sql = args["sql"].as_str().ok_or("Missing sql")?.to_string(); |
| 37 | |
| 38 | // SQL 安全检查 |
| 39 | let config = SqlSafetyConfig::from_env(); |
| 40 | match check_sql_safety(&sql, &config) { |
| 41 | SqlSafetyCheck::Allowed => {} |
| 42 | SqlSafetyCheck::Rejected(reason) => { |
| 43 | return Ok(ToolResult::error(reason)); |
| 44 | } |
| 45 | } |
| 46 | |
| 47 | // 执行查询 |
| 48 | let result = crate::commands::execute_with_retry_from_app_state( |
| 49 | state, |
| 50 | connection_id, |
| 51 | database, |
| 52 | |driver| { |
| 53 | let sql = sql.clone(); |
| 54 | async move { driver.execute_query(sql).await } |
| 55 | }, |
| 56 | ) |
| 57 | .await?; |
| 58 | |
| 59 | // 格式化结果 |
| 60 | if !result.success { |
| 61 | return Ok(ToolResult::error( |
| 62 | result.error.unwrap_or_else(|| "Query failed".to_string()), |
| 63 | )); |
| 64 | } |
| 65 | |
| 66 | // 截断行数 |
| 67 | let max_rows = config.max_rows; |
| 68 | let data = if result.data.len() > max_rows { |
| 69 | &result.data[..max_rows] |
| 70 | } else { |
| 71 | &result.data |
| 72 | }; |
| 73 | |
| 74 | // 构建 Markdown 表格 |
| 75 | let mut output = String::new(); |
| 76 | |
| 77 | if result.columns.is_empty() { |
| 78 | output.push_str(&format!( |
| 79 | "Query executed successfully. {} row(s) affected.\n", |
| 80 | result.row_count |
| 81 | )); |
| 82 | } else { |
| 83 | // 表头 |
| 84 | output.push_str("| "); |
| 85 | for col in &result.columns { |
| 86 | output.push_str(&col.name); |
| 87 | output.push_str(" | "); |
| 88 | } |
no test coverage detected