Skip to main content

corium_sql/
lib.rs

1//! Read-only SQL execution over immutable Corium database values.
2//!
3//! A [`SqlSession`] captures one [`corium_db::Db`] time view. Current,
4//! as-of, and since views expose one wide table per attribute namespace;
5//! history views expose normalized event relations only.
6
7mod catalog;
8mod value;
9
10use arrow::record_batch::RecordBatch;
11use corium_db::{Db, DbView};
12use datafusion::execution::context::SQLOptions;
13use datafusion::physical_plan::SendableRecordBatchStream;
14use datafusion::prelude::SessionContext;
15use futures::StreamExt as _;
16use thiserror::Error;
17
18pub use value::{SqlColumn, SqlRow, SqlType, SqlValue};
19
20/// SQL planning, catalog, or execution failure.
21#[derive(Debug, Error)]
22pub enum SqlError {
23    /// `DataFusion` rejected or failed the query.
24    #[error(transparent)]
25    DataFusion(#[from] datafusion::error::DataFusionError),
26    /// Arrow rejected a generated batch.
27    #[error(transparent)]
28    Arrow(#[from] arrow::error::ArrowError),
29    /// A Corium schema cannot be represented by the SQL projection.
30    #[error("SQL schema error: {0}")]
31    Schema(String),
32}
33
34/// One read-only SQL environment over a fixed immutable database view.
35pub struct SqlSession {
36    context: SessionContext,
37    basis_t: u64,
38    view: DbView,
39}
40
41impl SqlSession {
42    /// Builds the SQL catalog for `db`.
43    ///
44    /// Current, as-of, and since views get namespace-derived wide tables in
45    /// `corium`; every view gets normalized relations in `corium_sys`.
46    ///
47    /// # Errors
48    /// Returns [`SqlError`] when the database schema cannot be projected.
49    pub fn new(db: &Db) -> Result<Self, SqlError> {
50        let context = SessionContext::new();
51        catalog::register(&context, db)?;
52        Ok(Self {
53            context,
54            basis_t: db.basis_t(),
55            view: db.view(),
56        })
57    }
58
59    /// Transaction basis captured by this session.
60    #[must_use]
61    pub const fn basis_t(&self) -> u64 {
62        self.basis_t
63    }
64
65    /// Corium time view captured by this session.
66    #[must_use]
67    pub const fn view(&self) -> DbView {
68        self.view
69    }
70
71    /// Registered Corium relations as `schema.table` names.
72    #[must_use]
73    pub fn tables(&self) -> Vec<String> {
74        let Some(catalog) = self.context.catalog("datafusion") else {
75            return Vec::new();
76        };
77        let mut tables = Vec::new();
78        for schema_name in ["corium", "corium_sys"] {
79            if let Some(schema) = catalog.schema(schema_name) {
80                tables.extend(
81                    schema
82                        .table_names()
83                        .into_iter()
84                        .map(|table| format!("{schema_name}.{table}")),
85                );
86            }
87        }
88        tables.sort();
89        tables
90    }
91
92    /// Plans and starts a read-only SQL query.
93    ///
94    /// DDL, DML, and session-mutating statements are rejected. Dropping the
95    /// returned stream cancels unfinished execution.
96    ///
97    /// # Errors
98    /// Returns [`SqlError`] for SQL parsing, planning, or execution failure.
99    pub async fn query(&self, sql: &str) -> Result<SqlQuery, SqlError> {
100        let options = SQLOptions::new()
101            .with_allow_ddl(false)
102            .with_allow_dml(false)
103            .with_allow_statements(false);
104        let frame = self.context.sql_with_options(sql, options).await?;
105        let stream = frame.execute_stream().await?;
106        let columns = stream
107            .schema()
108            .fields()
109            .iter()
110            .map(|field| SqlColumn::from_arrow(field))
111            .collect();
112        Ok(SqlQuery {
113            columns,
114            stream,
115            batch: None,
116            row: 0,
117        })
118    }
119}
120
121/// Streaming result of one SQL statement.
122pub struct SqlQuery {
123    columns: Vec<SqlColumn>,
124    stream: SendableRecordBatchStream,
125    batch: Option<RecordBatch>,
126    row: usize,
127}
128
129impl SqlQuery {
130    /// Result columns in projection order.
131    #[must_use]
132    pub fn columns(&self) -> &[SqlColumn] {
133        &self.columns
134    }
135
136    /// Reads the next result row, or `None` at end of stream.
137    ///
138    /// # Errors
139    /// Returns [`SqlError`] when execution or value conversion fails.
140    pub async fn next_row(&mut self) -> Result<Option<SqlRow>, SqlError> {
141        loop {
142            if let Some(batch) = &self.batch
143                && self.row < batch.num_rows()
144            {
145                let row = batch
146                    .columns()
147                    .iter()
148                    .map(|array| {
149                        datafusion::common::ScalarValue::try_from_array(array.as_ref(), self.row)
150                            .map_err(SqlError::from)
151                            .and_then(SqlValue::from_scalar)
152                    })
153                    .collect::<Result<Vec<_>, _>>()?;
154                self.row += 1;
155                return Ok(Some(row));
156            }
157            match self.stream.next().await {
158                Some(batch) => {
159                    self.batch = Some(batch?);
160                    self.row = 0;
161                }
162                None => return Ok(None),
163            }
164        }
165    }
166
167    /// Collects all remaining rows.
168    ///
169    /// # Errors
170    /// Returns [`SqlError`] when execution or value conversion fails.
171    pub async fn collect(mut self) -> Result<Vec<SqlRow>, SqlError> {
172        let mut rows = Vec::new();
173        while let Some(row) = self.next_row().await? {
174            rows.push(row);
175        }
176        Ok(rows)
177    }
178}
179
180#[cfg(test)]
181mod tests {
182    use corium_core::{
183        Attribute, Cardinality, Datom, EntityId, Keyword, KeywordInterner, Schema, Value, ValueType,
184    };
185    use corium_db::{Db, Idents};
186
187    use super::*;
188
189    fn fixture() -> Db {
190        let name = EntityId::from_raw(10);
191        let tags = EntityId::from_raw(11);
192        let release_year = EntityId::from_raw(12);
193        let mut schema = Schema::default();
194        schema.insert(Attribute {
195            id: name,
196            value_type: ValueType::Str,
197            cardinality: Cardinality::One,
198            unique: None,
199            is_component: false,
200            indexed: false,
201            no_history: false,
202        });
203        schema.insert(Attribute {
204            id: tags,
205            value_type: ValueType::Str,
206            cardinality: Cardinality::Many,
207            unique: None,
208            is_component: false,
209            indexed: true,
210            no_history: false,
211        });
212        schema.insert(Attribute {
213            id: release_year,
214            value_type: ValueType::Long,
215            cardinality: Cardinality::One,
216            unique: None,
217            is_component: false,
218            indexed: true,
219            no_history: false,
220        });
221        let mut idents = Idents::default();
222        idents.insert(Keyword::parse("artist/name"), name);
223        idents.insert(Keyword::parse("artist/tags"), tags);
224        idents.insert(Keyword::parse("artist/release-year"), release_year);
225        let e = EntityId::from_raw(1_000);
226        let second = EntityId::from_raw(1_001);
227        let tx = EntityId::from_raw(1);
228        Db::new(schema)
229            .with_naming(idents, KeywordInterner::default())
230            .with_transaction(
231                1,
232                &[
233                    Datom {
234                        e,
235                        a: name,
236                        v: Value::Str("Boards of Canada".into()),
237                        tx,
238                        added: true,
239                    },
240                    Datom {
241                        e,
242                        a: tags,
243                        v: Value::Str("ambient".into()),
244                        tx,
245                        added: true,
246                    },
247                    Datom {
248                        e,
249                        a: tags,
250                        v: Value::Str("electronic".into()),
251                        tx,
252                        added: true,
253                    },
254                    Datom {
255                        e,
256                        a: release_year,
257                        v: Value::Long(1998),
258                        tx,
259                        added: true,
260                    },
261                    Datom {
262                        e: second,
263                        a: name,
264                        v: Value::Str("Tycho".into()),
265                        tx,
266                        added: true,
267                    },
268                    Datom {
269                        e: second,
270                        a: release_year,
271                        v: Value::Long(2011),
272                        tx,
273                        added: true,
274                    },
275                ],
276            )
277    }
278
279    #[tokio::test]
280    async fn missing_many_attribute_is_an_empty_list() {
281        let session = SqlSession::new(&fixture()).expect("session");
282        let rows = session
283            .query("SELECT tags FROM corium.artist WHERE name = 'Tycho'")
284            .await
285            .expect("query")
286            .collect()
287            .await
288            .expect("rows");
289        assert_eq!(rows, vec![vec![SqlValue::List(Vec::new())]]);
290    }
291
292    #[tokio::test]
293    async fn entity_equality_uses_the_wide_provider_lookup_path() {
294        let session = SqlSession::new(&fixture()).expect("session");
295        let rows = session
296            .query("SELECT name FROM corium.artist WHERE e = 1001")
297            .await
298            .expect("query")
299            .collect()
300            .await
301            .expect("rows");
302        assert_eq!(rows, vec![vec![SqlValue::Text("Tycho".into())]]);
303
304        let missing = session
305            .query("SELECT name FROM corium.artist WHERE e = 9999")
306            .await
307            .expect("query")
308            .collect()
309            .await
310            .expect("rows");
311        assert!(missing.is_empty());
312    }
313
314    #[tokio::test]
315    async fn exact_identifiers_and_indexed_ranges_are_supported() {
316        let session = SqlSession::new(&fixture()).expect("session");
317        let rows = session
318            .query(
319                "SELECT name FROM corium.artist \
320                 WHERE \"release-year\" >= 2000 ORDER BY name",
321            )
322            .await
323            .expect("query")
324            .collect()
325            .await
326            .expect("rows");
327        assert_eq!(rows, vec![vec![SqlValue::Text("Tycho".into())]]);
328    }
329
330    #[tokio::test]
331    async fn wide_table_exposes_lists_with_set_semantics() {
332        let session = SqlSession::new(&fixture()).expect("session");
333        let query = session
334            .query(
335                "SELECT e, name, tags FROM corium.artist \
336                 WHERE array_has(tags, 'ambient')",
337            )
338            .await
339            .expect("query");
340        let rows = query.collect().await.expect("rows");
341        assert_eq!(rows.len(), 1);
342        assert_eq!(rows[0][0], SqlValue::Unsigned(1_000));
343        assert_eq!(rows[0][1], SqlValue::Text("Boards of Canada".into()));
344        assert_eq!(
345            rows[0][2],
346            SqlValue::List(vec![
347                SqlValue::Text("ambient".into()),
348                SqlValue::Text("electronic".into()),
349            ])
350        );
351    }
352
353    #[tokio::test]
354    async fn explain_reports_attribute_filter_pushdown() {
355        let session = SqlSession::new(&fixture()).expect("session");
356        let rows = session
357            .query(
358                "EXPLAIN SELECT e FROM corium.artist \
359                 WHERE name >= 'Boards of Canada' AND array_has(tags, 'ambient')",
360            )
361            .await
362            .expect("explain")
363            .collect()
364            .await
365            .expect("rows");
366        let explanation = rows
367            .iter()
368            .flatten()
369            .map(ToString::to_string)
370            .collect::<Vec<_>>()
371            .join("\n");
372        assert!(explanation.contains("partial_filters="));
373        assert!(explanation.contains("name >="));
374        assert!(explanation.contains("array_has"));
375    }
376
377    #[tokio::test]
378    async fn history_session_exposes_events_but_not_wide_tables() {
379        let session = SqlSession::new(&fixture().history()).expect("session");
380        let rows = session
381            .query("SELECT count(*) FROM corium_sys.datoms")
382            .await
383            .expect("event query")
384            .collect()
385            .await
386            .expect("rows");
387        assert_eq!(rows, vec![vec![SqlValue::Integer(6)]]);
388        assert!(session.query("SELECT * FROM corium.artist").await.is_err());
389    }
390
391    #[tokio::test]
392    async fn data_definition_and_modification_are_rejected() {
393        let session = SqlSession::new(&fixture()).expect("session");
394        assert!(
395            session
396                .query("CREATE TABLE nope AS SELECT 1")
397                .await
398                .is_err()
399        );
400        assert!(
401            session
402                .query("INSERT INTO corium.artist (e) VALUES (42)")
403                .await
404                .is_err()
405        );
406    }
407}