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
44pub(super) async fn execute(
45    db: &PgPool,
46    catalog: &SchemaCatalog,
47    q: TableQuery,
48) -> Result<Vec<Value>, QueryError> {
49    let table = catalog.table(&q.table)?;
50
51    validate_columns(table, &q)?;
52
53    let mut query = build_query(table, &q)?;
54    let rows = timeout(
55        Duration::from_secs(QUERY_TIMEOUT_SECS),
56        query.build().fetch_all(db),
57    )
58    .await??;
59
60    Ok(rows.iter().map(|row| row_to_json(row, table)).collect())
61}
62
63pub(super) async fn execute_batch(
64    db: &PgPool,
65    catalog: &SchemaCatalog,
66    batch: TableBatchQuery,
67) -> Result<Vec<BatchQueryResult>, QueryError> {
68    if batch.queries.is_empty() {
69        return Err(QueryError::EmptyBatch);
70    }
71
72    if batch.queries.len() > MAX_BATCH_QUERIES {
73        return Err(QueryError::BatchTooLarge);
74    }
75
76    let futs = batch.queries.into_iter().map(|q| {
77        let table_name = q.table.clone();
78
79        async move { (table_name, execute(db, catalog, q).await) }
80    });
81
82    Ok(futures::future::join_all(futs)
83        .await
84        .into_iter()
85        .enumerate()
86        .map(|(index, (table, result))| match result {
87            Ok(rows) => BatchQueryResult {
88                index,
89                table,
90                row_count: rows.len(),
91                results: Some(rows),
92                error: None,
93            },
94            Err(err) => BatchQueryResult {
95                index,
96                table,
97                row_count: 0,
98                results: None,
99                error: Some(err.to_string()),
100            },
101        })
102        .collect())
103}
104
105pub(super) async fn execute_count(
106    db: &PgPool,
107    catalog: &SchemaCatalog,
108    q: CountQuery,
109) -> Result<CountResult, QueryError> {
110    let table = catalog.table(&q.table)?;
111
112    let mut qb = build_count_query(table, &q)?;
113    let row = timeout(
114        Duration::from_secs(QUERY_TIMEOUT_SECS),
115        qb.build().fetch_one(db),
116    )
117    .await??;
118
119    let count: i64 = row.try_get(0)?;
120
121    Ok(CountResult {
122        table: q.table,
123        count,
124    })
125}
126
127fn build_count_query<'a>(
128    table: &'a Table,
129    q: &'a CountQuery,
130) -> Result<QueryBuilder<'a, sqlx::Postgres>, QueryError> {
131    if let Some(filter) = q
132        .filters
133        .iter()
134        .find(|filter| get_column(table, &filter.column).is_none())
135    {
136        return Err(format!(
137            "Column '{}' not filterable on {}",
138            filter.column,
139            table.name()
140        )
141        .into());
142    }
143
144    let mut qb = QueryBuilder::new("SELECT COUNT(*) FROM ");
145
146    qb.push(table.name());
147
148    let mut first = true;
149
150    // #t(block: rust_alloc_in_loop) error message built only on lookup failure
151    for f in &q.filters {
152        let col_def = get_column(table, &f.column)
153            .ok_or_else(|| format!("Column '{}' not found after validation", f.column))?;
154
155        qb.push(if first { " WHERE " } else { " AND " });
156        first = false;
157        push_filter(&mut qb, col_def.name, col_def.typ, &f.op, f.value.as_ref())?;
158    }
159
160    Ok(qb)
161}
162
163fn row_to_json(row: &sqlx::postgres::PgRow, table: &Table) -> Value {
164    let mut map = Map::with_capacity(table.columns.len());
165
166    for col in table.columns {
167        let value = match col.typ {
168            ColType::Int => row
169                .try_get::<i64, _>(col.name)
170                .map(Value::from)
171                .or_else(|_| row.try_get::<i32, _>(col.name).map(Value::from))
172                .unwrap_or(Value::Null),
173            ColType::Float => row
174                .try_get::<f64, _>(col.name)
175                .map(Value::from)
176                .or_else(|_| row.try_get::<f32, _>(col.name).map(Value::from))
177                .unwrap_or(Value::Null),
178            ColType::Text => row
179                .try_get::<String, _>(col.name)
180                .map_or(Value::Null, Value::from),
181            ColType::Bool => row
182                .try_get::<bool, _>(col.name)
183                .map_or(Value::Null, Value::from),
184            ColType::Json => row.try_get::<Value, _>(col.name).unwrap_or(Value::Null),
185            ColType::Timestamp => row
186                .try_get::<chrono::DateTime<chrono::Utc>, _>(col.name)
187                .map_or(Value::Null, |v| Value::from(v.to_rfc3339())),
188            ColType::IntArray => row
189                .try_get::<Vec<i32>, _>(col.name)
190                .map_or(Value::Null, |v| {
191                    Value::Array(v.into_iter().map(Value::from).collect())
192                }),
193        };
194
195        // #t(rust_alloc_in_loop) column names are small strings built per row for JSON output
196        map.insert(col.name.to_string(), value);
197    }
198
199    Value::Object(map)
200}
wowlab_sentinel/mcp/tools/query/execute.rs

wowlab_sentinel/mcp/tools/query/
execute.rs
