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::{
13    BatchQueryResult, CountQuery, CountResult, TableBatchQuery, TableQuery,
14};
15
16pub(super) async fn handle_query_count(
17    db: &PgPool,
18    catalog: &crate::mcp::schema::SchemaCatalog,
19    params: CountQuery,
20) -> Result<rmcp::model::CallToolResult, rmcp::ErrorData> {
21    Ok(crate::mcp::tool_result(
22        crate::telemetry::McpTool::QueryCount,
23        execute_count(db, catalog, params).await,
24    ))
25}
26
27pub(super) async fn handle_query(
28    db: &PgPool,
29    catalog: &crate::mcp::schema::SchemaCatalog,
30    params: TableQuery,
31) -> Result<rmcp::model::CallToolResult, rmcp::ErrorData> {
32    let results = match execute(db, catalog, params).await {
33        Ok(results) => results,
34        Err(error) => {
35            return Ok(crate::mcp::mcp_error(
36                crate::telemetry::McpTool::Query,
37                &error,
38            ));
39        }
40    };
41
42    crate::telemetry::record_mcp_query_rows(results.len() as u64);
43
44    Ok(crate::mcp::json_result(results))
45}
46
47pub(super) async fn handle_query_batch(
48    db: &PgPool,
49    catalog: &crate::mcp::schema::SchemaCatalog,
50    params: TableBatchQuery,
51) -> Result<rmcp::model::CallToolResult, rmcp::ErrorData> {
52    let results = match execute_batch(db, catalog, params).await {
53        Ok(results) => results,
54        Err(error) => {
55            return Ok(crate::mcp::mcp_error(
56                crate::telemetry::McpTool::QueryBatch,
57                &error,
58            ));
59        }
60    };
61
62    let total_rows = record_batch_metrics(&results);
63
64    crate::telemetry::record_mcp_query_rows(total_rows);
65
66    Ok(crate::mcp::json_result(results))
67}
68
69fn record_batch_metrics(results: &[BatchQueryResult]) -> u64 {
70    results
71        .iter()
72        .map(|item| {
73            if let Some(error_type) = item.error_type {
74                crate::telemetry::record_mcp_error(
75                    crate::telemetry::McpTool::QueryBatch,
76                    error_type,
77                );
78            }
79
80            item.row_count as u64
81        })
82        .sum()
83}
wowlab_sentinel/mcp/tools/query/mod.rs

wowlab_sentinel/mcp/tools/query/
mod.rs
