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
24pub type DefaultProvider = Box<dyn Fn() -> Value + Send + Sync>;
26
27#[derive(Debug, Clone)]
30pub struct ExplainPlan {
31 pub index_accelerated: bool,
33 pub exact: bool,
36 pub pushed_conditions: Vec<String>,
38}
39
40#[derive(Debug, Clone)]
42pub struct SimilarRow {
43 pub row: crate::schema::Row,
44 pub similarity: f64,
45}
46
47fn parse_string_set(value: Option<&Value>) -> std::collections::HashSet<String> {
51 let arr = match value {
52 Some(Value::Array(a)) => Some(a.clone()),
53 Some(Value::String(s)) => serde_json::from_str::<Value>(s)
54 .ok()
55 .and_then(|v| v.as_array().cloned()),
56 _ => None,
57 };
58 arr.into_iter()
59 .flatten()
60 .filter_map(|v| match v {
61 Value::String(s) => Some(s),
62 Value::Number(n) => Some(n.to_string()),
63 Value::Bool(b) => Some(b.to_string()),
64 _ => None,
65 })
66 .collect()
67}
68
69#[derive(Debug, Clone, Copy, PartialEq, Eq)]
71pub enum IncrementalAggKind {
72 Count,
73 Sum,
74 Min,
75 Max,
76 Avg,
77}
78
79#[derive(Debug, Clone)]
81pub struct IncrementalAggregate {
82 pub value: Value,
85 pub incremental: bool,
89 pub delta_rows: u64,
91}
92
93fn incremental_cache_key(
97 table_id: u32,
98 column: Option<u16>,
99 agg: IncrementalAggKind,
100 conditions: &[mongreldb_core::query::Condition],
101) -> u64 {
102 use std::hash::{Hash, Hasher};
103 let mut h = std::collections::hash_map::DefaultHasher::new();
104 table_id.hash(&mut h);
105 column.hash(&mut h);
106 (agg as u8).hash(&mut h);
107 format!("{conditions:?}").hash(&mut h);
109 h.finish()
110}
111
112fn agg_state_value(s: &AggState) -> Value {
116 let num_f64 = |x: f64| {
117 serde_json::Number::from_f64(x)
118 .map(Value::Number)
119 .unwrap_or(Value::Null)
120 };
121 match s {
122 AggState::Count(n) => Value::from(*n),
123 AggState::SumI { sum, .. } => i64::try_from(*sum)
124 .map(Value::from)
125 .unwrap_or_else(|_| num_f64(*sum as f64)),
126 AggState::SumF { sum, .. } => num_f64(*sum),
127 AggState::AvgI { sum, count } if *count > 0 => num_f64(*sum as f64 / *count as f64),
128 AggState::AvgF { sum, count } if *count > 0 => num_f64(*sum / *count as f64),
129 AggState::AvgI { .. } | AggState::AvgF { .. } => Value::Null,
130 AggState::MinI(n) | AggState::MaxI(n) => Value::from(*n),
131 AggState::MinF(f) | AggState::MaxF(f) => num_f64(*f),
132 AggState::Empty => Value::Null,
133 }
134}
135
136#[derive(Debug, Clone, Copy, PartialEq, Eq)]
138pub enum ApproxAggKind {
139 Count,
140 Sum,
141 Avg,
142}
143
144#[derive(Debug, Clone)]
148pub struct ApproxAggregate {
149 pub point: f64,
150 pub ci_low: f64,
151 pub ci_high: f64,
152 pub n_population: u64,
153 pub n_sample_live: usize,
154 pub n_passing: usize,
155}
156
157fn condition_label(c: &mongreldb_core::query::Condition) -> String {
160 let dbg = format!("{c:?}");
161 dbg.split(['(', '{', ' ']).next().unwrap_or("").to_string()
162}
163
164pub struct Database {
169 pub(crate) inner: Arc<CoreDatabase>,
170 pub(crate) schema: KitSchema,
171 pub(crate) root: PathBuf,
172 pub(crate) default_providers: HashMap<String, DefaultProvider>,
174 pub(crate) session: parking_lot::Mutex<Option<mongreldb_query::MongrelSession>>,
181}
182
183impl Database {
184 pub fn open(path: &Path) -> Result<Self> {
186 let inner = Arc::new(CoreDatabase::open(path)?);
187 let schema = load_schema(path)?;
188 ensure_internal_tables(&inner)?;
190 reap_rotated_wal_segments(&inner);
191 Ok(Self {
192 inner,
193 schema,
194 root: path.to_path_buf(),
195 default_providers: HashMap::new(),
196 session: parking_lot::Mutex::new(None),
197 })
198 }
199
200 pub fn open_encrypted(path: &Path, passphrase: &str) -> Result<Self> {
202 let inner = Arc::new(CoreDatabase::open_encrypted(path, passphrase)?);
203 let schema = load_schema(path)?;
204 ensure_internal_tables(&inner)?;
205 reap_rotated_wal_segments(&inner);
206 Ok(Self {
207 inner,
208 schema,
209 root: path.to_path_buf(),
210 default_providers: HashMap::new(),
211 session: parking_lot::Mutex::new(None),
212 })
213 }
214
215 pub fn create_encrypted(path: &Path, schema: KitSchema, passphrase: &str) -> Result<Self> {
219 std::fs::create_dir_all(path)?;
220 let inner = Arc::new(CoreDatabase::create_encrypted(path, passphrase)?);
221 ensure_internal_tables(&inner)?;
222 store_schema(path, &schema)?;
223 for table in &schema.tables {
224 create_core_table(&inner, &table.name, to_core_schema(table))?;
225 }
226 Ok(Self {
227 inner,
228 schema,
229 root: path.to_path_buf(),
230 default_providers: HashMap::new(),
231 session: parking_lot::Mutex::new(None),
232 })
233 }
234
235 pub fn create(path: &Path, schema: KitSchema) -> Result<Self> {
237 std::fs::create_dir_all(path)?;
238 let inner = Arc::new(CoreDatabase::create(path)?);
239
240 ensure_internal_tables(&inner)?;
243
244 store_schema(path, &schema)?;
247
248 for table in &schema.tables {
250 create_core_table(&inner, &table.name, to_core_schema(table))?;
251 }
252
253 Ok(Self {
254 inner,
255 schema,
256 root: path.to_path_buf(),
257 default_providers: HashMap::new(),
258 session: parking_lot::Mutex::new(None),
259 })
260 }
261
262 pub fn register_default(
265 &mut self,
266 name: impl Into<String>,
267 provider: impl Fn() -> Value + Send + Sync + 'static,
268 ) {
269 self.default_providers
270 .insert(name.into(), Box::new(provider));
271 }
272
273 pub fn raw(&self) -> &CoreDatabase {
277 &self.inner
278 }
279
280 pub fn table_names(&self) -> Vec<String> {
282 self.schema
283 .tables
284 .iter()
285 .map(|t| t.name.clone())
286 .filter(|n| !n.starts_with("__kit_"))
287 .collect()
288 }
289
290 pub fn create_procedure(
291 &self,
292 spec: &ProcedureSpec,
293 ) -> Result<mongreldb_core::StoredProcedure> {
294 let procedure = core_procedure(spec)?;
295 self.inner
296 .create_procedure(procedure)
297 .map_err(KitError::from)
298 }
299
300 pub fn replace_procedure(
301 &self,
302 spec: &ProcedureSpec,
303 ) -> Result<mongreldb_core::StoredProcedure> {
304 let procedure = core_procedure(spec)?;
305 self.inner
306 .create_or_replace_procedure(procedure)
307 .map_err(KitError::from)
308 }
309
310 pub fn drop_procedure(&self, name: &str) -> Result<()> {
311 self.inner.drop_procedure(name).map_err(KitError::from)
312 }
313
314 pub fn call_procedure(
315 &self,
316 name: &str,
317 args: serde_json::Map<String, Value>,
318 ) -> Result<mongreldb_core::ProcedureCallResult> {
319 let args = args
320 .iter()
321 .map(|(key, value)| Ok((key.clone(), json_to_core_value(value)?)))
322 .collect::<Result<HashMap<_, _>>>()?;
323 self.inner
324 .call_procedure(name, args)
325 .map_err(KitError::from)
326 }
327
328 pub fn create_trigger(&self, spec: &TriggerSpec) -> Result<mongreldb_core::StoredTrigger> {
329 let trigger = core_trigger(spec)?;
330 self.inner.create_trigger(trigger).map_err(KitError::from)
331 }
332
333 pub fn replace_trigger(&self, spec: &TriggerSpec) -> Result<mongreldb_core::StoredTrigger> {
334 let trigger = core_trigger(spec)?;
335 self.inner
336 .create_or_replace_trigger(trigger)
337 .map_err(KitError::from)
338 }
339
340 pub fn drop_trigger(&self, name: &str) -> Result<()> {
341 self.inner.drop_trigger(name).map_err(KitError::from)
342 }
343
344 pub fn triggers(&self) -> Vec<mongreldb_core::StoredTrigger> {
345 self.inner.triggers()
346 }
347
348 pub fn trigger(&self, name: &str) -> Option<mongreldb_core::StoredTrigger> {
349 self.inner.trigger(name)
350 }
351
352 pub fn allocate_sequence(&self, name: &str, count: i64) -> Result<i64> {
358 use crate::internal::cols;
359 let mut attempt = 0;
360 loop {
361 let mut txn = self.inner.begin();
362 let snapshot = txn.read_snapshot();
363 let existing = self
364 .visible_core_rows_at(crate::internal::SEQUENCES, snapshot)?
365 .into_iter()
366 .find(|r| internal_bytes(r, cols::SEQ_NAME) == Some(name.to_string()));
367
368 let now = crate::internal::iso_now();
369 let (start, next, old_row_id) = match &existing {
373 Some(row) => {
374 let current = match row.columns.get(&cols::SEQ_NEXT) {
375 Some(CoreValue::Int64(i)) => *i,
376 _ => 1,
377 };
378 (current, current + count, Some(row.row_id))
379 }
380 None => (1, 1 + count, None),
381 };
382
383 if let Some(rid) = old_row_id {
384 txn.delete(crate::internal::SEQUENCES, rid)
385 .map_err(KitError::from)?;
386 }
387 txn.put(
388 crate::internal::SEQUENCES,
389 vec![
390 (cols::SEQ_NAME, CoreValue::Bytes(name.as_bytes().to_vec())),
391 (cols::SEQ_NEXT, CoreValue::Int64(next)),
392 (cols::SEQ_UPDATED, CoreValue::Bytes(now.into_bytes())),
393 ],
394 )
395 .map_err(KitError::from)?;
396 match txn.commit() {
397 Ok(_) => return Ok(start),
398 Err(mongreldb_core::MongrelError::Conflict(_)) if attempt < 10_000 => {
399 attempt += 1;
400 std::thread::yield_now();
401 continue;
402 }
403 Err(e) => return Err(KitError::from(e)),
404 }
405 }
406 }
407
408 pub fn transaction<T, F>(&self, max_retries: usize, mut f: F) -> Result<T>
411 where
412 F: FnMut(&mut crate::txn::Transaction<'_>) -> Result<T>,
413 {
414 let mut attempt = 0;
415 loop {
416 let mut txn = self.begin()?;
417 match f(&mut txn) {
418 Ok(value) => match txn.commit() {
419 Ok(()) => return Ok(value),
420 Err(KitError::Conflict(_)) if attempt < max_retries => {
421 attempt += 1;
422 continue;
423 }
424 Err(e) => return Err(e),
425 },
426 Err(KitError::Conflict(_)) if attempt < max_retries => {
427 txn.rollback();
428 attempt += 1;
429 continue;
430 }
431 Err(e) => {
432 txn.rollback();
433 return Err(e);
434 }
435 }
436 }
437 }
438
439 pub fn table(&self, name: &str) -> Option<&KitTable> {
441 self.schema.table(name)
442 }
443
444 pub fn schema(&self) -> &KitSchema {
446 &self.schema
447 }
448
449 pub fn begin(&self) -> Result<crate::txn::Transaction<'_>> {
451 let core_txn = self.inner.begin();
452 Ok(crate::txn::Transaction::new(self, core_txn))
453 }
454
455 pub fn set_schema(&mut self, schema: KitSchema) {
457 self.schema = schema;
458 }
459
460 pub fn check_internal_tables(&self) -> Result<()> {
463 let schema_file = self.root.join(SCHEMA_FILE);
464 if !schema_file.exists() {
465 return Err(KitError::Integrity(format!(
466 "schema file {} is missing",
467 schema_file.display()
468 )));
469 }
470 for (name, _) in internal_tables_core() {
471 if self.inner.table_id(name).is_err() {
472 return Err(KitError::Integrity(format!(
473 "internal table {name} is missing"
474 )));
475 }
476 }
477 Ok(())
478 }
479
480 pub fn gc(&self) -> Result<usize> {
483 self.inner.gc().map_err(KitError::from)
484 }
485
486 pub fn check(&self) -> Vec<serde_json::Value> {
489 self.inner
490 .check()
491 .into_iter()
492 .map(|i| {
493 serde_json::json!({
494 "table_id": i.table_id,
495 "table_name": i.table_name,
496 "severity": i.severity,
497 "description": i.description,
498 })
499 })
500 .collect()
501 }
502
503 pub fn doctor(&self) -> Result<Vec<u64>> {
505 self.inner.doctor().map_err(KitError::from)
506 }
507
508 pub fn snapshot_epoch(&self) -> u64 {
512 self.inner.snapshot().0.epoch.0
513 }
514
515 pub fn export_tsv(&self, table: &str) -> Result<String> {
519 let t = self
520 .schema
521 .tables
522 .iter()
523 .find(|t| t.name == table)
524 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?
525 .clone();
526 let tx = self.begin()?;
527 let rows = tx.all_rows(table)?;
528 Ok(crate::tsv::rows_to_tsv(&t, &rows))
529 }
530
531 pub fn import_tsv(&self, table: &str, text: &str) -> Result<usize> {
535 let t = self
536 .schema
537 .tables
538 .iter()
539 .find(|t| t.name == table)
540 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?
541 .clone();
542 let rows = crate::tsv::tsv_to_rows(&t, text)?;
543 let n = rows.len();
544 self.transaction(1, |tx| {
545 tx.insert_many(table, rows.clone())?;
546 Ok(())
547 })?;
548 Ok(n)
549 }
550
551 pub fn explain(
556 &self,
557 table: &str,
558 predicate: &mongreldb_kit_core::query::Expr,
559 ) -> Result<ExplainPlan> {
560 let t = self
561 .schema
562 .tables
563 .iter()
564 .find(|t| t.name == table)
565 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
566 Ok(match crate::pushdown::translate_predicate(t, predicate) {
567 Some(p) => ExplainPlan {
568 index_accelerated: p.can_push(),
569 exact: p.fully_translated,
570 pushed_conditions: p.conditions.iter().map(condition_label).collect(),
571 },
572 None => ExplainPlan {
573 index_accelerated: false,
574 exact: false,
575 pushed_conditions: Vec::new(),
576 },
577 })
578 }
579
580 pub fn rows_at_epoch(&self, table: &str, epoch: u64) -> Result<Vec<crate::schema::Row>> {
586 let t = self
587 .schema
588 .tables
589 .iter()
590 .find(|t| t.name == table)
591 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
592 let current = self.snapshot_epoch();
593 if epoch > current {
594 return Err(KitError::Validation(format!(
595 "epoch {epoch} is in the future (current committed epoch is {current})"
596 )));
597 }
598 let snap = Snapshot::at(mongreldb_core::epoch::Epoch(epoch));
599 let rows = self.visible_core_rows_at(table, snap)?;
600 rows.iter()
601 .map(|r| crate::schema::core_row_to_json(r, t))
602 .collect()
603 }
604
605 pub fn approx_aggregate(
611 &self,
612 table: &str,
613 column: Option<&str>,
614 agg: ApproxAggKind,
615 z: f64,
616 ) -> Result<Option<ApproxAggregate>> {
617 let t = self
618 .schema
619 .tables
620 .iter()
621 .find(|t| t.name == table)
622 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
623 if matches!(agg, ApproxAggKind::Sum | ApproxAggKind::Avg) && column.is_none() {
624 return Err(KitError::Validation(
625 "approx sum/avg requires a column".into(),
626 ));
627 }
628 let cid = match column {
629 Some(name) => Some(
630 t.columns
631 .iter()
632 .find(|c| c.name == name)
633 .ok_or_else(|| KitError::Validation(format!("unknown column '{name}'")))?
634 .id as u16,
635 ),
636 None => None,
637 };
638 let core_agg = match agg {
639 ApproxAggKind::Count => ApproxAgg::Count,
640 ApproxAggKind::Sum => ApproxAgg::Sum,
641 ApproxAggKind::Avg => ApproxAgg::Avg,
642 };
643 let handle = self.inner.table(table).map_err(KitError::from)?;
644 let mut guard = handle.lock();
645 let res = guard
646 .approx_aggregate(&[], cid, core_agg, z)
647 .map_err(KitError::from)?;
648 Ok(res.map(|r| ApproxAggregate {
649 point: r.point,
650 ci_low: r.ci_low,
651 ci_high: r.ci_high,
652 n_population: r.n_population,
653 n_sample_live: r.n_sample_live,
654 n_passing: r.n_passing,
655 }))
656 }
657
658 pub fn scan_batched<F>(&self, table: &str, batch_size: usize, mut f: F) -> Result<()>
664 where
665 F: FnMut(&[serde_json::Map<String, Value>]) -> Result<()>,
666 {
667 let kit_t = self
668 .schema
669 .tables
670 .iter()
671 .find(|t| t.name == table)
672 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
673 let batch_size = batch_size.max(1);
674 let (snapshot, _pin) = self.inner.snapshot();
677 let handle = self.inner.table(table).map_err(KitError::from)?;
678 let guard = handle.lock();
679
680 let mut projection: Vec<(u16, mongreldb_core::schema::TypeId)> = Vec::new();
682 let mut meta: Vec<(String, mongreldb_kit_core::schema::ColumnType)> = Vec::new();
683 for c in &guard.schema().columns {
684 if let Some(kc) = kit_t.columns.iter().find(|kc| kc.id as u16 == c.id) {
685 projection.push((c.id, c.ty));
686 meta.push((kc.name.clone(), kc.storage_type));
687 }
688 }
689
690 match guard
691 .scan_cursor(snapshot, projection, &[])
692 .map_err(KitError::from)?
693 {
694 Some(mut cursor) => {
695 let mut buf: Vec<serde_json::Map<String, Value>> = Vec::with_capacity(batch_size);
696 while let Some(batch) = cursor.next_batch().map_err(KitError::from)? {
697 let nrows = batch.first().map(|c| c.len()).unwrap_or(0);
698 for j in 0..nrows {
699 let mut m = serde_json::Map::new();
700 for (ci, (name, ty)) in meta.iter().enumerate() {
701 let cv = batch
702 .get(ci)
703 .and_then(|col| col.value_at(j))
704 .unwrap_or(CoreValue::Null);
705 m.insert(name.clone(), crate::schema::core_to_json(&cv, *ty)?);
706 }
707 buf.push(m);
708 if buf.len() >= batch_size {
709 f(&buf)?;
710 buf.clear();
711 }
712 }
713 }
714 if !buf.is_empty() {
715 f(&buf)?;
716 }
717 Ok(())
718 }
719 None => {
720 drop(guard);
721 let rows = self.visible_core_rows_at(table, snapshot)?;
722 let maps: Vec<serde_json::Map<String, Value>> = rows
723 .iter()
724 .map(|r| crate::schema::core_row_to_json(r, kit_t).map(|row| row.values))
725 .collect::<Result<Vec<_>>>()?;
726 for chunk in maps.chunks(batch_size) {
727 f(chunk)?;
728 }
729 Ok(())
730 }
731 }
732 }
733
734 pub fn set_similarity(
743 &self,
744 table: &str,
745 column: &str,
746 query: &[String],
747 k: usize,
748 ) -> Result<Vec<SimilarRow>> {
749 let t = self
750 .schema
751 .tables
752 .iter()
753 .find(|t| t.name == table)
754 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
755 let col = t.columns.iter().find(|c| c.name == column).ok_or_else(|| {
756 KitError::Validation(format!("unknown column '{column}' on table '{table}'"))
757 })?;
758 let query_set: std::collections::HashSet<String> = query.iter().cloned().collect();
759
760 let has_minhash = t.indexes.iter().any(|idx| {
761 idx.kind == KitIndexKind::MinHash && idx.columns.iter().any(|c| c == column)
762 });
763 let rows = if has_minhash {
764 let query_hashes: Vec<u64> = query
766 .iter()
767 .map(|s| mongreldb_core::index::minhash_token_hash(s))
768 .collect();
769 let cand_k = k.saturating_mul(8).max(k + 64);
771 let cond = mongreldb_core::query::Condition::MinHashSimilar {
772 column_id: col.id as u16,
773 query: query_hashes,
774 k: cand_k,
775 };
776 let (snapshot, _pin) = self.inner.snapshot();
777 let core_rows = self.query_core_rows_at(table, &[cond], snapshot)?;
778 core_rows
779 .iter()
780 .map(|r| crate::schema::core_row_to_json(r, t))
781 .collect::<Result<Vec<_>>>()?
782 } else {
783 let tx = self.begin()?;
784 tx.all_rows(table)?
785 };
786
787 let mut scored: Vec<SimilarRow> = Vec::new();
788 for row in rows {
789 let set = parse_string_set(row.values.get(column));
790 let inter = set.iter().filter(|x| query_set.contains(*x)).count();
791 let union = set.len() + query_set.len() - inter;
792 let sim = if union == 0 {
793 0.0
794 } else {
795 inter as f64 / union as f64
796 };
797 if sim > 0.0 {
798 scored.push(SimilarRow {
799 row,
800 similarity: sim,
801 });
802 }
803 }
804 scored.sort_by(|a, b| {
805 b.similarity
806 .partial_cmp(&a.similarity)
807 .unwrap_or(std::cmp::Ordering::Equal)
808 });
809 scored.truncate(k);
810 Ok(scored)
811 }
812
813 pub fn flush(&self) -> Result<()> {
817 for name in self.inner.table_names() {
818 let handle = self.inner.table(&name).map_err(KitError::from)?;
819 let mut guard = handle.lock();
820 guard.flush().map_err(KitError::from)?;
821 }
822 Ok(())
823 }
824
825 pub fn incremental_aggregate(
837 &self,
838 table: &str,
839 column: Option<&str>,
840 agg: IncrementalAggKind,
841 filter: Option<&mongreldb_kit_core::query::Expr>,
842 ) -> Result<IncrementalAggregate> {
843 let t = self
844 .schema
845 .tables
846 .iter()
847 .find(|t| t.name == table)
848 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
849 if !matches!(agg, IncrementalAggKind::Count) && column.is_none() {
850 return Err(KitError::Validation(
851 "sum/min/max/avg incremental aggregate requires a column".into(),
852 ));
853 }
854 let cid = match column {
855 Some(name) => Some(
856 t.columns
857 .iter()
858 .find(|c| c.name == name)
859 .ok_or_else(|| KitError::Validation(format!("unknown column '{name}'")))?
860 .id as u16,
861 ),
862 None => None,
863 };
864 let conditions = match filter {
865 Some(expr) => {
866 let plan = crate::pushdown::translate_predicate(t, expr).ok_or_else(|| {
867 KitError::Validation(
868 "filter is not index-translatable for an incremental aggregate".into(),
869 )
870 })?;
871 if !plan.fully_translated {
872 return Err(KitError::Validation(
873 "filter has a residual that an incremental aggregate cannot apply exactly"
874 .into(),
875 ));
876 }
877 plan.conditions
878 }
879 None => Vec::new(),
880 };
881 let core_agg = match agg {
882 IncrementalAggKind::Count => NativeAgg::Count,
883 IncrementalAggKind::Sum => NativeAgg::Sum,
884 IncrementalAggKind::Min => NativeAgg::Min,
885 IncrementalAggKind::Max => NativeAgg::Max,
886 IncrementalAggKind::Avg => NativeAgg::Avg,
887 };
888 let cache_key = incremental_cache_key(t.id, cid, agg, &conditions);
889 let handle = self.inner.table(table).map_err(KitError::from)?;
890 let mut guard = handle.lock();
891 let res = guard
892 .aggregate_incremental(cache_key, &conditions, cid, core_agg)
893 .map_err(KitError::from)?;
894 Ok(IncrementalAggregate {
895 value: agg_state_value(&res.state),
896 incremental: res.incremental,
897 delta_rows: res.delta_rows,
898 })
899 }
900
901 pub fn applied_migrations(&self) -> Result<Vec<mongreldb_kit_core::migrations::Migration>> {
903 crate::migrate::load_applied_migrations(&self.inner)
904 }
905
906 pub(crate) fn core_db(&self) -> &CoreDatabase {
907 &self.inner
908 }
909
910 pub(crate) fn core_arc(&self) -> Arc<CoreDatabase> {
913 Arc::clone(&self.inner)
914 }
915
916 pub fn close(&self) -> Result<()> {
921 self.inner.close().map_err(KitError::from)
922 }
923
924 pub fn compact_all(&self) -> Result<(usize, usize)> {
929 self.inner.compact().map_err(KitError::from)
930 }
931
932 pub fn compact_table(&self, name: &str) -> Result<bool> {
935 self.inner.compact_table(name).map_err(KitError::from)
936 }
937
938 pub fn rename_table(&mut self, from: &str, to: &str) -> Result<()> {
948 if from.starts_with("__kit_") || to.starts_with("__kit_") {
949 return Err(KitError::Validation(
950 "rename_table: names beginning with '__kit_' are reserved for internal tables"
951 .into(),
952 ));
953 }
954 self.inner.rename_table(from, to).map_err(KitError::from)?;
955 if !self.schema.rename_table(from, to) {
958 return Err(KitError::Integrity(format!(
961 "rename_table: kit schema has no table '{from}' (or '{to}' already exists)"
962 )));
963 }
964 for table in &mut self.schema.tables {
965 for fk in &mut table.foreign_keys {
966 if fk.references_table == from {
967 fk.references_table = to.to_string();
968 }
969 }
970 }
971 store_schema(&self.root, &self.schema)?;
972 Ok(())
973 }
974
975 pub fn analyze(&self) -> Result<()> {
980 for name in self.inner.table_names() {
981 let handle = self.inner.table(&name).map_err(KitError::from)?;
982 handle.lock().ensure_indexes_complete()?;
983 }
984 Ok(())
985 }
986
987 pub fn vacuum(&self) -> Result<usize> {
991 self.inner.compact().map_err(KitError::from)?;
992 self.inner.gc().map_err(KitError::from)
993 }
994
995 pub fn create_view(&self, spec: &ViewSpec) -> Result<()> {
1001 self.sql(&spec.create_sql())?;
1002 Ok(())
1003 }
1004
1005 pub fn drop_view(&self, name: &str) -> Result<()> {
1007 self.sql(&format!("DROP VIEW IF EXISTS {name}"))?;
1008 Ok(())
1009 }
1010
1011 pub fn reserve_auto_inc(&self, table: &str) -> Result<Option<i64>> {
1019 let handle = self.inner.table(table).map_err(KitError::from)?;
1020 let mut guard = handle.lock();
1021 guard.reserve_auto_inc().map_err(KitError::from)
1022 }
1023
1024 pub fn create_user(&self, username: &str, password: &str) -> Result<()> {
1028 self.inner.create_user(username, password).map_err(KitError::from)?;
1029 Ok(())
1030 }
1031
1032 pub fn drop_user(&self, username: &str) -> Result<()> {
1034 self.inner.drop_user(username).map_err(KitError::from)
1035 }
1036
1037 pub fn alter_user_password(&self, username: &str, new_password: &str) -> Result<()> {
1039 self.inner
1040 .alter_user_password(username, new_password)
1041 .map_err(KitError::from)
1042 }
1043
1044 pub fn verify_user(
1046 &self,
1047 username: &str,
1048 password: &str,
1049 ) -> Result<Option<mongreldb_core::auth::UserEntry>> {
1050 self.inner.verify_user(username, password).map_err(KitError::from)
1051 }
1052
1053 pub fn set_user_admin(&self, username: &str, is_admin: bool) -> Result<()> {
1055 self.inner.set_user_admin(username, is_admin).map_err(KitError::from)
1056 }
1057
1058 pub fn users(&self) -> Vec<String> {
1060 self.inner.users().into_iter().map(|u| u.username).collect()
1061 }
1062
1063 pub fn create_role(&self, name: &str) -> Result<()> {
1065 self.inner.create_role(name).map_err(KitError::from)?;
1066 Ok(())
1067 }
1068
1069 pub fn drop_role(&self, name: &str) -> Result<()> {
1071 self.inner.drop_role(name).map_err(KitError::from)
1072 }
1073
1074 pub fn roles(&self) -> Vec<String> {
1076 self.inner.roles().into_iter().map(|r| r.name).collect()
1077 }
1078
1079 pub fn grant_role(&self, username: &str, role_name: &str) -> Result<()> {
1081 self.inner.grant_role(username, role_name).map_err(KitError::from)
1082 }
1083
1084 pub fn revoke_role(&self, username: &str, role_name: &str) -> Result<()> {
1086 self.inner.revoke_role(username, role_name).map_err(KitError::from)
1087 }
1088
1089 pub fn grant_permission(
1091 &self,
1092 role_name: &str,
1093 permission: mongreldb_core::auth::Permission,
1094 ) -> Result<()> {
1095 self.inner.grant_permission(role_name, permission).map_err(KitError::from)
1096 }
1097
1098 pub fn revoke_permission(
1100 &self,
1101 role_name: &str,
1102 permission: mongreldb_core::auth::Permission,
1103 ) -> Result<()> {
1104 self.inner.revoke_permission(role_name, permission).map_err(KitError::from)
1105 }
1106
1107 pub fn set_spill_threshold(&self, bytes: u64) {
1113 self.inner.set_spill_threshold(bytes);
1114 }
1115
1116 pub fn set_recursive_triggers(&self, enabled: bool) {
1118 self.inner.set_recursive_triggers(enabled);
1119 }
1120
1121 pub fn trigger_config(&self) -> mongreldb_core::TriggerConfig {
1123 self.inner.trigger_config()
1124 }
1125
1126 pub fn set_trigger_config(&self, config: mongreldb_core::TriggerConfig) -> Result<()> {
1128 self.inner
1129 .set_trigger_config(config)
1130 .map_err(KitError::from)
1131 }
1132
1133 pub fn set_table_compaction_zstd_level(&self, table: &str, level: i32) -> Result<()> {
1135 let handle = self.inner.table(table).map_err(KitError::from)?;
1136 handle.lock().set_compaction_zstd_level(level);
1137 Ok(())
1138 }
1139
1140 pub fn set_table_result_cache_max_bytes(&self, table: &str, max_bytes: u64) -> Result<()> {
1142 let handle = self.inner.table(table).map_err(KitError::from)?;
1143 handle.lock().set_result_cache_max_bytes(max_bytes);
1144 Ok(())
1145 }
1146
1147 pub fn set_table_mutable_run_spill_bytes(&self, table: &str, bytes: u64) -> Result<()> {
1149 let handle = self.inner.table(table).map_err(KitError::from)?;
1150 handle.lock().set_mutable_run_spill_bytes(bytes);
1151 Ok(())
1152 }
1153
1154 pub fn set_table_sync_byte_threshold(&self, table: &str, threshold: u64) -> Result<()> {
1156 let handle = self.inner.table(table).map_err(KitError::from)?;
1157 handle.lock().set_sync_byte_threshold(threshold);
1158 Ok(())
1159 }
1160
1161 pub fn set_table_index_build_policy(
1164 &self,
1165 table: &str,
1166 policy: mongreldb_core::IndexBuildPolicy,
1167 ) -> Result<()> {
1168 let handle = self.inner.table(table).map_err(KitError::from)?;
1169 handle.lock().set_index_build_policy(policy);
1170 Ok(())
1171 }
1172
1173 pub fn table_page_cache_stats(&self, table: &str) -> Result<mongreldb_core::cache::CacheStats> {
1175 let handle = self.inner.table(table).map_err(KitError::from)?;
1176 let stats = handle.lock().page_cache_stats();
1177 Ok(stats)
1178 }
1179
1180 pub fn table_run_count(&self, table: &str) -> Result<usize> {
1182 let handle = self.inner.table(table).map_err(KitError::from)?;
1183 let n = handle.lock().run_count();
1184 Ok(n)
1185 }
1186
1187 pub fn table_memtable_len(&self, table: &str) -> Result<usize> {
1189 let handle = self.inner.table(table).map_err(KitError::from)?;
1190 let n = handle.lock().memtable_len();
1191 Ok(n)
1192 }
1193
1194 pub fn table_mutable_run_len(&self, table: &str) -> Result<usize> {
1196 let handle = self.inner.table(table).map_err(KitError::from)?;
1197 let n = handle.lock().mutable_run_len();
1198 Ok(n)
1199 }
1200
1201 pub fn table_page_cache_len(&self, table: &str) -> Result<usize> {
1203 let handle = self.inner.table(table).map_err(KitError::from)?;
1204 let n = handle.lock().page_cache_len();
1205 Ok(n)
1206 }
1207
1208 pub fn table_decoded_cache_len(&self, table: &str) -> Result<usize> {
1210 let handle = self.inner.table(table).map_err(KitError::from)?;
1211 let n = handle.lock().decoded_cache_len();
1212 Ok(n)
1213 }
1214
1215 pub fn sql(&self, statement: &str) -> Result<Vec<arrow::record_batch::RecordBatch>> {
1233 let session = match self.session.lock().take() {
1239 Some(s) => s,
1240 None => {
1241 mongreldb_query::MongrelSession::open(self.core_arc()).map_err(KitError::from)?
1242 }
1243 };
1244 let runtime = sql_runtime();
1245 let result = runtime
1246 .block_on(session.run(statement))
1247 .map_err(KitError::from);
1248 *self.session.lock() = Some(session);
1250 result
1251 }
1252
1253 pub fn refresh_sql_session(&self) -> Result<()> {
1258 let session =
1259 mongreldb_query::MongrelSession::open(self.core_arc()).map_err(KitError::from)?;
1260 *self.session.lock() = Some(session);
1261 Ok(())
1262 }
1263
1264 pub fn sql_arrow(&self, statement: &str) -> Result<Vec<u8>> {
1270 let batches = self.sql(statement)?;
1271 crate::arrow_util::batches_to_ipc(&batches)
1272 }
1273
1274 pub fn sql_rows(&self, statement: &str) -> Result<Vec<serde_json::Map<String, Value>>> {
1278 let batches = self.sql(statement)?;
1279 crate::arrow_util::batches_to_rows(&batches)
1280 }
1281
1282 pub(crate) fn lookup_row_id(&self, table: &str, key: &[u8]) -> Result<Option<RowId>> {
1286 let handle = self.inner.table(table).map_err(KitError::from)?;
1287 let mut guard = handle.lock();
1288 guard.ensure_indexes_complete()?;
1289 Ok(guard.lookup_pk(key))
1290 }
1291
1292 pub(crate) fn root(&self) -> &Path {
1293 &self.root
1294 }
1295
1296 pub(crate) fn visible_core_rows_at(
1300 &self,
1301 table_name: &str,
1302 snapshot: Snapshot,
1303 ) -> Result<Vec<CoreRow>> {
1304 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1305 let guard = handle.lock();
1306 guard.visible_rows(snapshot).map_err(KitError::from)
1307 }
1308
1309 pub(crate) fn query_core_rows_at(
1316 &self,
1317 table_name: &str,
1318 conditions: &[mongreldb_core::query::Condition],
1319 snapshot: Snapshot,
1320 ) -> Result<Vec<CoreRow>> {
1321 if conditions.is_empty() {
1322 return self.visible_core_rows_at(table_name, snapshot);
1323 }
1324 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1325 let mut guard = handle.lock();
1326 let q = mongreldb_core::query::Query {
1327 conditions: conditions.to_vec(),
1328 };
1329 guard.query(&q).map_err(KitError::from)
1330 }
1331
1332 pub(crate) fn flush_table(&self, table_name: &str) -> Result<()> {
1340 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1341 handle.lock().flush().map_err(KitError::from)?;
1342 Ok(())
1343 }
1344
1345 pub(crate) fn count_core_rows_at(
1355 &self,
1356 table_name: &str,
1357 conditions: &[mongreldb_core::query::Condition],
1358 snapshot: Snapshot,
1359 ) -> Result<Option<u64>> {
1360 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1361 let mut guard = handle.lock();
1362 if guard.snapshot().epoch != snapshot.epoch {
1363 return Ok(None); }
1365 guard
1366 .count_conditions(conditions, snapshot)
1367 .map_err(KitError::from)
1368 }
1369
1370 pub(crate) fn aggregate_core_at(
1379 &self,
1380 table_name: &str,
1381 column: Option<u16>,
1382 conditions: &[mongreldb_core::query::Condition],
1383 agg: NativeAgg,
1384 snapshot: Snapshot,
1385 ) -> Result<Option<NativeAggResult>> {
1386 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1387 let guard = handle.lock();
1388 if guard.snapshot().epoch != snapshot.epoch {
1389 return Ok(None); }
1391 guard
1392 .aggregate_native(snapshot, column, conditions, agg)
1393 .map_err(KitError::from)
1394 }
1395
1396 pub(crate) fn count_distinct_core_at(
1405 &self,
1406 table_name: &str,
1407 column_id: u16,
1408 snapshot: Snapshot,
1409 ) -> Result<Option<u64>> {
1410 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1411 let mut guard = handle.lock();
1412 if guard.snapshot().epoch != snapshot.epoch {
1413 return Ok(None); }
1415 guard
1416 .count_distinct_from_bitmap(column_id)
1417 .map_err(KitError::from)
1418 }
1419
1420 #[allow(dead_code)]
1422 pub(crate) fn get_core_row(&self, table_name: &str, row_id: u64) -> Result<Option<CoreRow>> {
1423 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1424 let guard = handle.lock();
1425 let snapshot = guard.snapshot();
1426 Ok(guard.get(mongreldb_core::RowId(row_id), snapshot))
1427 }
1428}
1429
1430pub(crate) fn create_core_table(db: &CoreDatabase, name: &str, schema: CoreSchema) -> Result<()> {
1431 if db.table_id(name).is_ok() {
1432 return Ok(());
1433 }
1434 db.create_table(name, schema).map_err(KitError::from)?;
1435 Ok(())
1436}
1437
1438fn sql_runtime() -> &'static tokio::runtime::Runtime {
1443 use std::sync::OnceLock;
1444 static RT: OnceLock<tokio::runtime::Runtime> = OnceLock::new();
1445 RT.get_or_init(|| {
1446 tokio::runtime::Builder::new_current_thread()
1447 .enable_all()
1448 .build()
1449 .expect("failed to build kit SQL tokio runtime")
1450 })
1451}
1452
1453fn core_procedure(spec: &ProcedureSpec) -> Result<mongreldb_core::StoredProcedure> {
1454 let parsed: mongreldb_core::StoredProcedure =
1455 serde_json::from_value(spec.json.clone()).map_err(KitError::from)?;
1456 mongreldb_core::StoredProcedure::new(parsed.name, parsed.mode, parsed.params, parsed.body, 0)
1457 .map_err(KitError::from)
1458}
1459
1460fn core_trigger(spec: &TriggerSpec) -> Result<mongreldb_core::StoredTrigger> {
1461 let parsed: mongreldb_core::StoredTrigger =
1462 serde_json::from_value(spec.json.clone()).map_err(KitError::from)?;
1463 mongreldb_core::StoredTrigger::new(
1464 parsed.name,
1465 mongreldb_core::TriggerDefinition {
1466 target: parsed.target,
1467 timing: parsed.timing,
1468 event: parsed.event,
1469 update_of: parsed.update_of,
1470 target_columns: parsed.target_columns,
1471 when: parsed.when,
1472 program: parsed.program,
1473 },
1474 0,
1475 )
1476 .map_err(KitError::from)
1477}
1478
1479fn json_to_core_value(value: &Value) -> Result<CoreValue> {
1480 match value {
1481 Value::Null => Ok(CoreValue::Null),
1482 Value::Bool(value) => Ok(CoreValue::Bool(*value)),
1483 Value::Number(value) => {
1484 if let Some(value) = value.as_i64() {
1485 Ok(CoreValue::Int64(value))
1486 } else if let Some(value) = value.as_f64() {
1487 Ok(CoreValue::Float64(value))
1488 } else {
1489 Err(KitError::Validation("unsupported JSON number".into()))
1490 }
1491 }
1492 Value::String(value) => Ok(CoreValue::Bytes(value.as_bytes().to_vec())),
1493 Value::Array(_) | Value::Object(_) => Err(KitError::Validation(
1494 "procedure args only support scalar JSON values".into(),
1495 )),
1496 }
1497}
1498
1499pub(crate) fn internal_bytes(row: &CoreRow, col_id: u16) -> Option<String> {
1501 match row.columns.get(&col_id) {
1502 Some(CoreValue::Bytes(b)) => String::from_utf8(b.clone()).ok(),
1503 _ => None,
1504 }
1505}
1506
1507fn reap_rotated_wal_segments(db: &CoreDatabase) {
1522 let _ = db.gc();
1523}
1524
1525pub(crate) fn load_schema(path: &Path) -> Result<KitSchema> {
1526 let file = path.join(SCHEMA_FILE);
1527 let json = std::fs::read_to_string(&file)
1528 .map_err(|e| KitError::Migration(format!("cannot read schema file: {e}")))?;
1529 let schema: KitSchema = serde_json::from_str(&json)?;
1530 Ok(schema)
1531}
1532
1533pub(crate) fn store_schema(path: &Path, schema: &KitSchema) -> Result<()> {
1534 let file = path.join(SCHEMA_FILE);
1535 let json = serde_json::to_string_pretty(schema)?;
1536 std::fs::write(&file, json)?;
1537 Ok(())
1538}
1539
1540pub(crate) fn persist_schema(db: &Database, schema: &KitSchema) -> Result<()> {
1542 store_schema(&db.root, schema)
1543}