Skip to main content

ontocore_query/
sql.rs

1use crate::QueryError;
2use ontocore_catalog::OntologyCatalog;
3use ontocore_core::{limits::MAX_QUERY_BYTES, limits::MAX_SQL_RESULT_ROWS, EntityKind};
4use serde::Serialize;
5use sqlparser::ast::{
6    Expr, GroupByExpr, Select, SelectItem, SetExpr, Statement, TableFactor, Value,
7};
8use sqlparser::dialect::GenericDialect;
9use sqlparser::parser::Parser;
10use std::collections::BTreeMap;
11
12pub type Result<T> = std::result::Result<T, QueryError>;
13
14type Row = BTreeMap<String, String>;
15
16#[derive(Debug, Clone, Serialize)]
17pub struct QueryResult {
18    pub columns: Vec<String>,
19    pub rows: Vec<Row>,
20    pub truncated: bool,
21}
22
23pub fn run_sql(catalog: &OntologyCatalog, sql: &str) -> Result<QueryResult> {
24    if sql.len() > MAX_QUERY_BYTES {
25        return Err(QueryError::Sql(format!(
26            "query exceeds maximum length of {MAX_QUERY_BYTES} bytes"
27        )));
28    }
29
30    let dialect = GenericDialect {};
31    let statements =
32        Parser::parse_sql(&dialect, sql).map_err(|e| QueryError::Sql(e.to_string()))?;
33
34    if statements.len() > 1 {
35        return Err(QueryError::Sql("only a single SQL statement is supported".to_string()));
36    }
37
38    let statement =
39        statements.into_iter().next().ok_or_else(|| QueryError::Sql("empty query".to_string()))?;
40
41    match statement {
42        Statement::Query(query) => {
43            if query.order_by.is_some() {
44                return Err(QueryError::Sql("ORDER BY is not supported".to_string()));
45            }
46            if query.limit.is_some() || query.offset.is_some() {
47                return Err(QueryError::Sql("LIMIT and OFFSET are not supported".to_string()));
48            }
49            let select = match *query.body {
50                SetExpr::Select(select) => select,
51                _ => return Err(QueryError::Sql("only SELECT queries are supported".to_string())),
52            };
53            execute_select(catalog, select)
54        }
55        _ => Err(QueryError::Sql("only SELECT queries are supported".to_string())),
56    }
57}
58
59fn execute_select(catalog: &OntologyCatalog, select: Box<Select>) -> Result<QueryResult> {
60    if select.distinct.is_some() {
61        return Err(QueryError::Sql("DISTINCT is not supported".to_string()));
62    }
63    if group_by_present(&select.group_by) {
64        return Err(QueryError::Sql("GROUP BY is not supported".to_string()));
65    }
66    if select.from.len() > 1 {
67        return Err(QueryError::Sql("JOIN is not supported".to_string()));
68    }
69    let table_name = table_name_from_select(&select)?;
70    if let Some(selection) = &select.selection {
71        if matches!(selection, Expr::Identifier(_)) {
72            return Err(QueryError::Sql(
73                "bare column names are not supported in WHERE; use column = 'value'".to_string(),
74            ));
75        }
76        validate_filter(selection)?;
77    }
78
79    let mut rows = Vec::new();
80    let mut truncated = false;
81    for row in table_row_iter(catalog, &table_name)? {
82        if let Some(selection) = &select.selection {
83            if !evaluate_filter(selection, &row)? {
84                continue;
85            }
86        }
87        if rows.len() >= MAX_SQL_RESULT_ROWS {
88            truncated = true;
89            break;
90        }
91        rows.push(row);
92    }
93
94    let (columns, projected_rows) = project_rows(&select.projection, &rows)?;
95    Ok(QueryResult { columns, rows: projected_rows, truncated })
96}
97
98fn table_name_from_select(select: &Select) -> Result<String> {
99    let from =
100        select.from.first().ok_or_else(|| QueryError::Sql("missing FROM clause".to_string()))?;
101
102    match &from.relation {
103        TableFactor::Table { name, .. } => Ok(name.to_string().to_ascii_lowercase()),
104        _ => Err(QueryError::Sql("unsupported table expression".to_string())),
105    }
106}
107
108fn table_row_iter<'a>(
109    catalog: &'a OntologyCatalog,
110    table: &str,
111) -> Result<Box<dyn Iterator<Item = Row> + 'a>> {
112    let data = catalog.data();
113    match table {
114        "ontologies" => Ok(Box::new(data.documents.iter().map(|doc| {
115            let mut row = BTreeMap::new();
116            row.insert("id".into(), doc.id.clone());
117            row.insert("path".into(), doc.path.display().to_string());
118            row.insert("format".into(), doc.format.as_str().to_string());
119            row.insert("base_iri".into(), doc.base_iri.clone().unwrap_or_default());
120            row.insert("parse_status".into(), doc.parse_status.as_str().to_string());
121            row.insert("content_hash".into(), doc.content_hash.clone());
122            row.insert("modified_time".into(), doc.modified_time.to_string());
123            row
124        }))),
125        "classes" => entity_row_iter(catalog, EntityKind::Class),
126        "object_properties" => entity_row_iter(catalog, EntityKind::ObjectProperty),
127        "data_properties" => entity_row_iter(catalog, EntityKind::DataProperty),
128        "annotation_properties" => entity_row_iter(catalog, EntityKind::AnnotationProperty),
129        "individuals" => entity_row_iter(catalog, EntityKind::Individual),
130        "entities" => Ok(Box::new(data.entities.iter().map(entity_to_row))),
131        "annotations" => Ok(Box::new(data.annotations.iter().map(|a| {
132            let mut row = BTreeMap::new();
133            row.insert("subject".into(), a.subject.clone());
134            row.insert("predicate".into(), a.predicate.clone());
135            row.insert("object".into(), a.object.clone());
136            row.insert("ontology_id".into(), a.ontology_id.clone());
137            row
138        }))),
139        "axioms" => Ok(Box::new(data.axioms.iter().map(|a| {
140            let mut row = BTreeMap::new();
141            row.insert("id".into(), a.id.clone());
142            row.insert("ontology_id".into(), a.ontology_id.clone());
143            row.insert("subject".into(), a.subject.clone());
144            row.insert("predicate".into(), a.predicate.clone());
145            row.insert("object".into(), a.object.clone());
146            row.insert("axiom_kind".into(), a.axiom_kind.clone());
147            row
148        }))),
149        "namespaces" => Ok(Box::new(data.namespaces.iter().map(|n| {
150            let mut row = BTreeMap::new();
151            row.insert("prefix".into(), n.prefix.clone());
152            row.insert("iri".into(), n.iri.clone());
153            row.insert("ontology_id".into(), n.ontology_id.clone());
154            row
155        }))),
156        "imports" => Ok(Box::new(data.imports.iter().map(|i| {
157            let mut row = BTreeMap::new();
158            row.insert("ontology_id".into(), i.ontology_id.clone());
159            row.insert("import_iri".into(), i.import_iri.clone());
160            row
161        }))),
162        "diagnostics" => Ok(Box::new(data.diagnostics.iter().map(|d| {
163            let mut row = BTreeMap::new();
164            row.insert("code".into(), d.code.as_str().to_string());
165            row.insert("severity".into(), d.severity.as_str().to_string());
166            row.insert("message".into(), d.message.clone());
167            row.insert("file".into(), d.file.display().to_string());
168            row.insert("line".into(), d.range.line.map(|l| l.to_string()).unwrap_or_default());
169            row.insert("column".into(), d.range.column.map(|c| c.to_string()).unwrap_or_default());
170            row.insert("entity_iri".into(), d.entity_iri.clone().unwrap_or_default());
171            row
172        }))),
173        "properties" => {
174            let mut iter: Box<dyn Iterator<Item = Row>> = Box::new(std::iter::empty());
175            for kind in [
176                EntityKind::ObjectProperty,
177                EntityKind::DataProperty,
178                EntityKind::AnnotationProperty,
179            ] {
180                let next = entity_row_iter(catalog, kind)?;
181                iter = Box::new(iter.chain(next));
182            }
183            Ok(iter)
184        }
185        other => Err(QueryError::Sql(format!("unknown table: {other}"))),
186    }
187}
188
189fn entity_row_iter(
190    catalog: &OntologyCatalog,
191    kind: EntityKind,
192) -> Result<Box<dyn Iterator<Item = Row> + '_>> {
193    Ok(Box::new(catalog.data().entities.iter().filter(move |e| e.kind == kind).map(entity_to_row)))
194}
195
196fn entity_to_row(entity: &ontocore_core::Entity) -> Row {
197    let mut row = BTreeMap::new();
198    row.insert("iri".into(), entity.iri.clone());
199    row.insert("short_name".into(), entity.short_name.clone());
200    row.insert("kind".into(), entity.kind.as_str().to_string());
201    row.insert("ontology_id".into(), entity.ontology_id.clone());
202    row.insert("labels".into(), entity.labels.join("; "));
203    row.insert("comments".into(), entity.comments.join("; "));
204    row.insert("deprecated".into(), entity.deprecated.to_string());
205    if let Some(ref obo_id) = entity.obo_id {
206        row.insert("obo_id".into(), obo_id.clone());
207    }
208    row
209}
210
211struct ProjectionCol {
212    name: String,
213    source: String,
214}
215
216fn projection_columns(projection: &[SelectItem]) -> Result<Vec<ProjectionCol>> {
217    if projection.len() == 1 && matches!(projection[0], SelectItem::Wildcard(_)) {
218        return Ok(Vec::new());
219    }
220
221    let mut columns = Vec::new();
222    for item in projection {
223        match item {
224            SelectItem::UnnamedExpr(Expr::Identifier(ident)) => {
225                let col = ident.value.to_ascii_lowercase();
226                columns.push(ProjectionCol { name: col.clone(), source: col });
227            }
228            SelectItem::ExprWithAlias { expr, alias, .. } => {
229                let source = match expr {
230                    Expr::Identifier(ident) => ident.value.to_ascii_lowercase(),
231                    _ => {
232                        return Err(QueryError::Sql(
233                            "only simple column projections are supported".to_string(),
234                        ));
235                    }
236                };
237                columns.push(ProjectionCol { name: alias.value.to_ascii_lowercase(), source });
238            }
239            SelectItem::Wildcard(_) => {
240                return Err(QueryError::Sql("wildcard projection must be used alone".to_string()));
241            }
242            _ => {
243                return Err(QueryError::Sql(
244                    "only simple column projections are supported".to_string(),
245                ));
246            }
247        }
248    }
249    Ok(columns)
250}
251
252fn project_rows(projection: &[SelectItem], rows: &[Row]) -> Result<(Vec<String>, Vec<Row>)> {
253    if projection.len() == 1 && matches!(projection[0], SelectItem::Wildcard(_)) {
254        let columns = rows.first().map(|r| r.keys().cloned().collect()).unwrap_or_default();
255        return Ok((columns, rows.to_vec()));
256    }
257
258    let columns_spec = projection_columns(projection)?;
259    let columns: Vec<String> = columns_spec.iter().map(|c| c.name.clone()).collect();
260
261    let projected = rows
262        .iter()
263        .map(|row| {
264            let mut out = BTreeMap::new();
265            for col in &columns_spec {
266                out.insert(col.name.clone(), row.get(&col.source).cloned().unwrap_or_default());
267            }
268            out
269        })
270        .collect();
271
272    Ok((columns, projected))
273}
274
275fn group_by_present(group_by: &GroupByExpr) -> bool {
276    match group_by {
277        GroupByExpr::All(_) => true,
278        GroupByExpr::Expressions(exprs, _) => !exprs.is_empty(),
279    }
280}
281
282fn validate_filter(expr: &Expr) -> Result<()> {
283    match expr {
284        Expr::BinaryOp { left, op, right } => {
285            use sqlparser::ast::BinaryOperator;
286            match op {
287                BinaryOperator::Eq
288                | BinaryOperator::NotEq
289                | BinaryOperator::And
290                | BinaryOperator::Or => {
291                    validate_filter(left)?;
292                    validate_filter(right)?;
293                    Ok(())
294                }
295                other => Err(QueryError::Sql(format!("unsupported WHERE operator: {other:?}"))),
296            }
297        }
298        Expr::Identifier(_) | Expr::Value(_) => Ok(()),
299        other => Err(QueryError::Sql(format!("unsupported WHERE expression: {other:?}"))),
300    }
301}
302
303fn evaluate_filter(expr: &Expr, row: &Row) -> Result<bool> {
304    match expr {
305        Expr::BinaryOp { left, op, right } => {
306            use sqlparser::ast::BinaryOperator;
307            match op {
308                BinaryOperator::Eq => Ok(eval_expr(left, row) == eval_expr(right, row)),
309                BinaryOperator::NotEq => Ok(eval_expr(left, row) != eval_expr(right, row)),
310                BinaryOperator::And => {
311                    Ok(evaluate_filter(left, row)? && evaluate_filter(right, row)?)
312                }
313                BinaryOperator::Or => {
314                    Ok(evaluate_filter(left, row)? || evaluate_filter(right, row)?)
315                }
316                other => Err(QueryError::Sql(format!("unsupported WHERE operator: {other:?}"))),
317            }
318        }
319        Expr::Identifier(ident) => {
320            Ok(row.get(&ident.value.to_ascii_lowercase()).map(|v| v == "true").unwrap_or(false))
321        }
322        Expr::Value(Value::Boolean(b)) => Ok(*b),
323        other => Err(QueryError::Sql(format!("unsupported WHERE expression: {other:?}"))),
324    }
325}
326
327fn eval_expr(expr: &Expr, row: &Row) -> String {
328    match expr {
329        Expr::Identifier(ident) => {
330            row.get(&ident.value.to_ascii_lowercase()).cloned().unwrap_or_default()
331        }
332        Expr::Value(Value::SingleQuotedString(s) | Value::DoubleQuotedString(s)) => s.clone(),
333        Expr::Value(Value::Boolean(b)) => b.to_string(),
334        Expr::Value(Value::Number(n, _)) => n.clone(),
335        _ => String::new(),
336    }
337}
338
339pub fn to_csv(result: &QueryResult) -> Result<String> {
340    let mut writer = csv::Writer::from_writer(Vec::new());
341    writer.write_record(&result.columns).map_err(|e| crate::QueryError::Export(e.to_string()))?;
342    for row in &result.rows {
343        let values: Vec<String> =
344            result.columns.iter().map(|c| row.get(c).cloned().unwrap_or_default()).collect();
345        writer.write_record(&values).map_err(|e| crate::QueryError::Export(e.to_string()))?;
346    }
347    let bytes = writer.into_inner().map_err(|e| crate::QueryError::Export(e.to_string()))?;
348    String::from_utf8(bytes).map_err(|e| crate::QueryError::Export(e.to_string()))
349}
350
351pub fn to_json(result: &QueryResult) -> Result<String> {
352    let rows: Vec<Vec<String>> = result
353        .rows
354        .iter()
355        .map(|row| result.columns.iter().map(|c| row.get(c).cloned().unwrap_or_default()).collect())
356        .collect();
357    serde_json::to_string_pretty(&serde_json::json!({
358        "columns": result.columns,
359        "rows": rows,
360        "truncated": result.truncated,
361    }))
362    .map_err(|e| crate::QueryError::Export(e.to_string()))
363}
364
365#[cfg(test)]
366mod tests {
367    use super::*;
368    use ontocore_catalog::OntologyCatalog;
369    use ontocore_core::limits::MAX_QUERY_BYTES;
370    use std::path::PathBuf;
371
372    fn fixture_catalog() -> OntologyCatalog {
373        let fixtures = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../fixtures");
374        ontocore_catalog::IndexBuilder::new().workspace(fixtures).build().expect("index fixtures")
375    }
376
377    #[test]
378    fn where_and_filters_rows() {
379        let catalog = fixture_catalog();
380        let result = run_sql(
381            &catalog,
382            "SELECT short_name FROM classes WHERE short_name = 'Person' AND deprecated = 'false'",
383        )
384        .expect("and filter");
385        assert_eq!(result.rows.len(), 1);
386        assert_eq!(result.rows[0].get("short_name").map(String::as_str), Some("Person"));
387    }
388
389    #[test]
390    fn where_or_filters_rows() {
391        let catalog = fixture_catalog();
392        let result = run_sql(
393            &catalog,
394            "SELECT short_name FROM classes WHERE short_name = 'Person' OR short_name = 'Thing'",
395        )
396        .expect("or filter");
397        let names: Vec<_> =
398            result.rows.iter().filter_map(|r| r.get("short_name").cloned()).collect();
399        assert!(names.contains(&"Person".to_string()));
400        assert!(names.contains(&"Thing".to_string()));
401    }
402
403    #[test]
404    fn where_not_eq_excludes_matching_row() {
405        let catalog = fixture_catalog();
406        let result =
407            run_sql(&catalog, "SELECT short_name FROM classes WHERE short_name != 'Person'")
408                .expect("not eq filter");
409        assert!(!result
410            .rows
411            .iter()
412            .any(|r| r.get("short_name").map(String::as_str) == Some("Person")));
413        assert!(!result.rows.is_empty());
414    }
415
416    #[test]
417    fn unsupported_like_returns_error() {
418        let catalog = fixture_catalog();
419        let err = run_sql(&catalog, "SELECT short_name FROM classes WHERE short_name LIKE 'Per%'")
420            .unwrap_err();
421        assert!(matches!(err, crate::QueryError::Sql(msg) if msg.contains("unsupported WHERE")));
422    }
423
424    #[test]
425    fn unsupported_limit_returns_error() {
426        let catalog = fixture_catalog();
427        let err = run_sql(&catalog, "SELECT short_name FROM classes LIMIT 1").unwrap_err();
428        assert!(matches!(err, crate::QueryError::Sql(msg) if msg.contains("LIMIT")));
429    }
430
431    #[test]
432    fn rejects_oversized_query() {
433        let catalog = fixture_catalog();
434        let padding = "x".repeat(MAX_QUERY_BYTES);
435        let sql = format!("SELECT short_name FROM classes WHERE short_name = '{padding}'");
436        let err = run_sql(&catalog, &sql).unwrap_err();
437        assert!(matches!(err, crate::QueryError::Sql(msg) if msg.contains("maximum length")));
438    }
439
440    #[test]
441    fn select_alias_uses_source_column() {
442        let catalog = fixture_catalog();
443        let result =
444            run_sql(&catalog, "SELECT short_name AS name FROM classes WHERE short_name = 'Person'")
445                .expect("alias projection");
446        assert_eq!(result.rows.len(), 1);
447        assert_eq!(result.rows[0].get("name").map(String::as_str), Some("Person"));
448    }
449}