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    if matches!(op, FilterOp::IsNull | FilterOp::IsNotNull) {
118        push_quoted(qb, col_name);
119        qb.push(if matches!(op, FilterOp::IsNull) {
120            " IS NULL"
121        } else {
122            " IS NOT NULL"
123        });
124
125        return Ok(());
126    }
127
128    let val = value.ok_or("Filter requires a value")?;
129
130    match op {
131        FilterOp::Eq | FilterOp::Ne => {
132            if col_type == ColType::Json {
133                return Err(
134                    "eq/ne not supported on json columns; use jsonb_contains or jsonb_has_key"
135                        .into(),
136                );
137            }
138
139            if col_type == ColType::IntArray {
140                return Err("eq/ne not supported on int_array columns; use array_contains".into());
141            }
142
143            let symbol = if matches!(op, FilterOp::Eq) {
144                " = "
145            } else {
146                " <> "
147            };
148
149            push_quoted(qb, col_name);
150            qb.push(symbol);
151            push_typed_value(qb, val, col_type)?;
152        }
153        FilterOp::Gt | FilterOp::Gte | FilterOp::Lt | FilterOp::Lte => {
154            if !matches!(col_type, ColType::Int | ColType::Float | ColType::Timestamp) {
155                return Err(format!("{op:?} only works on numeric or timestamp columns").into());
156            }
157
158            let symbol = match op {
159                FilterOp::Gt => " > ",
160                FilterOp::Gte => " >= ",
161                FilterOp::Lt => " < ",
162                FilterOp::Lte => " <= ",
163                _ => unreachable!("outer match restricts this branch to comparison operators"),
164            };
165
166            push_quoted(qb, col_name);
167            qb.push(symbol);
168            push_typed_value(qb, val, col_type)?;
169        }
170        FilterOp::Contains => {
171            if col_type != ColType::Text {
172                return Err("contains only works on text columns".into());
173            }
174
175            let s = val.as_str().ok_or("contains requires string value")?;
176
177            push_quoted(qb, col_name);
178            qb.push(" ILIKE ");
179            let escaped = escape_ilike(s);
180
181            qb.push_bind(format!("%{escaped}%"));
182        }
183        FilterOp::StartsWith => {
184            if col_type != ColType::Text {
185                return Err("starts_with only works on text columns".into());
186            }
187
188            let s = val.as_str().ok_or("starts_with requires string value")?;
189
190            qb.push("starts_with(");
191            push_quoted(qb, col_name);
192            qb.push(", ");
193            qb.push_bind(s.to_string()).push(")");
194        }
195        FilterOp::In => {
196            let values = val.as_array().ok_or("in requires array value")?;
197
198            if values.is_empty() {
199                return Err("in requires a non-empty array".into());
200            }
201
202            push_quoted(qb, col_name);
203            qb.push(" = ANY(");
204
205            match col_type {
206                ColType::Int => {
207                    let v: Vec<i64> = values
208                        .iter()
209                        .map(|v| v.as_i64().ok_or("in: expected integer values"))
210                        .collect::<Result<_, _>>()?;
211
212                    qb.push_bind(v);
213                }
214                ColType::Float => {
215                    let v: Vec<f64> = values
216                        .iter()
217                        .map(|v| v.as_f64().ok_or("in: expected float values"))
218                        .collect::<Result<_, _>>()?;
219
220                    qb.push_bind(v);
221                }
222                ColType::Text => {
223                    let v: Vec<String> = values
224                        .iter()
225                        .map(|v| {
226                            v.as_str()
227                                .map(String::from)
228                                .ok_or("in: expected string values")
229                        })
230                        .collect::<Result<_, _>>()?;
231
232                    qb.push_bind(v);
233                }
234                ColType::Bool => {
235                    let v: Vec<bool> = values
236                        .iter()
237                        .map(|v| v.as_bool().ok_or("in: expected bool values"))
238                        .collect::<Result<_, _>>()?;
239
240                    qb.push_bind(v);
241                }
242                ColType::Json => {
243                    return Err("in not supported on json columns".into());
244                }
245                ColType::IntArray => {
246                    return Err("in not supported on int_array columns; use array_contains".into());
247                }
248                ColType::Timestamp => {
249                    return Err("in not supported on timestamp columns; use gt/lt ranges".into());
250                }
251            }
252
253            qb.push(")");
254        }
255        FilterOp::JsonbContains => {
256            if col_type != ColType::Json {
257                return Err("jsonb_contains only works on json columns".into());
258            }
259
260            push_quoted(qb, col_name);
261            qb.push(" @> ");
262            qb.push_bind(val.clone());
263        }
264        FilterOp::JsonbHasKey => {
265            if col_type != ColType::Json {
266                return Err("jsonb_has_key only works on json columns".into());
267            }
268
269            let key = val.as_str().ok_or("jsonb_has_key requires a string key")?;
270
271            qb.push("jsonb_exists(");
272            push_quoted(qb, col_name);
273            qb.push(", ");
274            qb.push_bind(key.to_string());
275            qb.push(")");
276        }
277        FilterOp::JsonbArrayContains => {
278            if col_type != ColType::Json {
279                return Err("jsonb_array_contains only works on json columns".into());
280            }
281
282            if val.is_array() {
283                return Err(
284                    "jsonb_array_contains value must be an object, not an array (it is wrapped in [] automatically)".into(),
285                );
286            }
287
288            push_quoted(qb, col_name);
289            qb.push(" @> ");
290            qb.push_bind(Value::Array(vec![val.clone()]));
291        }
292        FilterOp::ArrayContains => {
293            if col_type != ColType::IntArray {
294                return Err("array_contains only works on int_array columns".into());
295            }
296
297            let n = val
298                .as_i64()
299                .ok_or("array_contains requires an integer value")?;
300            let n = i32::try_from(n)
301                .map_err(|error| format!("array_contains integer is outside i32 range: {error}"))?;
302
303            qb.push_bind(n);
304            qb.push(" = ANY(");
305            push_quoted(qb, col_name);
306            qb.push(")");
307        }
308        FilterOp::IsNull | FilterOp::IsNotNull => {
309            unreachable!("null operators return before value validation")
310        }
311    }
312
313    Ok(())
314}
315
316pub(super) fn push_typed_value(
317    qb: &mut SqlQuery<'_>,
318    val: &Value,
319    expected: ColType,
320) -> Result<(), super::execute::QueryError> {
321    match (expected, val) {
322        (ColType::Int, Value::Number(n)) => {
323            qb.push_bind(n.as_i64().ok_or("expected integer")?);
324        }
325        (ColType::Float, Value::Number(n)) => {
326            qb.push_bind(n.as_f64().ok_or("expected float")?);
327        }
328        (ColType::Text, Value::String(s)) => {
329            qb.push_bind(s.clone());
330        }
331        (ColType::Bool, Value::Bool(b)) => {
332            qb.push_bind(*b);
333        }
334        (ColType::Timestamp, Value::String(s)) => {
335            let ts = chrono::DateTime::parse_from_rfc3339(s)
336                .map_err(|e| format!("expected RFC 3339 timestamp: {e}"))?;
337
338            qb.push_bind(ts.with_timezone(&chrono::Utc));
339        }
340        _ => return Err(format!("type mismatch: expected {expected:?}").into()),
341    }
342
343    Ok(())
344}