1use crate::error::{KitError, Result};
4use crate::internal::{ensure_internal_tables, internal_tables_core};
5use crate::schema::to_core_schema;
6use mongreldb_core::epoch::Snapshot;
7use mongreldb_core::memtable::Row as CoreRow;
8use mongreldb_core::memtable::Value as CoreValue;
9use mongreldb_core::schema::Schema as CoreSchema;
10use mongreldb_core::Database as CoreDatabase;
11use mongreldb_core::{AggState, ApproxAgg, NativeAgg, NativeAggResult, RowId};
12use mongreldb_kit_core::schema::IndexKind as KitIndexKind;
13use mongreldb_kit_core::schema::Schema as KitSchema;
14use mongreldb_kit_core::schema::Table as KitTable;
15use mongreldb_kit_core::{ProcedureSpec, TriggerSpec, ViewSpec};
16use serde_json::Value;
17
18use std::collections::HashMap;
19use std::path::{Path, PathBuf};
20use std::sync::Arc;
21
22const SCHEMA_FILE: &str = "kit_schema.json";
23
24#[derive(Clone, Copy, Debug, Default)]
29pub struct OpenOptions {
30 pub lock_timeout_ms: u32,
36}
37
38impl OpenOptions {
39 pub fn new() -> Self {
41 Self::default()
42 }
43
44 pub fn with_lock_timeout_ms(mut self, ms: u32) -> Self {
47 self.lock_timeout_ms = ms;
48 self
49 }
50}
51
52#[derive(Debug, Clone, Default)]
53pub struct SqlOptions {
54 pub query_id: Option<mongreldb_query::QueryId>,
55 pub timeout: Option<std::time::Duration>,
56}
57
58#[derive(Debug, Clone, Copy, PartialEq, Eq)]
59pub struct SqlOutputLimits {
60 pub max_rows: usize,
61 pub max_bytes: usize,
62}
63
64impl Default for SqlOutputLimits {
65 fn default() -> Self {
66 Self {
67 max_rows: 1_000_000,
68 max_bytes: 64 * 1024 * 1024,
69 }
70 }
71}
72
73pub struct SqlQueryHandle {
74 query_id: mongreldb_query::QueryId,
75 session: Arc<mongreldb_query::MongrelSession>,
76 worker: Option<
77 std::thread::JoinHandle<mongreldb_query::Result<mongreldb_query::ManagedQueryBatches>>,
78 >,
79}
80
81impl SqlQueryHandle {
82 pub fn id(&self) -> mongreldb_query::QueryId {
83 self.query_id
84 }
85
86 pub fn cancel(&self) -> mongreldb_query::CancelOutcome {
87 self.session.cancel_query(self.query_id)
88 }
89
90 pub fn status(&self) -> Option<mongreldb_query::QueryStatus> {
91 self.session.query_registry().status(self.query_id)
92 }
93
94 pub fn wait(self) -> Result<Vec<arrow::record_batch::RecordBatch>> {
95 let output = self.wait_for_serialization()?;
96 let batches = output.batches().to_vec();
97 complete_sql_output(output)?;
98 Ok(batches)
99 }
100
101 pub fn wait_arrow(self) -> Result<Vec<u8>> {
102 self.wait_arrow_with_limits(SqlOutputLimits::default())
103 }
104
105 pub fn wait_arrow_with_limits(self, limits: SqlOutputLimits) -> Result<Vec<u8>> {
106 let output = self.wait_for_serialization()?;
107 let result = crate::arrow_util::batches_to_ipc_controlled_with_limits(
108 output.batches(),
109 output.query(),
110 limits,
111 );
112 match result {
113 Ok(bytes) => {
114 complete_sql_output(output)?;
115 Ok(bytes)
116 }
117 Err(error) => {
118 fail_sql_output(output, &error);
119 Err(error)
120 }
121 }
122 }
123
124 pub fn wait_rows(self) -> Result<Vec<serde_json::Map<String, Value>>> {
125 self.wait_rows_with_limits(SqlOutputLimits::default())
126 }
127
128 pub fn wait_rows_with_limits(
129 self,
130 limits: SqlOutputLimits,
131 ) -> Result<Vec<serde_json::Map<String, Value>>> {
132 let output = self.wait_for_serialization()?;
133 let result = crate::arrow_util::batches_to_rows_controlled_with_limits(
134 output.batches(),
135 output.query(),
136 limits,
137 );
138 match result {
139 Ok(rows) => {
140 complete_sql_output(output)?;
141 Ok(rows)
142 }
143 Err(error) => {
144 fail_sql_output(output, &error);
145 Err(error)
146 }
147 }
148 }
149
150 pub fn wait_for_serialization(mut self) -> Result<mongreldb_query::ManagedQueryBatches> {
151 let result = self
152 .worker
153 .take()
154 .expect("SQL worker is present")
155 .join()
156 .map_err(|_| KitError::Storage("SQL worker panicked".into()))?;
157 let output = result.map_err(|error| {
158 let status = self.session.query_registry().status(self.query_id);
159 crate::error::query_error_with_status(error, status.as_ref())
160 })?;
161 self.session
162 .fire_test_hook(mongreldb_query::SqlTestHookPoint::BeforeSerializationBatch);
163 Ok(output)
164 }
165}
166
167#[doc(hidden)]
169pub fn complete_sql_output(output: mongreldb_query::ManagedQueryBatches) -> Result<()> {
170 let query = output.query().clone();
171 loop {
172 if let Err(error) = query.checkpoint() {
173 let status = query.status();
174 output.fail();
175 return Err(crate::error::query_error_with_status(error, Some(&status)));
176 }
177 let phase = query.phase();
178 match phase {
179 mongreldb_query::SqlQueryPhase::Serializing
180 | mongreldb_query::SqlQueryPhase::CommitCritical => {
181 if query
182 .transition(phase, mongreldb_query::SqlQueryPhase::Completed)
183 .is_ok()
184 {
185 break;
186 }
187 }
188 mongreldb_query::SqlQueryPhase::Completed => break,
189 mongreldb_query::SqlQueryPhase::Cancelling => std::thread::yield_now(),
190 phase => {
191 let error = mongreldb_query::MongrelQueryError::InvalidQueryState(format!(
192 "query {} cannot complete output conversion from {phase:?}",
193 query.id()
194 ));
195 let status = query.status();
196 output.fail_with_error(
197 error.code(),
198 mongreldb_query::QueryTerminalErrorCategory::Execution,
199 );
200 return Err(crate::error::query_error_with_status(error, Some(&status)));
201 }
202 }
203 }
204 output.complete().map_err(|error| {
205 let status = query.status();
206 crate::error::query_error_with_status(error, Some(&status))
207 })
208}
209
210#[doc(hidden)]
212pub fn fail_sql_output(output: mongreldb_query::ManagedQueryBatches, error: &KitError) {
213 match error {
214 KitError::ResultLimitExceeded { .. } => output.fail_result_limit(),
215 KitError::SerializationFailed { .. } => output.fail_serialization(),
216 _ => output.fail(),
217 }
218}
219
220impl Drop for SqlQueryHandle {
221 fn drop(&mut self) {
222 if self.worker.is_some() {
223 let _ = self.session.cancel_query(self.query_id);
224 }
225 }
226}
227
228pub type DefaultProvider = Box<dyn Fn() -> Value + Send + Sync>;
230
231#[derive(Debug, Clone)]
234pub struct ExplainPlan {
235 pub index_accelerated: bool,
237 pub exact: bool,
240 pub pushed_conditions: Vec<String>,
242}
243
244#[derive(Debug, Clone)]
246pub struct SimilarRow {
247 pub row: crate::schema::Row,
248 pub similarity: f64,
249}
250
251fn parse_string_set(value: Option<&Value>) -> std::collections::HashSet<String> {
255 let arr = match value {
256 Some(Value::Array(a)) => Some(a.clone()),
257 Some(Value::String(s)) => serde_json::from_str::<Value>(s)
258 .ok()
259 .and_then(|v| v.as_array().cloned()),
260 _ => None,
261 };
262 arr.into_iter()
263 .flatten()
264 .filter_map(|v| match v {
265 Value::String(s) => Some(s),
266 Value::Number(n) => Some(n.to_string()),
267 Value::Bool(b) => Some(b.to_string()),
268 _ => None,
269 })
270 .collect()
271}
272
273#[derive(Debug, Clone, Copy, PartialEq, Eq)]
275pub enum IncrementalAggKind {
276 Count,
277 Sum,
278 Min,
279 Max,
280 Avg,
281}
282
283#[derive(Debug, Clone)]
285pub struct IncrementalAggregate {
286 pub value: Value,
289 pub incremental: bool,
293 pub delta_rows: u64,
295}
296
297fn incremental_cache_key(
301 table_id: u32,
302 column: Option<u16>,
303 agg: IncrementalAggKind,
304 conditions: &[mongreldb_core::query::Condition],
305) -> u64 {
306 use std::hash::{Hash, Hasher};
307 let mut h = std::collections::hash_map::DefaultHasher::new();
308 table_id.hash(&mut h);
309 column.hash(&mut h);
310 (agg as u8).hash(&mut h);
311 format!("{conditions:?}").hash(&mut h);
313 h.finish()
314}
315
316fn agg_state_value(s: &AggState) -> Value {
320 let num_f64 = |x: f64| {
321 serde_json::Number::from_f64(x)
322 .map(Value::Number)
323 .unwrap_or(Value::Null)
324 };
325 match s {
326 AggState::Count(n) => Value::from(*n),
327 AggState::SumI { sum, .. } => i64::try_from(*sum)
328 .map(Value::from)
329 .unwrap_or_else(|_| num_f64(*sum as f64)),
330 AggState::SumF { sum, .. } => num_f64(*sum),
331 AggState::AvgI { sum, count } if *count > 0 => num_f64(*sum as f64 / *count as f64),
332 AggState::AvgF { sum, count } if *count > 0 => num_f64(*sum / *count as f64),
333 AggState::AvgI { .. } | AggState::AvgF { .. } => Value::Null,
334 AggState::MinI(n) | AggState::MaxI(n) => Value::from(*n),
335 AggState::MinF(f) | AggState::MaxF(f) => num_f64(*f),
336 AggState::Empty => Value::Null,
337 }
338}
339
340#[derive(Debug, Clone, Copy, PartialEq, Eq)]
342pub enum ApproxAggKind {
343 Count,
344 Sum,
345 Avg,
346}
347
348#[derive(Debug, Clone)]
352pub struct ApproxAggregate {
353 pub point: f64,
354 pub ci_low: f64,
355 pub ci_high: f64,
356 pub n_population: u64,
357 pub n_sample_live: usize,
358 pub n_passing: usize,
359}
360
361fn condition_label(c: &mongreldb_core::query::Condition) -> String {
364 let dbg = format!("{c:?}");
365 dbg.split(['(', '{', ' ']).next().unwrap_or("").to_string()
366}
367
368fn open_core_with_retry<T>(
369 timeout_ms: u32,
370 mut open: impl FnMut() -> mongreldb_core::Result<T>,
371) -> mongreldb_core::Result<T> {
372 if timeout_ms == 0 {
373 return open();
374 }
375 let deadline = std::time::Instant::now() + std::time::Duration::from_millis(timeout_ms as u64);
376 let mut next_sleep = std::time::Duration::from_millis(1);
377 loop {
378 match open() {
379 Ok(db) => return Ok(db),
380 Err(err) if is_lock_contention(&err) => {
381 let now = std::time::Instant::now();
382 if now >= deadline {
383 return Err(err);
384 }
385 let sleep = next_sleep.min(deadline - now);
386 std::thread::sleep(sleep);
387 next_sleep = next_sleep
388 .saturating_mul(10)
389 .min(std::time::Duration::from_millis(50));
390 }
391 Err(err) => return Err(err),
392 }
393 }
394}
395
396fn is_lock_contention(err: &mongreldb_core::MongrelError) -> bool {
397 matches!(err, mongreldb_core::MongrelError::DatabaseLocked { .. })
398}
399
400pub struct Database {
405 pub(crate) inner: Arc<CoreDatabase>,
406 pub(crate) schema: KitSchema,
407 pub(crate) root: PathBuf,
408 pub(crate) default_providers: HashMap<String, DefaultProvider>,
410 pub(crate) session: parking_lot::RwLock<Option<Arc<mongreldb_query::MongrelSession>>>,
417}
418
419impl Database {
420 pub fn open(path: &Path) -> Result<Self> {
422 let inner = Arc::new(CoreDatabase::open(path)?);
423 let schema = load_schema(path)?;
424 ensure_internal_tables(&inner)?;
426 reap_rotated_wal_segments(&inner);
427 Ok(Self {
428 inner,
429 schema,
430 root: path.to_path_buf(),
431 default_providers: HashMap::new(),
432 session: parking_lot::RwLock::new(None),
433 })
434 }
435
436 pub fn open_with_options(path: &Path, opts: OpenOptions) -> Result<Self> {
444 let inner = Arc::new(open_core_with_retry(opts.lock_timeout_ms, || {
445 CoreDatabase::open(path)
446 })?);
447 let schema = load_schema(path)?;
448 ensure_internal_tables(&inner)?;
449 reap_rotated_wal_segments(&inner);
450 Ok(Self {
451 inner,
452 schema,
453 root: path.to_path_buf(),
454 default_providers: HashMap::new(),
455 session: parking_lot::RwLock::new(None),
456 })
457 }
458
459 pub fn open_encrypted(path: &Path, passphrase: &str) -> Result<Self> {
461 let inner = Arc::new(CoreDatabase::open_encrypted(path, passphrase)?);
462 let schema = load_schema(path)?;
463 ensure_internal_tables(&inner)?;
464 reap_rotated_wal_segments(&inner);
465 Ok(Self {
466 inner,
467 schema,
468 root: path.to_path_buf(),
469 default_providers: HashMap::new(),
470 session: parking_lot::RwLock::new(None),
471 })
472 }
473
474 pub fn open_encrypted_with_options(
478 path: &Path,
479 passphrase: &str,
480 opts: OpenOptions,
481 ) -> Result<Self> {
482 let inner = Arc::new(open_core_with_retry(opts.lock_timeout_ms, || {
483 CoreDatabase::open_encrypted(path, passphrase)
484 })?);
485 let schema = load_schema(path)?;
486 ensure_internal_tables(&inner)?;
487 reap_rotated_wal_segments(&inner);
488 Ok(Self {
489 inner,
490 schema,
491 root: path.to_path_buf(),
492 default_providers: HashMap::new(),
493 session: parking_lot::RwLock::new(None),
494 })
495 }
496
497 pub fn create_encrypted(path: &Path, schema: KitSchema, passphrase: &str) -> Result<Self> {
501 std::fs::create_dir_all(path)?;
502 let inner = Arc::new(CoreDatabase::create_encrypted(path, passphrase)?);
503 ensure_internal_tables(&inner)?;
504 store_schema(path, &schema)?;
505 for table in &schema.tables {
506 create_core_table(&inner, &table.name, to_core_schema(table)?)?;
507 }
508 Ok(Self {
509 inner,
510 schema,
511 root: path.to_path_buf(),
512 default_providers: HashMap::new(),
513 session: parking_lot::RwLock::new(None),
514 })
515 }
516
517 pub fn create(path: &Path, schema: KitSchema) -> Result<Self> {
519 std::fs::create_dir_all(path)?;
520 let inner = Arc::new(CoreDatabase::create(path)?);
521
522 ensure_internal_tables(&inner)?;
525
526 store_schema(path, &schema)?;
529
530 for table in &schema.tables {
532 create_core_table(&inner, &table.name, to_core_schema(table)?)?;
533 }
534
535 Ok(Self {
536 inner,
537 schema,
538 root: path.to_path_buf(),
539 default_providers: HashMap::new(),
540 session: parking_lot::RwLock::new(None),
541 })
542 }
543
544 pub fn open_with_credentials(path: &Path, username: &str, password: &str) -> Result<Self> {
554 let inner = Arc::new(CoreDatabase::open_with_credentials(
555 path, username, password,
556 )?);
557 let schema = load_schema(path)?;
558 ensure_internal_tables(&inner)?;
559 reap_rotated_wal_segments(&inner);
560 Ok(Self {
561 inner,
562 schema,
563 root: path.to_path_buf(),
564 default_providers: HashMap::new(),
565 session: parking_lot::RwLock::new(None),
566 })
567 }
568
569 pub fn open_with_credentials_and_options(
573 path: &Path,
574 username: &str,
575 password: &str,
576 opts: OpenOptions,
577 ) -> Result<Self> {
578 let inner = Arc::new(open_core_with_retry(opts.lock_timeout_ms, || {
579 CoreDatabase::open_with_credentials(path, username, password)
580 })?);
581 let schema = load_schema(path)?;
582 ensure_internal_tables(&inner)?;
583 reap_rotated_wal_segments(&inner);
584 Ok(Self {
585 inner,
586 schema,
587 root: path.to_path_buf(),
588 default_providers: HashMap::new(),
589 session: parking_lot::RwLock::new(None),
590 })
591 }
592
593 pub fn create_with_credentials(
599 path: &Path,
600 schema: KitSchema,
601 admin_username: &str,
602 admin_password: &str,
603 ) -> Result<Self> {
604 std::fs::create_dir_all(path)?;
605 let inner = Arc::new(CoreDatabase::create_with_credentials(
606 path,
607 admin_username,
608 admin_password,
609 )?);
610 ensure_internal_tables(&inner)?;
611 store_schema(path, &schema)?;
612 for table in &schema.tables {
613 create_core_table(&inner, &table.name, to_core_schema(table)?)?;
614 }
615 Ok(Self {
616 inner,
617 schema,
618 root: path.to_path_buf(),
619 default_providers: HashMap::new(),
620 session: parking_lot::RwLock::new(None),
621 })
622 }
623
624 pub fn open_encrypted_with_credentials(
627 path: &Path,
628 passphrase: &str,
629 username: &str,
630 password: &str,
631 ) -> Result<Self> {
632 let inner = Arc::new(CoreDatabase::open_encrypted_with_credentials(
633 path, passphrase, username, password,
634 )?);
635 let schema = load_schema(path)?;
636 ensure_internal_tables(&inner)?;
637 reap_rotated_wal_segments(&inner);
638 Ok(Self {
639 inner,
640 schema,
641 root: path.to_path_buf(),
642 default_providers: HashMap::new(),
643 session: parking_lot::RwLock::new(None),
644 })
645 }
646
647 pub fn open_encrypted_with_credentials_and_options(
651 path: &Path,
652 passphrase: &str,
653 username: &str,
654 password: &str,
655 opts: OpenOptions,
656 ) -> Result<Self> {
657 let inner = Arc::new(open_core_with_retry(opts.lock_timeout_ms, || {
658 CoreDatabase::open_encrypted_with_credentials(path, passphrase, username, password)
659 })?);
660 let schema = load_schema(path)?;
661 ensure_internal_tables(&inner)?;
662 reap_rotated_wal_segments(&inner);
663 Ok(Self {
664 inner,
665 schema,
666 root: path.to_path_buf(),
667 default_providers: HashMap::new(),
668 session: parking_lot::RwLock::new(None),
669 })
670 }
671
672 pub fn create_encrypted_with_credentials(
676 path: &Path,
677 schema: KitSchema,
678 passphrase: &str,
679 admin_username: &str,
680 admin_password: &str,
681 ) -> Result<Self> {
682 std::fs::create_dir_all(path)?;
683 let inner = Arc::new(CoreDatabase::create_encrypted_with_credentials(
684 path,
685 passphrase,
686 admin_username,
687 admin_password,
688 )?);
689 ensure_internal_tables(&inner)?;
690 store_schema(path, &schema)?;
691 for table in &schema.tables {
692 create_core_table(&inner, &table.name, to_core_schema(table)?)?;
693 }
694 Ok(Self {
695 inner,
696 schema,
697 root: path.to_path_buf(),
698 default_providers: HashMap::new(),
699 session: parking_lot::RwLock::new(None),
700 })
701 }
702
703 pub fn enable_auth(&self, admin_username: &str, admin_password: &str) -> Result<()> {
707 self.inner
708 .enable_auth(admin_username, admin_password)
709 .map_err(KitError::from)
710 }
711
712 pub fn disable_auth(&self) -> Result<()> {
716 self.inner.disable_auth().map_err(KitError::from)
717 }
718
719 pub fn require_auth_enabled(&self) -> bool {
721 self.inner.require_auth_enabled()
722 }
723
724 pub fn refresh_principal(&self) -> Result<()> {
728 self.inner.refresh_principal().map_err(KitError::from)?;
729 *self.session.write() = None;
732 Ok(())
733 }
734
735 pub fn register_default(
738 &mut self,
739 name: impl Into<String>,
740 provider: impl Fn() -> Value + Send + Sync + 'static,
741 ) {
742 self.default_providers
743 .insert(name.into(), Box::new(provider));
744 }
745
746 pub fn raw(&self) -> &CoreDatabase {
750 &self.inner
751 }
752
753 pub fn table_names(&self) -> Vec<String> {
755 self.schema
756 .tables
757 .iter()
758 .map(|t| t.name.clone())
759 .filter(|n| !n.starts_with("__kit_"))
760 .collect()
761 }
762
763 pub fn create_procedure(
764 &self,
765 spec: &ProcedureSpec,
766 ) -> Result<mongreldb_core::StoredProcedure> {
767 let procedure = core_procedure(spec)?;
768 self.inner
769 .create_procedure(procedure)
770 .map_err(KitError::from)
771 }
772
773 pub fn replace_procedure(
774 &self,
775 spec: &ProcedureSpec,
776 ) -> Result<mongreldb_core::StoredProcedure> {
777 let procedure = core_procedure(spec)?;
778 self.inner
779 .create_or_replace_procedure(procedure)
780 .map_err(KitError::from)
781 }
782
783 pub fn drop_procedure(&self, name: &str) -> Result<()> {
784 self.inner.drop_procedure(name).map_err(KitError::from)
785 }
786
787 pub fn call_procedure(
788 &self,
789 name: &str,
790 args: serde_json::Map<String, Value>,
791 ) -> Result<mongreldb_core::ProcedureCallResult> {
792 let args = args
793 .iter()
794 .map(|(key, value)| Ok((key.clone(), json_to_core_value(value)?)))
795 .collect::<Result<HashMap<_, _>>>()?;
796 self.inner
797 .call_procedure(name, args)
798 .map_err(KitError::from)
799 }
800
801 pub fn create_trigger(&self, spec: &TriggerSpec) -> Result<mongreldb_core::StoredTrigger> {
802 let trigger = core_trigger(spec)?;
803 self.inner.create_trigger(trigger).map_err(KitError::from)
804 }
805
806 pub fn replace_trigger(&self, spec: &TriggerSpec) -> Result<mongreldb_core::StoredTrigger> {
807 let trigger = core_trigger(spec)?;
808 self.inner
809 .create_or_replace_trigger(trigger)
810 .map_err(KitError::from)
811 }
812
813 pub fn drop_trigger(&self, name: &str) -> Result<()> {
814 self.inner.drop_trigger(name).map_err(KitError::from)
815 }
816
817 pub fn triggers(&self) -> Vec<mongreldb_core::StoredTrigger> {
818 self.inner.triggers()
819 }
820
821 pub fn trigger(&self, name: &str) -> Option<mongreldb_core::StoredTrigger> {
822 self.inner.trigger(name)
823 }
824
825 pub fn allocate_sequence(&self, name: &str, count: i64) -> Result<i64> {
831 use crate::internal::cols;
832 let mut attempt = 0;
833 loop {
834 let mut txn = self.inner.begin();
835 let snapshot = txn.read_snapshot();
836 let existing = self
837 .visible_core_rows_at(crate::internal::SEQUENCES, snapshot)?
838 .into_iter()
839 .find(|r| internal_bytes(r, cols::SEQ_NAME) == Some(name.to_string()));
840
841 let now = crate::internal::iso_now();
842 let (start, next, old_row_id) = match &existing {
846 Some(row) => {
847 let current = match row.columns.get(&cols::SEQ_NEXT) {
848 Some(CoreValue::Int64(i)) => *i,
849 _ => 1,
850 };
851 (current, current + count, Some(row.row_id))
852 }
853 None => (1, 1 + count, None),
854 };
855
856 if let Some(rid) = old_row_id {
857 txn.delete(crate::internal::SEQUENCES, rid)
858 .map_err(KitError::from)?;
859 }
860 txn.put(
861 crate::internal::SEQUENCES,
862 vec![
863 (cols::SEQ_NAME, CoreValue::Bytes(name.as_bytes().to_vec())),
864 (cols::SEQ_NEXT, CoreValue::Int64(next)),
865 (cols::SEQ_UPDATED, CoreValue::Bytes(now.into_bytes())),
866 ],
867 )
868 .map_err(KitError::from)?;
869 match txn.commit() {
870 Ok(_) => return Ok(start),
871 Err(mongreldb_core::MongrelError::Conflict(_)) if attempt < 10_000 => {
872 attempt += 1;
873 std::thread::yield_now();
874 continue;
875 }
876 Err(e) => return Err(KitError::from(e)),
877 }
878 }
879 }
880
881 pub fn transaction<T, F>(&self, max_retries: usize, mut f: F) -> Result<T>
884 where
885 F: FnMut(&mut crate::txn::Transaction<'_>) -> Result<T>,
886 {
887 let mut attempt = 0;
888 loop {
889 let mut txn = self.begin()?;
890 match f(&mut txn) {
891 Ok(value) => match txn.commit() {
892 Ok(()) => return Ok(value),
893 Err(KitError::Conflict(_)) if attempt < max_retries => {
894 attempt += 1;
895 continue;
896 }
897 Err(e) => return Err(e),
898 },
899 Err(KitError::Conflict(_)) if attempt < max_retries => {
900 txn.rollback();
901 attempt += 1;
902 continue;
903 }
904 Err(e) => {
905 txn.rollback();
906 return Err(e);
907 }
908 }
909 }
910 }
911
912 pub fn table(&self, name: &str) -> Option<&KitTable> {
914 self.schema.table(name)
915 }
916
917 pub fn schema(&self) -> &KitSchema {
919 &self.schema
920 }
921
922 pub fn begin(&self) -> Result<crate::txn::Transaction<'_>> {
924 let core_txn = self.inner.begin();
925 Ok(crate::txn::Transaction::new(self, core_txn))
926 }
927
928 pub fn set_schema(&mut self, schema: KitSchema) {
930 self.schema = schema;
931 }
932
933 pub fn check_internal_tables(&self) -> Result<()> {
936 let schema_file = self.root.join(SCHEMA_FILE);
937 if !schema_file.exists() {
938 return Err(KitError::Integrity(format!(
939 "schema file {} is missing",
940 schema_file.display()
941 )));
942 }
943 for (name, _) in internal_tables_core() {
944 if self.inner.table_id(name).is_err() {
945 return Err(KitError::Integrity(format!(
946 "internal table {name} is missing"
947 )));
948 }
949 }
950 Ok(())
951 }
952
953 pub fn gc(&self) -> Result<usize> {
956 self.inner.gc().map_err(KitError::from)
957 }
958
959 pub fn check(&self) -> Vec<serde_json::Value> {
962 self.inner
963 .check()
964 .into_iter()
965 .map(|i| {
966 serde_json::json!({
967 "table_id": i.table_id,
968 "table_name": i.table_name,
969 "severity": i.severity,
970 "description": i.description,
971 })
972 })
973 .collect()
974 }
975
976 pub fn doctor(&self) -> Result<Vec<u64>> {
978 self.inner.doctor().map_err(KitError::from)
979 }
980
981 pub fn snapshot_epoch(&self) -> u64 {
985 self.inner.snapshot().0.epoch.0
986 }
987
988 pub fn set_history_retention_epochs(&self, epochs: u64) -> Result<()> {
989 self.inner
990 .set_history_retention_epochs(epochs)
991 .map_err(KitError::from)
992 }
993
994 pub fn history_retention_epochs(&self) -> u64 {
995 self.inner.history_retention_epochs()
996 }
997
998 pub fn earliest_retained_epoch(&self) -> u64 {
999 self.inner.earliest_retained_epoch().0
1000 }
1001
1002 pub fn export_tsv(&self, table: &str) -> Result<String> {
1006 let t = self
1007 .schema
1008 .tables
1009 .iter()
1010 .find(|t| t.name == table)
1011 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?
1012 .clone();
1013 let tx = self.begin()?;
1014 let rows = tx.all_rows(table)?;
1015 Ok(crate::tsv::rows_to_tsv(&t, &rows))
1016 }
1017
1018 pub fn import_tsv(&self, table: &str, text: &str) -> Result<usize> {
1022 let t = self
1023 .schema
1024 .tables
1025 .iter()
1026 .find(|t| t.name == table)
1027 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?
1028 .clone();
1029 let rows = crate::tsv::tsv_to_rows(&t, text)?;
1030 let n = rows.len();
1031 self.transaction(1, |tx| {
1032 tx.insert_many(table, rows.clone())?;
1033 Ok(())
1034 })?;
1035 Ok(n)
1036 }
1037
1038 pub fn explain(
1043 &self,
1044 table: &str,
1045 predicate: &mongreldb_kit_core::query::Expr,
1046 ) -> Result<ExplainPlan> {
1047 let t = self
1048 .schema
1049 .tables
1050 .iter()
1051 .find(|t| t.name == table)
1052 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
1053 Ok(match crate::pushdown::translate_predicate(t, predicate) {
1054 Some(p) => ExplainPlan {
1055 index_accelerated: p.can_push(),
1056 exact: p.fully_translated,
1057 pushed_conditions: p.conditions.iter().map(condition_label).collect(),
1058 },
1059 None => ExplainPlan {
1060 index_accelerated: false,
1061 exact: false,
1062 pushed_conditions: Vec::new(),
1063 },
1064 })
1065 }
1066
1067 pub fn rows_at_epoch(&self, table: &str, epoch: u64) -> Result<Vec<crate::schema::Row>> {
1073 let t = self
1074 .schema
1075 .tables
1076 .iter()
1077 .find(|t| t.name == table)
1078 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
1079 let current = self.snapshot_epoch();
1080 if epoch > current {
1081 return Err(KitError::Validation(format!(
1082 "epoch {epoch} is in the future (current committed epoch is {current})"
1083 )));
1084 }
1085 let snap = Snapshot::at(mongreldb_core::epoch::Epoch(epoch));
1086 let rows = self.visible_core_rows_at(table, snap)?;
1087 rows.iter()
1088 .map(|r| crate::schema::core_row_to_json(r, t))
1089 .collect()
1090 }
1091
1092 pub fn approx_aggregate(
1098 &self,
1099 table: &str,
1100 column: Option<&str>,
1101 agg: ApproxAggKind,
1102 z: f64,
1103 ) -> Result<Option<ApproxAggregate>> {
1104 let t = self
1105 .schema
1106 .tables
1107 .iter()
1108 .find(|t| t.name == table)
1109 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
1110 if matches!(agg, ApproxAggKind::Sum | ApproxAggKind::Avg) && column.is_none() {
1111 return Err(KitError::Validation(
1112 "approx sum/avg requires a column".into(),
1113 ));
1114 }
1115 let cid = match column {
1116 Some(name) => Some(
1117 t.columns
1118 .iter()
1119 .find(|c| c.name == name)
1120 .ok_or_else(|| KitError::Validation(format!("unknown column '{name}'")))?
1121 .id as u16,
1122 ),
1123 None => None,
1124 };
1125 let core_agg = match agg {
1126 ApproxAggKind::Count => ApproxAgg::Count,
1127 ApproxAggKind::Sum => ApproxAgg::Sum,
1128 ApproxAggKind::Avg => ApproxAgg::Avg,
1129 };
1130 let handle = self.inner.table(table).map_err(KitError::from)?;
1131 let mut guard = handle.lock();
1132 let res = guard
1133 .approx_aggregate(&[], cid, core_agg, z)
1134 .map_err(KitError::from)?;
1135 Ok(res.map(|r| ApproxAggregate {
1136 point: r.point,
1137 ci_low: r.ci_low,
1138 ci_high: r.ci_high,
1139 n_population: r.n_population,
1140 n_sample_live: r.n_sample_live,
1141 n_passing: r.n_passing,
1142 }))
1143 }
1144
1145 pub fn scan_batched<F>(&self, table: &str, batch_size: usize, mut f: F) -> Result<()>
1151 where
1152 F: FnMut(&[serde_json::Map<String, Value>]) -> Result<()>,
1153 {
1154 let kit_t = self
1155 .schema
1156 .tables
1157 .iter()
1158 .find(|t| t.name == table)
1159 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
1160 let batch_size = batch_size.max(1);
1161 let (snapshot, _pin) = self.inner.snapshot();
1164 let handle = self.inner.table(table).map_err(KitError::from)?;
1165 let guard = handle.lock();
1166
1167 let mut projection: Vec<(u16, mongreldb_core::schema::TypeId)> = Vec::new();
1169 let mut meta: Vec<(String, mongreldb_kit_core::schema::ColumnType)> = Vec::new();
1170 for c in &guard.schema().columns {
1171 if let Some(kc) = kit_t.columns.iter().find(|kc| kc.id as u16 == c.id) {
1172 projection.push((c.id, c.ty.clone()));
1173 meta.push((kc.name.clone(), kc.storage_type));
1174 }
1175 }
1176
1177 match guard
1178 .scan_cursor(snapshot, projection, &[])
1179 .map_err(KitError::from)?
1180 {
1181 Some(mut cursor) => {
1182 let mut buf: Vec<serde_json::Map<String, Value>> = Vec::with_capacity(batch_size);
1183 while let Some(batch) = cursor.next_batch().map_err(KitError::from)? {
1184 let nrows = batch.first().map(|c| c.len()).unwrap_or(0);
1185 for j in 0..nrows {
1186 let mut m = serde_json::Map::new();
1187 for (ci, (name, ty)) in meta.iter().enumerate() {
1188 let cv = batch
1189 .get(ci)
1190 .and_then(|col| col.value_at(j))
1191 .unwrap_or(CoreValue::Null);
1192 m.insert(name.clone(), crate::schema::core_to_json(&cv, *ty)?);
1193 }
1194 buf.push(m);
1195 if buf.len() >= batch_size {
1196 f(&buf)?;
1197 buf.clear();
1198 }
1199 }
1200 }
1201 if !buf.is_empty() {
1202 f(&buf)?;
1203 }
1204 Ok(())
1205 }
1206 None => {
1207 drop(guard);
1208 let rows = self.visible_core_rows_at(table, snapshot)?;
1209 let maps: Vec<serde_json::Map<String, Value>> = rows
1210 .iter()
1211 .map(|r| crate::schema::core_row_to_json(r, kit_t).map(|row| row.values))
1212 .collect::<Result<Vec<_>>>()?;
1213 for chunk in maps.chunks(batch_size) {
1214 f(chunk)?;
1215 }
1216 Ok(())
1217 }
1218 }
1219 }
1220
1221 pub fn set_similarity(
1230 &self,
1231 table: &str,
1232 column: &str,
1233 query: &[String],
1234 k: usize,
1235 ) -> Result<Vec<SimilarRow>> {
1236 let t = self
1237 .schema
1238 .tables
1239 .iter()
1240 .find(|t| t.name == table)
1241 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
1242 let col = t.columns.iter().find(|c| c.name == column).ok_or_else(|| {
1243 KitError::Validation(format!("unknown column '{column}' on table '{table}'"))
1244 })?;
1245 let query_set: std::collections::HashSet<String> = query.iter().cloned().collect();
1246
1247 let has_minhash = t.indexes.iter().any(|idx| {
1248 idx.kind == KitIndexKind::MinHash && idx.columns.iter().any(|c| c == column)
1249 });
1250 let rows = if has_minhash {
1251 let query_hashes: Vec<u64> = query
1253 .iter()
1254 .map(|s| mongreldb_core::index::minhash_token_hash(s))
1255 .collect();
1256 let cand_k = k.saturating_mul(8).max(k + 64);
1258 let cond = mongreldb_core::query::Condition::MinHashSimilar {
1259 column_id: col.id as u16,
1260 query: query_hashes,
1261 k: cand_k,
1262 };
1263 let (snapshot, _pin) = self.inner.snapshot();
1264 let core_rows = self.query_core_rows_at(table, &[cond], snapshot)?;
1265 core_rows
1266 .iter()
1267 .map(|r| crate::schema::core_row_to_json(r, t))
1268 .collect::<Result<Vec<_>>>()?
1269 } else {
1270 let tx = self.begin()?;
1271 tx.all_rows(table)?
1272 };
1273
1274 let mut scored: Vec<SimilarRow> = Vec::new();
1275 for row in rows {
1276 let set = parse_string_set(row.values.get(column));
1277 let inter = set.iter().filter(|x| query_set.contains(*x)).count();
1278 let union = set.len() + query_set.len() - inter;
1279 let sim = if union == 0 {
1280 0.0
1281 } else {
1282 inter as f64 / union as f64
1283 };
1284 if sim > 0.0 {
1285 scored.push(SimilarRow {
1286 row,
1287 similarity: sim,
1288 });
1289 }
1290 }
1291 scored.sort_by(|a, b| {
1292 b.similarity
1293 .partial_cmp(&a.similarity)
1294 .unwrap_or(std::cmp::Ordering::Equal)
1295 });
1296 scored.truncate(k);
1297 Ok(scored)
1298 }
1299
1300 pub fn flush(&self) -> Result<()> {
1304 for name in self.inner.table_names() {
1305 let handle = self.inner.table(&name).map_err(KitError::from)?;
1306 let mut guard = handle.lock();
1307 guard.flush().map_err(KitError::from)?;
1308 }
1309 Ok(())
1310 }
1311
1312 pub fn incremental_aggregate(
1324 &self,
1325 table: &str,
1326 column: Option<&str>,
1327 agg: IncrementalAggKind,
1328 filter: Option<&mongreldb_kit_core::query::Expr>,
1329 ) -> Result<IncrementalAggregate> {
1330 let t = self
1331 .schema
1332 .tables
1333 .iter()
1334 .find(|t| t.name == table)
1335 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
1336 if !matches!(agg, IncrementalAggKind::Count) && column.is_none() {
1337 return Err(KitError::Validation(
1338 "sum/min/max/avg incremental aggregate requires a column".into(),
1339 ));
1340 }
1341 let cid = match column {
1342 Some(name) => Some(
1343 t.columns
1344 .iter()
1345 .find(|c| c.name == name)
1346 .ok_or_else(|| KitError::Validation(format!("unknown column '{name}'")))?
1347 .id as u16,
1348 ),
1349 None => None,
1350 };
1351 let conditions = match filter {
1352 Some(expr) => {
1353 let plan = crate::pushdown::translate_predicate(t, expr).ok_or_else(|| {
1354 KitError::Validation(
1355 "filter is not index-translatable for an incremental aggregate".into(),
1356 )
1357 })?;
1358 if !plan.fully_translated {
1359 return Err(KitError::Validation(
1360 "filter has a residual that an incremental aggregate cannot apply exactly"
1361 .into(),
1362 ));
1363 }
1364 plan.conditions
1365 }
1366 None => Vec::new(),
1367 };
1368 let core_agg = match agg {
1369 IncrementalAggKind::Count => NativeAgg::Count,
1370 IncrementalAggKind::Sum => NativeAgg::Sum,
1371 IncrementalAggKind::Min => NativeAgg::Min,
1372 IncrementalAggKind::Max => NativeAgg::Max,
1373 IncrementalAggKind::Avg => NativeAgg::Avg,
1374 };
1375 let cache_key = incremental_cache_key(t.id, cid, agg, &conditions);
1376 let handle = self.inner.table(table).map_err(KitError::from)?;
1377 let mut guard = handle.lock();
1378 let res = guard
1379 .aggregate_incremental(cache_key, &conditions, cid, core_agg)
1380 .map_err(KitError::from)?;
1381 Ok(IncrementalAggregate {
1382 value: agg_state_value(&res.state),
1383 incremental: res.incremental,
1384 delta_rows: res.delta_rows,
1385 })
1386 }
1387
1388 pub fn applied_migrations(&self) -> Result<Vec<mongreldb_kit_core::migrations::Migration>> {
1390 crate::migrate::load_applied_migrations(&self.inner)
1391 }
1392
1393 pub(crate) fn core_db(&self) -> &CoreDatabase {
1394 &self.inner
1395 }
1396
1397 pub(crate) fn core_arc(&self) -> Arc<CoreDatabase> {
1400 Arc::clone(&self.inner)
1401 }
1402
1403 pub fn close(&self) -> Result<()> {
1408 self.inner.close().map_err(KitError::from)
1409 }
1410
1411 pub fn compact_all(&self) -> Result<(usize, usize)> {
1416 self.inner.compact().map_err(KitError::from)
1417 }
1418
1419 pub fn compact_table(&self, name: &str) -> Result<bool> {
1422 self.inner.compact_table(name).map_err(KitError::from)
1423 }
1424
1425 pub fn rename_table(&mut self, from: &str, to: &str) -> Result<()> {
1435 if from.starts_with("__kit_") || to.starts_with("__kit_") {
1436 return Err(KitError::Validation(
1437 "rename_table: names beginning with '__kit_' are reserved for internal tables"
1438 .into(),
1439 ));
1440 }
1441 self.inner.rename_table(from, to).map_err(KitError::from)?;
1442 if !self.schema.rename_table(from, to) {
1445 return Err(KitError::Integrity(format!(
1448 "rename_table: kit schema has no table '{from}' (or '{to}' already exists)"
1449 )));
1450 }
1451 for table in &mut self.schema.tables {
1452 for fk in &mut table.foreign_keys {
1453 if fk.references_table == from {
1454 fk.references_table = to.to_string();
1455 }
1456 }
1457 }
1458 store_schema(&self.root, &self.schema)?;
1459 Ok(())
1460 }
1461
1462 pub fn analyze(&self) -> Result<()> {
1467 for name in self.inner.table_names() {
1468 let handle = self.inner.table(&name).map_err(KitError::from)?;
1469 handle.lock().ensure_indexes_complete()?;
1470 }
1471 Ok(())
1472 }
1473
1474 pub fn vacuum(&self) -> Result<usize> {
1478 self.inner.compact().map_err(KitError::from)?;
1479 self.inner.gc().map_err(KitError::from)
1480 }
1481
1482 pub fn create_view(&self, spec: &ViewSpec) -> Result<()> {
1488 self.sql(&spec.create_sql())?;
1489 Ok(())
1490 }
1491
1492 pub fn drop_view(&self, name: &str) -> Result<()> {
1494 self.sql(&format!("DROP VIEW IF EXISTS {name}"))?;
1495 Ok(())
1496 }
1497
1498 pub fn reserve_auto_inc(&self, table: &str) -> Result<Option<i64>> {
1506 let handle = self.inner.table(table).map_err(KitError::from)?;
1507 let mut guard = handle.lock();
1508 guard.reserve_auto_inc().map_err(KitError::from)
1509 }
1510
1511 pub fn create_user(&self, username: &str, password: &str) -> Result<()> {
1515 self.inner
1516 .create_user(username, password)
1517 .map_err(KitError::from)?;
1518 Ok(())
1519 }
1520
1521 pub fn drop_user(&self, username: &str) -> Result<()> {
1523 self.inner.drop_user(username).map_err(KitError::from)
1524 }
1525
1526 pub fn alter_user_password(&self, username: &str, new_password: &str) -> Result<()> {
1528 self.inner
1529 .alter_user_password(username, new_password)
1530 .map_err(KitError::from)
1531 }
1532
1533 pub fn verify_user(
1535 &self,
1536 username: &str,
1537 password: &str,
1538 ) -> Result<Option<mongreldb_core::auth::UserEntry>> {
1539 self.inner
1540 .verify_user(username, password)
1541 .map_err(KitError::from)
1542 }
1543
1544 pub fn set_user_admin(&self, username: &str, is_admin: bool) -> Result<()> {
1546 self.inner
1547 .set_user_admin(username, is_admin)
1548 .map_err(KitError::from)
1549 }
1550
1551 pub fn users(&self) -> Vec<String> {
1553 self.inner.users().into_iter().map(|u| u.username).collect()
1554 }
1555
1556 pub fn create_role(&self, name: &str) -> Result<()> {
1558 self.inner.create_role(name).map_err(KitError::from)?;
1559 Ok(())
1560 }
1561
1562 pub fn drop_role(&self, name: &str) -> Result<()> {
1564 self.inner.drop_role(name).map_err(KitError::from)
1565 }
1566
1567 pub fn roles(&self) -> Vec<String> {
1569 self.inner.roles().into_iter().map(|r| r.name).collect()
1570 }
1571
1572 pub fn grant_role(&self, username: &str, role_name: &str) -> Result<()> {
1574 self.inner
1575 .grant_role(username, role_name)
1576 .map_err(KitError::from)
1577 }
1578
1579 pub fn revoke_role(&self, username: &str, role_name: &str) -> Result<()> {
1581 self.inner
1582 .revoke_role(username, role_name)
1583 .map_err(KitError::from)
1584 }
1585
1586 pub fn grant_permission(
1588 &self,
1589 role_name: &str,
1590 permission: mongreldb_core::auth::Permission,
1591 ) -> Result<()> {
1592 self.inner
1593 .grant_permission(role_name, permission)
1594 .map_err(KitError::from)
1595 }
1596
1597 pub fn revoke_permission(
1599 &self,
1600 role_name: &str,
1601 permission: mongreldb_core::auth::Permission,
1602 ) -> Result<()> {
1603 self.inner
1604 .revoke_permission(role_name, permission)
1605 .map_err(KitError::from)
1606 }
1607
1608 pub fn set_spill_threshold(&self, bytes: u64) {
1614 self.inner.set_spill_threshold(bytes);
1615 }
1616
1617 pub fn set_recursive_triggers(&self, enabled: bool) {
1619 self.inner.set_recursive_triggers(enabled);
1620 }
1621
1622 pub fn trigger_config(&self) -> mongreldb_core::TriggerConfig {
1624 self.inner.trigger_config()
1625 }
1626
1627 pub fn set_trigger_config(&self, config: mongreldb_core::TriggerConfig) -> Result<()> {
1629 self.inner
1630 .set_trigger_config(config)
1631 .map_err(KitError::from)
1632 }
1633
1634 pub fn set_table_compaction_zstd_level(&self, table: &str, level: i32) -> Result<()> {
1636 let handle = self.inner.table(table).map_err(KitError::from)?;
1637 handle.lock().set_compaction_zstd_level(level);
1638 Ok(())
1639 }
1640
1641 pub fn set_table_result_cache_max_bytes(&self, table: &str, max_bytes: u64) -> Result<()> {
1643 let handle = self.inner.table(table).map_err(KitError::from)?;
1644 handle.lock().set_result_cache_max_bytes(max_bytes);
1645 Ok(())
1646 }
1647
1648 pub fn set_table_mutable_run_spill_bytes(&self, table: &str, bytes: u64) -> Result<()> {
1650 let handle = self.inner.table(table).map_err(KitError::from)?;
1651 handle.lock().set_mutable_run_spill_bytes(bytes);
1652 Ok(())
1653 }
1654
1655 pub fn set_table_sync_byte_threshold(&self, table: &str, threshold: u64) -> Result<()> {
1657 let handle = self.inner.table(table).map_err(KitError::from)?;
1658 handle.lock().set_sync_byte_threshold(threshold);
1659 Ok(())
1660 }
1661
1662 pub fn set_table_index_build_policy(
1665 &self,
1666 table: &str,
1667 policy: mongreldb_core::IndexBuildPolicy,
1668 ) -> Result<()> {
1669 let handle = self.inner.table(table).map_err(KitError::from)?;
1670 handle.lock().set_index_build_policy(policy);
1671 Ok(())
1672 }
1673
1674 pub fn table_page_cache_stats(&self, table: &str) -> Result<mongreldb_core::cache::CacheStats> {
1676 let handle = self.inner.table(table).map_err(KitError::from)?;
1677 let stats = handle.lock().page_cache_stats();
1678 Ok(stats)
1679 }
1680
1681 pub fn table_run_count(&self, table: &str) -> Result<usize> {
1683 let handle = self.inner.table(table).map_err(KitError::from)?;
1684 let n = handle.lock().run_count();
1685 Ok(n)
1686 }
1687
1688 pub fn table_memtable_len(&self, table: &str) -> Result<usize> {
1690 let handle = self.inner.table(table).map_err(KitError::from)?;
1691 let n = handle.lock().memtable_len();
1692 Ok(n)
1693 }
1694
1695 pub fn table_mutable_run_len(&self, table: &str) -> Result<usize> {
1697 let handle = self.inner.table(table).map_err(KitError::from)?;
1698 let n = handle.lock().mutable_run_len();
1699 Ok(n)
1700 }
1701
1702 pub fn table_page_cache_len(&self, table: &str) -> Result<usize> {
1704 let handle = self.inner.table(table).map_err(KitError::from)?;
1705 let n = handle.lock().page_cache_len();
1706 Ok(n)
1707 }
1708
1709 pub fn table_decoded_cache_len(&self, table: &str) -> Result<usize> {
1711 let handle = self.inner.table(table).map_err(KitError::from)?;
1712 let n = handle.lock().decoded_cache_len();
1713 Ok(n)
1714 }
1715
1716 pub fn sql(&self, statement: &str) -> Result<Vec<arrow::record_batch::RecordBatch>> {
1734 self.sql_with_options(statement, SqlOptions::default())
1735 }
1736
1737 fn sql_session(&self) -> Result<Arc<mongreldb_query::MongrelSession>> {
1738 if let Some(session) = self.session.read().as_ref() {
1739 return Ok(Arc::clone(session));
1740 }
1741 let session = Arc::new(
1742 mongreldb_query::MongrelSession::open(self.core_arc()).map_err(KitError::from)?,
1743 );
1744 let mut cached = self.session.write();
1745 Ok(Arc::clone(cached.get_or_insert(session)))
1746 }
1747
1748 #[doc(hidden)]
1749 pub fn set_sql_test_hook(&self, hook: Option<mongreldb_query::SqlTestHook>) -> Result<()> {
1750 self.sql_session()?.set_test_hook(hook);
1751 Ok(())
1752 }
1753
1754 pub fn start_sql(
1755 &self,
1756 statement: impl Into<String>,
1757 options: SqlOptions,
1758 ) -> Result<SqlQueryHandle> {
1759 let session = self.sql_session()?;
1760 let query = session
1761 .register_query(mongreldb_query::SqlQueryOptions {
1762 query_id: options.query_id,
1763 timeout: options.timeout,
1764 ..mongreldb_query::SqlQueryOptions::default()
1765 })
1766 .map_err(KitError::from)?;
1767 let query_id = query.id();
1768 let registration = mongreldb_query::RegisteredQueryGuard::new(query);
1769 let worker_session = Arc::clone(&session);
1770 let statement = statement.into();
1771 let worker = std::thread::Builder::new()
1772 .name(format!("mongreldb-kit-sql-{query_id}"))
1773 .spawn(move || {
1774 sql_runtime().block_on(
1775 worker_session
1776 .run_with_query_for_serialization(&statement, registration.into_query()),
1777 )
1778 })
1779 .map_err(|error| KitError::Storage(error.to_string()))?;
1780 Ok(SqlQueryHandle {
1781 query_id,
1782 session,
1783 worker: Some(worker),
1784 })
1785 }
1786
1787 pub fn sql_with_options(
1788 &self,
1789 statement: &str,
1790 options: SqlOptions,
1791 ) -> Result<Vec<arrow::record_batch::RecordBatch>> {
1792 self.start_sql(statement, options)?.wait()
1793 }
1794
1795 pub fn cancel_sql(&self, query_id: mongreldb_query::QueryId) -> mongreldb_query::CancelOutcome {
1796 self.session
1797 .read()
1798 .as_ref()
1799 .map_or(mongreldb_query::CancelOutcome::NotFound, |session| {
1800 session.cancel_query(query_id)
1801 })
1802 }
1803
1804 pub fn sql_query_status(
1805 &self,
1806 query_id: mongreldb_query::QueryId,
1807 ) -> Result<Option<mongreldb_query::QueryStatus>> {
1808 Ok(self.sql_session()?.query_registry().status(query_id))
1809 }
1810
1811 pub fn refresh_sql_session(&self) -> Result<()> {
1816 let session =
1817 mongreldb_query::MongrelSession::open(self.core_arc()).map_err(KitError::from)?;
1818 *self.session.write() = Some(Arc::new(session));
1819 Ok(())
1820 }
1821
1822 pub fn sql_arrow(&self, statement: &str) -> Result<Vec<u8>> {
1828 self.sql_arrow_with_options(statement, SqlOptions::default())
1829 }
1830
1831 pub fn sql_arrow_with_options(&self, statement: &str, options: SqlOptions) -> Result<Vec<u8>> {
1832 self.sql_serialized_with_options(statement, options, |output| {
1833 crate::arrow_util::batches_to_ipc_controlled(output.batches(), output.query())
1834 })
1835 }
1836
1837 pub fn sql_rows(&self, statement: &str) -> Result<Vec<serde_json::Map<String, Value>>> {
1841 self.sql_rows_with_options(statement, SqlOptions::default())
1842 }
1843
1844 pub fn sql_rows_with_options(
1845 &self,
1846 statement: &str,
1847 options: SqlOptions,
1848 ) -> Result<Vec<serde_json::Map<String, Value>>> {
1849 self.sql_serialized_with_options(statement, options, |output| {
1850 crate::arrow_util::batches_to_rows_controlled(output.batches(), output.query())
1851 })
1852 }
1853
1854 fn sql_serialized_with_options<T>(
1855 &self,
1856 statement: &str,
1857 options: SqlOptions,
1858 serialize: impl FnOnce(&mongreldb_query::ManagedQueryBatches) -> Result<T>,
1859 ) -> Result<T> {
1860 let session = self.sql_session()?;
1861 let query = session
1862 .register_query(mongreldb_query::SqlQueryOptions {
1863 query_id: options.query_id,
1864 timeout: options.timeout,
1865 ..mongreldb_query::SqlQueryOptions::default()
1866 })
1867 .map_err(KitError::from)?;
1868 let query_id = query.id();
1869 let output = sql_runtime()
1870 .block_on(session.run_with_query_for_serialization(statement, query))
1871 .map_err(|error| {
1872 let status = session.query_registry().status(query_id);
1873 crate::error::query_error_with_status(error, status.as_ref())
1874 })?;
1875 session.fire_test_hook(mongreldb_query::SqlTestHookPoint::BeforeSerializationBatch);
1876 match serialize(&output) {
1877 Ok(value) => {
1878 complete_sql_output(output)?;
1879 Ok(value)
1880 }
1881 Err(error) => {
1882 fail_sql_output(output, &error);
1883 Err(error)
1884 }
1885 }
1886 }
1887
1888 pub(crate) fn lookup_row_id(&self, table: &str, key: &[u8]) -> Result<Option<RowId>> {
1892 let handle = self.inner.table(table).map_err(KitError::from)?;
1893 let mut guard = handle.lock();
1894 guard.ensure_indexes_complete()?;
1895 Ok(guard.lookup_pk(key))
1896 }
1897
1898 pub(crate) fn root(&self) -> &Path {
1899 &self.root
1900 }
1901
1902 pub(crate) fn visible_core_rows_at(
1906 &self,
1907 table_name: &str,
1908 snapshot: Snapshot,
1909 ) -> Result<Vec<CoreRow>> {
1910 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1911 let guard = handle.lock();
1912 guard.visible_rows(snapshot).map_err(KitError::from)
1913 }
1914
1915 pub(crate) fn query_core_rows_at(
1922 &self,
1923 table_name: &str,
1924 conditions: &[mongreldb_core::query::Condition],
1925 snapshot: Snapshot,
1926 ) -> Result<Vec<CoreRow>> {
1927 if conditions.is_empty() {
1928 return self.visible_core_rows_at(table_name, snapshot);
1929 }
1930 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1931 let mut guard = handle.lock();
1932 let q = conditions
1933 .iter()
1934 .cloned()
1935 .fold(mongreldb_core::query::Query::new(), |query, condition| {
1936 query.and(condition)
1937 });
1938 guard.query(&q).map_err(KitError::from)
1939 }
1940
1941 pub(crate) fn flush_table(&self, table_name: &str) -> Result<()> {
1949 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1950 handle.lock().flush().map_err(KitError::from)?;
1951 Ok(())
1952 }
1953
1954 pub(crate) fn count_core_rows_at(
1964 &self,
1965 table_name: &str,
1966 conditions: &[mongreldb_core::query::Condition],
1967 snapshot: Snapshot,
1968 ) -> Result<Option<u64>> {
1969 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1970 let mut guard = handle.lock();
1971 if guard.snapshot().epoch != snapshot.epoch {
1972 return Ok(None); }
1974 guard
1975 .count_conditions(conditions, snapshot)
1976 .map_err(KitError::from)
1977 }
1978
1979 pub(crate) fn aggregate_core_at(
1988 &self,
1989 table_name: &str,
1990 column: Option<u16>,
1991 conditions: &[mongreldb_core::query::Condition],
1992 agg: NativeAgg,
1993 snapshot: Snapshot,
1994 ) -> Result<Option<NativeAggResult>> {
1995 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1996 let guard = handle.lock();
1997 if guard.snapshot().epoch != snapshot.epoch {
1998 return Ok(None); }
2000 guard
2001 .aggregate_native(snapshot, column, conditions, agg)
2002 .map_err(KitError::from)
2003 }
2004
2005 pub(crate) fn count_distinct_core_at(
2014 &self,
2015 table_name: &str,
2016 column_id: u16,
2017 snapshot: Snapshot,
2018 ) -> Result<Option<u64>> {
2019 let handle = self.inner.table(table_name).map_err(KitError::from)?;
2020 let mut guard = handle.lock();
2021 if guard.snapshot().epoch != snapshot.epoch {
2022 return Ok(None); }
2024 guard
2025 .count_distinct_from_bitmap(column_id)
2026 .map_err(KitError::from)
2027 }
2028
2029 #[allow(dead_code)]
2031 pub(crate) fn get_core_row(&self, table_name: &str, row_id: u64) -> Result<Option<CoreRow>> {
2032 let handle = self.inner.table(table_name).map_err(KitError::from)?;
2033 let guard = handle.lock();
2034 let snapshot = guard.snapshot();
2035 Ok(guard.get(mongreldb_core::RowId(row_id), snapshot))
2036 }
2037}
2038
2039pub(crate) fn create_core_table(db: &CoreDatabase, name: &str, schema: CoreSchema) -> Result<()> {
2040 if db.table_id(name).is_ok() {
2041 return Ok(());
2042 }
2043 db.create_table(name, schema).map_err(KitError::from)?;
2044 Ok(())
2045}
2046
2047fn sql_runtime() -> &'static tokio::runtime::Runtime {
2051 use std::sync::OnceLock;
2052 static RT: OnceLock<tokio::runtime::Runtime> = OnceLock::new();
2053 RT.get_or_init(|| {
2054 tokio::runtime::Builder::new_multi_thread()
2055 .worker_threads(4)
2056 .enable_all()
2057 .build()
2058 .expect("failed to build kit SQL tokio runtime")
2059 })
2060}
2061
2062fn core_procedure(spec: &ProcedureSpec) -> Result<mongreldb_core::StoredProcedure> {
2063 let parsed: mongreldb_core::StoredProcedure =
2064 serde_json::from_value(spec.json.clone()).map_err(KitError::from)?;
2065 mongreldb_core::StoredProcedure::new(parsed.name, parsed.mode, parsed.params, parsed.body, 0)
2066 .map_err(KitError::from)
2067}
2068
2069fn core_trigger(spec: &TriggerSpec) -> Result<mongreldb_core::StoredTrigger> {
2070 let parsed: mongreldb_core::StoredTrigger =
2071 serde_json::from_value(spec.json.clone()).map_err(KitError::from)?;
2072 mongreldb_core::StoredTrigger::new(
2073 parsed.name,
2074 mongreldb_core::TriggerDefinition {
2075 target: parsed.target,
2076 timing: parsed.timing,
2077 event: parsed.event,
2078 update_of: parsed.update_of,
2079 target_columns: parsed.target_columns,
2080 when: parsed.when,
2081 program: parsed.program,
2082 },
2083 0,
2084 )
2085 .map_err(KitError::from)
2086}
2087
2088fn json_to_core_value(value: &Value) -> Result<CoreValue> {
2089 match value {
2090 Value::Null => Ok(CoreValue::Null),
2091 Value::Bool(value) => Ok(CoreValue::Bool(*value)),
2092 Value::Number(value) => {
2093 if let Some(value) = value.as_i64() {
2094 Ok(CoreValue::Int64(value))
2095 } else if let Some(value) = value.as_f64() {
2096 Ok(CoreValue::Float64(value))
2097 } else {
2098 Err(KitError::Validation("unsupported JSON number".into()))
2099 }
2100 }
2101 Value::String(value) => Ok(CoreValue::Bytes(value.as_bytes().to_vec())),
2102 Value::Array(_) | Value::Object(_) => Err(KitError::Validation(
2103 "procedure args only support scalar JSON values".into(),
2104 )),
2105 }
2106}
2107
2108pub(crate) fn internal_bytes(row: &CoreRow, col_id: u16) -> Option<String> {
2110 match row.columns.get(&col_id) {
2111 Some(CoreValue::Bytes(b)) => String::from_utf8(b.clone()).ok(),
2112 _ => None,
2113 }
2114}
2115
2116fn reap_rotated_wal_segments(db: &CoreDatabase) {
2131 let _ = db.gc();
2132}
2133
2134pub(crate) fn load_schema(path: &Path) -> Result<KitSchema> {
2135 let file = path.join(SCHEMA_FILE);
2136 let json = std::fs::read_to_string(&file)
2137 .map_err(|e| KitError::Migration(format!("cannot read schema file: {e}")))?;
2138 let schema: KitSchema = serde_json::from_str(&json)?;
2139 Ok(schema)
2140}
2141
2142pub(crate) fn store_schema(path: &Path, schema: &KitSchema) -> Result<()> {
2143 let file = path.join(SCHEMA_FILE);
2144 let json = serde_json::to_string_pretty(schema)?;
2145 std::fs::write(&file, json)?;
2146 Ok(())
2147}
2148
2149pub(crate) fn persist_schema(db: &Database, schema: &KitSchema) -> Result<()> {
2151 store_schema(&db.root, schema)
2152}
2153
2154#[cfg(test)]
2155mod tests {
2156 use super::open_core_with_retry;
2157
2158 fn lock_error() -> mongreldb_core::MongrelError {
2159 mongreldb_core::MongrelError::DatabaseLocked {
2160 path: "/tmp/db".into(),
2161 message: "would block".into(),
2162 }
2163 }
2164
2165 #[test]
2166 fn open_retry_waits_for_lock_contention_only() {
2167 let mut calls = 0;
2168 let value = open_core_with_retry(50, || {
2169 calls += 1;
2170 if calls < 3 {
2171 Err(lock_error())
2172 } else {
2173 Ok(7)
2174 }
2175 })
2176 .unwrap();
2177 assert_eq!(value, 7);
2178 assert_eq!(calls, 3);
2179
2180 let mut non_lock_calls = 0;
2181 let err: mongreldb_core::Result<()> = open_core_with_retry(50, || {
2182 non_lock_calls += 1;
2183 Err(mongreldb_core::MongrelError::Other("nope".into()))
2184 });
2185 let err = err.unwrap_err();
2186 assert_eq!(non_lock_calls, 1);
2187 assert!(matches!(err, mongreldb_core::MongrelError::Other(_)));
2188 }
2189}