valence_backend_sql/
query.rs1use serde_json::{Map, Value};
4use valence_core::compiled_query::CompiledQuery;
5use valence_core::error::{Error, Result};
6
7pub 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
36pub 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
46pub 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
75pub 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)] {
85 f as i64
86 }
87 })
88 })
89 })
90 .unwrap_or(0)
91}
92
93pub 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
104pub 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
120pub 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
133pub 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
177pub 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
188pub 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 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, ¶ms);
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}