1use quarb::{AstAdapter, NodeId, Value};
16use quarb_relational::{RelationalModel, RowSpec, TableSpec};
17use rusqlite::Connection;
18use rusqlite::types::ValueRef;
19
20#[derive(Debug, thiserror::Error)]
22pub enum SqliteError {
23 #[error("sqlite: {0}")]
24 Sqlite(#[from] rusqlite::Error),
25 #[error("pushdown plan: {0}")]
26 Plan(String),
27}
28
29pub struct SqliteAdapter {
31 model: RelationalModel,
32}
33
34fn to_value(v: ValueRef<'_>) -> Value {
35 match v {
36 ValueRef::Null => Value::Null,
37 ValueRef::Integer(n) => Value::Int(n),
38 ValueRef::Real(f) => Value::Float(f),
39 ValueRef::Text(t) => Value::Str(String::from_utf8_lossy(t).into_owned()),
40 ValueRef::Blob(b) => Value::Str(format!("<blob {} bytes>", b.len())),
41 }
42}
43
44fn quote_ident(name: &str) -> String {
48 quarb_relational::quote_double(name)
49}
50
51fn specs(conn: &Connection, refs: &[(String, String)]) -> Result<Vec<TableSpec>, SqliteError> {
61 let named: Vec<(String, bool)> = conn
62 .prepare(
63 "SELECT name, type = 'view' FROM sqlite_master \
64 WHERE type IN ('table', 'view') \
65 AND name NOT LIKE 'sqlite_%' ORDER BY name",
66 )?
67 .query_map([], |r| Ok((r.get(0)?, r.get(1)?)))?
68 .collect::<Result<_, _>>()?;
69
70 let mut out = Vec::new();
71 for (name, is_view) in named {
72 let mut columns = Vec::new();
75 let mut json_columns = Vec::new();
76 let mut pk_cols: Vec<(i64, usize)> = Vec::new();
77 let mut column_types: Vec<String> = Vec::new();
78 {
79 let mut stmt = conn.prepare(&format!("PRAGMA table_info({})", quote_ident(&name)))?;
80 let mut rows = stmt.query([])?;
81 while let Some(r) = rows.next()? {
82 let col: String = r.get(1)?;
83 let declared: String = r.get(2)?;
84 let pk: i64 = r.get(5)?;
85 if pk > 0 {
86 pk_cols.push((pk, columns.len()));
87 }
88 if declared.trim().eq_ignore_ascii_case("json")
91 || declared.trim().eq_ignore_ascii_case("jsonb")
92 {
93 json_columns.push(columns.len());
94 }
95 columns.push(col);
96 column_types.push(declared);
97 }
98 }
99 pk_cols.sort();
100 let pk = (pk_cols.len() == 1)
106 .then(|| pk_cols[0].1)
107 .or_else(|| (is_view && columns.len() == 1).then_some(0));
108
109 let mut fks = Vec::new();
112 {
113 let mut stmt =
114 conn.prepare(&format!("PRAGMA foreign_key_list({})", quote_ident(&name)))?;
115 let mut rows = stmt.query([])?;
116 while let Some(r) = rows.next()? {
117 let target: String = r.get(2)?;
118 let from: String = r.get(3)?;
119 let to: Option<String> = r.get(4)?;
120 if let Some(idx) = columns.iter().position(|c| *c == from) {
121 fks.push((idx, target, to.unwrap_or_default()));
122 }
123 }
124 }
125 for (field, container) in refs {
129 let col = match field.split_once('.') {
130 Some((table, col)) if table == name => col,
131 Some(_) => continue,
132 None => field.as_str(),
133 };
134 if let Some(idx) = columns.iter().position(|c| c == col)
135 && !fks.iter().any(|(i, _, _)| *i == idx)
136 && !(*container == name && pk == Some(idx))
139 {
140 fks.push((idx, container.clone(), String::new()));
141 }
142 }
143 out.push(TableSpec {
144 name,
145 columns,
146 pk,
147 fks,
148 json_columns,
149 column_types,
150 });
151 }
152 Ok(out)
153}
154
155fn fetch_rows_where(
158 conn: &Connection,
159 spec: &TableSpec,
160 where_sql: Option<&str>,
161 limit: Option<usize>,
162) -> Result<Vec<RowSpec>, SqliteError> {
163 let cut = limit.map(|k| format!(" LIMIT {k}")).unwrap_or_default();
165 let cols = spec
166 .columns
167 .iter()
168 .map(|c| quote_ident(c))
169 .collect::<Vec<_>>()
170 .join(", ");
171 let table = quote_ident(&spec.name);
172 let filter = match where_sql {
173 Some(w) => format!(" WHERE {w}"),
174 None => String::new(),
175 };
176 let with_rowid = format!("SELECT rowid, {cols} FROM {table}{filter} ORDER BY rowid{cut}");
182 if let Ok(mut stmt) = conn.prepare(&with_rowid) {
183 let mut rows = stmt.query([])?;
184 let mut out = Vec::new();
185 while let Some(r) = rows.next()? {
186 let rowid: i64 = r.get(0)?;
187 let values: Vec<Value> = (0..spec.columns.len())
188 .map(|i| to_value(r.get_ref(i + 1).expect("column in range")))
189 .collect();
190 out.push(RowSpec {
191 rowid: Some(rowid),
192 values,
193 });
194 }
195 return Ok(out);
196 }
197 let order = match spec.pk {
201 Some(i) => quote_ident(&spec.columns[i]),
202 None => cols.clone(),
203 };
204 let mut stmt = conn.prepare(&format!(
205 "SELECT {cols} FROM {table}{filter} ORDER BY {order}{cut}"
206 ))?;
207 let mut rows = stmt.query([])?;
208 let mut out = Vec::new();
209 let mut rowid = 0i64;
210 while let Some(r) = rows.next()? {
211 rowid += 1;
212 let values: Vec<Value> = (0..spec.columns.len())
213 .map(|i| to_value(r.get_ref(i).expect("column in range")))
214 .collect();
215 out.push(RowSpec {
216 rowid: Some(rowid),
217 values,
218 });
219 }
220 Ok(out)
221}
222
223fn fetch_rows_keyed(
234 conn: &Connection,
235 spec: &TableSpec,
236 filter: Option<&str>,
237 col: Option<usize>,
238 values: &[Value],
239) -> Result<Vec<RowSpec>, String> {
240 use rusqlite::types::Value as Sql;
241 let column = match col.or(spec.pk) {
242 Some(i) => quote_ident(&spec.columns[i]),
243 None => "rowid".to_string(),
244 };
245 let mut binds: Vec<Sql> = Vec::new();
246 for v in values {
247 match v {
248 Value::Int(i) => {
249 binds.push(Sql::Integer(*i));
250 binds.push(Sql::Text(i.to_string()));
251 }
252 Value::Float(f) => {
253 binds.push(Sql::Real(*f));
254 binds.push(Sql::Text(v.to_string()));
255 }
256 Value::Bool(b) => binds.push(Sql::Integer(*b as i64)),
257 Value::Null => {}
258 other => {
259 let text = other.to_string();
260 if let Ok(i) = text.parse::<i64>() {
261 binds.push(Sql::Integer(i));
262 } else if let Ok(f) = text.parse::<f64>() {
263 binds.push(Sql::Real(f));
264 }
265 binds.push(Sql::Text(text));
266 }
267 }
268 }
269 if binds.is_empty() {
270 return Ok(Vec::new());
271 }
272 let cols = spec
273 .columns
274 .iter()
275 .map(|c| quote_ident(c))
276 .collect::<Vec<_>>()
277 .join(", ");
278 let marks = vec!["?"; binds.len()].join(", ");
279 let filter = match filter {
280 Some(w) => format!("({w}) AND "),
281 None => String::new(),
282 };
283 let sql = format!(
284 "SELECT rowid, {cols} FROM {} WHERE {filter}{column} IN ({marks}) ORDER BY rowid",
285 quote_ident(&spec.name)
286 );
287 let mut stmt = conn.prepare(&sql).map_err(|e| e.to_string())?;
288 let mut rows = stmt
289 .query(rusqlite::params_from_iter(binds.iter()))
290 .map_err(|e| e.to_string())?;
291 let mut out = Vec::new();
292 while let Some(r) = rows.next().map_err(|e| e.to_string())? {
293 let rowid: i64 = r.get(0).map_err(|e| e.to_string())?;
294 let values: Vec<Value> = (0..spec.columns.len())
295 .map(|i| to_value(r.get_ref(i + 1).expect("column in range")))
296 .collect();
297 out.push(RowSpec {
298 rowid: Some(rowid),
299 values,
300 });
301 }
302 Ok(out)
303}
304
305fn index_rows_per_key(conn: &Connection, table: &str, column: &str) -> Result<u64, String> {
310 let mut list = conn
311 .prepare(&format!("PRAGMA index_list({})", quote_ident(table)))
312 .map_err(|e| e.to_string())?;
313 let names: Vec<String> = list
314 .query_map([], |r| r.get::<_, String>(1))
315 .map_err(|e| e.to_string())?
316 .filter_map(Result::ok)
317 .collect();
318 for idx in names {
319 let first: Option<String> = conn
320 .query_row(
321 &format!("PRAGMA index_info({})", quote_ident(&idx)),
322 [],
323 |r| r.get(2),
324 )
325 .ok();
326 if first.as_deref() != Some(column) {
327 continue;
328 }
329 let stat: String = conn
330 .query_row(
331 "SELECT stat FROM sqlite_stat1 WHERE tbl = ?1 AND idx = ?2",
332 [table, idx.as_str()],
333 |r| r.get(0),
334 )
335 .map_err(|e| e.to_string())?;
336 return stat
337 .split_whitespace()
338 .nth(1)
339 .and_then(|t| t.parse().ok())
340 .ok_or_else(|| "no per-key figure in sqlite_stat1".to_string());
341 }
342 Err(format!("no index opens with {column}"))
343}
344
345impl SqliteAdapter {
346 pub fn open(path: &std::path::Path) -> Result<Self, SqliteError> {
350 Self::open_impl(path, None, &[])
351 }
352
353 pub fn open_with_refs(
357 path: &std::path::Path,
358 refs: &[(String, String)],
359 ) -> Result<Self, SqliteError> {
360 Self::open_impl(path, None, refs)
361 }
362
363 pub fn open_filtered(
366 path: &std::path::Path,
367 table: &str,
368 where_sql: &str,
369 ) -> Result<Self, SqliteError> {
370 Self::open_impl(path, Some((table.to_string(), where_sql.to_string())), &[])
371 }
372
373 pub fn open_filtered_with_refs(
375 path: &std::path::Path,
376 table: &str,
377 where_sql: &str,
378 refs: &[(String, String)],
379 ) -> Result<Self, SqliteError> {
380 Self::open_impl(path, Some((table.to_string(), where_sql.to_string())), refs)
381 }
382
383 fn open_impl(
384 path: &std::path::Path,
385 filter: Option<(String, String)>,
386 refs: &[(String, String)],
387 ) -> Result<Self, SqliteError> {
388 let conn = std::rc::Rc::new(Connection::open_with_flags(
389 path,
390 rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY,
391 )?);
392 let specs = specs(&conn, refs)?;
393 let (whole, keyed, stats, prefix, counting, listed) = (
394 conn.clone(),
395 conn.clone(),
396 conn.clone(),
397 conn.clone(),
398 conn.clone(),
399 conn,
400 );
401 let mut model = RelationalModel::lazy(
402 specs,
403 Box::new(move |_, spec| {
404 fetch_rows_where(&whole, spec, None, None).map_err(|e| e.to_string())
405 }),
406 )
407 .with_key_fetcher(Box::new(move |_, spec, col, values| {
409 fetch_rows_keyed(&keyed, spec, None, col, values)
410 }))
411 .with_prefix_fetcher(Box::new(move |_, spec, k| {
415 fetch_rows_where(&prefix, spec, None, Some(k)).map_err(|e| e.to_string())
416 }))
417 .with_counter(Box::new(move |_, spec| {
418 counting
419 .query_row(
420 &format!("SELECT count(*) FROM {}", quote_ident(&spec.name)),
421 [],
422 |r| r.get::<_, i64>(0),
423 )
424 .map(|n| n.max(0) as usize)
425 .map_err(|e| e.to_string())
426 }))
427 .with_estimator(std::rc::Rc::new(move |_, spec, equals| {
434 let stat: String = stats
435 .query_row(
436 "SELECT stat FROM sqlite_stat1 WHERE tbl = ?1 \
437 ORDER BY idx IS NULL DESC LIMIT 1",
438 [spec.name.as_str()],
439 |r| r.get(0),
440 )
441 .map_err(|e| e.to_string())?;
442 let n: u64 = stat
443 .split_whitespace()
444 .next()
445 .and_then(|t| t.parse().ok())
446 .ok_or("no row count in sqlite_stat1")?;
447 let mut per_key: Option<u64> = None;
452 for (col, _) in equals {
453 let column = &spec.columns[*col];
454 if let Ok(k) = index_rows_per_key(&stats, &spec.name, column) {
455 per_key = Some(per_key.map_or(k, |p| p.min(k)));
456 }
457 }
458 Ok(quarb::Estimate {
459 nodes: match (equals.is_empty(), per_key) {
460 (true, _) => (n, n),
461 (false, Some(k)) => (0, k.min(n)),
462 (false, None) => (0, n),
463 },
464 bytes: None,
465 requests: Some(0),
466 latency: quarb::Latency::Local,
467 confidence: quarb::Confidence::Guess,
468 })
469 }));
470 if let Some((table, where_sql)) = filter {
473 model = model.with_filtered_listing(
474 &table,
475 Box::new(move |_, spec| {
476 fetch_rows_where(&listed, spec, Some(&where_sql), None)
477 .map_err(|e| e.to_string())
478 }),
479 );
480 }
481 Ok(SqliteAdapter { model })
482 }
483
484 pub fn from_bytes(bytes: &[u8]) -> Result<Self, SqliteError> {
489 let mut conn = Connection::open_in_memory()?;
490 conn.deserialize_read_exact(rusqlite::MAIN_DB, bytes, bytes.len(), true)?;
491 Self::load(&conn)
492 }
493
494 pub fn load(conn: &Connection) -> Result<Self, SqliteError> {
497 Self::load_with_refs(conn, &[])
498 }
499
500 pub fn load_with_refs(
502 conn: &Connection,
503 refs: &[(String, String)],
504 ) -> Result<Self, SqliteError> {
505 let mut input = Vec::new();
506 for spec in specs(conn, refs)? {
507 let rows = fetch_rows_where(conn, &spec, None, None)?;
508 input.push((spec, rows));
509 }
510 Ok(SqliteAdapter {
511 model: RelationalModel::build(input),
512 })
513 }
514
515 pub fn prefetch(&self, table: &str) -> Result<(), String> {
520 self.model.prefetch(table)
521 }
522
523 pub fn locator(&self, node: NodeId) -> String {
524 self.model.locator(node)
525 }
526}
527
528fn unique_key(conn: &Connection, table: &str, cols: &[String]) -> Result<bool, SqliteError> {
536 use std::collections::HashSet;
537 let want: HashSet<&str> = cols.iter().map(|s| s.as_str()).collect();
538 if want.is_empty() {
539 return Ok(false);
540 }
541 let mut stmt = conn.prepare(&format!("PRAGMA table_info({})", quote_ident(table)))?;
543 let mut pk: HashSet<String> = HashSet::new();
544 let mut rows = stmt.query([])?;
545 while let Some(r) = rows.next()? {
546 let name: String = r.get(1)?;
547 let ord: i64 = r.get(5)?;
548 if ord > 0 {
549 pk.insert(name);
550 }
551 }
552 if !pk.is_empty() && pk.iter().all(|c| want.contains(c.as_str())) {
553 return Ok(true);
554 }
555 let mut stmt = conn.prepare(&format!("PRAGMA index_list({})", quote_ident(table)))?;
557 let mut uniques: Vec<String> = Vec::new();
558 let mut rows = stmt.query([])?;
559 while let Some(r) = rows.next()? {
560 let name: String = r.get(1)?;
561 let is_unique: i64 = r.get(2)?;
562 let partial: i64 = r.get(4).unwrap_or(0);
563 if is_unique == 1 && partial == 0 {
564 uniques.push(name);
565 }
566 }
567 for idx in uniques {
568 let mut stmt = conn.prepare(&format!("PRAGMA index_info({})", quote_ident(&idx)))?;
569 let mut cols_of: Vec<String> = Vec::new();
570 let mut rows = stmt.query([])?;
571 while let Some(r) = rows.next()? {
572 cols_of.push(r.get(2)?);
573 }
574 if !cols_of.is_empty() && cols_of.iter().all(|c| want.contains(c.as_str())) {
575 return Ok(true);
576 }
577 }
578 Ok(false)
579}
580
581pub fn column_catalog(
590 path: &std::path::Path,
591 tables: &[&str],
592) -> Result<Vec<(String, bool, bool)>, SqliteError> {
593 let conn = Connection::open_with_flags(path, rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY)?;
594 let mut catalog = Vec::new();
595 for table in tables {
596 let mut stmt = conn.prepare(&format!("PRAGMA table_info({})", quote_ident(table)))?;
597 let rows = stmt.query_map([], |r| Ok((r.get::<_, String>(1)?, r.get::<_, String>(2)?)))?;
598 for row in rows {
599 let (name, declared) = row?;
600 let json = matches!(
604 declared.trim().to_ascii_lowercase().as_str(),
605 "json" | "jsonb"
606 );
607 catalog.push((name, numeric_affinity(&declared), json));
608 }
609 }
610 Ok(catalog)
611}
612
613fn numeric_affinity(declared: &str) -> bool {
616 let d = declared.to_ascii_uppercase();
617 if d.contains("INT") {
618 true
619 } else if d.contains("CHAR") || d.contains("CLOB") || d.contains("TEXT") {
620 false
621 } else {
622 !(d.contains("BLOB") || d.trim().is_empty())
623 }
624}
625
626pub fn raw_query(
627 path: &std::path::Path,
628 sql: &str,
629 order_table: Option<&str>,
630 join_left: Option<(&str, &[String])>,
631 order_then: Option<&str>,
632) -> Result<(Vec<String>, Vec<Vec<Value>>), SqliteError> {
633 let conn = Connection::open_with_flags(path, rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY)?;
634 if let Some((table, cols)) = join_left
639 && !unique_key(&conn, table, cols)?
640 {
641 return Err(SqliteError::Plan(format!(
642 "join ON does not bind {table} by a unique key; \
643 the SQL JOIN would multiply rows"
644 )));
645 }
646 let sql = match order_table {
647 Some(t) => {
648 let catalog = specs(&conn, &[])?;
649 let key_of = |t: &str| {
650 catalog
651 .iter()
652 .find(|s| s.name == t)
653 .and_then(|s| s.pk.map(|i| s.columns[i].clone()))
654 .unwrap_or_else(|| "rowid".to_string())
655 };
656 match order_then {
659 Some(j) => format!("{sql} ORDER BY {t}.{}, {j}.{}", key_of(t), key_of(j)),
660 None => format!("{sql} ORDER BY {t}.{}", key_of(t)),
661 }
662 }
663 None => sql.to_string(),
664 };
665 quarb_relational::record_executed(&sql);
669 let mut stmt = conn.prepare(&sql)?;
670 let cols: Vec<String> = stmt.column_names().iter().map(|c| c.to_string()).collect();
671 let n = cols.len();
672 let mut rows = stmt.query([])?;
673 let mut out = Vec::new();
674 while let Some(r) = rows.next()? {
675 out.push(
676 (0..n)
677 .map(|i| to_value(r.get_ref(i).expect("column in range")))
678 .collect(),
679 );
680 }
681 Ok((cols, out))
682}
683
684impl AstAdapter for SqliteAdapter {
685 quarb::forward_ast_adapter!(model);
688}