Skip to main content

wowlab_sentinel/mcp/tools/query/
build.rs

1use serde_json::Value;
2use sqlx::QueryBuilder;
3
4use super::types::{FilterOp, TableQuery};
5use crate::mcp::schema::{ColType, Table, get_column};
6
7type SqlQuery<'a> = QueryBuilder<'a, sqlx::Postgres>;
8
9fn push_quoted(qb: &mut SqlQuery<'_>, name: &str) {
10    qb.push("\"");
11    qb.push(name);
12    qb.push("\"");
13}
14
15pub(super) fn escape_ilike(s: &str) -> String {
16    let mut out = String::with_capacity(s.len());
17
18    for c in s.chars() {
19        match c {
20            '%' | '_' | '\\' => {
21                out.push('\\');
22                out.push(c);
23            }
24            _ => out.push(c),
25        }
26    }
27
28    out
29}
30
31const MAX_ROW_LIMIT: u64 = 1000;
32
33pub(super) fn validate_columns(
34    table: &Table,
35    q: &TableQuery,
36) -> Result<(), super::execute::QueryError> {
37    for f in &q.filters {
38        if get_column(table, &f.column).is_none() {
39            // #t(rust_alloc_in_loop) error message built only on validation failure
40            return Err(format!("Column '{}' not filterable on {}", f.column, table.name()).into());
41        }
42    }
43
44    if let Some(col) = &q.order_by {
45        let col_def = get_column(table, col);
46
47        if col_def.is_none() {
48            return Err(format!("Column '{}' not sortable on {}", col, table.name()).into());
49        }
50
51        if col_def.is_some_and(|c| matches!(c.typ, ColType::Json | ColType::IntArray)) {
52            return Err(
53                format!("Column '{col}' is json/array and cannot be used for ordering").into(),
54            );
55        }
56    }
57
58    Ok(())
59}
60
61pub(super) fn build_query<'a>(
62    table: &'a Table,
63    q: &'a TableQuery,
64) -> Result<QueryBuilder<'a, sqlx::Postgres>, super::execute::QueryError> {
65    let mut qb: QueryBuilder<sqlx::Postgres> = QueryBuilder::new("SELECT ");
66
67    for (i, col) in table.columns.iter().enumerate() {
68        if i > 0 {
69            qb.push(", ");
70        }
71
72        push_quoted(&mut qb, col.name);
73    }
74
75    qb.push(" FROM ");
76    qb.push(table.name());
77
78    let mut first = true;
79
80    for f in &q.filters {
81        let col_def = get_column(table, &f.column)
82            .ok_or_else(|| format!("Column '{}' not found after validation", f.column))?;
83
84        qb.push(if first { " WHERE " } else { " AND " });
85        first = false;
86        push_filter(&mut qb, col_def.name, col_def.typ, &f.op, f.value.as_ref())?;
87    }
88
89    if let Some(user_col) = &q.order_by {
90        let col_def = get_column(table, user_col)
91            .ok_or_else(|| format!("Column '{user_col}' not found after validation"))?;
92
93        qb.push(" ORDER BY ");
94        push_quoted(&mut qb, col_def.name);
95        qb.push(if q.order_desc { " DESC" } else { " ASC" });
96    }
97
98    qb.push(" LIMIT ");
99    qb.push_bind(i64::try_from(q.limit.min(MAX_ROW_LIMIT)).unwrap_or(i64::MAX));
100
101    if let Some(off) = q.offset {
102        qb.push(" OFFSET ");
103        qb.push_bind(i64::try_from(off).unwrap_or(i64::MAX));
104    }
105
106    Ok(qb)
107}
108
109// #t(fn: rust_cyclomatic_complexity, rust_max_fn_lines) one-arm-per-FilterOp dispatch over all column types; splitting per variant would obscure shared error / column-type checks
110pub(super) fn push_filter(
111    qb: &mut SqlQuery<'_>,
112    col_name: &str,
113    col_type: ColType,
114    op: &FilterOp,
115    value: Option<&Value>,
116) -> Result<(), super::execute::QueryError> {
117    let val = value.ok_or("Filter requires a value")?;
118
119    match op {
120        FilterOp::Eq | FilterOp::Ne => {
121            if col_type == ColType::Json {
122                return Err(
123                    "eq/ne not supported on json columns; use jsonb_contains or jsonb_has_key"
124                        .into(),
125                );
126            }
127
128            if col_type == ColType::IntArray {
129                return Err("eq/ne not supported on int_array columns; use array_contains".into());
130            }
131
132            let symbol = if matches!(op, FilterOp::Eq) {
133                " = "
134            } else {
135                " <> "
136            };
137
138            push_quoted(qb, col_name);
139            qb.push(symbol);
140            push_typed_value(qb, val, col_type)?;
141        }
142        FilterOp::Gt | FilterOp::Gte | FilterOp::Lt | FilterOp::Lte => {
143            if !matches!(col_type, ColType::Int | ColType::Float | ColType::Timestamp) {
144                return Err(format!("{op:?} only works on numeric or timestamp columns").into());
145            }
146
147            let symbol = match op {
148                FilterOp::Gt => " > ",
149                FilterOp::Gte => " >= ",
150                FilterOp::Lt => " < ",
151                FilterOp::Lte => " <= ",
152                _ => unreachable!("outer match restricts this branch to comparison operators"),
153            };
154
155            push_quoted(qb, col_name);
156            qb.push(symbol);
157            push_typed_value(qb, val, col_type)?;
158        }
159        FilterOp::Contains => {
160            if col_type != ColType::Text {
161                return Err("contains only works on text columns".into());
162            }
163
164            let s = val.as_str().ok_or("contains requires string value")?;
165
166            push_quoted(qb, col_name);
167            qb.push(" ILIKE ");
168            let escaped = escape_ilike(s);
169
170            qb.push_bind(format!("%{escaped}%"));
171        }
172        FilterOp::StartsWith => {
173            if col_type != ColType::Text {
174                return Err("starts_with only works on text columns".into());
175            }
176
177            let s = val.as_str().ok_or("starts_with requires string value")?;
178
179            qb.push("starts_with(");
180            push_quoted(qb, col_name);
181            qb.push(", ");
182            qb.push_bind(s.to_string()).push(")");
183        }
184        FilterOp::In => {
185            let values = val.as_array().ok_or("in requires array value")?;
186
187            if values.is_empty() {
188                return Err("in requires a non-empty array".into());
189            }
190
191            push_quoted(qb, col_name);
192            qb.push(" = ANY(");
193
194            match col_type {
195                ColType::Int => {
196                    let v: Vec<i64> = values
197                        .iter()
198                        .map(|v| v.as_i64().ok_or("in: expected integer values"))
199                        .collect::<Result<_, _>>()?;
200
201                    qb.push_bind(v);
202                }
203                ColType::Float => {
204                    let v: Vec<f64> = values
205                        .iter()
206                        .map(|v| v.as_f64().ok_or("in: expected float values"))
207                        .collect::<Result<_, _>>()?;
208
209                    qb.push_bind(v);
210                }
211                ColType::Text => {
212                    let v: Vec<String> = values
213                        .iter()
214                        .map(|v| {
215                            v.as_str()
216                                .map(String::from)
217                                .ok_or("in: expected string values")
218                        })
219                        .collect::<Result<_, _>>()?;
220
221                    qb.push_bind(v);
222                }
223                ColType::Bool => {
224                    let v: Vec<bool> = values
225                        .iter()
226                        .map(|v| v.as_bool().ok_or("in: expected bool values"))
227                        .collect::<Result<_, _>>()?;
228
229                    qb.push_bind(v);
230                }
231                ColType::Json => {
232                    return Err("in not supported on json columns".into());
233                }
234                ColType::IntArray => {
235                    return Err("in not supported on int_array columns; use array_contains".into());
236                }
237                ColType::Timestamp => {
238                    return Err("in not supported on timestamp columns; use gt/lt ranges".into());
239                }
240            }
241
242            qb.push(")");
243        }
244        FilterOp::JsonbContains => {
245            if col_type != ColType::Json {
246                return Err("jsonb_contains only works on json columns".into());
247            }
248
249            push_quoted(qb, col_name);
250            qb.push(" @> ");
251            qb.push_bind(val.clone());
252        }
253        FilterOp::JsonbHasKey => {
254            if col_type != ColType::Json {
255                return Err("jsonb_has_key only works on json columns".into());
256            }
257
258            let key = val.as_str().ok_or("jsonb_has_key requires a string key")?;
259
260            qb.push("jsonb_exists(");
261            push_quoted(qb, col_name);
262            qb.push(", ");
263            qb.push_bind(key.to_string());
264            qb.push(")");
265        }
266        FilterOp::JsonbArrayContains => {
267            if col_type != ColType::Json {
268                return Err("jsonb_array_contains only works on json columns".into());
269            }
270
271            if val.is_array() {
272                return Err(
273                    "jsonb_array_contains value must be an object, not an array (it is wrapped in [] automatically)".into(),
274                );
275            }
276
277            push_quoted(qb, col_name);
278            qb.push(" @> ");
279            qb.push_bind(Value::Array(vec![val.clone()]));
280        }
281        FilterOp::ArrayContains => {
282            if col_type != ColType::IntArray {
283                return Err("array_contains only works on int_array columns".into());
284            }
285
286            let n = val
287                .as_i64()
288                .ok_or("array_contains requires an integer value")?;
289            let n = i32::try_from(n)
290                .map_err(|error| format!("array_contains integer is outside i32 range: {error}"))?;
291
292            qb.push_bind(n);
293            qb.push(" = ANY(");
294            push_quoted(qb, col_name);
295            qb.push(")");
296        }
297    }
298
299    Ok(())
300}
301
302pub(super) fn push_typed_value(
303    qb: &mut SqlQuery<'_>,
304    val: &Value,
305    expected: ColType,
306) -> Result<(), super::execute::QueryError> {
307    match (expected, val) {
308        (ColType::Int, Value::Number(n)) => {
309            qb.push_bind(n.as_i64().ok_or("expected integer")?);
310        }
311        (ColType::Float, Value::Number(n)) => {
312            qb.push_bind(n.as_f64().ok_or("expected float")?);
313        }
314        (ColType::Text, Value::String(s)) => {
315            qb.push_bind(s.clone());
316        }
317        (ColType::Bool, Value::Bool(b)) => {
318            qb.push_bind(*b);
319        }
320        (ColType::Timestamp, Value::String(s)) => {
321            let ts = chrono::DateTime::parse_from_rfc3339(s)
322                .map_err(|e| format!("expected RFC 3339 timestamp: {e}"))?;
323
324            qb.push_bind(ts.with_timezone(&chrono::Utc));
325        }
326        _ => return Err(format!("type mismatch: expected {expected:?}").into()),
327    }
328
329    Ok(())
330}