Skip to main content

valence_backend_sql/
query.rs

1//! Execute compiled SQL queries and map rows to JSON.
2
3use serde_json::{Map, Value};
4use valence_core::compiled_query::CompiledQuery;
5use valence_core::error::{Error, Result};
6
7/// Bind `$param_key` placeholders in SQL to `?` for SQLite positional binding.
8pub fn sql_with_positional_placeholders(
9    query: &str,
10    params: &[(String, Value)],
11) -> (String, Vec<Value>) {
12    let mut out = String::with_capacity(query.len());
13    let mut values = Vec::new();
14    let mut rest = query;
15    while let Some(dollar) = rest.find('$') {
16        out.push_str(&rest[..dollar]);
17        rest = &rest[dollar + 1..];
18        let key_len = rest
19            .chars()
20            .take_while(|c| c.is_ascii_alphanumeric() || *c == '_')
21            .count();
22        let key = &rest[..key_len];
23        rest = &rest[key_len..];
24        if let Some((_, value)) = params.iter().find(|(k, _)| k == key) {
25            out.push('?');
26            values.push(value.clone());
27        } else {
28            out.push('$');
29            out.push_str(key);
30        }
31    }
32    out.push_str(rest);
33    (out, values)
34}
35
36/// Decode a SQL row `(id, body)` into Valence JSON record shape.
37///
38/// # Errors
39///
40/// Currently infallible; returns `Ok` with an empty object body if JSON parse fails.
41pub fn row_to_json(table: &str, id: &str, body_text: &str) -> Result<Value> {
42    let body: Value = serde_json::from_str(body_text).unwrap_or_else(|_| Value::Object(Map::new()));
43    Ok(super::document::row_from_body(table, id, body))
44}
45
46/// Parse SELECT results from generic JSON rows returned by driver layer.
47///
48/// # Errors
49///
50/// Currently infallible; malformed row shapes are skipped or passed through.
51pub fn decode_select_rows(rows: Vec<Value>, default_table: &str) -> Result<Vec<Value>> {
52    let mut out = Vec::new();
53    for row in rows {
54        if let Some(obj) = row.as_object() {
55            if let (Some(id), Some(body)) = (obj.get("id"), obj.get("body")) {
56                let id_str = id.as_str().unwrap_or_default();
57                let body_val = if let Some(s) = body.as_str() {
58                    serde_json::from_str(s).unwrap_or_else(|_| Value::Object(Map::new()))
59                } else {
60                    body.clone()
61                };
62                out.push(super::document::row_from_body(
63                    default_table,
64                    id_str,
65                    body_val,
66                ));
67                continue;
68            }
69        }
70        out.push(row);
71    }
72    Ok(out)
73}
74
75/// Extract count from first row.
76pub fn first_count(rows: &[Value]) -> i64 {
77    rows.first()
78        .and_then(|v| {
79            v.as_i64()
80                .or_else(|| v.get("count").and_then(|c| c.as_i64()))
81                .or_else(|| {
82                    v.as_f64().map(|f| {
83                        #[allow(clippy::cast_possible_truncation)] // count aggregation
84                        {
85                            f as i64
86                        }
87                    })
88                })
89        })
90        .unwrap_or(0)
91}
92
93/// Extract id strings from SELECT id queries.
94pub fn extract_ids(rows: &[Value]) -> Vec<String> {
95    rows.iter()
96        .filter_map(|v| {
97            v.as_str()
98                .map(str::to_string)
99                .or_else(|| v.get("id").and_then(|id| id.as_str().map(str::to_string)))
100        })
101        .collect()
102}
103
104/// Validate compiled query is read-only SELECT.
105///
106/// # Errors
107///
108/// Returns [`Error::Internal`] when the query is not a `SELECT`.
109pub fn ensure_read_only(query: &str) -> Result<()> {
110    let upper = query.trim().to_uppercase();
111    if upper.starts_with("SELECT ") {
112        Ok(())
113    } else {
114        Err(Error::Internal(format!(
115            "unsupported SQL in execute_compiled_query: {query}"
116        )))
117    }
118}
119
120/// Normalize compiled query for SQLite execution (`?` placeholders).
121///
122/// # Errors
123///
124/// Returns [`Error::Internal`] when the compiled query is not read-only.
125pub fn prepare_compiled(compiled: &CompiledQuery) -> Result<(String, Vec<Value>)> {
126    ensure_read_only(&compiled.query_string)?;
127    Ok(sql_with_positional_placeholders(
128        &compiled.query_string,
129        &compiled.params,
130    ))
131}
132
133/// Translate SQLite-style `json_extract(expr, '$.a.b')` into Postgres jsonb operators.
134pub fn rewrite_json_extract_for_postgres(sql: &str) -> String {
135    let mut out = String::with_capacity(sql.len());
136    let mut rest = sql;
137    while let Some(start) = rest.find("json_extract(") {
138        out.push_str(&rest[..start]);
139        rest = &rest[start + "json_extract(".len()..];
140        let Some(comma) = rest.find(',') else {
141            out.push_str("json_extract(");
142            break;
143        };
144        let expr = rest[..comma].trim();
145        rest = rest[comma + 1..].trim_start();
146        let path = if let Some(stripped) = rest.strip_prefix("'$.") {
147            let Some(end_q) = stripped.find('\'') else {
148                out.push_str("json_extract(");
149                out.push_str(expr);
150                out.push_str(", ");
151                break;
152            };
153            let path = &stripped[..end_q];
154            rest = stripped[end_q + 1..].trim_start();
155            if let Some(r) = rest.strip_prefix(')') {
156                rest = r;
157            }
158            path
159        } else {
160            out.push_str("json_extract(");
161            out.push_str(expr);
162            out.push_str(", ");
163            continue;
164        };
165        let parts: Vec<&str> = path.split('.').filter(|p| !p.is_empty()).collect();
166        if parts.len() <= 1 {
167            let seg = parts.first().copied().unwrap_or("");
168            out.push_str(&format!("({expr}->>'{seg}')"));
169        } else {
170            out.push_str(&format!("({expr}#>>'{{{}}}')", parts.join(",")));
171        }
172    }
173    out.push_str(rest);
174    out
175}
176
177/// Normalize compiled query for Postgres execution (`$1`, `$2`, … placeholders).
178///
179/// # Errors
180///
181/// Returns [`Error::Internal`] when the compiled query is not read-only.
182pub fn prepare_compiled_postgres(compiled: &CompiledQuery) -> Result<(String, Vec<Value>)> {
183    ensure_read_only(&compiled.query_string)?;
184    let rewritten = rewrite_json_extract_for_postgres(&compiled.query_string);
185    Ok(sql_with_postgres_placeholders(&rewritten, &compiled.params))
186}
187
188/// Bind `$param_key` placeholders to Postgres numbered params.
189pub fn sql_with_postgres_placeholders(
190    query: &str,
191    params: &[(String, Value)],
192) -> (String, Vec<Value>) {
193    let mut out = String::with_capacity(query.len());
194    let mut values = Vec::new();
195    let mut rest = query;
196    let mut idx = 1usize;
197    while let Some(dollar) = rest.find('$') {
198        out.push_str(&rest[..dollar]);
199        rest = &rest[dollar + 1..];
200        let key_len = rest
201            .chars()
202            .take_while(|c| c.is_ascii_alphanumeric() || *c == '_')
203            .count();
204        let key = &rest[..key_len];
205        rest = &rest[key_len..];
206        if let Some((_, value)) = params.iter().find(|(k, _)| k == key) {
207            out.push_str(&format!("${idx}"));
208            values.push(value.clone());
209            idx += 1;
210        } else if key.chars().all(|c| c.is_ascii_digit()) {
211            // Already positional ($1) — keep as-is.
212            out.push('$');
213            out.push_str(key);
214        } else {
215            out.push('$');
216            out.push_str(key);
217        }
218    }
219    out.push_str(rest);
220    (out, values)
221}
222
223#[cfg(test)]
224mod tests {
225    #![allow(clippy::expect_used, clippy::unwrap_used)]
226
227    use super::*;
228    use serde_json::json;
229
230    #[test]
231    fn positional_preserves_json_path_and_binds_params() {
232        let q = "SELECT id, body FROM task WHERE (json_extract(body, '$.project') = $param_0 OR json_extract(body, '$.project') = $param_1 OR json_extract(body, '$.project.id') = $param_1)";
233        let params = vec![
234            ("param_0".into(), json!("project:mem-1")),
235            ("param_1".into(), json!("mem-1")),
236        ];
237        let (sql, values) = sql_with_positional_placeholders(q, &params);
238        assert_eq!(
239            sql,
240            "SELECT id, body FROM task WHERE (json_extract(body, '$.project') = ? OR json_extract(body, '$.project') = ? OR json_extract(body, '$.project.id') = ?)"
241        );
242        assert_eq!(values.len(), 3);
243        assert_eq!(values[0], json!("project:mem-1"));
244        assert_eq!(values[1], json!("mem-1"));
245        assert_eq!(values[2], json!("mem-1"));
246    }
247
248    #[test]
249    fn postgres_rewrites_json_extract_paths() {
250        let q = "SELECT id FROM task WHERE json_extract(body, '$.project') = $p OR json_extract(t.body, '$.project.id') = $p";
251        let out = rewrite_json_extract_for_postgres(q);
252        assert_eq!(
253            out,
254            "SELECT id FROM task WHERE (body->>'project') = $p OR (t.body#>>'{project,id}') = $p"
255        );
256    }
257
258    #[test]
259    fn prepare_compiled_postgres_rewrites_json_extract() {
260        let compiled = CompiledQuery {
261            query_string: "SELECT id, body FROM project WHERE json_extract(body, '$.name') = $param_0 LIMIT 10".into(),
262            params: vec![("param_0".into(), json!("alpha"))],
263        };
264        let (sql, values) = prepare_compiled_postgres(&compiled).expect("prepare");
265        assert!(
266            !sql.contains("json_extract"),
267            "raw json_extract must be rewritten: {sql}"
268        );
269        assert!(sql.contains("body->>'name'") || sql.contains("(body->>'name')"));
270        assert_eq!(values, vec![json!("alpha")]);
271    }
272}