Skip to main content

helios_sof/sqlquery/
engine.rs

1//! In-memory SQLite engine used by `$sqlquery-run`.
2//!
3//! One connection per request. Each depends-on ViewDefinition is materialized
4//! into a named table; the user's SQL then runs against those tables.
5
6use futures::Stream;
7use futures::StreamExt;
8use rusqlite::{Connection, ToSql, params_from_iter};
9use serde_json::Value;
10use std::pin::Pin;
11
12use super::{BoundParam, SqlQueryError};
13
14/// FHIR type code for a column. Mirrors the value-set used by
15/// `ViewDefinition.select.column.type` so we can pick the correct value[X]
16/// when rendering `_format=fhir`.
17#[derive(Debug, Clone, PartialEq, Eq)]
18pub enum ColumnFhirType {
19    Boolean,
20    Integer,
21    Integer64,
22    Decimal,
23    Date,
24    DateTime,
25    Instant,
26    Time,
27    Base64Binary,
28    /// Catch-all for `string`, `code`, `id`, `uri`, `canonical`, `url`,
29    /// `markdown`, `oid`, etc. The exact code is preserved so the FHIR
30    /// formatter can emit `valueCode` vs `valueString` correctly.
31    String(String),
32}
33
34impl ColumnFhirType {
35    pub fn from_code(code: &str) -> Self {
36        match code {
37            "boolean" => ColumnFhirType::Boolean,
38            "integer" | "positiveInt" | "unsignedInt" => ColumnFhirType::Integer,
39            "integer64" => ColumnFhirType::Integer64,
40            "decimal" => ColumnFhirType::Decimal,
41            "date" => ColumnFhirType::Date,
42            "dateTime" => ColumnFhirType::DateTime,
43            "instant" => ColumnFhirType::Instant,
44            "time" => ColumnFhirType::Time,
45            "base64Binary" => ColumnFhirType::Base64Binary,
46            other => ColumnFhirType::String(other.to_string()),
47        }
48    }
49
50    /// SQLite type-affinity declaration for `CREATE TABLE`.
51    pub fn sqlite_affinity(&self) -> &'static str {
52        match self {
53            ColumnFhirType::Boolean | ColumnFhirType::Integer | ColumnFhirType::Integer64 => {
54                "INTEGER"
55            }
56            ColumnFhirType::Decimal => "REAL",
57            _ => "TEXT",
58        }
59    }
60}
61
62/// One column in a materialized table.
63#[derive(Debug, Clone)]
64pub struct ColumnSchema {
65    pub name: String,
66    pub fhir_type: ColumnFhirType,
67}
68
69/// Per-table schema: the column list (order matters for INSERT).
70#[derive(Debug, Clone)]
71pub struct TableSchema {
72    pub columns: Vec<ColumnSchema>,
73}
74
75impl TableSchema {
76    /// Build a schema from a ViewDefinition's `select[].column[]` list.
77    /// Walks every `select` entry (including nested `select` under `forEach`)
78    /// and collects columns in document order.
79    pub fn from_view_definition(view: &Value) -> Self {
80        let mut columns = Vec::new();
81        if let Some(selects) = view.get("select").and_then(|v| v.as_array()) {
82            for s in selects {
83                collect_columns(s, &mut columns);
84            }
85        }
86        TableSchema { columns }
87    }
88}
89
90fn collect_columns(select: &Value, out: &mut Vec<ColumnSchema>) {
91    if let Some(cols) = select.get("column").and_then(|v| v.as_array()) {
92        for col in cols {
93            let Some(name) = col.get("name").and_then(|v| v.as_str()) else {
94                continue;
95            };
96            let type_code = col
97                .get("type")
98                .and_then(|v| v.as_str())
99                .unwrap_or("string")
100                .to_string();
101            out.push(ColumnSchema {
102                name: name.to_string(),
103                fhir_type: ColumnFhirType::from_code(&type_code),
104            });
105        }
106    }
107    if let Some(nested) = select.get("select").and_then(|v| v.as_array()) {
108        for s in nested {
109            collect_columns(s, out);
110        }
111    }
112    if let Some(union) = select.get("unionAll").and_then(|v| v.as_array()) {
113        for s in union {
114            collect_columns(s, out);
115        }
116    }
117}
118
119/// Result of running the user query.
120pub struct QueryResult {
121    pub columns: Vec<String>,
122    /// Column FHIR types, in `columns` order. Inferred from the rusqlite
123    /// declared column type plus a per-row check (NULL columns fall back to
124    /// `String`).
125    pub column_types: Vec<ColumnFhirType>,
126    /// Each row is a Vec of optional values in `columns` order.
127    pub rows: Vec<Vec<Option<Value>>>,
128}
129
130/// The in-memory SQLite engine.
131pub struct InMemorySqlEngine {
132    conn: Connection,
133}
134
135impl InMemorySqlEngine {
136    pub fn open() -> Result<Self, SqlQueryError> {
137        let conn = Connection::open_in_memory()?;
138        // Aggressive in-memory pragmas — we never persist this DB.
139        conn.execute_batch(
140            "PRAGMA journal_mode = MEMORY;
141             PRAGMA synchronous = OFF;
142             PRAGMA temp_store = MEMORY;
143             PRAGMA foreign_keys = OFF;",
144        )?;
145        Ok(Self { conn })
146    }
147
148    /// Returns an interrupt handle that can cancel a running statement from
149    /// another thread (used by the request-level timeout watchdog).
150    pub fn interrupt_handle(&self) -> rusqlite::InterruptHandle {
151        self.conn.get_interrupt_handle()
152    }
153
154    /// Create a table with the given label and schema.
155    pub fn create_table(&self, label: &str, schema: &TableSchema) -> Result<(), SqlQueryError> {
156        validate_identifier(label)?;
157        let mut columns_ddl = Vec::with_capacity(schema.columns.len());
158        for col in &schema.columns {
159            validate_identifier(&col.name)?;
160            columns_ddl.push(format!(
161                "\"{}\" {}",
162                col.name,
163                col.fhir_type.sqlite_affinity()
164            ));
165        }
166        let sql = if columns_ddl.is_empty() {
167            // SQLite needs at least one column.
168            format!("CREATE TABLE \"{label}\" (\"_empty\" TEXT)")
169        } else {
170            format!("CREATE TABLE \"{}\" ({})", label, columns_ddl.join(", "))
171        };
172        self.conn.execute(&sql, [])?;
173        Ok(())
174    }
175
176    /// Stream `rows` into `label`. Each row is a flat JSON object whose keys
177    /// match column names; missing or null keys become SQL NULL.
178    pub async fn insert_rows<S>(
179        &mut self,
180        label: &str,
181        schema: &TableSchema,
182        mut rows: Pin<Box<S>>,
183        max_rows: usize,
184    ) -> Result<usize, SqlQueryError>
185    where
186        S: Stream<Item = Result<Value, String>> + Send + ?Sized,
187    {
188        validate_identifier(label)?;
189        for col in &schema.columns {
190            validate_identifier(&col.name)?;
191        }
192        if schema.columns.is_empty() {
193            // Drain the stream without inserting; nothing to persist.
194            let mut n = 0usize;
195            while let Some(item) = rows.next().await {
196                item.map_err(SqlQueryError::MalformedLibrary)?;
197                n += 1;
198                if n > max_rows {
199                    return Err(SqlQueryError::RowCapExceeded { max: max_rows });
200                }
201            }
202            return Ok(n);
203        }
204
205        let placeholders = std::iter::repeat_n("?", schema.columns.len())
206            .collect::<Vec<_>>()
207            .join(", ");
208        let cols_quoted = schema
209            .columns
210            .iter()
211            .map(|c| format!("\"{}\"", c.name))
212            .collect::<Vec<_>>()
213            .join(", ");
214        let insert_sql = format!("INSERT INTO \"{label}\" ({cols_quoted}) VALUES ({placeholders})");
215
216        self.conn.execute("BEGIN", [])?;
217        let mut inserted = 0usize;
218        let result: Result<usize, SqlQueryError> = (|| {
219            let mut stmt = self.conn.prepare(&insert_sql)?;
220            while let Some(item) = futures::executor::block_on(rows.next()) {
221                let row = item.map_err(SqlQueryError::MalformedLibrary)?;
222                inserted += 1;
223                if inserted > max_rows {
224                    return Err(SqlQueryError::RowCapExceeded { max: max_rows });
225                }
226                let params: Vec<rusqlite::types::Value> = schema
227                    .columns
228                    .iter()
229                    .map(|c| json_to_sqlite_value(&row, c))
230                    .collect();
231                let param_refs: Vec<&dyn ToSql> = params.iter().map(|v| v as &dyn ToSql).collect();
232                stmt.execute(params_from_iter(param_refs))?;
233            }
234            Ok(inserted)
235        })();
236        match result {
237            Ok(n) => {
238                self.conn.execute("COMMIT", [])?;
239                Ok(n)
240            }
241            Err(e) => {
242                let _ = self.conn.execute("ROLLBACK", []);
243                Err(e)
244            }
245        }
246    }
247
248    /// Run a SELECT with named bindings and a row cap.
249    pub fn execute_select(
250        &self,
251        sql: &str,
252        bindings: &[BoundParam],
253        max_rows: usize,
254    ) -> Result<QueryResult, SqlQueryError> {
255        let mut stmt = self.conn.prepare(sql)?;
256
257        // Resolve each `:name` binding against the prepared statement's
258        // parameter index. Names not referenced by the SQL are silently
259        // ignored (the SQL may declare more params than it uses, or none).
260        for b in bindings {
261            let with_colon = format!(":{}", b.name);
262            if let Some(idx) = stmt.parameter_index(&with_colon)? {
263                stmt.raw_bind_parameter(idx, &b.value)?;
264            }
265        }
266
267        let columns: Vec<String> = stmt.column_names().into_iter().map(String::from).collect();
268        // Pre-seed with String to be overwritten per row.
269        let mut column_types: Vec<ColumnFhirType> = columns
270            .iter()
271            .map(|_| ColumnFhirType::String("string".to_string()))
272            .collect();
273        let mut rows_out: Vec<Vec<Option<Value>>> = Vec::new();
274
275        let mut rows_iter = stmt.raw_query();
276        while let Some(row) = rows_iter.next()? {
277            if rows_out.len() >= max_rows {
278                // SoF v2: the server's hard cap silently truncates the
279                // result set instead of erroring (spec PR #353: "Servers
280                // MAY enforce a maximum value, silently capping
281                // client-supplied limits at a smaller server-defined
282                // maximum"). Caller-supplied `_limit` is also a silent
283                // cap and is enforced at the handler. Source-row caps
284                // applied during `insert_rows` remain hard errors
285                // because truncating a depends-on table would silently
286                // change query semantics (JOINs, aggregates).
287                break;
288            }
289            let mut row_vals: Vec<Option<Value>> = Vec::with_capacity(columns.len());
290            for (i, _) in columns.iter().enumerate() {
291                let v: rusqlite::types::Value = row.get(i)?;
292                let (json_val, inferred) = sqlite_value_to_json(v);
293                if matches!(column_types[i], ColumnFhirType::String(_)) {
294                    if let Some(ft) = inferred {
295                        column_types[i] = ft;
296                    }
297                }
298                row_vals.push(json_val);
299            }
300            rows_out.push(row_vals);
301        }
302
303        Ok(QueryResult {
304            columns,
305            column_types,
306            rows: rows_out,
307        })
308    }
309}
310
311fn validate_identifier(name: &str) -> Result<(), SqlQueryError> {
312    if name.contains('"') || name.is_empty() {
313        return Err(SqlQueryError::InvalidIdentifier(name.to_string()));
314    }
315    Ok(())
316}
317
318fn json_to_sqlite_value(row: &Value, col: &ColumnSchema) -> rusqlite::types::Value {
319    use rusqlite::types::Value as RV;
320    let raw = row.get(&col.name).unwrap_or(&Value::Null);
321    match raw {
322        Value::Null => RV::Null,
323        Value::Bool(b) => RV::Integer(if *b { 1 } else { 0 }),
324        Value::Number(n) => {
325            if let Some(i) = n.as_i64() {
326                RV::Integer(i)
327            } else if let Some(f) = n.as_f64() {
328                RV::Real(f)
329            } else {
330                RV::Text(n.to_string())
331            }
332        }
333        Value::String(s) => match col.fhir_type {
334            ColumnFhirType::Integer | ColumnFhirType::Integer64 => s
335                .parse::<i64>()
336                .map(RV::Integer)
337                .unwrap_or(RV::Text(s.clone())),
338            ColumnFhirType::Decimal => s
339                .parse::<f64>()
340                .map(RV::Real)
341                .unwrap_or(RV::Text(s.clone())),
342            ColumnFhirType::Boolean => match s.as_str() {
343                "true" | "1" => RV::Integer(1),
344                "false" | "0" => RV::Integer(0),
345                _ => RV::Text(s.clone()),
346            },
347            _ => RV::Text(s.clone()),
348        },
349        Value::Array(_) | Value::Object(_) => RV::Text(raw.to_string()),
350    }
351}
352
353/// Maps a rusqlite value to JSON plus a best-guess `ColumnFhirType`. Useful
354/// for output columns the engine produced (e.g. `SELECT COUNT(*)`).
355fn sqlite_value_to_json(v: rusqlite::types::Value) -> (Option<Value>, Option<ColumnFhirType>) {
356    use rusqlite::types::Value as RV;
357    match v {
358        RV::Null => (None, None),
359        RV::Integer(i) => (Some(Value::Number(i.into())), Some(ColumnFhirType::Integer)),
360        RV::Real(f) => (
361            serde_json::Number::from_f64(f).map(Value::Number),
362            Some(ColumnFhirType::Decimal),
363        ),
364        RV::Text(s) => (Some(Value::String(s)), None),
365        RV::Blob(b) => (
366            Some(Value::String(
367                base64::engine::general_purpose::STANDARD.encode(b),
368            )),
369            Some(ColumnFhirType::Base64Binary),
370        ),
371    }
372}
373
374use base64::Engine as _;
375
376#[cfg(test)]
377mod tests {
378    use super::*;
379    use futures::stream;
380    use serde_json::json;
381
382    fn schema(cols: &[(&str, ColumnFhirType)]) -> TableSchema {
383        TableSchema {
384            columns: cols
385                .iter()
386                .map(|(n, t)| ColumnSchema {
387                    name: (*n).to_string(),
388                    fhir_type: t.clone(),
389                })
390                .collect(),
391        }
392    }
393
394    #[tokio::test]
395    async fn round_trip_basic() {
396        let mut engine = InMemorySqlEngine::open().unwrap();
397        let s = schema(&[
398            ("id", ColumnFhirType::String("id".into())),
399            ("n", ColumnFhirType::Integer),
400        ]);
401        engine.create_table("patients", &s).unwrap();
402        let rows = stream::iter(vec![
403            Ok(json!({"id": "a", "n": 1})),
404            Ok(json!({"id": "b", "n": 2})),
405        ]);
406        let inserted = engine
407            .insert_rows("patients", &s, Box::pin(rows), 10)
408            .await
409            .unwrap();
410        assert_eq!(inserted, 2);
411        let result = engine
412            .execute_select("SELECT id, n FROM patients ORDER BY n", &[], 10)
413            .unwrap();
414        assert_eq!(result.columns, vec!["id", "n"]);
415        assert_eq!(result.rows.len(), 2);
416        assert_eq!(result.rows[0][0], Some(Value::String("a".into())));
417        assert_eq!(result.rows[0][1], Some(Value::Number(1.into())));
418    }
419
420    #[tokio::test]
421    async fn null_handling() {
422        let mut engine = InMemorySqlEngine::open().unwrap();
423        let s = schema(&[
424            ("id", ColumnFhirType::String("id".into())),
425            ("age", ColumnFhirType::Integer),
426        ]);
427        engine.create_table("t", &s).unwrap();
428        let rows = stream::iter(vec![Ok(json!({"id": "a"}))]); // age missing
429        engine
430            .insert_rows("t", &s, Box::pin(rows), 10)
431            .await
432            .unwrap();
433        let result = engine
434            .execute_select("SELECT id, age FROM t", &[], 10)
435            .unwrap();
436        assert_eq!(result.rows[0][1], None);
437    }
438
439    #[tokio::test]
440    async fn row_cap_exceeded() {
441        let mut engine = InMemorySqlEngine::open().unwrap();
442        let s = schema(&[("n", ColumnFhirType::Integer)]);
443        engine.create_table("t", &s).unwrap();
444        let rows = stream::iter((0..10).map(|i| Ok(json!({"n": i}))));
445        let err = engine
446            .insert_rows("t", &s, Box::pin(rows), 3)
447            .await
448            .unwrap_err();
449        assert!(matches!(err, SqlQueryError::RowCapExceeded { max: 3 }));
450    }
451
452    #[tokio::test]
453    async fn execute_select_silently_truncates_at_max_rows() {
454        // SoF v2 PR #353: the server's hard cap silently truncates the
455        // result set; it must not error.
456        let mut engine = InMemorySqlEngine::open().unwrap();
457        let s = schema(&[("n", ColumnFhirType::Integer)]);
458        engine.create_table("t", &s).unwrap();
459        let rows = stream::iter((1..=10).map(|i| Ok(json!({"n": i}))));
460        engine
461            .insert_rows("t", &s, Box::pin(rows), 100)
462            .await
463            .unwrap();
464        let result = engine
465            .execute_select("SELECT n FROM t ORDER BY n", &[], 4)
466            .unwrap();
467        assert_eq!(result.rows.len(), 4);
468        assert_eq!(result.rows[0][0], Some(Value::Number(1.into())));
469        assert_eq!(result.rows[3][0], Some(Value::Number(4.into())));
470    }
471
472    #[test]
473    fn rejects_quote_in_identifier() {
474        let engine = InMemorySqlEngine::open().unwrap();
475        let s = schema(&[("a", ColumnFhirType::Integer)]);
476        let err = engine.create_table("bad\"name", &s).unwrap_err();
477        assert!(matches!(err, SqlQueryError::InvalidIdentifier(_)));
478    }
479
480    #[tokio::test]
481    async fn named_bindings_filter() {
482        let mut engine = InMemorySqlEngine::open().unwrap();
483        let s = schema(&[("n", ColumnFhirType::Integer)]);
484        engine.create_table("t", &s).unwrap();
485        let rows = stream::iter((1..=5).map(|i| Ok(json!({"n": i}))));
486        engine
487            .insert_rows("t", &s, Box::pin(rows), 100)
488            .await
489            .unwrap();
490        let bindings = vec![BoundParam {
491            name: "min".to_string(),
492            value: rusqlite::types::Value::Integer(3),
493        }];
494        let result = engine
495            .execute_select("SELECT n FROM t WHERE n >= :min ORDER BY n", &bindings, 100)
496            .unwrap();
497        assert_eq!(result.rows.len(), 3);
498    }
499
500    #[test]
501    fn schema_from_vd_select_columns() {
502        let vd = json!({
503            "select": [{
504                "column": [
505                    {"name": "id", "type": "id"},
506                    {"name": "n", "type": "integer"}
507                ]
508            }]
509        });
510        let s = TableSchema::from_view_definition(&vd);
511        assert_eq!(s.columns.len(), 2);
512        assert_eq!(s.columns[0].name, "id");
513        assert!(matches!(s.columns[1].fhir_type, ColumnFhirType::Integer));
514    }
515
516    #[test]
517    fn schema_walks_nested_selects_and_union() {
518        let vd = json!({
519            "select": [{
520                "column": [{"name": "a"}],
521                "select": [{"column": [{"name": "b"}]}],
522                "unionAll": [{"column": [{"name": "c"}]}]
523            }]
524        });
525        let s = TableSchema::from_view_definition(&vd);
526        assert_eq!(
527            s.columns.iter().map(|c| c.name.clone()).collect::<Vec<_>>(),
528            vec!["a", "b", "c"]
529        );
530    }
531}