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 open_with_credentials(path: &Path, username: &str, password: &str) -> Result<Self> {
272 let inner = Arc::new(CoreDatabase::open_with_credentials(
273 path, username, password,
274 )?);
275 let schema = load_schema(path)?;
276 ensure_internal_tables(&inner)?;
277 reap_rotated_wal_segments(&inner);
278 Ok(Self {
279 inner,
280 schema,
281 root: path.to_path_buf(),
282 default_providers: HashMap::new(),
283 session: parking_lot::Mutex::new(None),
284 })
285 }
286
287 pub fn create_with_credentials(
293 path: &Path,
294 schema: KitSchema,
295 admin_username: &str,
296 admin_password: &str,
297 ) -> Result<Self> {
298 std::fs::create_dir_all(path)?;
299 let inner = Arc::new(CoreDatabase::create_with_credentials(
300 path,
301 admin_username,
302 admin_password,
303 )?);
304 ensure_internal_tables(&inner)?;
305 store_schema(path, &schema)?;
306 for table in &schema.tables {
307 create_core_table(&inner, &table.name, to_core_schema(table))?;
308 }
309 Ok(Self {
310 inner,
311 schema,
312 root: path.to_path_buf(),
313 default_providers: HashMap::new(),
314 session: parking_lot::Mutex::new(None),
315 })
316 }
317
318 pub fn open_encrypted_with_credentials(
321 path: &Path,
322 passphrase: &str,
323 username: &str,
324 password: &str,
325 ) -> Result<Self> {
326 let inner = Arc::new(CoreDatabase::open_encrypted_with_credentials(
327 path, passphrase, username, password,
328 )?);
329 let schema = load_schema(path)?;
330 ensure_internal_tables(&inner)?;
331 reap_rotated_wal_segments(&inner);
332 Ok(Self {
333 inner,
334 schema,
335 root: path.to_path_buf(),
336 default_providers: HashMap::new(),
337 session: parking_lot::Mutex::new(None),
338 })
339 }
340
341 pub fn create_encrypted_with_credentials(
345 path: &Path,
346 schema: KitSchema,
347 passphrase: &str,
348 admin_username: &str,
349 admin_password: &str,
350 ) -> Result<Self> {
351 std::fs::create_dir_all(path)?;
352 let inner = Arc::new(CoreDatabase::create_encrypted_with_credentials(
353 path,
354 passphrase,
355 admin_username,
356 admin_password,
357 )?);
358 ensure_internal_tables(&inner)?;
359 store_schema(path, &schema)?;
360 for table in &schema.tables {
361 create_core_table(&inner, &table.name, to_core_schema(table))?;
362 }
363 Ok(Self {
364 inner,
365 schema,
366 root: path.to_path_buf(),
367 default_providers: HashMap::new(),
368 session: parking_lot::Mutex::new(None),
369 })
370 }
371
372 pub fn enable_auth(&self, admin_username: &str, admin_password: &str) -> Result<()> {
376 self.inner
377 .enable_auth(admin_username, admin_password)
378 .map_err(KitError::from)
379 }
380
381 pub fn disable_auth(&self) -> Result<()> {
385 self.inner.disable_auth().map_err(KitError::from)
386 }
387
388 pub fn require_auth_enabled(&self) -> bool {
390 self.inner.require_auth_enabled()
391 }
392
393 pub fn refresh_principal(&self) -> Result<()> {
397 self.inner.refresh_principal().map_err(KitError::from)
398 }
399
400 pub fn register_default(
403 &mut self,
404 name: impl Into<String>,
405 provider: impl Fn() -> Value + Send + Sync + 'static,
406 ) {
407 self.default_providers
408 .insert(name.into(), Box::new(provider));
409 }
410
411 pub fn raw(&self) -> &CoreDatabase {
415 &self.inner
416 }
417
418 pub fn table_names(&self) -> Vec<String> {
420 self.schema
421 .tables
422 .iter()
423 .map(|t| t.name.clone())
424 .filter(|n| !n.starts_with("__kit_"))
425 .collect()
426 }
427
428 pub fn create_procedure(
429 &self,
430 spec: &ProcedureSpec,
431 ) -> Result<mongreldb_core::StoredProcedure> {
432 let procedure = core_procedure(spec)?;
433 self.inner
434 .create_procedure(procedure)
435 .map_err(KitError::from)
436 }
437
438 pub fn replace_procedure(
439 &self,
440 spec: &ProcedureSpec,
441 ) -> Result<mongreldb_core::StoredProcedure> {
442 let procedure = core_procedure(spec)?;
443 self.inner
444 .create_or_replace_procedure(procedure)
445 .map_err(KitError::from)
446 }
447
448 pub fn drop_procedure(&self, name: &str) -> Result<()> {
449 self.inner.drop_procedure(name).map_err(KitError::from)
450 }
451
452 pub fn call_procedure(
453 &self,
454 name: &str,
455 args: serde_json::Map<String, Value>,
456 ) -> Result<mongreldb_core::ProcedureCallResult> {
457 let args = args
458 .iter()
459 .map(|(key, value)| Ok((key.clone(), json_to_core_value(value)?)))
460 .collect::<Result<HashMap<_, _>>>()?;
461 self.inner
462 .call_procedure(name, args)
463 .map_err(KitError::from)
464 }
465
466 pub fn create_trigger(&self, spec: &TriggerSpec) -> Result<mongreldb_core::StoredTrigger> {
467 let trigger = core_trigger(spec)?;
468 self.inner.create_trigger(trigger).map_err(KitError::from)
469 }
470
471 pub fn replace_trigger(&self, spec: &TriggerSpec) -> Result<mongreldb_core::StoredTrigger> {
472 let trigger = core_trigger(spec)?;
473 self.inner
474 .create_or_replace_trigger(trigger)
475 .map_err(KitError::from)
476 }
477
478 pub fn drop_trigger(&self, name: &str) -> Result<()> {
479 self.inner.drop_trigger(name).map_err(KitError::from)
480 }
481
482 pub fn triggers(&self) -> Vec<mongreldb_core::StoredTrigger> {
483 self.inner.triggers()
484 }
485
486 pub fn trigger(&self, name: &str) -> Option<mongreldb_core::StoredTrigger> {
487 self.inner.trigger(name)
488 }
489
490 pub fn allocate_sequence(&self, name: &str, count: i64) -> Result<i64> {
496 use crate::internal::cols;
497 let mut attempt = 0;
498 loop {
499 let mut txn = self.inner.begin();
500 let snapshot = txn.read_snapshot();
501 let existing = self
502 .visible_core_rows_at(crate::internal::SEQUENCES, snapshot)?
503 .into_iter()
504 .find(|r| internal_bytes(r, cols::SEQ_NAME) == Some(name.to_string()));
505
506 let now = crate::internal::iso_now();
507 let (start, next, old_row_id) = match &existing {
511 Some(row) => {
512 let current = match row.columns.get(&cols::SEQ_NEXT) {
513 Some(CoreValue::Int64(i)) => *i,
514 _ => 1,
515 };
516 (current, current + count, Some(row.row_id))
517 }
518 None => (1, 1 + count, None),
519 };
520
521 if let Some(rid) = old_row_id {
522 txn.delete(crate::internal::SEQUENCES, rid)
523 .map_err(KitError::from)?;
524 }
525 txn.put(
526 crate::internal::SEQUENCES,
527 vec![
528 (cols::SEQ_NAME, CoreValue::Bytes(name.as_bytes().to_vec())),
529 (cols::SEQ_NEXT, CoreValue::Int64(next)),
530 (cols::SEQ_UPDATED, CoreValue::Bytes(now.into_bytes())),
531 ],
532 )
533 .map_err(KitError::from)?;
534 match txn.commit() {
535 Ok(_) => return Ok(start),
536 Err(mongreldb_core::MongrelError::Conflict(_)) if attempt < 10_000 => {
537 attempt += 1;
538 std::thread::yield_now();
539 continue;
540 }
541 Err(e) => return Err(KitError::from(e)),
542 }
543 }
544 }
545
546 pub fn transaction<T, F>(&self, max_retries: usize, mut f: F) -> Result<T>
549 where
550 F: FnMut(&mut crate::txn::Transaction<'_>) -> Result<T>,
551 {
552 let mut attempt = 0;
553 loop {
554 let mut txn = self.begin()?;
555 match f(&mut txn) {
556 Ok(value) => match txn.commit() {
557 Ok(()) => return Ok(value),
558 Err(KitError::Conflict(_)) if attempt < max_retries => {
559 attempt += 1;
560 continue;
561 }
562 Err(e) => return Err(e),
563 },
564 Err(KitError::Conflict(_)) if attempt < max_retries => {
565 txn.rollback();
566 attempt += 1;
567 continue;
568 }
569 Err(e) => {
570 txn.rollback();
571 return Err(e);
572 }
573 }
574 }
575 }
576
577 pub fn table(&self, name: &str) -> Option<&KitTable> {
579 self.schema.table(name)
580 }
581
582 pub fn schema(&self) -> &KitSchema {
584 &self.schema
585 }
586
587 pub fn begin(&self) -> Result<crate::txn::Transaction<'_>> {
589 let core_txn = self.inner.begin();
590 Ok(crate::txn::Transaction::new(self, core_txn))
591 }
592
593 pub fn set_schema(&mut self, schema: KitSchema) {
595 self.schema = schema;
596 }
597
598 pub fn check_internal_tables(&self) -> Result<()> {
601 let schema_file = self.root.join(SCHEMA_FILE);
602 if !schema_file.exists() {
603 return Err(KitError::Integrity(format!(
604 "schema file {} is missing",
605 schema_file.display()
606 )));
607 }
608 for (name, _) in internal_tables_core() {
609 if self.inner.table_id(name).is_err() {
610 return Err(KitError::Integrity(format!(
611 "internal table {name} is missing"
612 )));
613 }
614 }
615 Ok(())
616 }
617
618 pub fn gc(&self) -> Result<usize> {
621 self.inner.gc().map_err(KitError::from)
622 }
623
624 pub fn check(&self) -> Vec<serde_json::Value> {
627 self.inner
628 .check()
629 .into_iter()
630 .map(|i| {
631 serde_json::json!({
632 "table_id": i.table_id,
633 "table_name": i.table_name,
634 "severity": i.severity,
635 "description": i.description,
636 })
637 })
638 .collect()
639 }
640
641 pub fn doctor(&self) -> Result<Vec<u64>> {
643 self.inner.doctor().map_err(KitError::from)
644 }
645
646 pub fn snapshot_epoch(&self) -> u64 {
650 self.inner.snapshot().0.epoch.0
651 }
652
653 pub fn export_tsv(&self, table: &str) -> Result<String> {
657 let t = self
658 .schema
659 .tables
660 .iter()
661 .find(|t| t.name == table)
662 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?
663 .clone();
664 let tx = self.begin()?;
665 let rows = tx.all_rows(table)?;
666 Ok(crate::tsv::rows_to_tsv(&t, &rows))
667 }
668
669 pub fn import_tsv(&self, table: &str, text: &str) -> Result<usize> {
673 let t = self
674 .schema
675 .tables
676 .iter()
677 .find(|t| t.name == table)
678 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?
679 .clone();
680 let rows = crate::tsv::tsv_to_rows(&t, text)?;
681 let n = rows.len();
682 self.transaction(1, |tx| {
683 tx.insert_many(table, rows.clone())?;
684 Ok(())
685 })?;
686 Ok(n)
687 }
688
689 pub fn explain(
694 &self,
695 table: &str,
696 predicate: &mongreldb_kit_core::query::Expr,
697 ) -> Result<ExplainPlan> {
698 let t = self
699 .schema
700 .tables
701 .iter()
702 .find(|t| t.name == table)
703 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
704 Ok(match crate::pushdown::translate_predicate(t, predicate) {
705 Some(p) => ExplainPlan {
706 index_accelerated: p.can_push(),
707 exact: p.fully_translated,
708 pushed_conditions: p.conditions.iter().map(condition_label).collect(),
709 },
710 None => ExplainPlan {
711 index_accelerated: false,
712 exact: false,
713 pushed_conditions: Vec::new(),
714 },
715 })
716 }
717
718 pub fn rows_at_epoch(&self, table: &str, epoch: u64) -> Result<Vec<crate::schema::Row>> {
724 let t = self
725 .schema
726 .tables
727 .iter()
728 .find(|t| t.name == table)
729 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
730 let current = self.snapshot_epoch();
731 if epoch > current {
732 return Err(KitError::Validation(format!(
733 "epoch {epoch} is in the future (current committed epoch is {current})"
734 )));
735 }
736 let snap = Snapshot::at(mongreldb_core::epoch::Epoch(epoch));
737 let rows = self.visible_core_rows_at(table, snap)?;
738 rows.iter()
739 .map(|r| crate::schema::core_row_to_json(r, t))
740 .collect()
741 }
742
743 pub fn approx_aggregate(
749 &self,
750 table: &str,
751 column: Option<&str>,
752 agg: ApproxAggKind,
753 z: f64,
754 ) -> Result<Option<ApproxAggregate>> {
755 let t = self
756 .schema
757 .tables
758 .iter()
759 .find(|t| t.name == table)
760 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
761 if matches!(agg, ApproxAggKind::Sum | ApproxAggKind::Avg) && column.is_none() {
762 return Err(KitError::Validation(
763 "approx sum/avg requires a column".into(),
764 ));
765 }
766 let cid = match column {
767 Some(name) => Some(
768 t.columns
769 .iter()
770 .find(|c| c.name == name)
771 .ok_or_else(|| KitError::Validation(format!("unknown column '{name}'")))?
772 .id as u16,
773 ),
774 None => None,
775 };
776 let core_agg = match agg {
777 ApproxAggKind::Count => ApproxAgg::Count,
778 ApproxAggKind::Sum => ApproxAgg::Sum,
779 ApproxAggKind::Avg => ApproxAgg::Avg,
780 };
781 let handle = self.inner.table(table).map_err(KitError::from)?;
782 let mut guard = handle.lock();
783 let res = guard
784 .approx_aggregate(&[], cid, core_agg, z)
785 .map_err(KitError::from)?;
786 Ok(res.map(|r| ApproxAggregate {
787 point: r.point,
788 ci_low: r.ci_low,
789 ci_high: r.ci_high,
790 n_population: r.n_population,
791 n_sample_live: r.n_sample_live,
792 n_passing: r.n_passing,
793 }))
794 }
795
796 pub fn scan_batched<F>(&self, table: &str, batch_size: usize, mut f: F) -> Result<()>
802 where
803 F: FnMut(&[serde_json::Map<String, Value>]) -> Result<()>,
804 {
805 let kit_t = self
806 .schema
807 .tables
808 .iter()
809 .find(|t| t.name == table)
810 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
811 let batch_size = batch_size.max(1);
812 let (snapshot, _pin) = self.inner.snapshot();
815 let handle = self.inner.table(table).map_err(KitError::from)?;
816 let guard = handle.lock();
817
818 let mut projection: Vec<(u16, mongreldb_core::schema::TypeId)> = Vec::new();
820 let mut meta: Vec<(String, mongreldb_kit_core::schema::ColumnType)> = Vec::new();
821 for c in &guard.schema().columns {
822 if let Some(kc) = kit_t.columns.iter().find(|kc| kc.id as u16 == c.id) {
823 projection.push((c.id, c.ty));
824 meta.push((kc.name.clone(), kc.storage_type));
825 }
826 }
827
828 match guard
829 .scan_cursor(snapshot, projection, &[])
830 .map_err(KitError::from)?
831 {
832 Some(mut cursor) => {
833 let mut buf: Vec<serde_json::Map<String, Value>> = Vec::with_capacity(batch_size);
834 while let Some(batch) = cursor.next_batch().map_err(KitError::from)? {
835 let nrows = batch.first().map(|c| c.len()).unwrap_or(0);
836 for j in 0..nrows {
837 let mut m = serde_json::Map::new();
838 for (ci, (name, ty)) in meta.iter().enumerate() {
839 let cv = batch
840 .get(ci)
841 .and_then(|col| col.value_at(j))
842 .unwrap_or(CoreValue::Null);
843 m.insert(name.clone(), crate::schema::core_to_json(&cv, *ty)?);
844 }
845 buf.push(m);
846 if buf.len() >= batch_size {
847 f(&buf)?;
848 buf.clear();
849 }
850 }
851 }
852 if !buf.is_empty() {
853 f(&buf)?;
854 }
855 Ok(())
856 }
857 None => {
858 drop(guard);
859 let rows = self.visible_core_rows_at(table, snapshot)?;
860 let maps: Vec<serde_json::Map<String, Value>> = rows
861 .iter()
862 .map(|r| crate::schema::core_row_to_json(r, kit_t).map(|row| row.values))
863 .collect::<Result<Vec<_>>>()?;
864 for chunk in maps.chunks(batch_size) {
865 f(chunk)?;
866 }
867 Ok(())
868 }
869 }
870 }
871
872 pub fn set_similarity(
881 &self,
882 table: &str,
883 column: &str,
884 query: &[String],
885 k: usize,
886 ) -> Result<Vec<SimilarRow>> {
887 let t = self
888 .schema
889 .tables
890 .iter()
891 .find(|t| t.name == table)
892 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
893 let col = t.columns.iter().find(|c| c.name == column).ok_or_else(|| {
894 KitError::Validation(format!("unknown column '{column}' on table '{table}'"))
895 })?;
896 let query_set: std::collections::HashSet<String> = query.iter().cloned().collect();
897
898 let has_minhash = t.indexes.iter().any(|idx| {
899 idx.kind == KitIndexKind::MinHash && idx.columns.iter().any(|c| c == column)
900 });
901 let rows = if has_minhash {
902 let query_hashes: Vec<u64> = query
904 .iter()
905 .map(|s| mongreldb_core::index::minhash_token_hash(s))
906 .collect();
907 let cand_k = k.saturating_mul(8).max(k + 64);
909 let cond = mongreldb_core::query::Condition::MinHashSimilar {
910 column_id: col.id as u16,
911 query: query_hashes,
912 k: cand_k,
913 };
914 let (snapshot, _pin) = self.inner.snapshot();
915 let core_rows = self.query_core_rows_at(table, &[cond], snapshot)?;
916 core_rows
917 .iter()
918 .map(|r| crate::schema::core_row_to_json(r, t))
919 .collect::<Result<Vec<_>>>()?
920 } else {
921 let tx = self.begin()?;
922 tx.all_rows(table)?
923 };
924
925 let mut scored: Vec<SimilarRow> = Vec::new();
926 for row in rows {
927 let set = parse_string_set(row.values.get(column));
928 let inter = set.iter().filter(|x| query_set.contains(*x)).count();
929 let union = set.len() + query_set.len() - inter;
930 let sim = if union == 0 {
931 0.0
932 } else {
933 inter as f64 / union as f64
934 };
935 if sim > 0.0 {
936 scored.push(SimilarRow {
937 row,
938 similarity: sim,
939 });
940 }
941 }
942 scored.sort_by(|a, b| {
943 b.similarity
944 .partial_cmp(&a.similarity)
945 .unwrap_or(std::cmp::Ordering::Equal)
946 });
947 scored.truncate(k);
948 Ok(scored)
949 }
950
951 pub fn flush(&self) -> Result<()> {
955 for name in self.inner.table_names() {
956 let handle = self.inner.table(&name).map_err(KitError::from)?;
957 let mut guard = handle.lock();
958 guard.flush().map_err(KitError::from)?;
959 }
960 Ok(())
961 }
962
963 pub fn incremental_aggregate(
975 &self,
976 table: &str,
977 column: Option<&str>,
978 agg: IncrementalAggKind,
979 filter: Option<&mongreldb_kit_core::query::Expr>,
980 ) -> Result<IncrementalAggregate> {
981 let t = self
982 .schema
983 .tables
984 .iter()
985 .find(|t| t.name == table)
986 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
987 if !matches!(agg, IncrementalAggKind::Count) && column.is_none() {
988 return Err(KitError::Validation(
989 "sum/min/max/avg incremental aggregate requires a column".into(),
990 ));
991 }
992 let cid = match column {
993 Some(name) => Some(
994 t.columns
995 .iter()
996 .find(|c| c.name == name)
997 .ok_or_else(|| KitError::Validation(format!("unknown column '{name}'")))?
998 .id as u16,
999 ),
1000 None => None,
1001 };
1002 let conditions = match filter {
1003 Some(expr) => {
1004 let plan = crate::pushdown::translate_predicate(t, expr).ok_or_else(|| {
1005 KitError::Validation(
1006 "filter is not index-translatable for an incremental aggregate".into(),
1007 )
1008 })?;
1009 if !plan.fully_translated {
1010 return Err(KitError::Validation(
1011 "filter has a residual that an incremental aggregate cannot apply exactly"
1012 .into(),
1013 ));
1014 }
1015 plan.conditions
1016 }
1017 None => Vec::new(),
1018 };
1019 let core_agg = match agg {
1020 IncrementalAggKind::Count => NativeAgg::Count,
1021 IncrementalAggKind::Sum => NativeAgg::Sum,
1022 IncrementalAggKind::Min => NativeAgg::Min,
1023 IncrementalAggKind::Max => NativeAgg::Max,
1024 IncrementalAggKind::Avg => NativeAgg::Avg,
1025 };
1026 let cache_key = incremental_cache_key(t.id, cid, agg, &conditions);
1027 let handle = self.inner.table(table).map_err(KitError::from)?;
1028 let mut guard = handle.lock();
1029 let res = guard
1030 .aggregate_incremental(cache_key, &conditions, cid, core_agg)
1031 .map_err(KitError::from)?;
1032 Ok(IncrementalAggregate {
1033 value: agg_state_value(&res.state),
1034 incremental: res.incremental,
1035 delta_rows: res.delta_rows,
1036 })
1037 }
1038
1039 pub fn applied_migrations(&self) -> Result<Vec<mongreldb_kit_core::migrations::Migration>> {
1041 crate::migrate::load_applied_migrations(&self.inner)
1042 }
1043
1044 pub(crate) fn core_db(&self) -> &CoreDatabase {
1045 &self.inner
1046 }
1047
1048 pub(crate) fn core_arc(&self) -> Arc<CoreDatabase> {
1051 Arc::clone(&self.inner)
1052 }
1053
1054 pub fn close(&self) -> Result<()> {
1059 self.inner.close().map_err(KitError::from)
1060 }
1061
1062 pub fn compact_all(&self) -> Result<(usize, usize)> {
1067 self.inner.compact().map_err(KitError::from)
1068 }
1069
1070 pub fn compact_table(&self, name: &str) -> Result<bool> {
1073 self.inner.compact_table(name).map_err(KitError::from)
1074 }
1075
1076 pub fn rename_table(&mut self, from: &str, to: &str) -> Result<()> {
1086 if from.starts_with("__kit_") || to.starts_with("__kit_") {
1087 return Err(KitError::Validation(
1088 "rename_table: names beginning with '__kit_' are reserved for internal tables"
1089 .into(),
1090 ));
1091 }
1092 self.inner.rename_table(from, to).map_err(KitError::from)?;
1093 if !self.schema.rename_table(from, to) {
1096 return Err(KitError::Integrity(format!(
1099 "rename_table: kit schema has no table '{from}' (or '{to}' already exists)"
1100 )));
1101 }
1102 for table in &mut self.schema.tables {
1103 for fk in &mut table.foreign_keys {
1104 if fk.references_table == from {
1105 fk.references_table = to.to_string();
1106 }
1107 }
1108 }
1109 store_schema(&self.root, &self.schema)?;
1110 Ok(())
1111 }
1112
1113 pub fn analyze(&self) -> Result<()> {
1118 for name in self.inner.table_names() {
1119 let handle = self.inner.table(&name).map_err(KitError::from)?;
1120 handle.lock().ensure_indexes_complete()?;
1121 }
1122 Ok(())
1123 }
1124
1125 pub fn vacuum(&self) -> Result<usize> {
1129 self.inner.compact().map_err(KitError::from)?;
1130 self.inner.gc().map_err(KitError::from)
1131 }
1132
1133 pub fn create_view(&self, spec: &ViewSpec) -> Result<()> {
1139 self.sql(&spec.create_sql())?;
1140 Ok(())
1141 }
1142
1143 pub fn drop_view(&self, name: &str) -> Result<()> {
1145 self.sql(&format!("DROP VIEW IF EXISTS {name}"))?;
1146 Ok(())
1147 }
1148
1149 pub fn reserve_auto_inc(&self, table: &str) -> Result<Option<i64>> {
1157 let handle = self.inner.table(table).map_err(KitError::from)?;
1158 let mut guard = handle.lock();
1159 guard.reserve_auto_inc().map_err(KitError::from)
1160 }
1161
1162 pub fn create_user(&self, username: &str, password: &str) -> Result<()> {
1166 self.inner
1167 .create_user(username, password)
1168 .map_err(KitError::from)?;
1169 Ok(())
1170 }
1171
1172 pub fn drop_user(&self, username: &str) -> Result<()> {
1174 self.inner.drop_user(username).map_err(KitError::from)
1175 }
1176
1177 pub fn alter_user_password(&self, username: &str, new_password: &str) -> Result<()> {
1179 self.inner
1180 .alter_user_password(username, new_password)
1181 .map_err(KitError::from)
1182 }
1183
1184 pub fn verify_user(
1186 &self,
1187 username: &str,
1188 password: &str,
1189 ) -> Result<Option<mongreldb_core::auth::UserEntry>> {
1190 self.inner
1191 .verify_user(username, password)
1192 .map_err(KitError::from)
1193 }
1194
1195 pub fn set_user_admin(&self, username: &str, is_admin: bool) -> Result<()> {
1197 self.inner
1198 .set_user_admin(username, is_admin)
1199 .map_err(KitError::from)
1200 }
1201
1202 pub fn users(&self) -> Vec<String> {
1204 self.inner.users().into_iter().map(|u| u.username).collect()
1205 }
1206
1207 pub fn create_role(&self, name: &str) -> Result<()> {
1209 self.inner.create_role(name).map_err(KitError::from)?;
1210 Ok(())
1211 }
1212
1213 pub fn drop_role(&self, name: &str) -> Result<()> {
1215 self.inner.drop_role(name).map_err(KitError::from)
1216 }
1217
1218 pub fn roles(&self) -> Vec<String> {
1220 self.inner.roles().into_iter().map(|r| r.name).collect()
1221 }
1222
1223 pub fn grant_role(&self, username: &str, role_name: &str) -> Result<()> {
1225 self.inner
1226 .grant_role(username, role_name)
1227 .map_err(KitError::from)
1228 }
1229
1230 pub fn revoke_role(&self, username: &str, role_name: &str) -> Result<()> {
1232 self.inner
1233 .revoke_role(username, role_name)
1234 .map_err(KitError::from)
1235 }
1236
1237 pub fn grant_permission(
1239 &self,
1240 role_name: &str,
1241 permission: mongreldb_core::auth::Permission,
1242 ) -> Result<()> {
1243 self.inner
1244 .grant_permission(role_name, permission)
1245 .map_err(KitError::from)
1246 }
1247
1248 pub fn revoke_permission(
1250 &self,
1251 role_name: &str,
1252 permission: mongreldb_core::auth::Permission,
1253 ) -> Result<()> {
1254 self.inner
1255 .revoke_permission(role_name, permission)
1256 .map_err(KitError::from)
1257 }
1258
1259 pub fn set_spill_threshold(&self, bytes: u64) {
1265 self.inner.set_spill_threshold(bytes);
1266 }
1267
1268 pub fn set_recursive_triggers(&self, enabled: bool) {
1270 self.inner.set_recursive_triggers(enabled);
1271 }
1272
1273 pub fn trigger_config(&self) -> mongreldb_core::TriggerConfig {
1275 self.inner.trigger_config()
1276 }
1277
1278 pub fn set_trigger_config(&self, config: mongreldb_core::TriggerConfig) -> Result<()> {
1280 self.inner
1281 .set_trigger_config(config)
1282 .map_err(KitError::from)
1283 }
1284
1285 pub fn set_table_compaction_zstd_level(&self, table: &str, level: i32) -> Result<()> {
1287 let handle = self.inner.table(table).map_err(KitError::from)?;
1288 handle.lock().set_compaction_zstd_level(level);
1289 Ok(())
1290 }
1291
1292 pub fn set_table_result_cache_max_bytes(&self, table: &str, max_bytes: u64) -> Result<()> {
1294 let handle = self.inner.table(table).map_err(KitError::from)?;
1295 handle.lock().set_result_cache_max_bytes(max_bytes);
1296 Ok(())
1297 }
1298
1299 pub fn set_table_mutable_run_spill_bytes(&self, table: &str, bytes: u64) -> Result<()> {
1301 let handle = self.inner.table(table).map_err(KitError::from)?;
1302 handle.lock().set_mutable_run_spill_bytes(bytes);
1303 Ok(())
1304 }
1305
1306 pub fn set_table_sync_byte_threshold(&self, table: &str, threshold: u64) -> Result<()> {
1308 let handle = self.inner.table(table).map_err(KitError::from)?;
1309 handle.lock().set_sync_byte_threshold(threshold);
1310 Ok(())
1311 }
1312
1313 pub fn set_table_index_build_policy(
1316 &self,
1317 table: &str,
1318 policy: mongreldb_core::IndexBuildPolicy,
1319 ) -> Result<()> {
1320 let handle = self.inner.table(table).map_err(KitError::from)?;
1321 handle.lock().set_index_build_policy(policy);
1322 Ok(())
1323 }
1324
1325 pub fn table_page_cache_stats(&self, table: &str) -> Result<mongreldb_core::cache::CacheStats> {
1327 let handle = self.inner.table(table).map_err(KitError::from)?;
1328 let stats = handle.lock().page_cache_stats();
1329 Ok(stats)
1330 }
1331
1332 pub fn table_run_count(&self, table: &str) -> Result<usize> {
1334 let handle = self.inner.table(table).map_err(KitError::from)?;
1335 let n = handle.lock().run_count();
1336 Ok(n)
1337 }
1338
1339 pub fn table_memtable_len(&self, table: &str) -> Result<usize> {
1341 let handle = self.inner.table(table).map_err(KitError::from)?;
1342 let n = handle.lock().memtable_len();
1343 Ok(n)
1344 }
1345
1346 pub fn table_mutable_run_len(&self, table: &str) -> Result<usize> {
1348 let handle = self.inner.table(table).map_err(KitError::from)?;
1349 let n = handle.lock().mutable_run_len();
1350 Ok(n)
1351 }
1352
1353 pub fn table_page_cache_len(&self, table: &str) -> Result<usize> {
1355 let handle = self.inner.table(table).map_err(KitError::from)?;
1356 let n = handle.lock().page_cache_len();
1357 Ok(n)
1358 }
1359
1360 pub fn table_decoded_cache_len(&self, table: &str) -> Result<usize> {
1362 let handle = self.inner.table(table).map_err(KitError::from)?;
1363 let n = handle.lock().decoded_cache_len();
1364 Ok(n)
1365 }
1366
1367 pub fn sql(&self, statement: &str) -> Result<Vec<arrow::record_batch::RecordBatch>> {
1385 let session = match self.session.lock().take() {
1391 Some(s) => s,
1392 None => {
1393 mongreldb_query::MongrelSession::open(self.core_arc()).map_err(KitError::from)?
1394 }
1395 };
1396 let runtime = sql_runtime();
1397 let result = runtime
1398 .block_on(session.run(statement))
1399 .map_err(KitError::from);
1400 *self.session.lock() = Some(session);
1402 result
1403 }
1404
1405 pub fn refresh_sql_session(&self) -> Result<()> {
1410 let session =
1411 mongreldb_query::MongrelSession::open(self.core_arc()).map_err(KitError::from)?;
1412 *self.session.lock() = Some(session);
1413 Ok(())
1414 }
1415
1416 pub fn sql_arrow(&self, statement: &str) -> Result<Vec<u8>> {
1422 let batches = self.sql(statement)?;
1423 crate::arrow_util::batches_to_ipc(&batches)
1424 }
1425
1426 pub fn sql_rows(&self, statement: &str) -> Result<Vec<serde_json::Map<String, Value>>> {
1430 let batches = self.sql(statement)?;
1431 crate::arrow_util::batches_to_rows(&batches)
1432 }
1433
1434 pub(crate) fn lookup_row_id(&self, table: &str, key: &[u8]) -> Result<Option<RowId>> {
1438 let handle = self.inner.table(table).map_err(KitError::from)?;
1439 let mut guard = handle.lock();
1440 guard.ensure_indexes_complete()?;
1441 Ok(guard.lookup_pk(key))
1442 }
1443
1444 pub(crate) fn root(&self) -> &Path {
1445 &self.root
1446 }
1447
1448 pub(crate) fn visible_core_rows_at(
1452 &self,
1453 table_name: &str,
1454 snapshot: Snapshot,
1455 ) -> Result<Vec<CoreRow>> {
1456 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1457 let guard = handle.lock();
1458 guard.visible_rows(snapshot).map_err(KitError::from)
1459 }
1460
1461 pub(crate) fn query_core_rows_at(
1468 &self,
1469 table_name: &str,
1470 conditions: &[mongreldb_core::query::Condition],
1471 snapshot: Snapshot,
1472 ) -> Result<Vec<CoreRow>> {
1473 if conditions.is_empty() {
1474 return self.visible_core_rows_at(table_name, snapshot);
1475 }
1476 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1477 let mut guard = handle.lock();
1478 let q = mongreldb_core::query::Query {
1479 conditions: conditions.to_vec(),
1480 };
1481 guard.query(&q).map_err(KitError::from)
1482 }
1483
1484 pub(crate) fn flush_table(&self, table_name: &str) -> Result<()> {
1492 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1493 handle.lock().flush().map_err(KitError::from)?;
1494 Ok(())
1495 }
1496
1497 pub(crate) fn count_core_rows_at(
1507 &self,
1508 table_name: &str,
1509 conditions: &[mongreldb_core::query::Condition],
1510 snapshot: Snapshot,
1511 ) -> Result<Option<u64>> {
1512 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1513 let mut guard = handle.lock();
1514 if guard.snapshot().epoch != snapshot.epoch {
1515 return Ok(None); }
1517 guard
1518 .count_conditions(conditions, snapshot)
1519 .map_err(KitError::from)
1520 }
1521
1522 pub(crate) fn aggregate_core_at(
1531 &self,
1532 table_name: &str,
1533 column: Option<u16>,
1534 conditions: &[mongreldb_core::query::Condition],
1535 agg: NativeAgg,
1536 snapshot: Snapshot,
1537 ) -> Result<Option<NativeAggResult>> {
1538 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1539 let guard = handle.lock();
1540 if guard.snapshot().epoch != snapshot.epoch {
1541 return Ok(None); }
1543 guard
1544 .aggregate_native(snapshot, column, conditions, agg)
1545 .map_err(KitError::from)
1546 }
1547
1548 pub(crate) fn count_distinct_core_at(
1557 &self,
1558 table_name: &str,
1559 column_id: u16,
1560 snapshot: Snapshot,
1561 ) -> Result<Option<u64>> {
1562 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1563 let mut guard = handle.lock();
1564 if guard.snapshot().epoch != snapshot.epoch {
1565 return Ok(None); }
1567 guard
1568 .count_distinct_from_bitmap(column_id)
1569 .map_err(KitError::from)
1570 }
1571
1572 #[allow(dead_code)]
1574 pub(crate) fn get_core_row(&self, table_name: &str, row_id: u64) -> Result<Option<CoreRow>> {
1575 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1576 let guard = handle.lock();
1577 let snapshot = guard.snapshot();
1578 Ok(guard.get(mongreldb_core::RowId(row_id), snapshot))
1579 }
1580}
1581
1582pub(crate) fn create_core_table(db: &CoreDatabase, name: &str, schema: CoreSchema) -> Result<()> {
1583 if db.table_id(name).is_ok() {
1584 return Ok(());
1585 }
1586 db.create_table(name, schema).map_err(KitError::from)?;
1587 Ok(())
1588}
1589
1590fn sql_runtime() -> &'static tokio::runtime::Runtime {
1595 use std::sync::OnceLock;
1596 static RT: OnceLock<tokio::runtime::Runtime> = OnceLock::new();
1597 RT.get_or_init(|| {
1598 tokio::runtime::Builder::new_current_thread()
1599 .enable_all()
1600 .build()
1601 .expect("failed to build kit SQL tokio runtime")
1602 })
1603}
1604
1605fn core_procedure(spec: &ProcedureSpec) -> Result<mongreldb_core::StoredProcedure> {
1606 let parsed: mongreldb_core::StoredProcedure =
1607 serde_json::from_value(spec.json.clone()).map_err(KitError::from)?;
1608 mongreldb_core::StoredProcedure::new(parsed.name, parsed.mode, parsed.params, parsed.body, 0)
1609 .map_err(KitError::from)
1610}
1611
1612fn core_trigger(spec: &TriggerSpec) -> Result<mongreldb_core::StoredTrigger> {
1613 let parsed: mongreldb_core::StoredTrigger =
1614 serde_json::from_value(spec.json.clone()).map_err(KitError::from)?;
1615 mongreldb_core::StoredTrigger::new(
1616 parsed.name,
1617 mongreldb_core::TriggerDefinition {
1618 target: parsed.target,
1619 timing: parsed.timing,
1620 event: parsed.event,
1621 update_of: parsed.update_of,
1622 target_columns: parsed.target_columns,
1623 when: parsed.when,
1624 program: parsed.program,
1625 },
1626 0,
1627 )
1628 .map_err(KitError::from)
1629}
1630
1631fn json_to_core_value(value: &Value) -> Result<CoreValue> {
1632 match value {
1633 Value::Null => Ok(CoreValue::Null),
1634 Value::Bool(value) => Ok(CoreValue::Bool(*value)),
1635 Value::Number(value) => {
1636 if let Some(value) = value.as_i64() {
1637 Ok(CoreValue::Int64(value))
1638 } else if let Some(value) = value.as_f64() {
1639 Ok(CoreValue::Float64(value))
1640 } else {
1641 Err(KitError::Validation("unsupported JSON number".into()))
1642 }
1643 }
1644 Value::String(value) => Ok(CoreValue::Bytes(value.as_bytes().to_vec())),
1645 Value::Array(_) | Value::Object(_) => Err(KitError::Validation(
1646 "procedure args only support scalar JSON values".into(),
1647 )),
1648 }
1649}
1650
1651pub(crate) fn internal_bytes(row: &CoreRow, col_id: u16) -> Option<String> {
1653 match row.columns.get(&col_id) {
1654 Some(CoreValue::Bytes(b)) => String::from_utf8(b.clone()).ok(),
1655 _ => None,
1656 }
1657}
1658
1659fn reap_rotated_wal_segments(db: &CoreDatabase) {
1674 let _ = db.gc();
1675}
1676
1677pub(crate) fn load_schema(path: &Path) -> Result<KitSchema> {
1678 let file = path.join(SCHEMA_FILE);
1679 let json = std::fs::read_to_string(&file)
1680 .map_err(|e| KitError::Migration(format!("cannot read schema file: {e}")))?;
1681 let schema: KitSchema = serde_json::from_str(&json)?;
1682 Ok(schema)
1683}
1684
1685pub(crate) fn store_schema(path: &Path, schema: &KitSchema) -> Result<()> {
1686 let file = path.join(SCHEMA_FILE);
1687 let json = serde_json::to_string_pretty(schema)?;
1688 std::fs::write(&file, json)?;
1689 Ok(())
1690}
1691
1692pub(crate) fn persist_schema(db: &Database, schema: &KitSchema) -> Result<()> {
1694 store_schema(&db.root, schema)
1695}