1use crate::error::{KitError, Result};
4use crate::internal::{ensure_internal_tables, internal_tables_core};
5use crate::schema::to_core_schema;
6use mongreldb_core::epoch::Snapshot;
7use mongreldb_core::memtable::Row as CoreRow;
8use mongreldb_core::memtable::Value as CoreValue;
9use mongreldb_core::schema::Schema as CoreSchema;
10use mongreldb_core::Database as CoreDatabase;
11use mongreldb_core::{AggState, ApproxAgg, NativeAgg, NativeAggResult, RowId};
12use mongreldb_kit_core::schema::IndexKind as KitIndexKind;
13use mongreldb_kit_core::schema::Schema as KitSchema;
14use mongreldb_kit_core::schema::Table as KitTable;
15use mongreldb_kit_core::{ProcedureSpec, TriggerSpec, ViewSpec};
16use serde_json::Value;
17
18use std::collections::HashMap;
19use std::path::{Path, PathBuf};
20use std::sync::Arc;
21
22const SCHEMA_FILE: &str = "kit_schema.json";
23
24#[derive(Clone, Copy, Debug, Default)]
29pub struct OpenOptions {
30 pub lock_timeout_ms: u32,
36}
37
38impl OpenOptions {
39 pub fn new() -> Self {
41 Self::default()
42 }
43
44 pub fn with_lock_timeout_ms(mut self, ms: u32) -> Self {
47 self.lock_timeout_ms = ms;
48 self
49 }
50}
51
52#[derive(Debug, Clone, Default)]
53pub struct SqlOptions {
54 pub query_id: Option<mongreldb_query::QueryId>,
55 pub timeout: Option<std::time::Duration>,
56}
57
58pub struct SqlQueryHandle {
59 query_id: mongreldb_query::QueryId,
60 session: Arc<mongreldb_query::MongrelSession>,
61 worker: Option<std::thread::JoinHandle<Result<Vec<arrow::record_batch::RecordBatch>>>>,
62}
63
64impl SqlQueryHandle {
65 pub fn id(&self) -> mongreldb_query::QueryId {
66 self.query_id
67 }
68
69 pub fn cancel(&self) -> mongreldb_query::CancelOutcome {
70 self.session.cancel_query(self.query_id)
71 }
72
73 pub fn wait(mut self) -> Result<Vec<arrow::record_batch::RecordBatch>> {
74 self.worker
75 .take()
76 .expect("SQL worker is present")
77 .join()
78 .map_err(|_| KitError::Storage("SQL worker panicked".into()))?
79 }
80}
81
82impl Drop for SqlQueryHandle {
83 fn drop(&mut self) {
84 if self.worker.is_some() {
85 let _ = self.session.cancel_query(self.query_id);
86 }
87 }
88}
89
90pub type DefaultProvider = Box<dyn Fn() -> Value + Send + Sync>;
92
93#[derive(Debug, Clone)]
96pub struct ExplainPlan {
97 pub index_accelerated: bool,
99 pub exact: bool,
102 pub pushed_conditions: Vec<String>,
104}
105
106#[derive(Debug, Clone)]
108pub struct SimilarRow {
109 pub row: crate::schema::Row,
110 pub similarity: f64,
111}
112
113fn parse_string_set(value: Option<&Value>) -> std::collections::HashSet<String> {
117 let arr = match value {
118 Some(Value::Array(a)) => Some(a.clone()),
119 Some(Value::String(s)) => serde_json::from_str::<Value>(s)
120 .ok()
121 .and_then(|v| v.as_array().cloned()),
122 _ => None,
123 };
124 arr.into_iter()
125 .flatten()
126 .filter_map(|v| match v {
127 Value::String(s) => Some(s),
128 Value::Number(n) => Some(n.to_string()),
129 Value::Bool(b) => Some(b.to_string()),
130 _ => None,
131 })
132 .collect()
133}
134
135#[derive(Debug, Clone, Copy, PartialEq, Eq)]
137pub enum IncrementalAggKind {
138 Count,
139 Sum,
140 Min,
141 Max,
142 Avg,
143}
144
145#[derive(Debug, Clone)]
147pub struct IncrementalAggregate {
148 pub value: Value,
151 pub incremental: bool,
155 pub delta_rows: u64,
157}
158
159fn incremental_cache_key(
163 table_id: u32,
164 column: Option<u16>,
165 agg: IncrementalAggKind,
166 conditions: &[mongreldb_core::query::Condition],
167) -> u64 {
168 use std::hash::{Hash, Hasher};
169 let mut h = std::collections::hash_map::DefaultHasher::new();
170 table_id.hash(&mut h);
171 column.hash(&mut h);
172 (agg as u8).hash(&mut h);
173 format!("{conditions:?}").hash(&mut h);
175 h.finish()
176}
177
178fn agg_state_value(s: &AggState) -> Value {
182 let num_f64 = |x: f64| {
183 serde_json::Number::from_f64(x)
184 .map(Value::Number)
185 .unwrap_or(Value::Null)
186 };
187 match s {
188 AggState::Count(n) => Value::from(*n),
189 AggState::SumI { sum, .. } => i64::try_from(*sum)
190 .map(Value::from)
191 .unwrap_or_else(|_| num_f64(*sum as f64)),
192 AggState::SumF { sum, .. } => num_f64(*sum),
193 AggState::AvgI { sum, count } if *count > 0 => num_f64(*sum as f64 / *count as f64),
194 AggState::AvgF { sum, count } if *count > 0 => num_f64(*sum / *count as f64),
195 AggState::AvgI { .. } | AggState::AvgF { .. } => Value::Null,
196 AggState::MinI(n) | AggState::MaxI(n) => Value::from(*n),
197 AggState::MinF(f) | AggState::MaxF(f) => num_f64(*f),
198 AggState::Empty => Value::Null,
199 }
200}
201
202#[derive(Debug, Clone, Copy, PartialEq, Eq)]
204pub enum ApproxAggKind {
205 Count,
206 Sum,
207 Avg,
208}
209
210#[derive(Debug, Clone)]
214pub struct ApproxAggregate {
215 pub point: f64,
216 pub ci_low: f64,
217 pub ci_high: f64,
218 pub n_population: u64,
219 pub n_sample_live: usize,
220 pub n_passing: usize,
221}
222
223fn condition_label(c: &mongreldb_core::query::Condition) -> String {
226 let dbg = format!("{c:?}");
227 dbg.split(['(', '{', ' ']).next().unwrap_or("").to_string()
228}
229
230fn open_core_with_retry<T>(
231 timeout_ms: u32,
232 mut open: impl FnMut() -> mongreldb_core::Result<T>,
233) -> mongreldb_core::Result<T> {
234 if timeout_ms == 0 {
235 return open();
236 }
237 let deadline = std::time::Instant::now() + std::time::Duration::from_millis(timeout_ms as u64);
238 let mut next_sleep = std::time::Duration::from_millis(1);
239 loop {
240 match open() {
241 Ok(db) => return Ok(db),
242 Err(err) if is_lock_contention(&err) => {
243 let now = std::time::Instant::now();
244 if now >= deadline {
245 return Err(mongreldb_core::MongrelError::Io(std::io::Error::other(
246 format!("database lock timeout after {timeout_ms}ms: {err}"),
247 )));
248 }
249 let sleep = next_sleep.min(deadline - now);
250 std::thread::sleep(sleep);
251 next_sleep = next_sleep
252 .saturating_mul(10)
253 .min(std::time::Duration::from_millis(50));
254 }
255 Err(err) => return Err(err),
256 }
257 }
258}
259
260fn is_lock_contention(err: &mongreldb_core::MongrelError) -> bool {
261 matches!(
262 err,
263 mongreldb_core::MongrelError::Io(io)
264 if io.to_string().contains("locked by another process")
265 )
266}
267
268pub struct Database {
273 pub(crate) inner: Arc<CoreDatabase>,
274 pub(crate) schema: KitSchema,
275 pub(crate) root: PathBuf,
276 pub(crate) default_providers: HashMap<String, DefaultProvider>,
278 pub(crate) session: parking_lot::RwLock<Option<Arc<mongreldb_query::MongrelSession>>>,
285}
286
287impl Database {
288 pub fn open(path: &Path) -> Result<Self> {
290 let inner = Arc::new(CoreDatabase::open(path)?);
291 let schema = load_schema(path)?;
292 ensure_internal_tables(&inner)?;
294 reap_rotated_wal_segments(&inner);
295 Ok(Self {
296 inner,
297 schema,
298 root: path.to_path_buf(),
299 default_providers: HashMap::new(),
300 session: parking_lot::RwLock::new(None),
301 })
302 }
303
304 pub fn open_with_options(path: &Path, opts: OpenOptions) -> Result<Self> {
312 let inner = Arc::new(open_core_with_retry(opts.lock_timeout_ms, || {
313 CoreDatabase::open(path)
314 })?);
315 let schema = load_schema(path)?;
316 ensure_internal_tables(&inner)?;
317 reap_rotated_wal_segments(&inner);
318 Ok(Self {
319 inner,
320 schema,
321 root: path.to_path_buf(),
322 default_providers: HashMap::new(),
323 session: parking_lot::RwLock::new(None),
324 })
325 }
326
327 pub fn open_encrypted(path: &Path, passphrase: &str) -> Result<Self> {
329 let inner = Arc::new(CoreDatabase::open_encrypted(path, passphrase)?);
330 let schema = load_schema(path)?;
331 ensure_internal_tables(&inner)?;
332 reap_rotated_wal_segments(&inner);
333 Ok(Self {
334 inner,
335 schema,
336 root: path.to_path_buf(),
337 default_providers: HashMap::new(),
338 session: parking_lot::RwLock::new(None),
339 })
340 }
341
342 pub fn open_encrypted_with_options(
346 path: &Path,
347 passphrase: &str,
348 opts: OpenOptions,
349 ) -> Result<Self> {
350 let inner = Arc::new(open_core_with_retry(opts.lock_timeout_ms, || {
351 CoreDatabase::open_encrypted(path, passphrase)
352 })?);
353 let schema = load_schema(path)?;
354 ensure_internal_tables(&inner)?;
355 reap_rotated_wal_segments(&inner);
356 Ok(Self {
357 inner,
358 schema,
359 root: path.to_path_buf(),
360 default_providers: HashMap::new(),
361 session: parking_lot::RwLock::new(None),
362 })
363 }
364
365 pub fn create_encrypted(path: &Path, schema: KitSchema, passphrase: &str) -> Result<Self> {
369 std::fs::create_dir_all(path)?;
370 let inner = Arc::new(CoreDatabase::create_encrypted(path, passphrase)?);
371 ensure_internal_tables(&inner)?;
372 store_schema(path, &schema)?;
373 for table in &schema.tables {
374 create_core_table(&inner, &table.name, to_core_schema(table)?)?;
375 }
376 Ok(Self {
377 inner,
378 schema,
379 root: path.to_path_buf(),
380 default_providers: HashMap::new(),
381 session: parking_lot::RwLock::new(None),
382 })
383 }
384
385 pub fn create(path: &Path, schema: KitSchema) -> Result<Self> {
387 std::fs::create_dir_all(path)?;
388 let inner = Arc::new(CoreDatabase::create(path)?);
389
390 ensure_internal_tables(&inner)?;
393
394 store_schema(path, &schema)?;
397
398 for table in &schema.tables {
400 create_core_table(&inner, &table.name, to_core_schema(table)?)?;
401 }
402
403 Ok(Self {
404 inner,
405 schema,
406 root: path.to_path_buf(),
407 default_providers: HashMap::new(),
408 session: parking_lot::RwLock::new(None),
409 })
410 }
411
412 pub fn open_with_credentials(path: &Path, username: &str, password: &str) -> Result<Self> {
422 let inner = Arc::new(CoreDatabase::open_with_credentials(
423 path, username, password,
424 )?);
425 let schema = load_schema(path)?;
426 ensure_internal_tables(&inner)?;
427 reap_rotated_wal_segments(&inner);
428 Ok(Self {
429 inner,
430 schema,
431 root: path.to_path_buf(),
432 default_providers: HashMap::new(),
433 session: parking_lot::RwLock::new(None),
434 })
435 }
436
437 pub fn open_with_credentials_and_options(
441 path: &Path,
442 username: &str,
443 password: &str,
444 opts: OpenOptions,
445 ) -> Result<Self> {
446 let inner = Arc::new(open_core_with_retry(opts.lock_timeout_ms, || {
447 CoreDatabase::open_with_credentials(path, username, password)
448 })?);
449 let schema = load_schema(path)?;
450 ensure_internal_tables(&inner)?;
451 reap_rotated_wal_segments(&inner);
452 Ok(Self {
453 inner,
454 schema,
455 root: path.to_path_buf(),
456 default_providers: HashMap::new(),
457 session: parking_lot::RwLock::new(None),
458 })
459 }
460
461 pub fn create_with_credentials(
467 path: &Path,
468 schema: KitSchema,
469 admin_username: &str,
470 admin_password: &str,
471 ) -> Result<Self> {
472 std::fs::create_dir_all(path)?;
473 let inner = Arc::new(CoreDatabase::create_with_credentials(
474 path,
475 admin_username,
476 admin_password,
477 )?);
478 ensure_internal_tables(&inner)?;
479 store_schema(path, &schema)?;
480 for table in &schema.tables {
481 create_core_table(&inner, &table.name, to_core_schema(table)?)?;
482 }
483 Ok(Self {
484 inner,
485 schema,
486 root: path.to_path_buf(),
487 default_providers: HashMap::new(),
488 session: parking_lot::RwLock::new(None),
489 })
490 }
491
492 pub fn open_encrypted_with_credentials(
495 path: &Path,
496 passphrase: &str,
497 username: &str,
498 password: &str,
499 ) -> Result<Self> {
500 let inner = Arc::new(CoreDatabase::open_encrypted_with_credentials(
501 path, passphrase, username, password,
502 )?);
503 let schema = load_schema(path)?;
504 ensure_internal_tables(&inner)?;
505 reap_rotated_wal_segments(&inner);
506 Ok(Self {
507 inner,
508 schema,
509 root: path.to_path_buf(),
510 default_providers: HashMap::new(),
511 session: parking_lot::RwLock::new(None),
512 })
513 }
514
515 pub fn open_encrypted_with_credentials_and_options(
519 path: &Path,
520 passphrase: &str,
521 username: &str,
522 password: &str,
523 opts: OpenOptions,
524 ) -> Result<Self> {
525 let inner = Arc::new(open_core_with_retry(opts.lock_timeout_ms, || {
526 CoreDatabase::open_encrypted_with_credentials(path, passphrase, username, password)
527 })?);
528 let schema = load_schema(path)?;
529 ensure_internal_tables(&inner)?;
530 reap_rotated_wal_segments(&inner);
531 Ok(Self {
532 inner,
533 schema,
534 root: path.to_path_buf(),
535 default_providers: HashMap::new(),
536 session: parking_lot::RwLock::new(None),
537 })
538 }
539
540 pub fn create_encrypted_with_credentials(
544 path: &Path,
545 schema: KitSchema,
546 passphrase: &str,
547 admin_username: &str,
548 admin_password: &str,
549 ) -> Result<Self> {
550 std::fs::create_dir_all(path)?;
551 let inner = Arc::new(CoreDatabase::create_encrypted_with_credentials(
552 path,
553 passphrase,
554 admin_username,
555 admin_password,
556 )?);
557 ensure_internal_tables(&inner)?;
558 store_schema(path, &schema)?;
559 for table in &schema.tables {
560 create_core_table(&inner, &table.name, to_core_schema(table)?)?;
561 }
562 Ok(Self {
563 inner,
564 schema,
565 root: path.to_path_buf(),
566 default_providers: HashMap::new(),
567 session: parking_lot::RwLock::new(None),
568 })
569 }
570
571 pub fn enable_auth(&self, admin_username: &str, admin_password: &str) -> Result<()> {
575 self.inner
576 .enable_auth(admin_username, admin_password)
577 .map_err(KitError::from)
578 }
579
580 pub fn disable_auth(&self) -> Result<()> {
584 self.inner.disable_auth().map_err(KitError::from)
585 }
586
587 pub fn require_auth_enabled(&self) -> bool {
589 self.inner.require_auth_enabled()
590 }
591
592 pub fn refresh_principal(&self) -> Result<()> {
596 self.inner.refresh_principal().map_err(KitError::from)?;
597 *self.session.write() = None;
600 Ok(())
601 }
602
603 pub fn register_default(
606 &mut self,
607 name: impl Into<String>,
608 provider: impl Fn() -> Value + Send + Sync + 'static,
609 ) {
610 self.default_providers
611 .insert(name.into(), Box::new(provider));
612 }
613
614 pub fn raw(&self) -> &CoreDatabase {
618 &self.inner
619 }
620
621 pub fn table_names(&self) -> Vec<String> {
623 self.schema
624 .tables
625 .iter()
626 .map(|t| t.name.clone())
627 .filter(|n| !n.starts_with("__kit_"))
628 .collect()
629 }
630
631 pub fn create_procedure(
632 &self,
633 spec: &ProcedureSpec,
634 ) -> Result<mongreldb_core::StoredProcedure> {
635 let procedure = core_procedure(spec)?;
636 self.inner
637 .create_procedure(procedure)
638 .map_err(KitError::from)
639 }
640
641 pub fn replace_procedure(
642 &self,
643 spec: &ProcedureSpec,
644 ) -> Result<mongreldb_core::StoredProcedure> {
645 let procedure = core_procedure(spec)?;
646 self.inner
647 .create_or_replace_procedure(procedure)
648 .map_err(KitError::from)
649 }
650
651 pub fn drop_procedure(&self, name: &str) -> Result<()> {
652 self.inner.drop_procedure(name).map_err(KitError::from)
653 }
654
655 pub fn call_procedure(
656 &self,
657 name: &str,
658 args: serde_json::Map<String, Value>,
659 ) -> Result<mongreldb_core::ProcedureCallResult> {
660 let args = args
661 .iter()
662 .map(|(key, value)| Ok((key.clone(), json_to_core_value(value)?)))
663 .collect::<Result<HashMap<_, _>>>()?;
664 self.inner
665 .call_procedure(name, args)
666 .map_err(KitError::from)
667 }
668
669 pub fn create_trigger(&self, spec: &TriggerSpec) -> Result<mongreldb_core::StoredTrigger> {
670 let trigger = core_trigger(spec)?;
671 self.inner.create_trigger(trigger).map_err(KitError::from)
672 }
673
674 pub fn replace_trigger(&self, spec: &TriggerSpec) -> Result<mongreldb_core::StoredTrigger> {
675 let trigger = core_trigger(spec)?;
676 self.inner
677 .create_or_replace_trigger(trigger)
678 .map_err(KitError::from)
679 }
680
681 pub fn drop_trigger(&self, name: &str) -> Result<()> {
682 self.inner.drop_trigger(name).map_err(KitError::from)
683 }
684
685 pub fn triggers(&self) -> Vec<mongreldb_core::StoredTrigger> {
686 self.inner.triggers()
687 }
688
689 pub fn trigger(&self, name: &str) -> Option<mongreldb_core::StoredTrigger> {
690 self.inner.trigger(name)
691 }
692
693 pub fn allocate_sequence(&self, name: &str, count: i64) -> Result<i64> {
699 use crate::internal::cols;
700 let mut attempt = 0;
701 loop {
702 let mut txn = self.inner.begin();
703 let snapshot = txn.read_snapshot();
704 let existing = self
705 .visible_core_rows_at(crate::internal::SEQUENCES, snapshot)?
706 .into_iter()
707 .find(|r| internal_bytes(r, cols::SEQ_NAME) == Some(name.to_string()));
708
709 let now = crate::internal::iso_now();
710 let (start, next, old_row_id) = match &existing {
714 Some(row) => {
715 let current = match row.columns.get(&cols::SEQ_NEXT) {
716 Some(CoreValue::Int64(i)) => *i,
717 _ => 1,
718 };
719 (current, current + count, Some(row.row_id))
720 }
721 None => (1, 1 + count, None),
722 };
723
724 if let Some(rid) = old_row_id {
725 txn.delete(crate::internal::SEQUENCES, rid)
726 .map_err(KitError::from)?;
727 }
728 txn.put(
729 crate::internal::SEQUENCES,
730 vec![
731 (cols::SEQ_NAME, CoreValue::Bytes(name.as_bytes().to_vec())),
732 (cols::SEQ_NEXT, CoreValue::Int64(next)),
733 (cols::SEQ_UPDATED, CoreValue::Bytes(now.into_bytes())),
734 ],
735 )
736 .map_err(KitError::from)?;
737 match txn.commit() {
738 Ok(_) => return Ok(start),
739 Err(mongreldb_core::MongrelError::Conflict(_)) if attempt < 10_000 => {
740 attempt += 1;
741 std::thread::yield_now();
742 continue;
743 }
744 Err(e) => return Err(KitError::from(e)),
745 }
746 }
747 }
748
749 pub fn transaction<T, F>(&self, max_retries: usize, mut f: F) -> Result<T>
752 where
753 F: FnMut(&mut crate::txn::Transaction<'_>) -> Result<T>,
754 {
755 let mut attempt = 0;
756 loop {
757 let mut txn = self.begin()?;
758 match f(&mut txn) {
759 Ok(value) => match txn.commit() {
760 Ok(()) => return Ok(value),
761 Err(KitError::Conflict(_)) if attempt < max_retries => {
762 attempt += 1;
763 continue;
764 }
765 Err(e) => return Err(e),
766 },
767 Err(KitError::Conflict(_)) if attempt < max_retries => {
768 txn.rollback();
769 attempt += 1;
770 continue;
771 }
772 Err(e) => {
773 txn.rollback();
774 return Err(e);
775 }
776 }
777 }
778 }
779
780 pub fn table(&self, name: &str) -> Option<&KitTable> {
782 self.schema.table(name)
783 }
784
785 pub fn schema(&self) -> &KitSchema {
787 &self.schema
788 }
789
790 pub fn begin(&self) -> Result<crate::txn::Transaction<'_>> {
792 let core_txn = self.inner.begin();
793 Ok(crate::txn::Transaction::new(self, core_txn))
794 }
795
796 pub fn set_schema(&mut self, schema: KitSchema) {
798 self.schema = schema;
799 }
800
801 pub fn check_internal_tables(&self) -> Result<()> {
804 let schema_file = self.root.join(SCHEMA_FILE);
805 if !schema_file.exists() {
806 return Err(KitError::Integrity(format!(
807 "schema file {} is missing",
808 schema_file.display()
809 )));
810 }
811 for (name, _) in internal_tables_core() {
812 if self.inner.table_id(name).is_err() {
813 return Err(KitError::Integrity(format!(
814 "internal table {name} is missing"
815 )));
816 }
817 }
818 Ok(())
819 }
820
821 pub fn gc(&self) -> Result<usize> {
824 self.inner.gc().map_err(KitError::from)
825 }
826
827 pub fn check(&self) -> Vec<serde_json::Value> {
830 self.inner
831 .check()
832 .into_iter()
833 .map(|i| {
834 serde_json::json!({
835 "table_id": i.table_id,
836 "table_name": i.table_name,
837 "severity": i.severity,
838 "description": i.description,
839 })
840 })
841 .collect()
842 }
843
844 pub fn doctor(&self) -> Result<Vec<u64>> {
846 self.inner.doctor().map_err(KitError::from)
847 }
848
849 pub fn snapshot_epoch(&self) -> u64 {
853 self.inner.snapshot().0.epoch.0
854 }
855
856 pub fn set_history_retention_epochs(&self, epochs: u64) -> Result<()> {
857 self.inner
858 .set_history_retention_epochs(epochs)
859 .map_err(KitError::from)
860 }
861
862 pub fn history_retention_epochs(&self) -> u64 {
863 self.inner.history_retention_epochs()
864 }
865
866 pub fn earliest_retained_epoch(&self) -> u64 {
867 self.inner.earliest_retained_epoch().0
868 }
869
870 pub fn export_tsv(&self, table: &str) -> Result<String> {
874 let t = self
875 .schema
876 .tables
877 .iter()
878 .find(|t| t.name == table)
879 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?
880 .clone();
881 let tx = self.begin()?;
882 let rows = tx.all_rows(table)?;
883 Ok(crate::tsv::rows_to_tsv(&t, &rows))
884 }
885
886 pub fn import_tsv(&self, table: &str, text: &str) -> Result<usize> {
890 let t = self
891 .schema
892 .tables
893 .iter()
894 .find(|t| t.name == table)
895 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?
896 .clone();
897 let rows = crate::tsv::tsv_to_rows(&t, text)?;
898 let n = rows.len();
899 self.transaction(1, |tx| {
900 tx.insert_many(table, rows.clone())?;
901 Ok(())
902 })?;
903 Ok(n)
904 }
905
906 pub fn explain(
911 &self,
912 table: &str,
913 predicate: &mongreldb_kit_core::query::Expr,
914 ) -> Result<ExplainPlan> {
915 let t = self
916 .schema
917 .tables
918 .iter()
919 .find(|t| t.name == table)
920 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
921 Ok(match crate::pushdown::translate_predicate(t, predicate) {
922 Some(p) => ExplainPlan {
923 index_accelerated: p.can_push(),
924 exact: p.fully_translated,
925 pushed_conditions: p.conditions.iter().map(condition_label).collect(),
926 },
927 None => ExplainPlan {
928 index_accelerated: false,
929 exact: false,
930 pushed_conditions: Vec::new(),
931 },
932 })
933 }
934
935 pub fn rows_at_epoch(&self, table: &str, epoch: u64) -> Result<Vec<crate::schema::Row>> {
941 let t = self
942 .schema
943 .tables
944 .iter()
945 .find(|t| t.name == table)
946 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
947 let current = self.snapshot_epoch();
948 if epoch > current {
949 return Err(KitError::Validation(format!(
950 "epoch {epoch} is in the future (current committed epoch is {current})"
951 )));
952 }
953 let snap = Snapshot::at(mongreldb_core::epoch::Epoch(epoch));
954 let rows = self.visible_core_rows_at(table, snap)?;
955 rows.iter()
956 .map(|r| crate::schema::core_row_to_json(r, t))
957 .collect()
958 }
959
960 pub fn approx_aggregate(
966 &self,
967 table: &str,
968 column: Option<&str>,
969 agg: ApproxAggKind,
970 z: f64,
971 ) -> Result<Option<ApproxAggregate>> {
972 let t = self
973 .schema
974 .tables
975 .iter()
976 .find(|t| t.name == table)
977 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
978 if matches!(agg, ApproxAggKind::Sum | ApproxAggKind::Avg) && column.is_none() {
979 return Err(KitError::Validation(
980 "approx sum/avg requires a column".into(),
981 ));
982 }
983 let cid = match column {
984 Some(name) => Some(
985 t.columns
986 .iter()
987 .find(|c| c.name == name)
988 .ok_or_else(|| KitError::Validation(format!("unknown column '{name}'")))?
989 .id as u16,
990 ),
991 None => None,
992 };
993 let core_agg = match agg {
994 ApproxAggKind::Count => ApproxAgg::Count,
995 ApproxAggKind::Sum => ApproxAgg::Sum,
996 ApproxAggKind::Avg => ApproxAgg::Avg,
997 };
998 let handle = self.inner.table(table).map_err(KitError::from)?;
999 let mut guard = handle.lock();
1000 let res = guard
1001 .approx_aggregate(&[], cid, core_agg, z)
1002 .map_err(KitError::from)?;
1003 Ok(res.map(|r| ApproxAggregate {
1004 point: r.point,
1005 ci_low: r.ci_low,
1006 ci_high: r.ci_high,
1007 n_population: r.n_population,
1008 n_sample_live: r.n_sample_live,
1009 n_passing: r.n_passing,
1010 }))
1011 }
1012
1013 pub fn scan_batched<F>(&self, table: &str, batch_size: usize, mut f: F) -> Result<()>
1019 where
1020 F: FnMut(&[serde_json::Map<String, Value>]) -> Result<()>,
1021 {
1022 let kit_t = self
1023 .schema
1024 .tables
1025 .iter()
1026 .find(|t| t.name == table)
1027 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
1028 let batch_size = batch_size.max(1);
1029 let (snapshot, _pin) = self.inner.snapshot();
1032 let handle = self.inner.table(table).map_err(KitError::from)?;
1033 let guard = handle.lock();
1034
1035 let mut projection: Vec<(u16, mongreldb_core::schema::TypeId)> = Vec::new();
1037 let mut meta: Vec<(String, mongreldb_kit_core::schema::ColumnType)> = Vec::new();
1038 for c in &guard.schema().columns {
1039 if let Some(kc) = kit_t.columns.iter().find(|kc| kc.id as u16 == c.id) {
1040 projection.push((c.id, c.ty.clone()));
1041 meta.push((kc.name.clone(), kc.storage_type));
1042 }
1043 }
1044
1045 match guard
1046 .scan_cursor(snapshot, projection, &[])
1047 .map_err(KitError::from)?
1048 {
1049 Some(mut cursor) => {
1050 let mut buf: Vec<serde_json::Map<String, Value>> = Vec::with_capacity(batch_size);
1051 while let Some(batch) = cursor.next_batch().map_err(KitError::from)? {
1052 let nrows = batch.first().map(|c| c.len()).unwrap_or(0);
1053 for j in 0..nrows {
1054 let mut m = serde_json::Map::new();
1055 for (ci, (name, ty)) in meta.iter().enumerate() {
1056 let cv = batch
1057 .get(ci)
1058 .and_then(|col| col.value_at(j))
1059 .unwrap_or(CoreValue::Null);
1060 m.insert(name.clone(), crate::schema::core_to_json(&cv, *ty)?);
1061 }
1062 buf.push(m);
1063 if buf.len() >= batch_size {
1064 f(&buf)?;
1065 buf.clear();
1066 }
1067 }
1068 }
1069 if !buf.is_empty() {
1070 f(&buf)?;
1071 }
1072 Ok(())
1073 }
1074 None => {
1075 drop(guard);
1076 let rows = self.visible_core_rows_at(table, snapshot)?;
1077 let maps: Vec<serde_json::Map<String, Value>> = rows
1078 .iter()
1079 .map(|r| crate::schema::core_row_to_json(r, kit_t).map(|row| row.values))
1080 .collect::<Result<Vec<_>>>()?;
1081 for chunk in maps.chunks(batch_size) {
1082 f(chunk)?;
1083 }
1084 Ok(())
1085 }
1086 }
1087 }
1088
1089 pub fn set_similarity(
1098 &self,
1099 table: &str,
1100 column: &str,
1101 query: &[String],
1102 k: usize,
1103 ) -> Result<Vec<SimilarRow>> {
1104 let t = self
1105 .schema
1106 .tables
1107 .iter()
1108 .find(|t| t.name == table)
1109 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
1110 let col = t.columns.iter().find(|c| c.name == column).ok_or_else(|| {
1111 KitError::Validation(format!("unknown column '{column}' on table '{table}'"))
1112 })?;
1113 let query_set: std::collections::HashSet<String> = query.iter().cloned().collect();
1114
1115 let has_minhash = t.indexes.iter().any(|idx| {
1116 idx.kind == KitIndexKind::MinHash && idx.columns.iter().any(|c| c == column)
1117 });
1118 let rows = if has_minhash {
1119 let query_hashes: Vec<u64> = query
1121 .iter()
1122 .map(|s| mongreldb_core::index::minhash_token_hash(s))
1123 .collect();
1124 let cand_k = k.saturating_mul(8).max(k + 64);
1126 let cond = mongreldb_core::query::Condition::MinHashSimilar {
1127 column_id: col.id as u16,
1128 query: query_hashes,
1129 k: cand_k,
1130 };
1131 let (snapshot, _pin) = self.inner.snapshot();
1132 let core_rows = self.query_core_rows_at(table, &[cond], snapshot)?;
1133 core_rows
1134 .iter()
1135 .map(|r| crate::schema::core_row_to_json(r, t))
1136 .collect::<Result<Vec<_>>>()?
1137 } else {
1138 let tx = self.begin()?;
1139 tx.all_rows(table)?
1140 };
1141
1142 let mut scored: Vec<SimilarRow> = Vec::new();
1143 for row in rows {
1144 let set = parse_string_set(row.values.get(column));
1145 let inter = set.iter().filter(|x| query_set.contains(*x)).count();
1146 let union = set.len() + query_set.len() - inter;
1147 let sim = if union == 0 {
1148 0.0
1149 } else {
1150 inter as f64 / union as f64
1151 };
1152 if sim > 0.0 {
1153 scored.push(SimilarRow {
1154 row,
1155 similarity: sim,
1156 });
1157 }
1158 }
1159 scored.sort_by(|a, b| {
1160 b.similarity
1161 .partial_cmp(&a.similarity)
1162 .unwrap_or(std::cmp::Ordering::Equal)
1163 });
1164 scored.truncate(k);
1165 Ok(scored)
1166 }
1167
1168 pub fn flush(&self) -> Result<()> {
1172 for name in self.inner.table_names() {
1173 let handle = self.inner.table(&name).map_err(KitError::from)?;
1174 let mut guard = handle.lock();
1175 guard.flush().map_err(KitError::from)?;
1176 }
1177 Ok(())
1178 }
1179
1180 pub fn incremental_aggregate(
1192 &self,
1193 table: &str,
1194 column: Option<&str>,
1195 agg: IncrementalAggKind,
1196 filter: Option<&mongreldb_kit_core::query::Expr>,
1197 ) -> Result<IncrementalAggregate> {
1198 let t = self
1199 .schema
1200 .tables
1201 .iter()
1202 .find(|t| t.name == table)
1203 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
1204 if !matches!(agg, IncrementalAggKind::Count) && column.is_none() {
1205 return Err(KitError::Validation(
1206 "sum/min/max/avg incremental aggregate requires a column".into(),
1207 ));
1208 }
1209 let cid = match column {
1210 Some(name) => Some(
1211 t.columns
1212 .iter()
1213 .find(|c| c.name == name)
1214 .ok_or_else(|| KitError::Validation(format!("unknown column '{name}'")))?
1215 .id as u16,
1216 ),
1217 None => None,
1218 };
1219 let conditions = match filter {
1220 Some(expr) => {
1221 let plan = crate::pushdown::translate_predicate(t, expr).ok_or_else(|| {
1222 KitError::Validation(
1223 "filter is not index-translatable for an incremental aggregate".into(),
1224 )
1225 })?;
1226 if !plan.fully_translated {
1227 return Err(KitError::Validation(
1228 "filter has a residual that an incremental aggregate cannot apply exactly"
1229 .into(),
1230 ));
1231 }
1232 plan.conditions
1233 }
1234 None => Vec::new(),
1235 };
1236 let core_agg = match agg {
1237 IncrementalAggKind::Count => NativeAgg::Count,
1238 IncrementalAggKind::Sum => NativeAgg::Sum,
1239 IncrementalAggKind::Min => NativeAgg::Min,
1240 IncrementalAggKind::Max => NativeAgg::Max,
1241 IncrementalAggKind::Avg => NativeAgg::Avg,
1242 };
1243 let cache_key = incremental_cache_key(t.id, cid, agg, &conditions);
1244 let handle = self.inner.table(table).map_err(KitError::from)?;
1245 let mut guard = handle.lock();
1246 let res = guard
1247 .aggregate_incremental(cache_key, &conditions, cid, core_agg)
1248 .map_err(KitError::from)?;
1249 Ok(IncrementalAggregate {
1250 value: agg_state_value(&res.state),
1251 incremental: res.incremental,
1252 delta_rows: res.delta_rows,
1253 })
1254 }
1255
1256 pub fn applied_migrations(&self) -> Result<Vec<mongreldb_kit_core::migrations::Migration>> {
1258 crate::migrate::load_applied_migrations(&self.inner)
1259 }
1260
1261 pub(crate) fn core_db(&self) -> &CoreDatabase {
1262 &self.inner
1263 }
1264
1265 pub(crate) fn core_arc(&self) -> Arc<CoreDatabase> {
1268 Arc::clone(&self.inner)
1269 }
1270
1271 pub fn close(&self) -> Result<()> {
1276 self.inner.close().map_err(KitError::from)
1277 }
1278
1279 pub fn compact_all(&self) -> Result<(usize, usize)> {
1284 self.inner.compact().map_err(KitError::from)
1285 }
1286
1287 pub fn compact_table(&self, name: &str) -> Result<bool> {
1290 self.inner.compact_table(name).map_err(KitError::from)
1291 }
1292
1293 pub fn rename_table(&mut self, from: &str, to: &str) -> Result<()> {
1303 if from.starts_with("__kit_") || to.starts_with("__kit_") {
1304 return Err(KitError::Validation(
1305 "rename_table: names beginning with '__kit_' are reserved for internal tables"
1306 .into(),
1307 ));
1308 }
1309 self.inner.rename_table(from, to).map_err(KitError::from)?;
1310 if !self.schema.rename_table(from, to) {
1313 return Err(KitError::Integrity(format!(
1316 "rename_table: kit schema has no table '{from}' (or '{to}' already exists)"
1317 )));
1318 }
1319 for table in &mut self.schema.tables {
1320 for fk in &mut table.foreign_keys {
1321 if fk.references_table == from {
1322 fk.references_table = to.to_string();
1323 }
1324 }
1325 }
1326 store_schema(&self.root, &self.schema)?;
1327 Ok(())
1328 }
1329
1330 pub fn analyze(&self) -> Result<()> {
1335 for name in self.inner.table_names() {
1336 let handle = self.inner.table(&name).map_err(KitError::from)?;
1337 handle.lock().ensure_indexes_complete()?;
1338 }
1339 Ok(())
1340 }
1341
1342 pub fn vacuum(&self) -> Result<usize> {
1346 self.inner.compact().map_err(KitError::from)?;
1347 self.inner.gc().map_err(KitError::from)
1348 }
1349
1350 pub fn create_view(&self, spec: &ViewSpec) -> Result<()> {
1356 self.sql(&spec.create_sql())?;
1357 Ok(())
1358 }
1359
1360 pub fn drop_view(&self, name: &str) -> Result<()> {
1362 self.sql(&format!("DROP VIEW IF EXISTS {name}"))?;
1363 Ok(())
1364 }
1365
1366 pub fn reserve_auto_inc(&self, table: &str) -> Result<Option<i64>> {
1374 let handle = self.inner.table(table).map_err(KitError::from)?;
1375 let mut guard = handle.lock();
1376 guard.reserve_auto_inc().map_err(KitError::from)
1377 }
1378
1379 pub fn create_user(&self, username: &str, password: &str) -> Result<()> {
1383 self.inner
1384 .create_user(username, password)
1385 .map_err(KitError::from)?;
1386 Ok(())
1387 }
1388
1389 pub fn drop_user(&self, username: &str) -> Result<()> {
1391 self.inner.drop_user(username).map_err(KitError::from)
1392 }
1393
1394 pub fn alter_user_password(&self, username: &str, new_password: &str) -> Result<()> {
1396 self.inner
1397 .alter_user_password(username, new_password)
1398 .map_err(KitError::from)
1399 }
1400
1401 pub fn verify_user(
1403 &self,
1404 username: &str,
1405 password: &str,
1406 ) -> Result<Option<mongreldb_core::auth::UserEntry>> {
1407 self.inner
1408 .verify_user(username, password)
1409 .map_err(KitError::from)
1410 }
1411
1412 pub fn set_user_admin(&self, username: &str, is_admin: bool) -> Result<()> {
1414 self.inner
1415 .set_user_admin(username, is_admin)
1416 .map_err(KitError::from)
1417 }
1418
1419 pub fn users(&self) -> Vec<String> {
1421 self.inner.users().into_iter().map(|u| u.username).collect()
1422 }
1423
1424 pub fn create_role(&self, name: &str) -> Result<()> {
1426 self.inner.create_role(name).map_err(KitError::from)?;
1427 Ok(())
1428 }
1429
1430 pub fn drop_role(&self, name: &str) -> Result<()> {
1432 self.inner.drop_role(name).map_err(KitError::from)
1433 }
1434
1435 pub fn roles(&self) -> Vec<String> {
1437 self.inner.roles().into_iter().map(|r| r.name).collect()
1438 }
1439
1440 pub fn grant_role(&self, username: &str, role_name: &str) -> Result<()> {
1442 self.inner
1443 .grant_role(username, role_name)
1444 .map_err(KitError::from)
1445 }
1446
1447 pub fn revoke_role(&self, username: &str, role_name: &str) -> Result<()> {
1449 self.inner
1450 .revoke_role(username, role_name)
1451 .map_err(KitError::from)
1452 }
1453
1454 pub fn grant_permission(
1456 &self,
1457 role_name: &str,
1458 permission: mongreldb_core::auth::Permission,
1459 ) -> Result<()> {
1460 self.inner
1461 .grant_permission(role_name, permission)
1462 .map_err(KitError::from)
1463 }
1464
1465 pub fn revoke_permission(
1467 &self,
1468 role_name: &str,
1469 permission: mongreldb_core::auth::Permission,
1470 ) -> Result<()> {
1471 self.inner
1472 .revoke_permission(role_name, permission)
1473 .map_err(KitError::from)
1474 }
1475
1476 pub fn set_spill_threshold(&self, bytes: u64) {
1482 self.inner.set_spill_threshold(bytes);
1483 }
1484
1485 pub fn set_recursive_triggers(&self, enabled: bool) {
1487 self.inner.set_recursive_triggers(enabled);
1488 }
1489
1490 pub fn trigger_config(&self) -> mongreldb_core::TriggerConfig {
1492 self.inner.trigger_config()
1493 }
1494
1495 pub fn set_trigger_config(&self, config: mongreldb_core::TriggerConfig) -> Result<()> {
1497 self.inner
1498 .set_trigger_config(config)
1499 .map_err(KitError::from)
1500 }
1501
1502 pub fn set_table_compaction_zstd_level(&self, table: &str, level: i32) -> Result<()> {
1504 let handle = self.inner.table(table).map_err(KitError::from)?;
1505 handle.lock().set_compaction_zstd_level(level);
1506 Ok(())
1507 }
1508
1509 pub fn set_table_result_cache_max_bytes(&self, table: &str, max_bytes: u64) -> Result<()> {
1511 let handle = self.inner.table(table).map_err(KitError::from)?;
1512 handle.lock().set_result_cache_max_bytes(max_bytes);
1513 Ok(())
1514 }
1515
1516 pub fn set_table_mutable_run_spill_bytes(&self, table: &str, bytes: u64) -> Result<()> {
1518 let handle = self.inner.table(table).map_err(KitError::from)?;
1519 handle.lock().set_mutable_run_spill_bytes(bytes);
1520 Ok(())
1521 }
1522
1523 pub fn set_table_sync_byte_threshold(&self, table: &str, threshold: u64) -> Result<()> {
1525 let handle = self.inner.table(table).map_err(KitError::from)?;
1526 handle.lock().set_sync_byte_threshold(threshold);
1527 Ok(())
1528 }
1529
1530 pub fn set_table_index_build_policy(
1533 &self,
1534 table: &str,
1535 policy: mongreldb_core::IndexBuildPolicy,
1536 ) -> Result<()> {
1537 let handle = self.inner.table(table).map_err(KitError::from)?;
1538 handle.lock().set_index_build_policy(policy);
1539 Ok(())
1540 }
1541
1542 pub fn table_page_cache_stats(&self, table: &str) -> Result<mongreldb_core::cache::CacheStats> {
1544 let handle = self.inner.table(table).map_err(KitError::from)?;
1545 let stats = handle.lock().page_cache_stats();
1546 Ok(stats)
1547 }
1548
1549 pub fn table_run_count(&self, table: &str) -> Result<usize> {
1551 let handle = self.inner.table(table).map_err(KitError::from)?;
1552 let n = handle.lock().run_count();
1553 Ok(n)
1554 }
1555
1556 pub fn table_memtable_len(&self, table: &str) -> Result<usize> {
1558 let handle = self.inner.table(table).map_err(KitError::from)?;
1559 let n = handle.lock().memtable_len();
1560 Ok(n)
1561 }
1562
1563 pub fn table_mutable_run_len(&self, table: &str) -> Result<usize> {
1565 let handle = self.inner.table(table).map_err(KitError::from)?;
1566 let n = handle.lock().mutable_run_len();
1567 Ok(n)
1568 }
1569
1570 pub fn table_page_cache_len(&self, table: &str) -> Result<usize> {
1572 let handle = self.inner.table(table).map_err(KitError::from)?;
1573 let n = handle.lock().page_cache_len();
1574 Ok(n)
1575 }
1576
1577 pub fn table_decoded_cache_len(&self, table: &str) -> Result<usize> {
1579 let handle = self.inner.table(table).map_err(KitError::from)?;
1580 let n = handle.lock().decoded_cache_len();
1581 Ok(n)
1582 }
1583
1584 pub fn sql(&self, statement: &str) -> Result<Vec<arrow::record_batch::RecordBatch>> {
1602 self.sql_with_options(statement, SqlOptions::default())
1603 }
1604
1605 fn sql_session(&self) -> Result<Arc<mongreldb_query::MongrelSession>> {
1606 if let Some(session) = self.session.read().as_ref() {
1607 return Ok(Arc::clone(session));
1608 }
1609 let session = Arc::new(
1610 mongreldb_query::MongrelSession::open(self.core_arc()).map_err(KitError::from)?,
1611 );
1612 let mut cached = self.session.write();
1613 Ok(Arc::clone(cached.get_or_insert(session)))
1614 }
1615
1616 #[doc(hidden)]
1617 pub fn set_sql_test_hook(&self, hook: Option<mongreldb_query::SqlTestHook>) -> Result<()> {
1618 self.sql_session()?.set_test_hook(hook);
1619 Ok(())
1620 }
1621
1622 pub fn start_sql(
1623 &self,
1624 statement: impl Into<String>,
1625 options: SqlOptions,
1626 ) -> Result<SqlQueryHandle> {
1627 let session = self.sql_session()?;
1628 let query = session
1629 .register_query(mongreldb_query::SqlQueryOptions {
1630 query_id: options.query_id,
1631 timeout: options.timeout,
1632 ..mongreldb_query::SqlQueryOptions::default()
1633 })
1634 .map_err(KitError::from)?;
1635 let query_id = query.id();
1636 let registration = mongreldb_query::RegisteredQueryGuard::new(query);
1637 let worker_session = Arc::clone(&session);
1638 let statement = statement.into();
1639 let worker = std::thread::Builder::new()
1640 .name(format!("mongreldb-kit-sql-{query_id}"))
1641 .spawn(move || {
1642 sql_runtime()
1643 .block_on(worker_session.run_with_query(&statement, registration.into_query()))
1644 .map_err(KitError::from)
1645 })
1646 .map_err(|error| KitError::Storage(error.to_string()))?;
1647 Ok(SqlQueryHandle {
1648 query_id,
1649 session,
1650 worker: Some(worker),
1651 })
1652 }
1653
1654 pub fn sql_with_options(
1655 &self,
1656 statement: &str,
1657 options: SqlOptions,
1658 ) -> Result<Vec<arrow::record_batch::RecordBatch>> {
1659 self.start_sql(statement, options)?.wait()
1660 }
1661
1662 pub fn cancel_sql(&self, query_id: mongreldb_query::QueryId) -> mongreldb_query::CancelOutcome {
1663 self.session
1664 .read()
1665 .as_ref()
1666 .map_or(mongreldb_query::CancelOutcome::NotFound, |session| {
1667 session.cancel_query(query_id)
1668 })
1669 }
1670
1671 pub fn refresh_sql_session(&self) -> Result<()> {
1676 let session =
1677 mongreldb_query::MongrelSession::open(self.core_arc()).map_err(KitError::from)?;
1678 *self.session.write() = Some(Arc::new(session));
1679 Ok(())
1680 }
1681
1682 pub fn sql_arrow(&self, statement: &str) -> Result<Vec<u8>> {
1688 self.sql_arrow_with_options(statement, SqlOptions::default())
1689 }
1690
1691 pub fn sql_arrow_with_options(&self, statement: &str, options: SqlOptions) -> Result<Vec<u8>> {
1692 self.sql_serialized_with_options(statement, options, |output| {
1693 crate::arrow_util::batches_to_ipc_controlled(output.batches(), output.query())
1694 })
1695 }
1696
1697 pub fn sql_rows(&self, statement: &str) -> Result<Vec<serde_json::Map<String, Value>>> {
1701 self.sql_rows_with_options(statement, SqlOptions::default())
1702 }
1703
1704 pub fn sql_rows_with_options(
1705 &self,
1706 statement: &str,
1707 options: SqlOptions,
1708 ) -> Result<Vec<serde_json::Map<String, Value>>> {
1709 self.sql_serialized_with_options(statement, options, |output| {
1710 crate::arrow_util::batches_to_rows_controlled(output.batches(), output.query())
1711 })
1712 }
1713
1714 fn sql_serialized_with_options<T>(
1715 &self,
1716 statement: &str,
1717 options: SqlOptions,
1718 serialize: impl FnOnce(&mongreldb_query::ManagedQueryBatches) -> Result<T>,
1719 ) -> Result<T> {
1720 let session = self.sql_session()?;
1721 let query = session
1722 .register_query(mongreldb_query::SqlQueryOptions {
1723 query_id: options.query_id,
1724 timeout: options.timeout,
1725 ..mongreldb_query::SqlQueryOptions::default()
1726 })
1727 .map_err(KitError::from)?;
1728 let output = sql_runtime()
1729 .block_on(session.run_with_query_for_serialization(statement, query))
1730 .map_err(KitError::from)?;
1731 session.fire_test_hook(mongreldb_query::SqlTestHookPoint::BeforeSerializationBatch);
1732 match serialize(&output) {
1733 Ok(value) => {
1734 output.complete();
1735 Ok(value)
1736 }
1737 Err(error) => {
1738 output.fail();
1739 Err(error)
1740 }
1741 }
1742 }
1743
1744 pub(crate) fn lookup_row_id(&self, table: &str, key: &[u8]) -> Result<Option<RowId>> {
1748 let handle = self.inner.table(table).map_err(KitError::from)?;
1749 let mut guard = handle.lock();
1750 guard.ensure_indexes_complete()?;
1751 Ok(guard.lookup_pk(key))
1752 }
1753
1754 pub(crate) fn root(&self) -> &Path {
1755 &self.root
1756 }
1757
1758 pub(crate) fn visible_core_rows_at(
1762 &self,
1763 table_name: &str,
1764 snapshot: Snapshot,
1765 ) -> Result<Vec<CoreRow>> {
1766 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1767 let guard = handle.lock();
1768 guard.visible_rows(snapshot).map_err(KitError::from)
1769 }
1770
1771 pub(crate) fn query_core_rows_at(
1778 &self,
1779 table_name: &str,
1780 conditions: &[mongreldb_core::query::Condition],
1781 snapshot: Snapshot,
1782 ) -> Result<Vec<CoreRow>> {
1783 if conditions.is_empty() {
1784 return self.visible_core_rows_at(table_name, snapshot);
1785 }
1786 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1787 let mut guard = handle.lock();
1788 let q = conditions
1789 .iter()
1790 .cloned()
1791 .fold(mongreldb_core::query::Query::new(), |query, condition| {
1792 query.and(condition)
1793 });
1794 guard.query(&q).map_err(KitError::from)
1795 }
1796
1797 pub(crate) fn flush_table(&self, table_name: &str) -> Result<()> {
1805 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1806 handle.lock().flush().map_err(KitError::from)?;
1807 Ok(())
1808 }
1809
1810 pub(crate) fn count_core_rows_at(
1820 &self,
1821 table_name: &str,
1822 conditions: &[mongreldb_core::query::Condition],
1823 snapshot: Snapshot,
1824 ) -> Result<Option<u64>> {
1825 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1826 let mut guard = handle.lock();
1827 if guard.snapshot().epoch != snapshot.epoch {
1828 return Ok(None); }
1830 guard
1831 .count_conditions(conditions, snapshot)
1832 .map_err(KitError::from)
1833 }
1834
1835 pub(crate) fn aggregate_core_at(
1844 &self,
1845 table_name: &str,
1846 column: Option<u16>,
1847 conditions: &[mongreldb_core::query::Condition],
1848 agg: NativeAgg,
1849 snapshot: Snapshot,
1850 ) -> Result<Option<NativeAggResult>> {
1851 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1852 let guard = handle.lock();
1853 if guard.snapshot().epoch != snapshot.epoch {
1854 return Ok(None); }
1856 guard
1857 .aggregate_native(snapshot, column, conditions, agg)
1858 .map_err(KitError::from)
1859 }
1860
1861 pub(crate) fn count_distinct_core_at(
1870 &self,
1871 table_name: &str,
1872 column_id: u16,
1873 snapshot: Snapshot,
1874 ) -> Result<Option<u64>> {
1875 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1876 let mut guard = handle.lock();
1877 if guard.snapshot().epoch != snapshot.epoch {
1878 return Ok(None); }
1880 guard
1881 .count_distinct_from_bitmap(column_id)
1882 .map_err(KitError::from)
1883 }
1884
1885 #[allow(dead_code)]
1887 pub(crate) fn get_core_row(&self, table_name: &str, row_id: u64) -> Result<Option<CoreRow>> {
1888 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1889 let guard = handle.lock();
1890 let snapshot = guard.snapshot();
1891 Ok(guard.get(mongreldb_core::RowId(row_id), snapshot))
1892 }
1893}
1894
1895pub(crate) fn create_core_table(db: &CoreDatabase, name: &str, schema: CoreSchema) -> Result<()> {
1896 if db.table_id(name).is_ok() {
1897 return Ok(());
1898 }
1899 db.create_table(name, schema).map_err(KitError::from)?;
1900 Ok(())
1901}
1902
1903fn sql_runtime() -> &'static tokio::runtime::Runtime {
1907 use std::sync::OnceLock;
1908 static RT: OnceLock<tokio::runtime::Runtime> = OnceLock::new();
1909 RT.get_or_init(|| {
1910 tokio::runtime::Builder::new_multi_thread()
1911 .worker_threads(4)
1912 .enable_all()
1913 .build()
1914 .expect("failed to build kit SQL tokio runtime")
1915 })
1916}
1917
1918fn core_procedure(spec: &ProcedureSpec) -> Result<mongreldb_core::StoredProcedure> {
1919 let parsed: mongreldb_core::StoredProcedure =
1920 serde_json::from_value(spec.json.clone()).map_err(KitError::from)?;
1921 mongreldb_core::StoredProcedure::new(parsed.name, parsed.mode, parsed.params, parsed.body, 0)
1922 .map_err(KitError::from)
1923}
1924
1925fn core_trigger(spec: &TriggerSpec) -> Result<mongreldb_core::StoredTrigger> {
1926 let parsed: mongreldb_core::StoredTrigger =
1927 serde_json::from_value(spec.json.clone()).map_err(KitError::from)?;
1928 mongreldb_core::StoredTrigger::new(
1929 parsed.name,
1930 mongreldb_core::TriggerDefinition {
1931 target: parsed.target,
1932 timing: parsed.timing,
1933 event: parsed.event,
1934 update_of: parsed.update_of,
1935 target_columns: parsed.target_columns,
1936 when: parsed.when,
1937 program: parsed.program,
1938 },
1939 0,
1940 )
1941 .map_err(KitError::from)
1942}
1943
1944fn json_to_core_value(value: &Value) -> Result<CoreValue> {
1945 match value {
1946 Value::Null => Ok(CoreValue::Null),
1947 Value::Bool(value) => Ok(CoreValue::Bool(*value)),
1948 Value::Number(value) => {
1949 if let Some(value) = value.as_i64() {
1950 Ok(CoreValue::Int64(value))
1951 } else if let Some(value) = value.as_f64() {
1952 Ok(CoreValue::Float64(value))
1953 } else {
1954 Err(KitError::Validation("unsupported JSON number".into()))
1955 }
1956 }
1957 Value::String(value) => Ok(CoreValue::Bytes(value.as_bytes().to_vec())),
1958 Value::Array(_) | Value::Object(_) => Err(KitError::Validation(
1959 "procedure args only support scalar JSON values".into(),
1960 )),
1961 }
1962}
1963
1964pub(crate) fn internal_bytes(row: &CoreRow, col_id: u16) -> Option<String> {
1966 match row.columns.get(&col_id) {
1967 Some(CoreValue::Bytes(b)) => String::from_utf8(b.clone()).ok(),
1968 _ => None,
1969 }
1970}
1971
1972fn reap_rotated_wal_segments(db: &CoreDatabase) {
1987 let _ = db.gc();
1988}
1989
1990pub(crate) fn load_schema(path: &Path) -> Result<KitSchema> {
1991 let file = path.join(SCHEMA_FILE);
1992 let json = std::fs::read_to_string(&file)
1993 .map_err(|e| KitError::Migration(format!("cannot read schema file: {e}")))?;
1994 let schema: KitSchema = serde_json::from_str(&json)?;
1995 Ok(schema)
1996}
1997
1998pub(crate) fn store_schema(path: &Path, schema: &KitSchema) -> Result<()> {
1999 let file = path.join(SCHEMA_FILE);
2000 let json = serde_json::to_string_pretty(schema)?;
2001 std::fs::write(&file, json)?;
2002 Ok(())
2003}
2004
2005pub(crate) fn persist_schema(db: &Database, schema: &KitSchema) -> Result<()> {
2007 store_schema(&db.root, schema)
2008}
2009
2010#[cfg(test)]
2011mod tests {
2012 use super::open_core_with_retry;
2013
2014 fn lock_error() -> mongreldb_core::MongrelError {
2015 mongreldb_core::MongrelError::Io(std::io::Error::other(
2016 "database at /tmp/db is locked by another process: would block",
2017 ))
2018 }
2019
2020 #[test]
2021 fn open_retry_waits_for_lock_contention_only() {
2022 let mut calls = 0;
2023 let value = open_core_with_retry(50, || {
2024 calls += 1;
2025 if calls < 3 {
2026 Err(lock_error())
2027 } else {
2028 Ok(7)
2029 }
2030 })
2031 .unwrap();
2032 assert_eq!(value, 7);
2033 assert_eq!(calls, 3);
2034
2035 let mut non_lock_calls = 0;
2036 let err: mongreldb_core::Result<()> = open_core_with_retry(50, || {
2037 non_lock_calls += 1;
2038 Err(mongreldb_core::MongrelError::Other("nope".into()))
2039 });
2040 let err = err.unwrap_err();
2041 assert_eq!(non_lock_calls, 1);
2042 assert!(matches!(err, mongreldb_core::MongrelError::Other(_)));
2043 }
2044}