wowlab_sentinel/mcp/tools/query/
mod.rs

1mod build;
2mod execute;
3#[cfg(test)]
4mod tests;
5
6mod types;
7
8use execute::{execute, execute_batch, execute_count};
9use sqlx::PgPool;
10
11#[rustfmt::skip]
12pub(super) use types::{CountQuery, TableBatchQuery, TableQuery};
13
14pub(super) async fn handle_query_count(
15    db: &PgPool,
16    catalog: &crate::mcp::schema::SchemaCatalog,
17    params: CountQuery,
18) -> Result<rmcp::model::CallToolResult, rmcp::ErrorData> {
19    let result = execute_count(db, catalog, params)
20        .await
21        .map_err(|e| crate::mcp::mcp_error(crate::telemetry::McpTool::QueryCount, &e))?;
22
23    Ok(crate::mcp::json_result(result))
24}
25
26pub(super) async fn handle_query(
27    db: &PgPool,
28    catalog: &crate::mcp::schema::SchemaCatalog,
29    params: TableQuery,
30) -> Result<rmcp::model::CallToolResult, rmcp::ErrorData> {
31    let results = execute(db, catalog, params)
32        .await
33        .map_err(|e| crate::mcp::mcp_error(crate::telemetry::McpTool::Query, &e))?;
34
35    crate::telemetry::record_mcp_query_rows(results.len() as u64);
36
37    Ok(if results.is_empty() {
38        rmcp::model::CallToolResult::success(vec![rmcp::model::Content::text("No results")])
39    } else {
40        crate::mcp::json_result(results)
41    })
42}
43
44pub(super) async fn handle_query_batch(
45    db: &PgPool,
46    catalog: &crate::mcp::schema::SchemaCatalog,
47    params: TableBatchQuery,
48) -> Result<rmcp::model::CallToolResult, rmcp::ErrorData> {
49    let results = execute_batch(db, catalog, params)
50        .await
51        .map_err(|e| crate::mcp::mcp_error(crate::telemetry::McpTool::QueryBatch, &e))?;
52
53    let total_rows = record_batch_metrics(&results);
54
55    crate::telemetry::record_mcp_query_rows(total_rows);
56
57    Ok(crate::mcp::json_result(results))
58}
59
60fn record_batch_metrics(results: &[types::BatchQueryResult]) -> u64 {
61    results
62        .iter()
63        .map(|item| {
64            if let Some(err) = &item.error {
65                crate::telemetry::record_mcp_error(
66                    crate::telemetry::McpTool::QueryBatch,
67                    crate::mcp::classify_error(err),
68                );
69            }
70
71            item.row_count as u64
72        })
73        .sum()
74}
wowlab_sentinel/mcp/tools/query/mod.rs

wowlab_sentinel/mcp/tools/query/
mod.rs
