wowlab_sentinel/mcp/tools/query/
execute.rs

1use std::time::Duration;
2
3use serde_json::{Map, Value};
4use sqlx::{PgPool, QueryBuilder, Row};
5use tokio::time::timeout;
6
7use super::{
8    build::{build_query, push_filter, validate_columns},
9    types::{BatchQueryResult, CountQuery, CountResult, TableBatchQuery, TableQuery},
10};
11use crate::mcp::schema::{ColType, SchemaCatalog, Table, UnknownTableError, get_column};
12
13const QUERY_TIMEOUT_SECS: u64 = 30;
14const MAX_BATCH_QUERIES: usize = 50;
15
16#[derive(Debug, thiserror::Error)]
17pub(super) enum QueryError {
18    #[error(transparent)]
19    UnknownTable(#[from] UnknownTableError),
20    #[error("{0}")]
21    Invalid(String),
22    #[error("Query timeout: {0}")]
23    Timeout(#[from] tokio::time::error::Elapsed),
24    #[error(transparent)]
25    Database(#[from] sqlx::Error),
26    #[error("Batch requires at least one query")]
27    EmptyBatch,
28    #[error("Batch supports at most 50 queries")]
29    BatchTooLarge,
30}
31
32impl From<String> for QueryError {
33    fn from(message: String) -> Self {
34        Self::Invalid(message)
35    }
36}
37
38impl From<&str> for QueryError {
39    fn from(message: &str) -> Self {
40        Self::Invalid(message.to_owned())
41    }
42}
43
44impl crate::mcp::McpFailure for QueryError {
45    fn error_type(&self) -> crate::telemetry::McpErrorType {
46        match self {
47            Self::UnknownTable(_) | Self::Invalid(_) | Self::EmptyBatch | Self::BatchTooLarge => {
48                crate::telemetry::McpErrorType::InvalidParams
49            }
50            Self::Timeout(_) => crate::telemetry::McpErrorType::Timeout,
51            Self::Database(_) => crate::telemetry::McpErrorType::Database,
52        }
53    }
54}
55
56pub(super) async fn execute(
57    db: &PgPool,
58    catalog: &SchemaCatalog,
59    q: TableQuery,
60) -> Result<Vec<Value>, QueryError> {
61    let table = catalog.table(&q.table)?;
62
63    validate_columns(table, &q)?;
64
65    let mut query = build_query(table, &q)?;
66    let rows = timeout(
67        Duration::from_secs(QUERY_TIMEOUT_SECS),
68        query.build().fetch_all(db),
69    )
70    .await??;
71
72    Ok(rows.iter().map(|row| row_to_json(row, table)).collect())
73}
74
75pub(super) async fn execute_batch(
76    db: &PgPool,
77    catalog: &SchemaCatalog,
78    batch: TableBatchQuery,
79) -> Result<Vec<BatchQueryResult>, QueryError> {
80    if batch.queries.is_empty() {
81        return Err(QueryError::EmptyBatch);
82    }
83
84    if batch.queries.len() > MAX_BATCH_QUERIES {
85        return Err(QueryError::BatchTooLarge);
86    }
87
88    let futs = batch.queries.into_iter().map(|q| {
89        let table_name = q.table.clone();
90
91        async move { (table_name, execute(db, catalog, q).await) }
92    });
93
94    Ok(futures::future::join_all(futs)
95        .await
96        .into_iter()
97        .enumerate()
98        .map(|(index, (table, result))| match result {
99            Ok(rows) => BatchQueryResult {
100                index,
101                table,
102                row_count: rows.len(),
103                results: Some(rows),
104                error: None,
105                error_type: None,
106            },
107            Err(err) => BatchQueryResult {
108                index,
109                table,
110                row_count: 0,
111                results: None,
112                error: Some(err.to_string()),
113                error_type: Some(crate::mcp::McpFailure::error_type(&err)),
114            },
115        })
116        .collect())
117}
118
119pub(super) async fn execute_count(
120    db: &PgPool,
121    catalog: &SchemaCatalog,
122    q: CountQuery,
123) -> Result<CountResult, QueryError> {
124    let table = catalog.table(&q.table)?;
125
126    let mut qb = build_count_query(table, &q)?;
127    let row = timeout(
128        Duration::from_secs(QUERY_TIMEOUT_SECS),
129        qb.build().fetch_one(db),
130    )
131    .await??;
132
133    let count: i64 = row.try_get(0)?;
134
135    Ok(CountResult {
136        table: q.table,
137        count,
138    })
139}
140
141fn build_count_query<'a>(
142    table: &'a Table,
143    q: &'a CountQuery,
144) -> Result<QueryBuilder<'a, sqlx::Postgres>, QueryError> {
145    if let Some(filter) = q
146        .filters
147        .iter()
148        .find(|filter| get_column(table, &filter.column).is_none())
149    {
150        return Err(format!(
151            "Column '{}' not filterable on {}",
152            filter.column,
153            table.name()
154        )
155        .into());
156    }
157
158    let mut qb = QueryBuilder::new("SELECT COUNT(*) FROM ");
159
160    qb.push(table.name());
161
162    let mut first = true;
163
164    // #t(block: rust_alloc_in_loop) error message built only on lookup failure
165    for f in &q.filters {
166        let col_def = get_column(table, &f.column)
167            .ok_or_else(|| format!("Column '{}' not found after validation", f.column))?;
168
169        qb.push(if first { " WHERE " } else { " AND " });
170        first = false;
171        push_filter(&mut qb, col_def.name, col_def.typ, &f.op, f.value.as_ref())?;
172    }
173
174    Ok(qb)
175}
176
177fn row_to_json(row: &sqlx::postgres::PgRow, table: &Table) -> Value {
178    let mut map = Map::with_capacity(table.columns.len());
179
180    for col in table.columns {
181        let value = match col.typ {
182            ColType::Int => row
183                .try_get::<i64, _>(col.name)
184                .map(Value::from)
185                .or_else(|_| row.try_get::<i32, _>(col.name).map(Value::from))
186                .unwrap_or(Value::Null),
187            ColType::Float => row
188                .try_get::<f64, _>(col.name)
189                .map(Value::from)
190                .or_else(|_| row.try_get::<f32, _>(col.name).map(Value::from))
191                .unwrap_or(Value::Null),
192            ColType::Text => row
193                .try_get::<String, _>(col.name)
194                .map_or(Value::Null, Value::from),
195            ColType::Bool => row
196                .try_get::<bool, _>(col.name)
197                .map_or(Value::Null, Value::from),
198            ColType::Json => row.try_get::<Value, _>(col.name).unwrap_or(Value::Null),
199            ColType::Timestamp => row
200                .try_get::<chrono::DateTime<chrono::Utc>, _>(col.name)
201                .map_or(Value::Null, |v| Value::from(v.to_rfc3339())),
202            ColType::IntArray => row
203                .try_get::<Vec<i32>, _>(col.name)
204                .map_or(Value::Null, |v| {
205                    Value::Array(v.into_iter().map(Value::from).collect())
206                }),
207        };
208
209        // #t(rust_alloc_in_loop) column names are small strings built per row for JSON output
210        map.insert(col.name.to_string(), value);
211    }
212
213    Value::Object(map)
214}
wowlab_sentinel/mcp/tools/query/execute.rs

wowlab_sentinel/mcp/tools/query/
execute.rs
