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 *self.session.lock() = None;
401 Ok(())
402 }
403
404 pub fn register_default(
407 &mut self,
408 name: impl Into<String>,
409 provider: impl Fn() -> Value + Send + Sync + 'static,
410 ) {
411 self.default_providers
412 .insert(name.into(), Box::new(provider));
413 }
414
415 pub fn raw(&self) -> &CoreDatabase {
419 &self.inner
420 }
421
422 pub fn table_names(&self) -> Vec<String> {
424 self.schema
425 .tables
426 .iter()
427 .map(|t| t.name.clone())
428 .filter(|n| !n.starts_with("__kit_"))
429 .collect()
430 }
431
432 pub fn create_procedure(
433 &self,
434 spec: &ProcedureSpec,
435 ) -> Result<mongreldb_core::StoredProcedure> {
436 let procedure = core_procedure(spec)?;
437 self.inner
438 .create_procedure(procedure)
439 .map_err(KitError::from)
440 }
441
442 pub fn replace_procedure(
443 &self,
444 spec: &ProcedureSpec,
445 ) -> Result<mongreldb_core::StoredProcedure> {
446 let procedure = core_procedure(spec)?;
447 self.inner
448 .create_or_replace_procedure(procedure)
449 .map_err(KitError::from)
450 }
451
452 pub fn drop_procedure(&self, name: &str) -> Result<()> {
453 self.inner.drop_procedure(name).map_err(KitError::from)
454 }
455
456 pub fn call_procedure(
457 &self,
458 name: &str,
459 args: serde_json::Map<String, Value>,
460 ) -> Result<mongreldb_core::ProcedureCallResult> {
461 let args = args
462 .iter()
463 .map(|(key, value)| Ok((key.clone(), json_to_core_value(value)?)))
464 .collect::<Result<HashMap<_, _>>>()?;
465 self.inner
466 .call_procedure(name, args)
467 .map_err(KitError::from)
468 }
469
470 pub fn create_trigger(&self, spec: &TriggerSpec) -> Result<mongreldb_core::StoredTrigger> {
471 let trigger = core_trigger(spec)?;
472 self.inner.create_trigger(trigger).map_err(KitError::from)
473 }
474
475 pub fn replace_trigger(&self, spec: &TriggerSpec) -> Result<mongreldb_core::StoredTrigger> {
476 let trigger = core_trigger(spec)?;
477 self.inner
478 .create_or_replace_trigger(trigger)
479 .map_err(KitError::from)
480 }
481
482 pub fn drop_trigger(&self, name: &str) -> Result<()> {
483 self.inner.drop_trigger(name).map_err(KitError::from)
484 }
485
486 pub fn triggers(&self) -> Vec<mongreldb_core::StoredTrigger> {
487 self.inner.triggers()
488 }
489
490 pub fn trigger(&self, name: &str) -> Option<mongreldb_core::StoredTrigger> {
491 self.inner.trigger(name)
492 }
493
494 pub fn allocate_sequence(&self, name: &str, count: i64) -> Result<i64> {
500 use crate::internal::cols;
501 let mut attempt = 0;
502 loop {
503 let mut txn = self.inner.begin();
504 let snapshot = txn.read_snapshot();
505 let existing = self
506 .visible_core_rows_at(crate::internal::SEQUENCES, snapshot)?
507 .into_iter()
508 .find(|r| internal_bytes(r, cols::SEQ_NAME) == Some(name.to_string()));
509
510 let now = crate::internal::iso_now();
511 let (start, next, old_row_id) = match &existing {
515 Some(row) => {
516 let current = match row.columns.get(&cols::SEQ_NEXT) {
517 Some(CoreValue::Int64(i)) => *i,
518 _ => 1,
519 };
520 (current, current + count, Some(row.row_id))
521 }
522 None => (1, 1 + count, None),
523 };
524
525 if let Some(rid) = old_row_id {
526 txn.delete(crate::internal::SEQUENCES, rid)
527 .map_err(KitError::from)?;
528 }
529 txn.put(
530 crate::internal::SEQUENCES,
531 vec![
532 (cols::SEQ_NAME, CoreValue::Bytes(name.as_bytes().to_vec())),
533 (cols::SEQ_NEXT, CoreValue::Int64(next)),
534 (cols::SEQ_UPDATED, CoreValue::Bytes(now.into_bytes())),
535 ],
536 )
537 .map_err(KitError::from)?;
538 match txn.commit() {
539 Ok(_) => return Ok(start),
540 Err(mongreldb_core::MongrelError::Conflict(_)) if attempt < 10_000 => {
541 attempt += 1;
542 std::thread::yield_now();
543 continue;
544 }
545 Err(e) => return Err(KitError::from(e)),
546 }
547 }
548 }
549
550 pub fn transaction<T, F>(&self, max_retries: usize, mut f: F) -> Result<T>
553 where
554 F: FnMut(&mut crate::txn::Transaction<'_>) -> Result<T>,
555 {
556 let mut attempt = 0;
557 loop {
558 let mut txn = self.begin()?;
559 match f(&mut txn) {
560 Ok(value) => match txn.commit() {
561 Ok(()) => return Ok(value),
562 Err(KitError::Conflict(_)) if attempt < max_retries => {
563 attempt += 1;
564 continue;
565 }
566 Err(e) => return Err(e),
567 },
568 Err(KitError::Conflict(_)) if attempt < max_retries => {
569 txn.rollback();
570 attempt += 1;
571 continue;
572 }
573 Err(e) => {
574 txn.rollback();
575 return Err(e);
576 }
577 }
578 }
579 }
580
581 pub fn table(&self, name: &str) -> Option<&KitTable> {
583 self.schema.table(name)
584 }
585
586 pub fn schema(&self) -> &KitSchema {
588 &self.schema
589 }
590
591 pub fn begin(&self) -> Result<crate::txn::Transaction<'_>> {
593 let core_txn = self.inner.begin();
594 Ok(crate::txn::Transaction::new(self, core_txn))
595 }
596
597 pub fn set_schema(&mut self, schema: KitSchema) {
599 self.schema = schema;
600 }
601
602 pub fn check_internal_tables(&self) -> Result<()> {
605 let schema_file = self.root.join(SCHEMA_FILE);
606 if !schema_file.exists() {
607 return Err(KitError::Integrity(format!(
608 "schema file {} is missing",
609 schema_file.display()
610 )));
611 }
612 for (name, _) in internal_tables_core() {
613 if self.inner.table_id(name).is_err() {
614 return Err(KitError::Integrity(format!(
615 "internal table {name} is missing"
616 )));
617 }
618 }
619 Ok(())
620 }
621
622 pub fn gc(&self) -> Result<usize> {
625 self.inner.gc().map_err(KitError::from)
626 }
627
628 pub fn check(&self) -> Vec<serde_json::Value> {
631 self.inner
632 .check()
633 .into_iter()
634 .map(|i| {
635 serde_json::json!({
636 "table_id": i.table_id,
637 "table_name": i.table_name,
638 "severity": i.severity,
639 "description": i.description,
640 })
641 })
642 .collect()
643 }
644
645 pub fn doctor(&self) -> Result<Vec<u64>> {
647 self.inner.doctor().map_err(KitError::from)
648 }
649
650 pub fn snapshot_epoch(&self) -> u64 {
654 self.inner.snapshot().0.epoch.0
655 }
656
657 pub fn export_tsv(&self, table: &str) -> Result<String> {
661 let t = self
662 .schema
663 .tables
664 .iter()
665 .find(|t| t.name == table)
666 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?
667 .clone();
668 let tx = self.begin()?;
669 let rows = tx.all_rows(table)?;
670 Ok(crate::tsv::rows_to_tsv(&t, &rows))
671 }
672
673 pub fn import_tsv(&self, table: &str, text: &str) -> Result<usize> {
677 let t = self
678 .schema
679 .tables
680 .iter()
681 .find(|t| t.name == table)
682 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?
683 .clone();
684 let rows = crate::tsv::tsv_to_rows(&t, text)?;
685 let n = rows.len();
686 self.transaction(1, |tx| {
687 tx.insert_many(table, rows.clone())?;
688 Ok(())
689 })?;
690 Ok(n)
691 }
692
693 pub fn explain(
698 &self,
699 table: &str,
700 predicate: &mongreldb_kit_core::query::Expr,
701 ) -> Result<ExplainPlan> {
702 let t = self
703 .schema
704 .tables
705 .iter()
706 .find(|t| t.name == table)
707 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
708 Ok(match crate::pushdown::translate_predicate(t, predicate) {
709 Some(p) => ExplainPlan {
710 index_accelerated: p.can_push(),
711 exact: p.fully_translated,
712 pushed_conditions: p.conditions.iter().map(condition_label).collect(),
713 },
714 None => ExplainPlan {
715 index_accelerated: false,
716 exact: false,
717 pushed_conditions: Vec::new(),
718 },
719 })
720 }
721
722 pub fn rows_at_epoch(&self, table: &str, epoch: u64) -> Result<Vec<crate::schema::Row>> {
728 let t = self
729 .schema
730 .tables
731 .iter()
732 .find(|t| t.name == table)
733 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
734 let current = self.snapshot_epoch();
735 if epoch > current {
736 return Err(KitError::Validation(format!(
737 "epoch {epoch} is in the future (current committed epoch is {current})"
738 )));
739 }
740 let snap = Snapshot::at(mongreldb_core::epoch::Epoch(epoch));
741 let rows = self.visible_core_rows_at(table, snap)?;
742 rows.iter()
743 .map(|r| crate::schema::core_row_to_json(r, t))
744 .collect()
745 }
746
747 pub fn approx_aggregate(
753 &self,
754 table: &str,
755 column: Option<&str>,
756 agg: ApproxAggKind,
757 z: f64,
758 ) -> Result<Option<ApproxAggregate>> {
759 let t = self
760 .schema
761 .tables
762 .iter()
763 .find(|t| t.name == table)
764 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
765 if matches!(agg, ApproxAggKind::Sum | ApproxAggKind::Avg) && column.is_none() {
766 return Err(KitError::Validation(
767 "approx sum/avg requires a column".into(),
768 ));
769 }
770 let cid = match column {
771 Some(name) => Some(
772 t.columns
773 .iter()
774 .find(|c| c.name == name)
775 .ok_or_else(|| KitError::Validation(format!("unknown column '{name}'")))?
776 .id as u16,
777 ),
778 None => None,
779 };
780 let core_agg = match agg {
781 ApproxAggKind::Count => ApproxAgg::Count,
782 ApproxAggKind::Sum => ApproxAgg::Sum,
783 ApproxAggKind::Avg => ApproxAgg::Avg,
784 };
785 let handle = self.inner.table(table).map_err(KitError::from)?;
786 let mut guard = handle.lock();
787 let res = guard
788 .approx_aggregate(&[], cid, core_agg, z)
789 .map_err(KitError::from)?;
790 Ok(res.map(|r| ApproxAggregate {
791 point: r.point,
792 ci_low: r.ci_low,
793 ci_high: r.ci_high,
794 n_population: r.n_population,
795 n_sample_live: r.n_sample_live,
796 n_passing: r.n_passing,
797 }))
798 }
799
800 pub fn scan_batched<F>(&self, table: &str, batch_size: usize, mut f: F) -> Result<()>
806 where
807 F: FnMut(&[serde_json::Map<String, Value>]) -> Result<()>,
808 {
809 let kit_t = self
810 .schema
811 .tables
812 .iter()
813 .find(|t| t.name == table)
814 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
815 let batch_size = batch_size.max(1);
816 let (snapshot, _pin) = self.inner.snapshot();
819 let handle = self.inner.table(table).map_err(KitError::from)?;
820 let guard = handle.lock();
821
822 let mut projection: Vec<(u16, mongreldb_core::schema::TypeId)> = Vec::new();
824 let mut meta: Vec<(String, mongreldb_kit_core::schema::ColumnType)> = Vec::new();
825 for c in &guard.schema().columns {
826 if let Some(kc) = kit_t.columns.iter().find(|kc| kc.id as u16 == c.id) {
827 projection.push((c.id, c.ty));
828 meta.push((kc.name.clone(), kc.storage_type));
829 }
830 }
831
832 match guard
833 .scan_cursor(snapshot, projection, &[])
834 .map_err(KitError::from)?
835 {
836 Some(mut cursor) => {
837 let mut buf: Vec<serde_json::Map<String, Value>> = Vec::with_capacity(batch_size);
838 while let Some(batch) = cursor.next_batch().map_err(KitError::from)? {
839 let nrows = batch.first().map(|c| c.len()).unwrap_or(0);
840 for j in 0..nrows {
841 let mut m = serde_json::Map::new();
842 for (ci, (name, ty)) in meta.iter().enumerate() {
843 let cv = batch
844 .get(ci)
845 .and_then(|col| col.value_at(j))
846 .unwrap_or(CoreValue::Null);
847 m.insert(name.clone(), crate::schema::core_to_json(&cv, *ty)?);
848 }
849 buf.push(m);
850 if buf.len() >= batch_size {
851 f(&buf)?;
852 buf.clear();
853 }
854 }
855 }
856 if !buf.is_empty() {
857 f(&buf)?;
858 }
859 Ok(())
860 }
861 None => {
862 drop(guard);
863 let rows = self.visible_core_rows_at(table, snapshot)?;
864 let maps: Vec<serde_json::Map<String, Value>> = rows
865 .iter()
866 .map(|r| crate::schema::core_row_to_json(r, kit_t).map(|row| row.values))
867 .collect::<Result<Vec<_>>>()?;
868 for chunk in maps.chunks(batch_size) {
869 f(chunk)?;
870 }
871 Ok(())
872 }
873 }
874 }
875
876 pub fn set_similarity(
885 &self,
886 table: &str,
887 column: &str,
888 query: &[String],
889 k: usize,
890 ) -> Result<Vec<SimilarRow>> {
891 let t = self
892 .schema
893 .tables
894 .iter()
895 .find(|t| t.name == table)
896 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
897 let col = t.columns.iter().find(|c| c.name == column).ok_or_else(|| {
898 KitError::Validation(format!("unknown column '{column}' on table '{table}'"))
899 })?;
900 let query_set: std::collections::HashSet<String> = query.iter().cloned().collect();
901
902 let has_minhash = t.indexes.iter().any(|idx| {
903 idx.kind == KitIndexKind::MinHash && idx.columns.iter().any(|c| c == column)
904 });
905 let rows = if has_minhash {
906 let query_hashes: Vec<u64> = query
908 .iter()
909 .map(|s| mongreldb_core::index::minhash_token_hash(s))
910 .collect();
911 let cand_k = k.saturating_mul(8).max(k + 64);
913 let cond = mongreldb_core::query::Condition::MinHashSimilar {
914 column_id: col.id as u16,
915 query: query_hashes,
916 k: cand_k,
917 };
918 let (snapshot, _pin) = self.inner.snapshot();
919 let core_rows = self.query_core_rows_at(table, &[cond], snapshot)?;
920 core_rows
921 .iter()
922 .map(|r| crate::schema::core_row_to_json(r, t))
923 .collect::<Result<Vec<_>>>()?
924 } else {
925 let tx = self.begin()?;
926 tx.all_rows(table)?
927 };
928
929 let mut scored: Vec<SimilarRow> = Vec::new();
930 for row in rows {
931 let set = parse_string_set(row.values.get(column));
932 let inter = set.iter().filter(|x| query_set.contains(*x)).count();
933 let union = set.len() + query_set.len() - inter;
934 let sim = if union == 0 {
935 0.0
936 } else {
937 inter as f64 / union as f64
938 };
939 if sim > 0.0 {
940 scored.push(SimilarRow {
941 row,
942 similarity: sim,
943 });
944 }
945 }
946 scored.sort_by(|a, b| {
947 b.similarity
948 .partial_cmp(&a.similarity)
949 .unwrap_or(std::cmp::Ordering::Equal)
950 });
951 scored.truncate(k);
952 Ok(scored)
953 }
954
955 pub fn flush(&self) -> Result<()> {
959 for name in self.inner.table_names() {
960 let handle = self.inner.table(&name).map_err(KitError::from)?;
961 let mut guard = handle.lock();
962 guard.flush().map_err(KitError::from)?;
963 }
964 Ok(())
965 }
966
967 pub fn incremental_aggregate(
979 &self,
980 table: &str,
981 column: Option<&str>,
982 agg: IncrementalAggKind,
983 filter: Option<&mongreldb_kit_core::query::Expr>,
984 ) -> Result<IncrementalAggregate> {
985 let t = self
986 .schema
987 .tables
988 .iter()
989 .find(|t| t.name == table)
990 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
991 if !matches!(agg, IncrementalAggKind::Count) && column.is_none() {
992 return Err(KitError::Validation(
993 "sum/min/max/avg incremental aggregate requires a column".into(),
994 ));
995 }
996 let cid = match column {
997 Some(name) => Some(
998 t.columns
999 .iter()
1000 .find(|c| c.name == name)
1001 .ok_or_else(|| KitError::Validation(format!("unknown column '{name}'")))?
1002 .id as u16,
1003 ),
1004 None => None,
1005 };
1006 let conditions = match filter {
1007 Some(expr) => {
1008 let plan = crate::pushdown::translate_predicate(t, expr).ok_or_else(|| {
1009 KitError::Validation(
1010 "filter is not index-translatable for an incremental aggregate".into(),
1011 )
1012 })?;
1013 if !plan.fully_translated {
1014 return Err(KitError::Validation(
1015 "filter has a residual that an incremental aggregate cannot apply exactly"
1016 .into(),
1017 ));
1018 }
1019 plan.conditions
1020 }
1021 None => Vec::new(),
1022 };
1023 let core_agg = match agg {
1024 IncrementalAggKind::Count => NativeAgg::Count,
1025 IncrementalAggKind::Sum => NativeAgg::Sum,
1026 IncrementalAggKind::Min => NativeAgg::Min,
1027 IncrementalAggKind::Max => NativeAgg::Max,
1028 IncrementalAggKind::Avg => NativeAgg::Avg,
1029 };
1030 let cache_key = incremental_cache_key(t.id, cid, agg, &conditions);
1031 let handle = self.inner.table(table).map_err(KitError::from)?;
1032 let mut guard = handle.lock();
1033 let res = guard
1034 .aggregate_incremental(cache_key, &conditions, cid, core_agg)
1035 .map_err(KitError::from)?;
1036 Ok(IncrementalAggregate {
1037 value: agg_state_value(&res.state),
1038 incremental: res.incremental,
1039 delta_rows: res.delta_rows,
1040 })
1041 }
1042
1043 pub fn applied_migrations(&self) -> Result<Vec<mongreldb_kit_core::migrations::Migration>> {
1045 crate::migrate::load_applied_migrations(&self.inner)
1046 }
1047
1048 pub(crate) fn core_db(&self) -> &CoreDatabase {
1049 &self.inner
1050 }
1051
1052 pub(crate) fn core_arc(&self) -> Arc<CoreDatabase> {
1055 Arc::clone(&self.inner)
1056 }
1057
1058 pub fn close(&self) -> Result<()> {
1063 self.inner.close().map_err(KitError::from)
1064 }
1065
1066 pub fn compact_all(&self) -> Result<(usize, usize)> {
1071 self.inner.compact().map_err(KitError::from)
1072 }
1073
1074 pub fn compact_table(&self, name: &str) -> Result<bool> {
1077 self.inner.compact_table(name).map_err(KitError::from)
1078 }
1079
1080 pub fn rename_table(&mut self, from: &str, to: &str) -> Result<()> {
1090 if from.starts_with("__kit_") || to.starts_with("__kit_") {
1091 return Err(KitError::Validation(
1092 "rename_table: names beginning with '__kit_' are reserved for internal tables"
1093 .into(),
1094 ));
1095 }
1096 self.inner.rename_table(from, to).map_err(KitError::from)?;
1097 if !self.schema.rename_table(from, to) {
1100 return Err(KitError::Integrity(format!(
1103 "rename_table: kit schema has no table '{from}' (or '{to}' already exists)"
1104 )));
1105 }
1106 for table in &mut self.schema.tables {
1107 for fk in &mut table.foreign_keys {
1108 if fk.references_table == from {
1109 fk.references_table = to.to_string();
1110 }
1111 }
1112 }
1113 store_schema(&self.root, &self.schema)?;
1114 Ok(())
1115 }
1116
1117 pub fn analyze(&self) -> Result<()> {
1122 for name in self.inner.table_names() {
1123 let handle = self.inner.table(&name).map_err(KitError::from)?;
1124 handle.lock().ensure_indexes_complete()?;
1125 }
1126 Ok(())
1127 }
1128
1129 pub fn vacuum(&self) -> Result<usize> {
1133 self.inner.compact().map_err(KitError::from)?;
1134 self.inner.gc().map_err(KitError::from)
1135 }
1136
1137 pub fn create_view(&self, spec: &ViewSpec) -> Result<()> {
1143 self.sql(&spec.create_sql())?;
1144 Ok(())
1145 }
1146
1147 pub fn drop_view(&self, name: &str) -> Result<()> {
1149 self.sql(&format!("DROP VIEW IF EXISTS {name}"))?;
1150 Ok(())
1151 }
1152
1153 pub fn reserve_auto_inc(&self, table: &str) -> Result<Option<i64>> {
1161 let handle = self.inner.table(table).map_err(KitError::from)?;
1162 let mut guard = handle.lock();
1163 guard.reserve_auto_inc().map_err(KitError::from)
1164 }
1165
1166 pub fn create_user(&self, username: &str, password: &str) -> Result<()> {
1170 self.inner
1171 .create_user(username, password)
1172 .map_err(KitError::from)?;
1173 Ok(())
1174 }
1175
1176 pub fn drop_user(&self, username: &str) -> Result<()> {
1178 self.inner.drop_user(username).map_err(KitError::from)
1179 }
1180
1181 pub fn alter_user_password(&self, username: &str, new_password: &str) -> Result<()> {
1183 self.inner
1184 .alter_user_password(username, new_password)
1185 .map_err(KitError::from)
1186 }
1187
1188 pub fn verify_user(
1190 &self,
1191 username: &str,
1192 password: &str,
1193 ) -> Result<Option<mongreldb_core::auth::UserEntry>> {
1194 self.inner
1195 .verify_user(username, password)
1196 .map_err(KitError::from)
1197 }
1198
1199 pub fn set_user_admin(&self, username: &str, is_admin: bool) -> Result<()> {
1201 self.inner
1202 .set_user_admin(username, is_admin)
1203 .map_err(KitError::from)
1204 }
1205
1206 pub fn users(&self) -> Vec<String> {
1208 self.inner.users().into_iter().map(|u| u.username).collect()
1209 }
1210
1211 pub fn create_role(&self, name: &str) -> Result<()> {
1213 self.inner.create_role(name).map_err(KitError::from)?;
1214 Ok(())
1215 }
1216
1217 pub fn drop_role(&self, name: &str) -> Result<()> {
1219 self.inner.drop_role(name).map_err(KitError::from)
1220 }
1221
1222 pub fn roles(&self) -> Vec<String> {
1224 self.inner.roles().into_iter().map(|r| r.name).collect()
1225 }
1226
1227 pub fn grant_role(&self, username: &str, role_name: &str) -> Result<()> {
1229 self.inner
1230 .grant_role(username, role_name)
1231 .map_err(KitError::from)
1232 }
1233
1234 pub fn revoke_role(&self, username: &str, role_name: &str) -> Result<()> {
1236 self.inner
1237 .revoke_role(username, role_name)
1238 .map_err(KitError::from)
1239 }
1240
1241 pub fn grant_permission(
1243 &self,
1244 role_name: &str,
1245 permission: mongreldb_core::auth::Permission,
1246 ) -> Result<()> {
1247 self.inner
1248 .grant_permission(role_name, permission)
1249 .map_err(KitError::from)
1250 }
1251
1252 pub fn revoke_permission(
1254 &self,
1255 role_name: &str,
1256 permission: mongreldb_core::auth::Permission,
1257 ) -> Result<()> {
1258 self.inner
1259 .revoke_permission(role_name, permission)
1260 .map_err(KitError::from)
1261 }
1262
1263 pub fn set_spill_threshold(&self, bytes: u64) {
1269 self.inner.set_spill_threshold(bytes);
1270 }
1271
1272 pub fn set_recursive_triggers(&self, enabled: bool) {
1274 self.inner.set_recursive_triggers(enabled);
1275 }
1276
1277 pub fn trigger_config(&self) -> mongreldb_core::TriggerConfig {
1279 self.inner.trigger_config()
1280 }
1281
1282 pub fn set_trigger_config(&self, config: mongreldb_core::TriggerConfig) -> Result<()> {
1284 self.inner
1285 .set_trigger_config(config)
1286 .map_err(KitError::from)
1287 }
1288
1289 pub fn set_table_compaction_zstd_level(&self, table: &str, level: i32) -> Result<()> {
1291 let handle = self.inner.table(table).map_err(KitError::from)?;
1292 handle.lock().set_compaction_zstd_level(level);
1293 Ok(())
1294 }
1295
1296 pub fn set_table_result_cache_max_bytes(&self, table: &str, max_bytes: u64) -> Result<()> {
1298 let handle = self.inner.table(table).map_err(KitError::from)?;
1299 handle.lock().set_result_cache_max_bytes(max_bytes);
1300 Ok(())
1301 }
1302
1303 pub fn set_table_mutable_run_spill_bytes(&self, table: &str, bytes: u64) -> Result<()> {
1305 let handle = self.inner.table(table).map_err(KitError::from)?;
1306 handle.lock().set_mutable_run_spill_bytes(bytes);
1307 Ok(())
1308 }
1309
1310 pub fn set_table_sync_byte_threshold(&self, table: &str, threshold: u64) -> Result<()> {
1312 let handle = self.inner.table(table).map_err(KitError::from)?;
1313 handle.lock().set_sync_byte_threshold(threshold);
1314 Ok(())
1315 }
1316
1317 pub fn set_table_index_build_policy(
1320 &self,
1321 table: &str,
1322 policy: mongreldb_core::IndexBuildPolicy,
1323 ) -> Result<()> {
1324 let handle = self.inner.table(table).map_err(KitError::from)?;
1325 handle.lock().set_index_build_policy(policy);
1326 Ok(())
1327 }
1328
1329 pub fn table_page_cache_stats(&self, table: &str) -> Result<mongreldb_core::cache::CacheStats> {
1331 let handle = self.inner.table(table).map_err(KitError::from)?;
1332 let stats = handle.lock().page_cache_stats();
1333 Ok(stats)
1334 }
1335
1336 pub fn table_run_count(&self, table: &str) -> Result<usize> {
1338 let handle = self.inner.table(table).map_err(KitError::from)?;
1339 let n = handle.lock().run_count();
1340 Ok(n)
1341 }
1342
1343 pub fn table_memtable_len(&self, table: &str) -> Result<usize> {
1345 let handle = self.inner.table(table).map_err(KitError::from)?;
1346 let n = handle.lock().memtable_len();
1347 Ok(n)
1348 }
1349
1350 pub fn table_mutable_run_len(&self, table: &str) -> Result<usize> {
1352 let handle = self.inner.table(table).map_err(KitError::from)?;
1353 let n = handle.lock().mutable_run_len();
1354 Ok(n)
1355 }
1356
1357 pub fn table_page_cache_len(&self, table: &str) -> Result<usize> {
1359 let handle = self.inner.table(table).map_err(KitError::from)?;
1360 let n = handle.lock().page_cache_len();
1361 Ok(n)
1362 }
1363
1364 pub fn table_decoded_cache_len(&self, table: &str) -> Result<usize> {
1366 let handle = self.inner.table(table).map_err(KitError::from)?;
1367 let n = handle.lock().decoded_cache_len();
1368 Ok(n)
1369 }
1370
1371 pub fn sql(&self, statement: &str) -> Result<Vec<arrow::record_batch::RecordBatch>> {
1389 let session = match self.session.lock().take() {
1390 Some(s) => s,
1391 None => {
1392 mongreldb_query::MongrelSession::open(self.core_arc()).map_err(KitError::from)?
1393 }
1394 };
1395 let runtime = sql_runtime();
1396 let result = runtime
1397 .block_on(session.run(statement))
1398 .map_err(KitError::from);
1399 *self.session.lock() = Some(session);
1401 result
1402 }
1403
1404 pub fn refresh_sql_session(&self) -> Result<()> {
1409 let session =
1410 mongreldb_query::MongrelSession::open(self.core_arc()).map_err(KitError::from)?;
1411 *self.session.lock() = Some(session);
1412 Ok(())
1413 }
1414
1415 pub fn sql_arrow(&self, statement: &str) -> Result<Vec<u8>> {
1421 let batches = self.sql(statement)?;
1422 crate::arrow_util::batches_to_ipc(&batches)
1423 }
1424
1425 pub fn sql_rows(&self, statement: &str) -> Result<Vec<serde_json::Map<String, Value>>> {
1429 let batches = self.sql(statement)?;
1430 crate::arrow_util::batches_to_rows(&batches)
1431 }
1432
1433 pub(crate) fn lookup_row_id(&self, table: &str, key: &[u8]) -> Result<Option<RowId>> {
1437 let handle = self.inner.table(table).map_err(KitError::from)?;
1438 let mut guard = handle.lock();
1439 guard.ensure_indexes_complete()?;
1440 Ok(guard.lookup_pk(key))
1441 }
1442
1443 pub(crate) fn root(&self) -> &Path {
1444 &self.root
1445 }
1446
1447 pub(crate) fn visible_core_rows_at(
1451 &self,
1452 table_name: &str,
1453 snapshot: Snapshot,
1454 ) -> Result<Vec<CoreRow>> {
1455 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1456 let guard = handle.lock();
1457 guard.visible_rows(snapshot).map_err(KitError::from)
1458 }
1459
1460 pub(crate) fn query_core_rows_at(
1467 &self,
1468 table_name: &str,
1469 conditions: &[mongreldb_core::query::Condition],
1470 snapshot: Snapshot,
1471 ) -> Result<Vec<CoreRow>> {
1472 if conditions.is_empty() {
1473 return self.visible_core_rows_at(table_name, snapshot);
1474 }
1475 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1476 let mut guard = handle.lock();
1477 let q = mongreldb_core::query::Query {
1478 conditions: conditions.to_vec(),
1479 };
1480 guard.query(&q).map_err(KitError::from)
1481 }
1482
1483 pub(crate) fn flush_table(&self, table_name: &str) -> Result<()> {
1491 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1492 handle.lock().flush().map_err(KitError::from)?;
1493 Ok(())
1494 }
1495
1496 pub(crate) fn count_core_rows_at(
1506 &self,
1507 table_name: &str,
1508 conditions: &[mongreldb_core::query::Condition],
1509 snapshot: Snapshot,
1510 ) -> Result<Option<u64>> {
1511 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1512 let mut guard = handle.lock();
1513 if guard.snapshot().epoch != snapshot.epoch {
1514 return Ok(None); }
1516 guard
1517 .count_conditions(conditions, snapshot)
1518 .map_err(KitError::from)
1519 }
1520
1521 pub(crate) fn aggregate_core_at(
1530 &self,
1531 table_name: &str,
1532 column: Option<u16>,
1533 conditions: &[mongreldb_core::query::Condition],
1534 agg: NativeAgg,
1535 snapshot: Snapshot,
1536 ) -> Result<Option<NativeAggResult>> {
1537 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1538 let guard = handle.lock();
1539 if guard.snapshot().epoch != snapshot.epoch {
1540 return Ok(None); }
1542 guard
1543 .aggregate_native(snapshot, column, conditions, agg)
1544 .map_err(KitError::from)
1545 }
1546
1547 pub(crate) fn count_distinct_core_at(
1556 &self,
1557 table_name: &str,
1558 column_id: u16,
1559 snapshot: Snapshot,
1560 ) -> Result<Option<u64>> {
1561 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1562 let mut guard = handle.lock();
1563 if guard.snapshot().epoch != snapshot.epoch {
1564 return Ok(None); }
1566 guard
1567 .count_distinct_from_bitmap(column_id)
1568 .map_err(KitError::from)
1569 }
1570
1571 #[allow(dead_code)]
1573 pub(crate) fn get_core_row(&self, table_name: &str, row_id: u64) -> Result<Option<CoreRow>> {
1574 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1575 let guard = handle.lock();
1576 let snapshot = guard.snapshot();
1577 Ok(guard.get(mongreldb_core::RowId(row_id), snapshot))
1578 }
1579}
1580
1581pub(crate) fn create_core_table(db: &CoreDatabase, name: &str, schema: CoreSchema) -> Result<()> {
1582 if db.table_id(name).is_ok() {
1583 return Ok(());
1584 }
1585 db.create_table(name, schema).map_err(KitError::from)?;
1586 Ok(())
1587}
1588
1589fn sql_runtime() -> &'static tokio::runtime::Runtime {
1594 use std::sync::OnceLock;
1595 static RT: OnceLock<tokio::runtime::Runtime> = OnceLock::new();
1596 RT.get_or_init(|| {
1597 tokio::runtime::Builder::new_current_thread()
1598 .enable_all()
1599 .build()
1600 .expect("failed to build kit SQL tokio runtime")
1601 })
1602}
1603
1604fn core_procedure(spec: &ProcedureSpec) -> Result<mongreldb_core::StoredProcedure> {
1605 let parsed: mongreldb_core::StoredProcedure =
1606 serde_json::from_value(spec.json.clone()).map_err(KitError::from)?;
1607 mongreldb_core::StoredProcedure::new(parsed.name, parsed.mode, parsed.params, parsed.body, 0)
1608 .map_err(KitError::from)
1609}
1610
1611fn core_trigger(spec: &TriggerSpec) -> Result<mongreldb_core::StoredTrigger> {
1612 let parsed: mongreldb_core::StoredTrigger =
1613 serde_json::from_value(spec.json.clone()).map_err(KitError::from)?;
1614 mongreldb_core::StoredTrigger::new(
1615 parsed.name,
1616 mongreldb_core::TriggerDefinition {
1617 target: parsed.target,
1618 timing: parsed.timing,
1619 event: parsed.event,
1620 update_of: parsed.update_of,
1621 target_columns: parsed.target_columns,
1622 when: parsed.when,
1623 program: parsed.program,
1624 },
1625 0,
1626 )
1627 .map_err(KitError::from)
1628}
1629
1630fn json_to_core_value(value: &Value) -> Result<CoreValue> {
1631 match value {
1632 Value::Null => Ok(CoreValue::Null),
1633 Value::Bool(value) => Ok(CoreValue::Bool(*value)),
1634 Value::Number(value) => {
1635 if let Some(value) = value.as_i64() {
1636 Ok(CoreValue::Int64(value))
1637 } else if let Some(value) = value.as_f64() {
1638 Ok(CoreValue::Float64(value))
1639 } else {
1640 Err(KitError::Validation("unsupported JSON number".into()))
1641 }
1642 }
1643 Value::String(value) => Ok(CoreValue::Bytes(value.as_bytes().to_vec())),
1644 Value::Array(_) | Value::Object(_) => Err(KitError::Validation(
1645 "procedure args only support scalar JSON values".into(),
1646 )),
1647 }
1648}
1649
1650pub(crate) fn internal_bytes(row: &CoreRow, col_id: u16) -> Option<String> {
1652 match row.columns.get(&col_id) {
1653 Some(CoreValue::Bytes(b)) => String::from_utf8(b.clone()).ok(),
1654 _ => None,
1655 }
1656}
1657
1658fn reap_rotated_wal_segments(db: &CoreDatabase) {
1673 let _ = db.gc();
1674}
1675
1676pub(crate) fn load_schema(path: &Path) -> Result<KitSchema> {
1677 let file = path.join(SCHEMA_FILE);
1678 let json = std::fs::read_to_string(&file)
1679 .map_err(|e| KitError::Migration(format!("cannot read schema file: {e}")))?;
1680 let schema: KitSchema = serde_json::from_str(&json)?;
1681 Ok(schema)
1682}
1683
1684pub(crate) fn store_schema(path: &Path, schema: &KitSchema) -> Result<()> {
1685 let file = path.join(SCHEMA_FILE);
1686 let json = serde_json::to_string_pretty(schema)?;
1687 std::fs::write(&file, json)?;
1688 Ok(())
1689}
1690
1691pub(crate) fn persist_schema(db: &Database, schema: &KitSchema) -> Result<()> {
1693 store_schema(&db.root, schema)
1694}