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 sequence_lock: parking_lot::Mutex<()>,
418}
419
420impl Database {
421 pub fn open(path: &Path) -> Result<Self> {
423 let inner = Arc::new(CoreDatabase::open(path)?);
424 let schema = load_schema(path)?;
425 ensure_internal_tables(&inner)?;
427 reap_rotated_wal_segments(&inner);
428 Ok(Self {
429 inner,
430 schema,
431 root: path.to_path_buf(),
432 default_providers: HashMap::new(),
433 session: parking_lot::RwLock::new(None),
434 sequence_lock: parking_lot::Mutex::new(()),
435 })
436 }
437
438 pub fn open_with_options(path: &Path, opts: OpenOptions) -> Result<Self> {
446 let inner = Arc::new(open_core_with_retry(opts.lock_timeout_ms, || {
447 CoreDatabase::open(path)
448 })?);
449 let schema = load_schema(path)?;
450 ensure_internal_tables(&inner)?;
451 reap_rotated_wal_segments(&inner);
452 Ok(Self {
453 inner,
454 schema,
455 root: path.to_path_buf(),
456 default_providers: HashMap::new(),
457 session: parking_lot::RwLock::new(None),
458 sequence_lock: parking_lot::Mutex::new(()),
459 })
460 }
461
462 pub fn open_encrypted(path: &Path, passphrase: &str) -> Result<Self> {
464 let inner = Arc::new(CoreDatabase::open_encrypted(path, passphrase)?);
465 let schema = load_schema(path)?;
466 ensure_internal_tables(&inner)?;
467 reap_rotated_wal_segments(&inner);
468 Ok(Self {
469 inner,
470 schema,
471 root: path.to_path_buf(),
472 default_providers: HashMap::new(),
473 session: parking_lot::RwLock::new(None),
474 sequence_lock: parking_lot::Mutex::new(()),
475 })
476 }
477
478 pub fn open_encrypted_with_options(
482 path: &Path,
483 passphrase: &str,
484 opts: OpenOptions,
485 ) -> Result<Self> {
486 let inner = Arc::new(open_core_with_retry(opts.lock_timeout_ms, || {
487 CoreDatabase::open_encrypted(path, passphrase)
488 })?);
489 let schema = load_schema(path)?;
490 ensure_internal_tables(&inner)?;
491 reap_rotated_wal_segments(&inner);
492 Ok(Self {
493 inner,
494 schema,
495 root: path.to_path_buf(),
496 default_providers: HashMap::new(),
497 session: parking_lot::RwLock::new(None),
498 sequence_lock: parking_lot::Mutex::new(()),
499 })
500 }
501
502 pub fn create_encrypted(path: &Path, schema: KitSchema, passphrase: &str) -> Result<Self> {
506 std::fs::create_dir_all(path)?;
507 let inner = Arc::new(CoreDatabase::create_encrypted(path, passphrase)?);
508 ensure_internal_tables(&inner)?;
509 store_schema(path, &schema)?;
510 for table in &schema.tables {
511 create_core_table(&inner, &table.name, to_core_schema(table)?)?;
512 }
513 Ok(Self {
514 inner,
515 schema,
516 root: path.to_path_buf(),
517 default_providers: HashMap::new(),
518 session: parking_lot::RwLock::new(None),
519 sequence_lock: parking_lot::Mutex::new(()),
520 })
521 }
522
523 pub fn create(path: &Path, schema: KitSchema) -> Result<Self> {
525 std::fs::create_dir_all(path)?;
526 let inner = Arc::new(CoreDatabase::create(path)?);
527
528 ensure_internal_tables(&inner)?;
531
532 store_schema(path, &schema)?;
535
536 for table in &schema.tables {
538 create_core_table(&inner, &table.name, to_core_schema(table)?)?;
539 }
540
541 Ok(Self {
542 inner,
543 schema,
544 root: path.to_path_buf(),
545 default_providers: HashMap::new(),
546 session: parking_lot::RwLock::new(None),
547 sequence_lock: parking_lot::Mutex::new(()),
548 })
549 }
550
551 pub fn open_with_credentials(path: &Path, username: &str, password: &str) -> Result<Self> {
561 let inner = Arc::new(CoreDatabase::open_with_credentials(
562 path, username, password,
563 )?);
564 let schema = load_schema(path)?;
565 ensure_internal_tables(&inner)?;
566 reap_rotated_wal_segments(&inner);
567 Ok(Self {
568 inner,
569 schema,
570 root: path.to_path_buf(),
571 default_providers: HashMap::new(),
572 session: parking_lot::RwLock::new(None),
573 sequence_lock: parking_lot::Mutex::new(()),
574 })
575 }
576
577 pub fn open_with_credentials_and_options(
581 path: &Path,
582 username: &str,
583 password: &str,
584 opts: OpenOptions,
585 ) -> Result<Self> {
586 let inner = Arc::new(open_core_with_retry(opts.lock_timeout_ms, || {
587 CoreDatabase::open_with_credentials(path, username, password)
588 })?);
589 let schema = load_schema(path)?;
590 ensure_internal_tables(&inner)?;
591 reap_rotated_wal_segments(&inner);
592 Ok(Self {
593 inner,
594 schema,
595 root: path.to_path_buf(),
596 default_providers: HashMap::new(),
597 session: parking_lot::RwLock::new(None),
598 sequence_lock: parking_lot::Mutex::new(()),
599 })
600 }
601
602 pub fn create_with_credentials(
608 path: &Path,
609 schema: KitSchema,
610 admin_username: &str,
611 admin_password: &str,
612 ) -> Result<Self> {
613 std::fs::create_dir_all(path)?;
614 let inner = Arc::new(CoreDatabase::create_with_credentials(
615 path,
616 admin_username,
617 admin_password,
618 )?);
619 ensure_internal_tables(&inner)?;
620 store_schema(path, &schema)?;
621 for table in &schema.tables {
622 create_core_table(&inner, &table.name, to_core_schema(table)?)?;
623 }
624 Ok(Self {
625 inner,
626 schema,
627 root: path.to_path_buf(),
628 default_providers: HashMap::new(),
629 session: parking_lot::RwLock::new(None),
630 sequence_lock: parking_lot::Mutex::new(()),
631 })
632 }
633
634 pub fn open_encrypted_with_credentials(
637 path: &Path,
638 passphrase: &str,
639 username: &str,
640 password: &str,
641 ) -> Result<Self> {
642 let inner = Arc::new(CoreDatabase::open_encrypted_with_credentials(
643 path, passphrase, username, password,
644 )?);
645 let schema = load_schema(path)?;
646 ensure_internal_tables(&inner)?;
647 reap_rotated_wal_segments(&inner);
648 Ok(Self {
649 inner,
650 schema,
651 root: path.to_path_buf(),
652 default_providers: HashMap::new(),
653 session: parking_lot::RwLock::new(None),
654 sequence_lock: parking_lot::Mutex::new(()),
655 })
656 }
657
658 pub fn open_encrypted_with_credentials_and_options(
662 path: &Path,
663 passphrase: &str,
664 username: &str,
665 password: &str,
666 opts: OpenOptions,
667 ) -> Result<Self> {
668 let inner = Arc::new(open_core_with_retry(opts.lock_timeout_ms, || {
669 CoreDatabase::open_encrypted_with_credentials(path, passphrase, username, password)
670 })?);
671 let schema = load_schema(path)?;
672 ensure_internal_tables(&inner)?;
673 reap_rotated_wal_segments(&inner);
674 Ok(Self {
675 inner,
676 schema,
677 root: path.to_path_buf(),
678 default_providers: HashMap::new(),
679 session: parking_lot::RwLock::new(None),
680 sequence_lock: parking_lot::Mutex::new(()),
681 })
682 }
683
684 pub fn create_encrypted_with_credentials(
688 path: &Path,
689 schema: KitSchema,
690 passphrase: &str,
691 admin_username: &str,
692 admin_password: &str,
693 ) -> Result<Self> {
694 std::fs::create_dir_all(path)?;
695 let inner = Arc::new(CoreDatabase::create_encrypted_with_credentials(
696 path,
697 passphrase,
698 admin_username,
699 admin_password,
700 )?);
701 ensure_internal_tables(&inner)?;
702 store_schema(path, &schema)?;
703 for table in &schema.tables {
704 create_core_table(&inner, &table.name, to_core_schema(table)?)?;
705 }
706 Ok(Self {
707 inner,
708 schema,
709 root: path.to_path_buf(),
710 default_providers: HashMap::new(),
711 session: parking_lot::RwLock::new(None),
712 sequence_lock: parking_lot::Mutex::new(()),
713 })
714 }
715
716 pub fn enable_auth(&self, admin_username: &str, admin_password: &str) -> Result<()> {
720 self.inner
721 .enable_auth(admin_username, admin_password)
722 .map_err(KitError::from)
723 }
724
725 pub fn disable_auth(&self) -> Result<()> {
729 self.inner.disable_auth().map_err(KitError::from)
730 }
731
732 pub fn require_auth_enabled(&self) -> bool {
734 self.inner.require_auth_enabled()
735 }
736
737 pub fn refresh_principal(&self) -> Result<()> {
741 self.inner.refresh_principal().map_err(KitError::from)?;
742 *self.session.write() = None;
745 Ok(())
746 }
747
748 pub fn register_default(
751 &mut self,
752 name: impl Into<String>,
753 provider: impl Fn() -> Value + Send + Sync + 'static,
754 ) {
755 self.default_providers
756 .insert(name.into(), Box::new(provider));
757 }
758
759 pub fn embedding_providers(&self) -> &mongreldb_core::EmbeddingProviderRegistry {
766 self.inner.embedding_providers()
767 }
768
769 pub fn register_embedding_provider(
771 &self,
772 provider: Arc<dyn mongreldb_core::EmbeddingProvider>,
773 ) {
774 self.embedding_providers().register(provider);
775 }
776
777 pub fn embed_texts(
784 &self,
785 source: &mongreldb_kit_core::schema::EmbeddingSource,
786 texts: &[&str],
787 expected_dim: u32,
788 ) -> Result<Vec<Vec<f32>>> {
789 let core_source = crate::schema::to_core_embedding_source(source);
790 self.inner
791 .embedding_providers()
792 .embed(&core_source, texts, expected_dim)
793 .map_err(|e| KitError::Validation(e.to_string()))
794 }
795
796 pub fn raw(&self) -> &CoreDatabase {
800 &self.inner
801 }
802
803 pub fn table_names(&self) -> Vec<String> {
805 self.schema
806 .tables
807 .iter()
808 .map(|t| t.name.clone())
809 .filter(|n| !n.starts_with("__kit_"))
810 .collect()
811 }
812
813 pub fn create_procedure(
814 &self,
815 spec: &ProcedureSpec,
816 ) -> Result<mongreldb_core::StoredProcedure> {
817 let procedure = core_procedure(spec)?;
818 self.inner
819 .create_procedure(procedure)
820 .map_err(KitError::from)
821 }
822
823 pub fn replace_procedure(
824 &self,
825 spec: &ProcedureSpec,
826 ) -> Result<mongreldb_core::StoredProcedure> {
827 let procedure = core_procedure(spec)?;
828 self.inner
829 .create_or_replace_procedure(procedure)
830 .map_err(KitError::from)
831 }
832
833 pub fn drop_procedure(&self, name: &str) -> Result<()> {
834 self.inner.drop_procedure(name).map_err(KitError::from)
835 }
836
837 pub fn call_procedure(
838 &self,
839 name: &str,
840 args: serde_json::Map<String, Value>,
841 ) -> Result<mongreldb_core::ProcedureCallResult> {
842 let args = args
843 .iter()
844 .map(|(key, value)| Ok((key.clone(), json_to_core_value(value)?)))
845 .collect::<Result<HashMap<_, _>>>()?;
846 self.inner
847 .call_procedure(name, args)
848 .map_err(KitError::from)
849 }
850
851 pub fn create_trigger(&self, spec: &TriggerSpec) -> Result<mongreldb_core::StoredTrigger> {
852 let trigger = core_trigger(spec)?;
853 self.inner.create_trigger(trigger).map_err(KitError::from)
854 }
855
856 pub fn replace_trigger(&self, spec: &TriggerSpec) -> Result<mongreldb_core::StoredTrigger> {
857 let trigger = core_trigger(spec)?;
858 self.inner
859 .create_or_replace_trigger(trigger)
860 .map_err(KitError::from)
861 }
862
863 pub fn drop_trigger(&self, name: &str) -> Result<()> {
864 self.inner.drop_trigger(name).map_err(KitError::from)
865 }
866
867 pub fn triggers(&self) -> Vec<mongreldb_core::StoredTrigger> {
868 self.inner.triggers()
869 }
870
871 pub fn trigger(&self, name: &str) -> Option<mongreldb_core::StoredTrigger> {
872 self.inner.trigger(name)
873 }
874
875 pub fn allocate_sequence(&self, name: &str, count: i64) -> Result<i64> {
881 use crate::internal::cols;
882 let _guard = self.sequence_lock.lock();
883 let mut attempt = 0;
884 loop {
885 let mut txn = self.inner.begin();
886 let snapshot = txn.read_snapshot();
887 let existing = self
888 .visible_core_rows_at(crate::internal::SEQUENCES, snapshot)?
889 .into_iter()
890 .find(|r| internal_bytes(r, cols::SEQ_NAME) == Some(name.to_string()));
891
892 let now = crate::internal::iso_now();
893 let (start, next, old_row_id) = match &existing {
897 Some(row) => {
898 let current = match row.columns.get(&cols::SEQ_NEXT) {
899 Some(CoreValue::Int64(i)) => *i,
900 _ => 1,
901 };
902 (current, current + count, Some(row.row_id))
903 }
904 None => (1, 1 + count, None),
905 };
906
907 if let Some(rid) = old_row_id {
908 txn.delete(crate::internal::SEQUENCES, rid)
909 .map_err(KitError::from)?;
910 }
911 txn.put(
912 crate::internal::SEQUENCES,
913 vec![
914 (cols::SEQ_NAME, CoreValue::Bytes(name.as_bytes().to_vec())),
915 (cols::SEQ_NEXT, CoreValue::Int64(next)),
916 (cols::SEQ_UPDATED, CoreValue::Bytes(now.into_bytes())),
917 ],
918 )
919 .map_err(KitError::from)?;
920 match txn.commit() {
921 Ok(_) => return Ok(start),
922 Err(mongreldb_core::MongrelError::Conflict(_)) if attempt < 10_000 => {
923 attempt += 1;
924 std::thread::yield_now();
925 continue;
926 }
927 Err(e) => return Err(KitError::from(e)),
928 }
929 }
930 }
931
932 pub fn transaction<T, F>(&self, max_retries: usize, mut f: F) -> Result<T>
935 where
936 F: FnMut(&mut crate::txn::Transaction<'_>) -> Result<T>,
937 {
938 let mut attempt = 0;
939 loop {
940 let mut txn = self.begin()?;
941 match f(&mut txn) {
942 Ok(value) => match txn.commit() {
943 Ok(()) => return Ok(value),
944 Err(KitError::Conflict(_)) if attempt < max_retries => {
945 attempt += 1;
946 continue;
947 }
948 Err(e) => return Err(e),
949 },
950 Err(KitError::Conflict(_)) if attempt < max_retries => {
951 txn.rollback();
952 attempt += 1;
953 continue;
954 }
955 Err(e) => {
956 txn.rollback();
957 return Err(e);
958 }
959 }
960 }
961 }
962
963 pub fn table(&self, name: &str) -> Option<&KitTable> {
965 self.schema.table(name)
966 }
967
968 pub fn schema(&self) -> &KitSchema {
970 &self.schema
971 }
972
973 pub fn begin(&self) -> Result<crate::txn::Transaction<'_>> {
975 let core_txn = self.inner.begin();
976 Ok(crate::txn::Transaction::new(self, core_txn))
977 }
978
979 pub fn set_schema(&mut self, schema: KitSchema) {
981 self.schema = schema;
982 }
983
984 pub fn check_internal_tables(&self) -> Result<()> {
987 let schema_file = self.root.join(SCHEMA_FILE);
988 if !schema_file.exists() {
989 return Err(KitError::Integrity(format!(
990 "schema file {} is missing",
991 schema_file.display()
992 )));
993 }
994 for (name, _) in internal_tables_core() {
995 if self.inner.table_id(name).is_err() {
996 return Err(KitError::Integrity(format!(
997 "internal table {name} is missing"
998 )));
999 }
1000 }
1001 Ok(())
1002 }
1003
1004 pub fn gc(&self) -> Result<usize> {
1007 self.inner.gc().map_err(KitError::from)
1008 }
1009
1010 pub fn check(&self) -> Vec<serde_json::Value> {
1013 self.inner
1014 .check()
1015 .into_iter()
1016 .map(|i| {
1017 serde_json::json!({
1018 "table_id": i.table_id,
1019 "table_name": i.table_name,
1020 "severity": i.severity,
1021 "description": i.description,
1022 })
1023 })
1024 .collect()
1025 }
1026
1027 pub fn doctor(&self) -> Result<Vec<u64>> {
1029 self.inner.doctor().map_err(KitError::from)
1030 }
1031
1032 pub fn snapshot_epoch(&self) -> u64 {
1036 self.inner.snapshot().0.epoch.0
1037 }
1038
1039 pub fn set_history_retention_epochs(&self, epochs: u64) -> Result<()> {
1040 self.inner
1041 .set_history_retention_epochs(epochs)
1042 .map_err(KitError::from)
1043 }
1044
1045 pub fn history_retention_epochs(&self) -> u64 {
1046 self.inner.history_retention_epochs()
1047 }
1048
1049 pub fn earliest_retained_epoch(&self) -> u64 {
1050 self.inner.earliest_retained_epoch().0
1051 }
1052
1053 pub fn export_tsv(&self, table: &str) -> Result<String> {
1057 let t = self
1058 .schema
1059 .tables
1060 .iter()
1061 .find(|t| t.name == table)
1062 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?
1063 .clone();
1064 let tx = self.begin()?;
1065 let rows = tx.all_rows(table)?;
1066 Ok(crate::tsv::rows_to_tsv(&t, &rows))
1067 }
1068
1069 pub fn import_tsv(&self, table: &str, text: &str) -> Result<usize> {
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 .clone();
1080 let rows = crate::tsv::tsv_to_rows(&t, text)?;
1081 let n = rows.len();
1082 self.transaction(1, |tx| {
1083 tx.insert_many(table, rows.clone())?;
1084 Ok(())
1085 })?;
1086 Ok(n)
1087 }
1088
1089 pub fn explain(
1094 &self,
1095 table: &str,
1096 predicate: &mongreldb_kit_core::query::Expr,
1097 ) -> Result<ExplainPlan> {
1098 let t = self
1099 .schema
1100 .tables
1101 .iter()
1102 .find(|t| t.name == table)
1103 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
1104 Ok(match crate::pushdown::translate_predicate(t, predicate) {
1105 Some(p) => ExplainPlan {
1106 index_accelerated: p.can_push(),
1107 exact: p.fully_translated,
1108 pushed_conditions: p.conditions.iter().map(condition_label).collect(),
1109 },
1110 None => ExplainPlan {
1111 index_accelerated: false,
1112 exact: false,
1113 pushed_conditions: Vec::new(),
1114 },
1115 })
1116 }
1117
1118 pub fn rows_at_epoch(&self, table: &str, epoch: u64) -> Result<Vec<crate::schema::Row>> {
1124 let t = self
1125 .schema
1126 .tables
1127 .iter()
1128 .find(|t| t.name == table)
1129 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
1130 let current = self.snapshot_epoch();
1131 if epoch > current {
1132 return Err(KitError::Validation(format!(
1133 "epoch {epoch} is in the future (current committed epoch is {current})"
1134 )));
1135 }
1136 let snap = Snapshot::at(mongreldb_core::epoch::Epoch(epoch));
1137 let rows = self.visible_core_rows_at(table, snap)?;
1138 rows.iter()
1139 .map(|r| crate::schema::core_row_to_json(r, t))
1140 .collect()
1141 }
1142
1143 pub fn approx_aggregate(
1149 &self,
1150 table: &str,
1151 column: Option<&str>,
1152 agg: ApproxAggKind,
1153 z: f64,
1154 ) -> Result<Option<ApproxAggregate>> {
1155 let t = self
1156 .schema
1157 .tables
1158 .iter()
1159 .find(|t| t.name == table)
1160 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
1161 if matches!(agg, ApproxAggKind::Sum | ApproxAggKind::Avg) && column.is_none() {
1162 return Err(KitError::Validation(
1163 "approx sum/avg requires a column".into(),
1164 ));
1165 }
1166 let cid = match column {
1167 Some(name) => Some(
1168 t.columns
1169 .iter()
1170 .find(|c| c.name == name)
1171 .ok_or_else(|| KitError::Validation(format!("unknown column '{name}'")))?
1172 .id as u16,
1173 ),
1174 None => None,
1175 };
1176 let core_agg = match agg {
1177 ApproxAggKind::Count => ApproxAgg::Count,
1178 ApproxAggKind::Sum => ApproxAgg::Sum,
1179 ApproxAggKind::Avg => ApproxAgg::Avg,
1180 };
1181 let handle = self.inner.table(table).map_err(KitError::from)?;
1182 let mut guard = handle.lock();
1183 let res = guard
1184 .approx_aggregate(&[], cid, core_agg, z)
1185 .map_err(KitError::from)?;
1186 Ok(res.map(|r| ApproxAggregate {
1187 point: r.point,
1188 ci_low: r.ci_low,
1189 ci_high: r.ci_high,
1190 n_population: r.n_population,
1191 n_sample_live: r.n_sample_live,
1192 n_passing: r.n_passing,
1193 }))
1194 }
1195
1196 pub fn scan_batched<F>(&self, table: &str, batch_size: usize, mut f: F) -> Result<()>
1202 where
1203 F: FnMut(&[serde_json::Map<String, Value>]) -> Result<()>,
1204 {
1205 let kit_t = self
1206 .schema
1207 .tables
1208 .iter()
1209 .find(|t| t.name == table)
1210 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
1211 let batch_size = batch_size.max(1);
1212 let (snapshot, _pin) = self.inner.snapshot();
1215 let handle = self.inner.table(table).map_err(KitError::from)?;
1216 let guard = handle.lock();
1217
1218 let mut projection: Vec<(u16, mongreldb_core::schema::TypeId)> = Vec::new();
1220 let mut meta: Vec<(String, mongreldb_kit_core::schema::ColumnType)> = Vec::new();
1221 for c in &guard.schema().columns {
1222 if let Some(kc) = kit_t.columns.iter().find(|kc| kc.id as u16 == c.id) {
1223 projection.push((c.id, c.ty.clone()));
1224 meta.push((kc.name.clone(), kc.storage_type));
1225 }
1226 }
1227
1228 match guard
1229 .scan_cursor(snapshot, projection, &[])
1230 .map_err(KitError::from)?
1231 {
1232 Some(mut cursor) => {
1233 let mut buf: Vec<serde_json::Map<String, Value>> = Vec::with_capacity(batch_size);
1234 while let Some(batch) = cursor.next_batch().map_err(KitError::from)? {
1235 let nrows = batch.first().map(|c| c.len()).unwrap_or(0);
1236 for j in 0..nrows {
1237 let mut m = serde_json::Map::new();
1238 for (ci, (name, ty)) in meta.iter().enumerate() {
1239 let cv = batch
1240 .get(ci)
1241 .and_then(|col| col.value_at(j))
1242 .unwrap_or(CoreValue::Null);
1243 m.insert(name.clone(), crate::schema::core_to_json(&cv, *ty)?);
1244 }
1245 buf.push(m);
1246 if buf.len() >= batch_size {
1247 f(&buf)?;
1248 buf.clear();
1249 }
1250 }
1251 }
1252 if !buf.is_empty() {
1253 f(&buf)?;
1254 }
1255 Ok(())
1256 }
1257 None => {
1258 drop(guard);
1259 let rows = self.visible_core_rows_at(table, snapshot)?;
1260 let maps: Vec<serde_json::Map<String, Value>> = rows
1261 .iter()
1262 .map(|r| crate::schema::core_row_to_json(r, kit_t).map(|row| row.values))
1263 .collect::<Result<Vec<_>>>()?;
1264 for chunk in maps.chunks(batch_size) {
1265 f(chunk)?;
1266 }
1267 Ok(())
1268 }
1269 }
1270 }
1271
1272 pub fn set_similarity(
1281 &self,
1282 table: &str,
1283 column: &str,
1284 query: &[String],
1285 k: usize,
1286 ) -> Result<Vec<SimilarRow>> {
1287 let t = self
1288 .schema
1289 .tables
1290 .iter()
1291 .find(|t| t.name == table)
1292 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
1293 let col = t.columns.iter().find(|c| c.name == column).ok_or_else(|| {
1294 KitError::Validation(format!("unknown column '{column}' on table '{table}'"))
1295 })?;
1296 let query_set: std::collections::HashSet<String> = query.iter().cloned().collect();
1297
1298 let has_minhash = t.indexes.iter().any(|idx| {
1299 idx.kind == KitIndexKind::MinHash && idx.columns.iter().any(|c| c == column)
1300 });
1301 let rows = if has_minhash {
1302 let query_hashes: Vec<u64> = query
1304 .iter()
1305 .map(|s| mongreldb_core::index::minhash_token_hash(s))
1306 .collect();
1307 let cand_k = k.saturating_mul(8).max(k + 64);
1309 let cond = mongreldb_core::query::Condition::MinHashSimilar {
1310 column_id: col.id as u16,
1311 query: query_hashes,
1312 k: cand_k,
1313 };
1314 let (snapshot, _pin) = self.inner.snapshot();
1315 let core_rows = self.query_core_rows_at(table, &[cond], snapshot)?;
1316 core_rows
1317 .iter()
1318 .map(|r| crate::schema::core_row_to_json(r, t))
1319 .collect::<Result<Vec<_>>>()?
1320 } else {
1321 let tx = self.begin()?;
1322 tx.all_rows(table)?
1323 };
1324
1325 let mut scored: Vec<SimilarRow> = Vec::new();
1326 for row in rows {
1327 let set = parse_string_set(row.values.get(column));
1328 let inter = set.iter().filter(|x| query_set.contains(*x)).count();
1329 let union = set.len() + query_set.len() - inter;
1330 let sim = if union == 0 {
1331 0.0
1332 } else {
1333 inter as f64 / union as f64
1334 };
1335 if sim > 0.0 {
1336 scored.push(SimilarRow {
1337 row,
1338 similarity: sim,
1339 });
1340 }
1341 }
1342 scored.sort_by(|a, b| {
1343 b.similarity
1344 .partial_cmp(&a.similarity)
1345 .unwrap_or(std::cmp::Ordering::Equal)
1346 });
1347 scored.truncate(k);
1348 Ok(scored)
1349 }
1350
1351 pub fn flush(&self) -> Result<()> {
1355 for name in self.inner.table_names() {
1356 let handle = self.inner.table(&name).map_err(KitError::from)?;
1357 let mut guard = handle.lock();
1358 guard.flush().map_err(KitError::from)?;
1359 }
1360 Ok(())
1361 }
1362
1363 pub fn incremental_aggregate(
1375 &self,
1376 table: &str,
1377 column: Option<&str>,
1378 agg: IncrementalAggKind,
1379 filter: Option<&mongreldb_kit_core::query::Expr>,
1380 ) -> Result<IncrementalAggregate> {
1381 let t = self
1382 .schema
1383 .tables
1384 .iter()
1385 .find(|t| t.name == table)
1386 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
1387 if !matches!(agg, IncrementalAggKind::Count) && column.is_none() {
1388 return Err(KitError::Validation(
1389 "sum/min/max/avg incremental aggregate requires a column".into(),
1390 ));
1391 }
1392 let cid = match column {
1393 Some(name) => Some(
1394 t.columns
1395 .iter()
1396 .find(|c| c.name == name)
1397 .ok_or_else(|| KitError::Validation(format!("unknown column '{name}'")))?
1398 .id as u16,
1399 ),
1400 None => None,
1401 };
1402 let conditions = match filter {
1403 Some(expr) => {
1404 let plan = crate::pushdown::translate_predicate(t, expr).ok_or_else(|| {
1405 KitError::Validation(
1406 "filter is not index-translatable for an incremental aggregate".into(),
1407 )
1408 })?;
1409 if !plan.fully_translated {
1410 return Err(KitError::Validation(
1411 "filter has a residual that an incremental aggregate cannot apply exactly"
1412 .into(),
1413 ));
1414 }
1415 plan.conditions
1416 }
1417 None => Vec::new(),
1418 };
1419 let core_agg = match agg {
1420 IncrementalAggKind::Count => NativeAgg::Count,
1421 IncrementalAggKind::Sum => NativeAgg::Sum,
1422 IncrementalAggKind::Min => NativeAgg::Min,
1423 IncrementalAggKind::Max => NativeAgg::Max,
1424 IncrementalAggKind::Avg => NativeAgg::Avg,
1425 };
1426 let cache_key = incremental_cache_key(t.id, cid, agg, &conditions);
1427 let handle = self.inner.table(table).map_err(KitError::from)?;
1428 let mut guard = handle.lock();
1429 let res = guard
1430 .aggregate_incremental(cache_key, &conditions, cid, core_agg)
1431 .map_err(KitError::from)?;
1432 Ok(IncrementalAggregate {
1433 value: agg_state_value(&res.state),
1434 incremental: res.incremental,
1435 delta_rows: res.delta_rows,
1436 })
1437 }
1438
1439 pub fn applied_migrations(&self) -> Result<Vec<mongreldb_kit_core::migrations::Migration>> {
1441 crate::migrate::load_applied_migrations(&self.inner)
1442 }
1443
1444 pub(crate) fn core_db(&self) -> &CoreDatabase {
1445 &self.inner
1446 }
1447
1448 pub(crate) fn core_arc(&self) -> Arc<CoreDatabase> {
1451 Arc::clone(&self.inner)
1452 }
1453
1454 pub fn close(&self) -> Result<()> {
1459 self.inner.close().map_err(KitError::from)
1460 }
1461
1462 pub fn compact_all(&self) -> Result<(usize, usize)> {
1467 self.inner.compact().map_err(KitError::from)
1468 }
1469
1470 pub fn compact_table(&self, name: &str) -> Result<bool> {
1473 self.inner.compact_table(name).map_err(KitError::from)
1474 }
1475
1476 pub fn rename_table(&mut self, from: &str, to: &str) -> Result<()> {
1486 if from.starts_with("__kit_") || to.starts_with("__kit_") {
1487 return Err(KitError::Validation(
1488 "rename_table: names beginning with '__kit_' are reserved for internal tables"
1489 .into(),
1490 ));
1491 }
1492 self.inner.rename_table(from, to).map_err(KitError::from)?;
1493 if !self.schema.rename_table(from, to) {
1496 return Err(KitError::Integrity(format!(
1499 "rename_table: kit schema has no table '{from}' (or '{to}' already exists)"
1500 )));
1501 }
1502 for table in &mut self.schema.tables {
1503 for fk in &mut table.foreign_keys {
1504 if fk.references_table == from {
1505 fk.references_table = to.to_string();
1506 }
1507 }
1508 }
1509 store_schema(&self.root, &self.schema)?;
1510 Ok(())
1511 }
1512
1513 pub fn analyze(&self) -> Result<()> {
1518 for name in self.inner.table_names() {
1519 let handle = self.inner.table(&name).map_err(KitError::from)?;
1520 handle.lock().ensure_indexes_complete()?;
1521 }
1522 Ok(())
1523 }
1524
1525 pub fn vacuum(&self) -> Result<usize> {
1529 self.inner.compact().map_err(KitError::from)?;
1530 self.inner.gc().map_err(KitError::from)
1531 }
1532
1533 pub fn create_view(&self, spec: &ViewSpec) -> Result<()> {
1539 self.sql(&spec.create_sql())?;
1540 Ok(())
1541 }
1542
1543 pub fn drop_view(&self, name: &str) -> Result<()> {
1545 self.sql(&format!("DROP VIEW IF EXISTS {name}"))?;
1546 Ok(())
1547 }
1548
1549 pub fn reserve_auto_inc(&self, table: &str) -> Result<Option<i64>> {
1557 let handle = self.inner.table(table).map_err(KitError::from)?;
1558 let mut guard = handle.lock();
1559 guard.reserve_auto_inc().map_err(KitError::from)
1560 }
1561
1562 pub fn create_user(&self, username: &str, password: &str) -> Result<()> {
1566 self.inner
1567 .create_user(username, password)
1568 .map_err(KitError::from)?;
1569 Ok(())
1570 }
1571
1572 pub fn drop_user(&self, username: &str) -> Result<()> {
1574 self.inner.drop_user(username).map_err(KitError::from)
1575 }
1576
1577 pub fn alter_user_password(&self, username: &str, new_password: &str) -> Result<()> {
1579 self.inner
1580 .alter_user_password(username, new_password)
1581 .map_err(KitError::from)
1582 }
1583
1584 pub fn verify_user(
1586 &self,
1587 username: &str,
1588 password: &str,
1589 ) -> Result<Option<mongreldb_core::auth::UserEntry>> {
1590 self.inner
1591 .verify_user(username, password)
1592 .map_err(KitError::from)
1593 }
1594
1595 pub fn set_user_admin(&self, username: &str, is_admin: bool) -> Result<()> {
1597 self.inner
1598 .set_user_admin(username, is_admin)
1599 .map_err(KitError::from)
1600 }
1601
1602 pub fn users(&self) -> Vec<String> {
1604 self.inner.users().into_iter().map(|u| u.username).collect()
1605 }
1606
1607 pub fn create_role(&self, name: &str) -> Result<()> {
1609 self.inner.create_role(name).map_err(KitError::from)?;
1610 Ok(())
1611 }
1612
1613 pub fn drop_role(&self, name: &str) -> Result<()> {
1615 self.inner.drop_role(name).map_err(KitError::from)
1616 }
1617
1618 pub fn roles(&self) -> Vec<String> {
1620 self.inner.roles().into_iter().map(|r| r.name).collect()
1621 }
1622
1623 pub fn grant_role(&self, username: &str, role_name: &str) -> Result<()> {
1625 self.inner
1626 .grant_role(username, role_name)
1627 .map_err(KitError::from)
1628 }
1629
1630 pub fn revoke_role(&self, username: &str, role_name: &str) -> Result<()> {
1632 self.inner
1633 .revoke_role(username, role_name)
1634 .map_err(KitError::from)
1635 }
1636
1637 pub fn grant_permission(
1639 &self,
1640 role_name: &str,
1641 permission: mongreldb_core::auth::Permission,
1642 ) -> Result<()> {
1643 self.inner
1644 .grant_permission(role_name, permission)
1645 .map_err(KitError::from)
1646 }
1647
1648 pub fn revoke_permission(
1650 &self,
1651 role_name: &str,
1652 permission: mongreldb_core::auth::Permission,
1653 ) -> Result<()> {
1654 self.inner
1655 .revoke_permission(role_name, permission)
1656 .map_err(KitError::from)
1657 }
1658
1659 pub fn set_spill_threshold(&self, bytes: u64) {
1665 self.inner.set_spill_threshold(bytes);
1666 }
1667
1668 pub fn set_recursive_triggers(&self, enabled: bool) {
1670 self.inner.set_recursive_triggers(enabled);
1671 }
1672
1673 pub fn trigger_config(&self) -> mongreldb_core::TriggerConfig {
1675 self.inner.trigger_config()
1676 }
1677
1678 pub fn set_trigger_config(&self, config: mongreldb_core::TriggerConfig) -> Result<()> {
1680 self.inner
1681 .set_trigger_config(config)
1682 .map_err(KitError::from)
1683 }
1684
1685 pub fn set_table_compaction_zstd_level(&self, table: &str, level: i32) -> Result<()> {
1687 let handle = self.inner.table(table).map_err(KitError::from)?;
1688 handle.lock().set_compaction_zstd_level(level);
1689 Ok(())
1690 }
1691
1692 pub fn set_table_result_cache_max_bytes(&self, table: &str, max_bytes: u64) -> Result<()> {
1694 let handle = self.inner.table(table).map_err(KitError::from)?;
1695 handle.lock().set_result_cache_max_bytes(max_bytes);
1696 Ok(())
1697 }
1698
1699 pub fn set_table_mutable_run_spill_bytes(&self, table: &str, bytes: u64) -> Result<()> {
1701 let handle = self.inner.table(table).map_err(KitError::from)?;
1702 handle.lock().set_mutable_run_spill_bytes(bytes);
1703 Ok(())
1704 }
1705
1706 pub fn set_table_sync_byte_threshold(&self, table: &str, threshold: u64) -> Result<()> {
1708 let handle = self.inner.table(table).map_err(KitError::from)?;
1709 handle.lock().set_sync_byte_threshold(threshold);
1710 Ok(())
1711 }
1712
1713 pub fn set_table_index_build_policy(
1716 &self,
1717 table: &str,
1718 policy: mongreldb_core::IndexBuildPolicy,
1719 ) -> Result<()> {
1720 let handle = self.inner.table(table).map_err(KitError::from)?;
1721 handle.lock().set_index_build_policy(policy);
1722 Ok(())
1723 }
1724
1725 pub fn table_page_cache_stats(&self, table: &str) -> Result<mongreldb_core::cache::CacheStats> {
1727 let handle = self.inner.table(table).map_err(KitError::from)?;
1728 let stats = handle.lock().page_cache_stats();
1729 Ok(stats)
1730 }
1731
1732 pub fn table_run_count(&self, table: &str) -> Result<usize> {
1734 let handle = self.inner.table(table).map_err(KitError::from)?;
1735 let n = handle.lock().run_count();
1736 Ok(n)
1737 }
1738
1739 pub fn table_memtable_len(&self, table: &str) -> Result<usize> {
1741 let handle = self.inner.table(table).map_err(KitError::from)?;
1742 let n = handle.lock().memtable_len();
1743 Ok(n)
1744 }
1745
1746 pub fn table_mutable_run_len(&self, table: &str) -> Result<usize> {
1748 let handle = self.inner.table(table).map_err(KitError::from)?;
1749 let n = handle.lock().mutable_run_len();
1750 Ok(n)
1751 }
1752
1753 pub fn table_page_cache_len(&self, table: &str) -> Result<usize> {
1755 let handle = self.inner.table(table).map_err(KitError::from)?;
1756 let n = handle.lock().page_cache_len();
1757 Ok(n)
1758 }
1759
1760 pub fn table_decoded_cache_len(&self, table: &str) -> Result<usize> {
1762 let handle = self.inner.table(table).map_err(KitError::from)?;
1763 let n = handle.lock().decoded_cache_len();
1764 Ok(n)
1765 }
1766
1767 pub fn sql(&self, statement: &str) -> Result<Vec<arrow::record_batch::RecordBatch>> {
1785 self.sql_with_options(statement, SqlOptions::default())
1786 }
1787
1788 fn sql_session(&self) -> Result<Arc<mongreldb_query::MongrelSession>> {
1789 if let Some(session) = self.session.read().as_ref() {
1790 return Ok(Arc::clone(session));
1791 }
1792 let session = Arc::new(
1793 mongreldb_query::MongrelSession::open(self.core_arc()).map_err(KitError::from)?,
1794 );
1795 let mut cached = self.session.write();
1796 Ok(Arc::clone(cached.get_or_insert(session)))
1797 }
1798
1799 #[doc(hidden)]
1800 pub fn set_sql_test_hook(&self, hook: Option<mongreldb_query::SqlTestHook>) -> Result<()> {
1801 self.sql_session()?.set_test_hook(hook);
1802 Ok(())
1803 }
1804
1805 pub fn start_sql(
1806 &self,
1807 statement: impl Into<String>,
1808 options: SqlOptions,
1809 ) -> Result<SqlQueryHandle> {
1810 let session = self.sql_session()?;
1811 let query = session
1812 .register_query(mongreldb_query::SqlQueryOptions {
1813 query_id: options.query_id,
1814 timeout: options.timeout,
1815 ..mongreldb_query::SqlQueryOptions::default()
1816 })
1817 .map_err(KitError::from)?;
1818 let query_id = query.id();
1819 let registration = mongreldb_query::RegisteredQueryGuard::new(query);
1820 let worker_session = Arc::clone(&session);
1821 let statement = statement.into();
1822 let worker = std::thread::Builder::new()
1823 .name(format!("mongreldb-kit-sql-{query_id}"))
1824 .spawn(move || {
1825 sql_runtime().block_on(
1826 worker_session
1827 .run_with_query_for_serialization(&statement, registration.into_query()),
1828 )
1829 })
1830 .map_err(|error| KitError::Storage(error.to_string()))?;
1831 Ok(SqlQueryHandle {
1832 query_id,
1833 session,
1834 worker: Some(worker),
1835 })
1836 }
1837
1838 pub fn sql_with_options(
1839 &self,
1840 statement: &str,
1841 options: SqlOptions,
1842 ) -> Result<Vec<arrow::record_batch::RecordBatch>> {
1843 self.start_sql(statement, options)?.wait()
1844 }
1845
1846 pub fn cancel_sql(&self, query_id: mongreldb_query::QueryId) -> mongreldb_query::CancelOutcome {
1847 self.session
1848 .read()
1849 .as_ref()
1850 .map_or(mongreldb_query::CancelOutcome::NotFound, |session| {
1851 session.cancel_query(query_id)
1852 })
1853 }
1854
1855 pub fn sql_query_status(
1856 &self,
1857 query_id: mongreldb_query::QueryId,
1858 ) -> Result<Option<mongreldb_query::QueryStatus>> {
1859 Ok(self.sql_session()?.query_registry().status(query_id))
1860 }
1861
1862 pub fn refresh_sql_session(&self) -> Result<()> {
1867 let session =
1868 mongreldb_query::MongrelSession::open(self.core_arc()).map_err(KitError::from)?;
1869 *self.session.write() = Some(Arc::new(session));
1870 Ok(())
1871 }
1872
1873 pub fn sql_arrow(&self, statement: &str) -> Result<Vec<u8>> {
1879 self.sql_arrow_with_options(statement, SqlOptions::default())
1880 }
1881
1882 pub fn sql_arrow_with_options(&self, statement: &str, options: SqlOptions) -> Result<Vec<u8>> {
1883 self.sql_serialized_with_options(statement, options, |output| {
1884 crate::arrow_util::batches_to_ipc_controlled(output.batches(), output.query())
1885 })
1886 }
1887
1888 pub fn sql_rows(&self, statement: &str) -> Result<Vec<serde_json::Map<String, Value>>> {
1892 self.sql_rows_with_options(statement, SqlOptions::default())
1893 }
1894
1895 pub fn sql_rows_with_options(
1896 &self,
1897 statement: &str,
1898 options: SqlOptions,
1899 ) -> Result<Vec<serde_json::Map<String, Value>>> {
1900 self.sql_serialized_with_options(statement, options, |output| {
1901 crate::arrow_util::batches_to_rows_controlled(output.batches(), output.query())
1902 })
1903 }
1904
1905 fn sql_serialized_with_options<T>(
1906 &self,
1907 statement: &str,
1908 options: SqlOptions,
1909 serialize: impl FnOnce(&mongreldb_query::ManagedQueryBatches) -> Result<T>,
1910 ) -> Result<T> {
1911 let session = self.sql_session()?;
1912 let query = session
1913 .register_query(mongreldb_query::SqlQueryOptions {
1914 query_id: options.query_id,
1915 timeout: options.timeout,
1916 ..mongreldb_query::SqlQueryOptions::default()
1917 })
1918 .map_err(KitError::from)?;
1919 let query_id = query.id();
1920 let output = sql_runtime()
1921 .block_on(session.run_with_query_for_serialization(statement, query))
1922 .map_err(|error| {
1923 let status = session.query_registry().status(query_id);
1924 crate::error::query_error_with_status(error, status.as_ref())
1925 })?;
1926 session.fire_test_hook(mongreldb_query::SqlTestHookPoint::BeforeSerializationBatch);
1927 match serialize(&output) {
1928 Ok(value) => {
1929 complete_sql_output(output)?;
1930 Ok(value)
1931 }
1932 Err(error) => {
1933 fail_sql_output(output, &error);
1934 Err(error)
1935 }
1936 }
1937 }
1938
1939 pub(crate) fn lookup_row_id(&self, table: &str, key: &[u8]) -> Result<Option<RowId>> {
1943 let handle = self.inner.table(table).map_err(KitError::from)?;
1944 let mut guard = handle.lock();
1945 guard.ensure_indexes_complete()?;
1946 Ok(guard.lookup_pk(key))
1947 }
1948
1949 pub(crate) fn root(&self) -> &Path {
1950 &self.root
1951 }
1952
1953 pub(crate) fn visible_core_rows_at(
1957 &self,
1958 table_name: &str,
1959 snapshot: Snapshot,
1960 ) -> Result<Vec<CoreRow>> {
1961 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1962 let guard = handle.lock();
1963 guard.visible_rows(snapshot).map_err(KitError::from)
1964 }
1965
1966 pub(crate) fn query_core_rows_at(
1973 &self,
1974 table_name: &str,
1975 conditions: &[mongreldb_core::query::Condition],
1976 snapshot: Snapshot,
1977 ) -> Result<Vec<CoreRow>> {
1978 if conditions.is_empty() {
1979 return self.visible_core_rows_at(table_name, snapshot);
1980 }
1981 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1982 let mut guard = handle.lock();
1983 let q = conditions
1984 .iter()
1985 .cloned()
1986 .fold(mongreldb_core::query::Query::new(), |query, condition| {
1987 query.and(condition)
1988 });
1989 guard.query(&q).map_err(KitError::from)
1990 }
1991
1992 pub(crate) fn flush_table(&self, table_name: &str) -> Result<()> {
2000 let handle = self.inner.table(table_name).map_err(KitError::from)?;
2001 handle.lock().flush().map_err(KitError::from)?;
2002 Ok(())
2003 }
2004
2005 pub(crate) fn count_core_rows_at(
2015 &self,
2016 table_name: &str,
2017 conditions: &[mongreldb_core::query::Condition],
2018 snapshot: Snapshot,
2019 ) -> Result<Option<u64>> {
2020 let handle = self.inner.table(table_name).map_err(KitError::from)?;
2021 let mut guard = handle.lock();
2022 if guard.snapshot().epoch != snapshot.epoch {
2023 return Ok(None); }
2025 guard
2026 .count_conditions(conditions, snapshot)
2027 .map_err(KitError::from)
2028 }
2029
2030 pub(crate) fn aggregate_core_at(
2039 &self,
2040 table_name: &str,
2041 column: Option<u16>,
2042 conditions: &[mongreldb_core::query::Condition],
2043 agg: NativeAgg,
2044 snapshot: Snapshot,
2045 ) -> Result<Option<NativeAggResult>> {
2046 let handle = self.inner.table(table_name).map_err(KitError::from)?;
2047 let guard = handle.lock();
2048 if guard.snapshot().epoch != snapshot.epoch {
2049 return Ok(None); }
2051 guard
2052 .aggregate_native(snapshot, column, conditions, agg)
2053 .map_err(KitError::from)
2054 }
2055
2056 pub(crate) fn count_distinct_core_at(
2065 &self,
2066 table_name: &str,
2067 column_id: u16,
2068 snapshot: Snapshot,
2069 ) -> Result<Option<u64>> {
2070 let handle = self.inner.table(table_name).map_err(KitError::from)?;
2071 let mut guard = handle.lock();
2072 if guard.snapshot().epoch != snapshot.epoch {
2073 return Ok(None); }
2075 guard
2076 .count_distinct_from_bitmap(column_id)
2077 .map_err(KitError::from)
2078 }
2079
2080 #[allow(dead_code)]
2082 pub(crate) fn get_core_row(&self, table_name: &str, row_id: u64) -> Result<Option<CoreRow>> {
2083 let handle = self.inner.table(table_name).map_err(KitError::from)?;
2084 let guard = handle.lock();
2085 let snapshot = guard.snapshot();
2086 Ok(guard.get(mongreldb_core::RowId(row_id), snapshot))
2087 }
2088}
2089
2090pub(crate) fn create_core_table(db: &CoreDatabase, name: &str, schema: CoreSchema) -> Result<()> {
2091 if db.table_id(name).is_ok() {
2092 return Ok(());
2093 }
2094 db.create_table(name, schema).map_err(KitError::from)?;
2095 Ok(())
2096}
2097
2098fn sql_runtime() -> &'static tokio::runtime::Runtime {
2102 use std::sync::OnceLock;
2103 static RT: OnceLock<tokio::runtime::Runtime> = OnceLock::new();
2104 RT.get_or_init(|| {
2105 tokio::runtime::Builder::new_multi_thread()
2106 .worker_threads(4)
2107 .enable_all()
2108 .build()
2109 .expect("failed to build kit SQL tokio runtime")
2110 })
2111}
2112
2113fn core_procedure(spec: &ProcedureSpec) -> Result<mongreldb_core::StoredProcedure> {
2114 let parsed: mongreldb_core::StoredProcedure =
2115 serde_json::from_value(spec.json.clone()).map_err(KitError::from)?;
2116 mongreldb_core::StoredProcedure::new(parsed.name, parsed.mode, parsed.params, parsed.body, 0)
2117 .map_err(KitError::from)
2118}
2119
2120fn core_trigger(spec: &TriggerSpec) -> Result<mongreldb_core::StoredTrigger> {
2121 let parsed: mongreldb_core::StoredTrigger =
2122 serde_json::from_value(spec.json.clone()).map_err(KitError::from)?;
2123 mongreldb_core::StoredTrigger::new(
2124 parsed.name,
2125 mongreldb_core::TriggerDefinition {
2126 target: parsed.target,
2127 timing: parsed.timing,
2128 event: parsed.event,
2129 update_of: parsed.update_of,
2130 target_columns: parsed.target_columns,
2131 when: parsed.when,
2132 program: parsed.program,
2133 },
2134 0,
2135 )
2136 .map_err(KitError::from)
2137}
2138
2139fn json_to_core_value(value: &Value) -> Result<CoreValue> {
2140 match value {
2141 Value::Null => Ok(CoreValue::Null),
2142 Value::Bool(value) => Ok(CoreValue::Bool(*value)),
2143 Value::Number(value) => {
2144 if let Some(value) = value.as_i64() {
2145 Ok(CoreValue::Int64(value))
2146 } else if let Some(value) = value.as_f64() {
2147 Ok(CoreValue::Float64(value))
2148 } else {
2149 Err(KitError::Validation("unsupported JSON number".into()))
2150 }
2151 }
2152 Value::String(value) => Ok(CoreValue::Bytes(value.as_bytes().to_vec())),
2153 Value::Array(_) | Value::Object(_) => Err(KitError::Validation(
2154 "procedure args only support scalar JSON values".into(),
2155 )),
2156 }
2157}
2158
2159pub(crate) fn internal_bytes(row: &CoreRow, col_id: u16) -> Option<String> {
2161 match row.columns.get(&col_id) {
2162 Some(CoreValue::Bytes(b)) => String::from_utf8(b.clone()).ok(),
2163 _ => None,
2164 }
2165}
2166
2167fn reap_rotated_wal_segments(db: &CoreDatabase) {
2182 let _ = db.gc();
2183}
2184
2185pub(crate) fn load_schema(path: &Path) -> Result<KitSchema> {
2186 let file = path.join(SCHEMA_FILE);
2187 let json = std::fs::read_to_string(&file)
2188 .map_err(|e| KitError::Migration(format!("cannot read schema file: {e}")))?;
2189 let schema: KitSchema = serde_json::from_str(&json)?;
2190 Ok(schema)
2191}
2192
2193pub(crate) fn store_schema(path: &Path, schema: &KitSchema) -> Result<()> {
2194 let file = path.join(SCHEMA_FILE);
2195 let json = serde_json::to_string_pretty(schema)?;
2196 std::fs::write(&file, json)?;
2197 Ok(())
2198}
2199
2200pub(crate) fn persist_schema(db: &Database, schema: &KitSchema) -> Result<()> {
2202 store_schema(&db.root, schema)
2203}
2204
2205#[cfg(test)]
2206mod tests {
2207 use super::open_core_with_retry;
2208
2209 fn lock_error() -> mongreldb_core::MongrelError {
2210 mongreldb_core::MongrelError::DatabaseLocked {
2211 path: "/tmp/db".into(),
2212 message: "would block".into(),
2213 }
2214 }
2215
2216 #[test]
2217 fn open_retry_waits_for_lock_contention_only() {
2218 let mut calls = 0;
2219 let value = open_core_with_retry(50, || {
2220 calls += 1;
2221 if calls < 3 {
2222 Err(lock_error())
2223 } else {
2224 Ok(7)
2225 }
2226 })
2227 .unwrap();
2228 assert_eq!(value, 7);
2229 assert_eq!(calls, 3);
2230
2231 let mut non_lock_calls = 0;
2232 let err: mongreldb_core::Result<()> = open_core_with_retry(50, || {
2233 non_lock_calls += 1;
2234 Err(mongreldb_core::MongrelError::Other("nope".into()))
2235 });
2236 let err = err.unwrap_err();
2237 assert_eq!(non_lock_calls, 1);
2238 assert!(matches!(err, mongreldb_core::MongrelError::Other(_)));
2239 }
2240}