1mod 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#[derive(Debug, Error)]
22pub enum SqlError {
23 #[error(transparent)]
25 DataFusion(#[from] datafusion::error::DataFusionError),
26 #[error(transparent)]
28 Arrow(#[from] arrow::error::ArrowError),
29 #[error("SQL schema error: {0}")]
31 Schema(String),
32}
33
34pub struct SqlSession {
36 context: SessionContext,
37 basis_t: u64,
38 view: DbView,
39}
40
41impl SqlSession {
42 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 #[must_use]
61 pub const fn basis_t(&self) -> u64 {
62 self.basis_t
63 }
64
65 #[must_use]
67 pub const fn view(&self) -> DbView {
68 self.view
69 }
70
71 #[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 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
121pub struct SqlQuery {
123 columns: Vec<SqlColumn>,
124 stream: SendableRecordBatchStream,
125 batch: Option<RecordBatch>,
126 row: usize,
127}
128
129impl SqlQuery {
130 #[must_use]
132 pub fn columns(&self) -> &[SqlColumn] {
133 &self.columns
134 }
135
136 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 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}