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 raw(&self) -> &CoreDatabase {
763 &self.inner
764 }
765
766 pub fn table_names(&self) -> Vec<String> {
768 self.schema
769 .tables
770 .iter()
771 .map(|t| t.name.clone())
772 .filter(|n| !n.starts_with("__kit_"))
773 .collect()
774 }
775
776 pub fn create_procedure(
777 &self,
778 spec: &ProcedureSpec,
779 ) -> Result<mongreldb_core::StoredProcedure> {
780 let procedure = core_procedure(spec)?;
781 self.inner
782 .create_procedure(procedure)
783 .map_err(KitError::from)
784 }
785
786 pub fn replace_procedure(
787 &self,
788 spec: &ProcedureSpec,
789 ) -> Result<mongreldb_core::StoredProcedure> {
790 let procedure = core_procedure(spec)?;
791 self.inner
792 .create_or_replace_procedure(procedure)
793 .map_err(KitError::from)
794 }
795
796 pub fn drop_procedure(&self, name: &str) -> Result<()> {
797 self.inner.drop_procedure(name).map_err(KitError::from)
798 }
799
800 pub fn call_procedure(
801 &self,
802 name: &str,
803 args: serde_json::Map<String, Value>,
804 ) -> Result<mongreldb_core::ProcedureCallResult> {
805 let args = args
806 .iter()
807 .map(|(key, value)| Ok((key.clone(), json_to_core_value(value)?)))
808 .collect::<Result<HashMap<_, _>>>()?;
809 self.inner
810 .call_procedure(name, args)
811 .map_err(KitError::from)
812 }
813
814 pub fn create_trigger(&self, spec: &TriggerSpec) -> Result<mongreldb_core::StoredTrigger> {
815 let trigger = core_trigger(spec)?;
816 self.inner.create_trigger(trigger).map_err(KitError::from)
817 }
818
819 pub fn replace_trigger(&self, spec: &TriggerSpec) -> Result<mongreldb_core::StoredTrigger> {
820 let trigger = core_trigger(spec)?;
821 self.inner
822 .create_or_replace_trigger(trigger)
823 .map_err(KitError::from)
824 }
825
826 pub fn drop_trigger(&self, name: &str) -> Result<()> {
827 self.inner.drop_trigger(name).map_err(KitError::from)
828 }
829
830 pub fn triggers(&self) -> Vec<mongreldb_core::StoredTrigger> {
831 self.inner.triggers()
832 }
833
834 pub fn trigger(&self, name: &str) -> Option<mongreldb_core::StoredTrigger> {
835 self.inner.trigger(name)
836 }
837
838 pub fn allocate_sequence(&self, name: &str, count: i64) -> Result<i64> {
844 use crate::internal::cols;
845 let _guard = self.sequence_lock.lock();
846 let mut attempt = 0;
847 loop {
848 let mut txn = self.inner.begin();
849 let snapshot = txn.read_snapshot();
850 let existing = self
851 .visible_core_rows_at(crate::internal::SEQUENCES, snapshot)?
852 .into_iter()
853 .find(|r| internal_bytes(r, cols::SEQ_NAME) == Some(name.to_string()));
854
855 let now = crate::internal::iso_now();
856 let (start, next, old_row_id) = match &existing {
860 Some(row) => {
861 let current = match row.columns.get(&cols::SEQ_NEXT) {
862 Some(CoreValue::Int64(i)) => *i,
863 _ => 1,
864 };
865 (current, current + count, Some(row.row_id))
866 }
867 None => (1, 1 + count, None),
868 };
869
870 if let Some(rid) = old_row_id {
871 txn.delete(crate::internal::SEQUENCES, rid)
872 .map_err(KitError::from)?;
873 }
874 txn.put(
875 crate::internal::SEQUENCES,
876 vec![
877 (cols::SEQ_NAME, CoreValue::Bytes(name.as_bytes().to_vec())),
878 (cols::SEQ_NEXT, CoreValue::Int64(next)),
879 (cols::SEQ_UPDATED, CoreValue::Bytes(now.into_bytes())),
880 ],
881 )
882 .map_err(KitError::from)?;
883 match txn.commit() {
884 Ok(_) => return Ok(start),
885 Err(mongreldb_core::MongrelError::Conflict(_)) if attempt < 10_000 => {
886 attempt += 1;
887 std::thread::yield_now();
888 continue;
889 }
890 Err(e) => return Err(KitError::from(e)),
891 }
892 }
893 }
894
895 pub fn transaction<T, F>(&self, max_retries: usize, mut f: F) -> Result<T>
898 where
899 F: FnMut(&mut crate::txn::Transaction<'_>) -> Result<T>,
900 {
901 let mut attempt = 0;
902 loop {
903 let mut txn = self.begin()?;
904 match f(&mut txn) {
905 Ok(value) => match txn.commit() {
906 Ok(()) => return Ok(value),
907 Err(KitError::Conflict(_)) if attempt < max_retries => {
908 attempt += 1;
909 continue;
910 }
911 Err(e) => return Err(e),
912 },
913 Err(KitError::Conflict(_)) if attempt < max_retries => {
914 txn.rollback();
915 attempt += 1;
916 continue;
917 }
918 Err(e) => {
919 txn.rollback();
920 return Err(e);
921 }
922 }
923 }
924 }
925
926 pub fn table(&self, name: &str) -> Option<&KitTable> {
928 self.schema.table(name)
929 }
930
931 pub fn schema(&self) -> &KitSchema {
933 &self.schema
934 }
935
936 pub fn begin(&self) -> Result<crate::txn::Transaction<'_>> {
938 let core_txn = self.inner.begin();
939 Ok(crate::txn::Transaction::new(self, core_txn))
940 }
941
942 pub fn set_schema(&mut self, schema: KitSchema) {
944 self.schema = schema;
945 }
946
947 pub fn check_internal_tables(&self) -> Result<()> {
950 let schema_file = self.root.join(SCHEMA_FILE);
951 if !schema_file.exists() {
952 return Err(KitError::Integrity(format!(
953 "schema file {} is missing",
954 schema_file.display()
955 )));
956 }
957 for (name, _) in internal_tables_core() {
958 if self.inner.table_id(name).is_err() {
959 return Err(KitError::Integrity(format!(
960 "internal table {name} is missing"
961 )));
962 }
963 }
964 Ok(())
965 }
966
967 pub fn gc(&self) -> Result<usize> {
970 self.inner.gc().map_err(KitError::from)
971 }
972
973 pub fn check(&self) -> Vec<serde_json::Value> {
976 self.inner
977 .check()
978 .into_iter()
979 .map(|i| {
980 serde_json::json!({
981 "table_id": i.table_id,
982 "table_name": i.table_name,
983 "severity": i.severity,
984 "description": i.description,
985 })
986 })
987 .collect()
988 }
989
990 pub fn doctor(&self) -> Result<Vec<u64>> {
992 self.inner.doctor().map_err(KitError::from)
993 }
994
995 pub fn snapshot_epoch(&self) -> u64 {
999 self.inner.snapshot().0.epoch.0
1000 }
1001
1002 pub fn set_history_retention_epochs(&self, epochs: u64) -> Result<()> {
1003 self.inner
1004 .set_history_retention_epochs(epochs)
1005 .map_err(KitError::from)
1006 }
1007
1008 pub fn history_retention_epochs(&self) -> u64 {
1009 self.inner.history_retention_epochs()
1010 }
1011
1012 pub fn earliest_retained_epoch(&self) -> u64 {
1013 self.inner.earliest_retained_epoch().0
1014 }
1015
1016 pub fn export_tsv(&self, table: &str) -> Result<String> {
1020 let t = self
1021 .schema
1022 .tables
1023 .iter()
1024 .find(|t| t.name == table)
1025 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?
1026 .clone();
1027 let tx = self.begin()?;
1028 let rows = tx.all_rows(table)?;
1029 Ok(crate::tsv::rows_to_tsv(&t, &rows))
1030 }
1031
1032 pub fn import_tsv(&self, table: &str, text: &str) -> Result<usize> {
1036 let t = self
1037 .schema
1038 .tables
1039 .iter()
1040 .find(|t| t.name == table)
1041 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?
1042 .clone();
1043 let rows = crate::tsv::tsv_to_rows(&t, text)?;
1044 let n = rows.len();
1045 self.transaction(1, |tx| {
1046 tx.insert_many(table, rows.clone())?;
1047 Ok(())
1048 })?;
1049 Ok(n)
1050 }
1051
1052 pub fn explain(
1057 &self,
1058 table: &str,
1059 predicate: &mongreldb_kit_core::query::Expr,
1060 ) -> Result<ExplainPlan> {
1061 let t = self
1062 .schema
1063 .tables
1064 .iter()
1065 .find(|t| t.name == table)
1066 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
1067 Ok(match crate::pushdown::translate_predicate(t, predicate) {
1068 Some(p) => ExplainPlan {
1069 index_accelerated: p.can_push(),
1070 exact: p.fully_translated,
1071 pushed_conditions: p.conditions.iter().map(condition_label).collect(),
1072 },
1073 None => ExplainPlan {
1074 index_accelerated: false,
1075 exact: false,
1076 pushed_conditions: Vec::new(),
1077 },
1078 })
1079 }
1080
1081 pub fn rows_at_epoch(&self, table: &str, epoch: u64) -> Result<Vec<crate::schema::Row>> {
1087 let t = self
1088 .schema
1089 .tables
1090 .iter()
1091 .find(|t| t.name == table)
1092 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
1093 let current = self.snapshot_epoch();
1094 if epoch > current {
1095 return Err(KitError::Validation(format!(
1096 "epoch {epoch} is in the future (current committed epoch is {current})"
1097 )));
1098 }
1099 let snap = Snapshot::at(mongreldb_core::epoch::Epoch(epoch));
1100 let rows = self.visible_core_rows_at(table, snap)?;
1101 rows.iter()
1102 .map(|r| crate::schema::core_row_to_json(r, t))
1103 .collect()
1104 }
1105
1106 pub fn approx_aggregate(
1112 &self,
1113 table: &str,
1114 column: Option<&str>,
1115 agg: ApproxAggKind,
1116 z: f64,
1117 ) -> Result<Option<ApproxAggregate>> {
1118 let t = self
1119 .schema
1120 .tables
1121 .iter()
1122 .find(|t| t.name == table)
1123 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
1124 if matches!(agg, ApproxAggKind::Sum | ApproxAggKind::Avg) && column.is_none() {
1125 return Err(KitError::Validation(
1126 "approx sum/avg requires a column".into(),
1127 ));
1128 }
1129 let cid = match column {
1130 Some(name) => Some(
1131 t.columns
1132 .iter()
1133 .find(|c| c.name == name)
1134 .ok_or_else(|| KitError::Validation(format!("unknown column '{name}'")))?
1135 .id as u16,
1136 ),
1137 None => None,
1138 };
1139 let core_agg = match agg {
1140 ApproxAggKind::Count => ApproxAgg::Count,
1141 ApproxAggKind::Sum => ApproxAgg::Sum,
1142 ApproxAggKind::Avg => ApproxAgg::Avg,
1143 };
1144 let handle = self.inner.table(table).map_err(KitError::from)?;
1145 let mut guard = handle.lock();
1146 let res = guard
1147 .approx_aggregate(&[], cid, core_agg, z)
1148 .map_err(KitError::from)?;
1149 Ok(res.map(|r| ApproxAggregate {
1150 point: r.point,
1151 ci_low: r.ci_low,
1152 ci_high: r.ci_high,
1153 n_population: r.n_population,
1154 n_sample_live: r.n_sample_live,
1155 n_passing: r.n_passing,
1156 }))
1157 }
1158
1159 pub fn scan_batched<F>(&self, table: &str, batch_size: usize, mut f: F) -> Result<()>
1165 where
1166 F: FnMut(&[serde_json::Map<String, Value>]) -> Result<()>,
1167 {
1168 let kit_t = self
1169 .schema
1170 .tables
1171 .iter()
1172 .find(|t| t.name == table)
1173 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
1174 let batch_size = batch_size.max(1);
1175 let (snapshot, _pin) = self.inner.snapshot();
1178 let handle = self.inner.table(table).map_err(KitError::from)?;
1179 let guard = handle.lock();
1180
1181 let mut projection: Vec<(u16, mongreldb_core::schema::TypeId)> = Vec::new();
1183 let mut meta: Vec<(String, mongreldb_kit_core::schema::ColumnType)> = Vec::new();
1184 for c in &guard.schema().columns {
1185 if let Some(kc) = kit_t.columns.iter().find(|kc| kc.id as u16 == c.id) {
1186 projection.push((c.id, c.ty.clone()));
1187 meta.push((kc.name.clone(), kc.storage_type));
1188 }
1189 }
1190
1191 match guard
1192 .scan_cursor(snapshot, projection, &[])
1193 .map_err(KitError::from)?
1194 {
1195 Some(mut cursor) => {
1196 let mut buf: Vec<serde_json::Map<String, Value>> = Vec::with_capacity(batch_size);
1197 while let Some(batch) = cursor.next_batch().map_err(KitError::from)? {
1198 let nrows = batch.first().map(|c| c.len()).unwrap_or(0);
1199 for j in 0..nrows {
1200 let mut m = serde_json::Map::new();
1201 for (ci, (name, ty)) in meta.iter().enumerate() {
1202 let cv = batch
1203 .get(ci)
1204 .and_then(|col| col.value_at(j))
1205 .unwrap_or(CoreValue::Null);
1206 m.insert(name.clone(), crate::schema::core_to_json(&cv, *ty)?);
1207 }
1208 buf.push(m);
1209 if buf.len() >= batch_size {
1210 f(&buf)?;
1211 buf.clear();
1212 }
1213 }
1214 }
1215 if !buf.is_empty() {
1216 f(&buf)?;
1217 }
1218 Ok(())
1219 }
1220 None => {
1221 drop(guard);
1222 let rows = self.visible_core_rows_at(table, snapshot)?;
1223 let maps: Vec<serde_json::Map<String, Value>> = rows
1224 .iter()
1225 .map(|r| crate::schema::core_row_to_json(r, kit_t).map(|row| row.values))
1226 .collect::<Result<Vec<_>>>()?;
1227 for chunk in maps.chunks(batch_size) {
1228 f(chunk)?;
1229 }
1230 Ok(())
1231 }
1232 }
1233 }
1234
1235 pub fn set_similarity(
1244 &self,
1245 table: &str,
1246 column: &str,
1247 query: &[String],
1248 k: usize,
1249 ) -> Result<Vec<SimilarRow>> {
1250 let t = self
1251 .schema
1252 .tables
1253 .iter()
1254 .find(|t| t.name == table)
1255 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
1256 let col = t.columns.iter().find(|c| c.name == column).ok_or_else(|| {
1257 KitError::Validation(format!("unknown column '{column}' on table '{table}'"))
1258 })?;
1259 let query_set: std::collections::HashSet<String> = query.iter().cloned().collect();
1260
1261 let has_minhash = t.indexes.iter().any(|idx| {
1262 idx.kind == KitIndexKind::MinHash && idx.columns.iter().any(|c| c == column)
1263 });
1264 let rows = if has_minhash {
1265 let query_hashes: Vec<u64> = query
1267 .iter()
1268 .map(|s| mongreldb_core::index::minhash_token_hash(s))
1269 .collect();
1270 let cand_k = k.saturating_mul(8).max(k + 64);
1272 let cond = mongreldb_core::query::Condition::MinHashSimilar {
1273 column_id: col.id as u16,
1274 query: query_hashes,
1275 k: cand_k,
1276 };
1277 let (snapshot, _pin) = self.inner.snapshot();
1278 let core_rows = self.query_core_rows_at(table, &[cond], snapshot)?;
1279 core_rows
1280 .iter()
1281 .map(|r| crate::schema::core_row_to_json(r, t))
1282 .collect::<Result<Vec<_>>>()?
1283 } else {
1284 let tx = self.begin()?;
1285 tx.all_rows(table)?
1286 };
1287
1288 let mut scored: Vec<SimilarRow> = Vec::new();
1289 for row in rows {
1290 let set = parse_string_set(row.values.get(column));
1291 let inter = set.iter().filter(|x| query_set.contains(*x)).count();
1292 let union = set.len() + query_set.len() - inter;
1293 let sim = if union == 0 {
1294 0.0
1295 } else {
1296 inter as f64 / union as f64
1297 };
1298 if sim > 0.0 {
1299 scored.push(SimilarRow {
1300 row,
1301 similarity: sim,
1302 });
1303 }
1304 }
1305 scored.sort_by(|a, b| {
1306 b.similarity
1307 .partial_cmp(&a.similarity)
1308 .unwrap_or(std::cmp::Ordering::Equal)
1309 });
1310 scored.truncate(k);
1311 Ok(scored)
1312 }
1313
1314 pub fn flush(&self) -> Result<()> {
1318 for name in self.inner.table_names() {
1319 let handle = self.inner.table(&name).map_err(KitError::from)?;
1320 let mut guard = handle.lock();
1321 guard.flush().map_err(KitError::from)?;
1322 }
1323 Ok(())
1324 }
1325
1326 pub fn incremental_aggregate(
1338 &self,
1339 table: &str,
1340 column: Option<&str>,
1341 agg: IncrementalAggKind,
1342 filter: Option<&mongreldb_kit_core::query::Expr>,
1343 ) -> Result<IncrementalAggregate> {
1344 let t = self
1345 .schema
1346 .tables
1347 .iter()
1348 .find(|t| t.name == table)
1349 .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
1350 if !matches!(agg, IncrementalAggKind::Count) && column.is_none() {
1351 return Err(KitError::Validation(
1352 "sum/min/max/avg incremental aggregate requires a column".into(),
1353 ));
1354 }
1355 let cid = match column {
1356 Some(name) => Some(
1357 t.columns
1358 .iter()
1359 .find(|c| c.name == name)
1360 .ok_or_else(|| KitError::Validation(format!("unknown column '{name}'")))?
1361 .id as u16,
1362 ),
1363 None => None,
1364 };
1365 let conditions = match filter {
1366 Some(expr) => {
1367 let plan = crate::pushdown::translate_predicate(t, expr).ok_or_else(|| {
1368 KitError::Validation(
1369 "filter is not index-translatable for an incremental aggregate".into(),
1370 )
1371 })?;
1372 if !plan.fully_translated {
1373 return Err(KitError::Validation(
1374 "filter has a residual that an incremental aggregate cannot apply exactly"
1375 .into(),
1376 ));
1377 }
1378 plan.conditions
1379 }
1380 None => Vec::new(),
1381 };
1382 let core_agg = match agg {
1383 IncrementalAggKind::Count => NativeAgg::Count,
1384 IncrementalAggKind::Sum => NativeAgg::Sum,
1385 IncrementalAggKind::Min => NativeAgg::Min,
1386 IncrementalAggKind::Max => NativeAgg::Max,
1387 IncrementalAggKind::Avg => NativeAgg::Avg,
1388 };
1389 let cache_key = incremental_cache_key(t.id, cid, agg, &conditions);
1390 let handle = self.inner.table(table).map_err(KitError::from)?;
1391 let mut guard = handle.lock();
1392 let res = guard
1393 .aggregate_incremental(cache_key, &conditions, cid, core_agg)
1394 .map_err(KitError::from)?;
1395 Ok(IncrementalAggregate {
1396 value: agg_state_value(&res.state),
1397 incremental: res.incremental,
1398 delta_rows: res.delta_rows,
1399 })
1400 }
1401
1402 pub fn applied_migrations(&self) -> Result<Vec<mongreldb_kit_core::migrations::Migration>> {
1404 crate::migrate::load_applied_migrations(&self.inner)
1405 }
1406
1407 pub(crate) fn core_db(&self) -> &CoreDatabase {
1408 &self.inner
1409 }
1410
1411 pub(crate) fn core_arc(&self) -> Arc<CoreDatabase> {
1414 Arc::clone(&self.inner)
1415 }
1416
1417 pub fn close(&self) -> Result<()> {
1422 self.inner.close().map_err(KitError::from)
1423 }
1424
1425 pub fn compact_all(&self) -> Result<(usize, usize)> {
1430 self.inner.compact().map_err(KitError::from)
1431 }
1432
1433 pub fn compact_table(&self, name: &str) -> Result<bool> {
1436 self.inner.compact_table(name).map_err(KitError::from)
1437 }
1438
1439 pub fn rename_table(&mut self, from: &str, to: &str) -> Result<()> {
1449 if from.starts_with("__kit_") || to.starts_with("__kit_") {
1450 return Err(KitError::Validation(
1451 "rename_table: names beginning with '__kit_' are reserved for internal tables"
1452 .into(),
1453 ));
1454 }
1455 self.inner.rename_table(from, to).map_err(KitError::from)?;
1456 if !self.schema.rename_table(from, to) {
1459 return Err(KitError::Integrity(format!(
1462 "rename_table: kit schema has no table '{from}' (or '{to}' already exists)"
1463 )));
1464 }
1465 for table in &mut self.schema.tables {
1466 for fk in &mut table.foreign_keys {
1467 if fk.references_table == from {
1468 fk.references_table = to.to_string();
1469 }
1470 }
1471 }
1472 store_schema(&self.root, &self.schema)?;
1473 Ok(())
1474 }
1475
1476 pub fn analyze(&self) -> Result<()> {
1481 for name in self.inner.table_names() {
1482 let handle = self.inner.table(&name).map_err(KitError::from)?;
1483 handle.lock().ensure_indexes_complete()?;
1484 }
1485 Ok(())
1486 }
1487
1488 pub fn vacuum(&self) -> Result<usize> {
1492 self.inner.compact().map_err(KitError::from)?;
1493 self.inner.gc().map_err(KitError::from)
1494 }
1495
1496 pub fn create_view(&self, spec: &ViewSpec) -> Result<()> {
1502 self.sql(&spec.create_sql())?;
1503 Ok(())
1504 }
1505
1506 pub fn drop_view(&self, name: &str) -> Result<()> {
1508 self.sql(&format!("DROP VIEW IF EXISTS {name}"))?;
1509 Ok(())
1510 }
1511
1512 pub fn reserve_auto_inc(&self, table: &str) -> Result<Option<i64>> {
1520 let handle = self.inner.table(table).map_err(KitError::from)?;
1521 let mut guard = handle.lock();
1522 guard.reserve_auto_inc().map_err(KitError::from)
1523 }
1524
1525 pub fn create_user(&self, username: &str, password: &str) -> Result<()> {
1529 self.inner
1530 .create_user(username, password)
1531 .map_err(KitError::from)?;
1532 Ok(())
1533 }
1534
1535 pub fn drop_user(&self, username: &str) -> Result<()> {
1537 self.inner.drop_user(username).map_err(KitError::from)
1538 }
1539
1540 pub fn alter_user_password(&self, username: &str, new_password: &str) -> Result<()> {
1542 self.inner
1543 .alter_user_password(username, new_password)
1544 .map_err(KitError::from)
1545 }
1546
1547 pub fn verify_user(
1549 &self,
1550 username: &str,
1551 password: &str,
1552 ) -> Result<Option<mongreldb_core::auth::UserEntry>> {
1553 self.inner
1554 .verify_user(username, password)
1555 .map_err(KitError::from)
1556 }
1557
1558 pub fn set_user_admin(&self, username: &str, is_admin: bool) -> Result<()> {
1560 self.inner
1561 .set_user_admin(username, is_admin)
1562 .map_err(KitError::from)
1563 }
1564
1565 pub fn users(&self) -> Vec<String> {
1567 self.inner.users().into_iter().map(|u| u.username).collect()
1568 }
1569
1570 pub fn create_role(&self, name: &str) -> Result<()> {
1572 self.inner.create_role(name).map_err(KitError::from)?;
1573 Ok(())
1574 }
1575
1576 pub fn drop_role(&self, name: &str) -> Result<()> {
1578 self.inner.drop_role(name).map_err(KitError::from)
1579 }
1580
1581 pub fn roles(&self) -> Vec<String> {
1583 self.inner.roles().into_iter().map(|r| r.name).collect()
1584 }
1585
1586 pub fn grant_role(&self, username: &str, role_name: &str) -> Result<()> {
1588 self.inner
1589 .grant_role(username, role_name)
1590 .map_err(KitError::from)
1591 }
1592
1593 pub fn revoke_role(&self, username: &str, role_name: &str) -> Result<()> {
1595 self.inner
1596 .revoke_role(username, role_name)
1597 .map_err(KitError::from)
1598 }
1599
1600 pub fn grant_permission(
1602 &self,
1603 role_name: &str,
1604 permission: mongreldb_core::auth::Permission,
1605 ) -> Result<()> {
1606 self.inner
1607 .grant_permission(role_name, permission)
1608 .map_err(KitError::from)
1609 }
1610
1611 pub fn revoke_permission(
1613 &self,
1614 role_name: &str,
1615 permission: mongreldb_core::auth::Permission,
1616 ) -> Result<()> {
1617 self.inner
1618 .revoke_permission(role_name, permission)
1619 .map_err(KitError::from)
1620 }
1621
1622 pub fn set_spill_threshold(&self, bytes: u64) {
1628 self.inner.set_spill_threshold(bytes);
1629 }
1630
1631 pub fn set_recursive_triggers(&self, enabled: bool) {
1633 self.inner.set_recursive_triggers(enabled);
1634 }
1635
1636 pub fn trigger_config(&self) -> mongreldb_core::TriggerConfig {
1638 self.inner.trigger_config()
1639 }
1640
1641 pub fn set_trigger_config(&self, config: mongreldb_core::TriggerConfig) -> Result<()> {
1643 self.inner
1644 .set_trigger_config(config)
1645 .map_err(KitError::from)
1646 }
1647
1648 pub fn set_table_compaction_zstd_level(&self, table: &str, level: i32) -> Result<()> {
1650 let handle = self.inner.table(table).map_err(KitError::from)?;
1651 handle.lock().set_compaction_zstd_level(level);
1652 Ok(())
1653 }
1654
1655 pub fn set_table_result_cache_max_bytes(&self, table: &str, max_bytes: u64) -> Result<()> {
1657 let handle = self.inner.table(table).map_err(KitError::from)?;
1658 handle.lock().set_result_cache_max_bytes(max_bytes);
1659 Ok(())
1660 }
1661
1662 pub fn set_table_mutable_run_spill_bytes(&self, table: &str, bytes: u64) -> Result<()> {
1664 let handle = self.inner.table(table).map_err(KitError::from)?;
1665 handle.lock().set_mutable_run_spill_bytes(bytes);
1666 Ok(())
1667 }
1668
1669 pub fn set_table_sync_byte_threshold(&self, table: &str, threshold: u64) -> Result<()> {
1671 let handle = self.inner.table(table).map_err(KitError::from)?;
1672 handle.lock().set_sync_byte_threshold(threshold);
1673 Ok(())
1674 }
1675
1676 pub fn set_table_index_build_policy(
1679 &self,
1680 table: &str,
1681 policy: mongreldb_core::IndexBuildPolicy,
1682 ) -> Result<()> {
1683 let handle = self.inner.table(table).map_err(KitError::from)?;
1684 handle.lock().set_index_build_policy(policy);
1685 Ok(())
1686 }
1687
1688 pub fn table_page_cache_stats(&self, table: &str) -> Result<mongreldb_core::cache::CacheStats> {
1690 let handle = self.inner.table(table).map_err(KitError::from)?;
1691 let stats = handle.lock().page_cache_stats();
1692 Ok(stats)
1693 }
1694
1695 pub fn table_run_count(&self, table: &str) -> Result<usize> {
1697 let handle = self.inner.table(table).map_err(KitError::from)?;
1698 let n = handle.lock().run_count();
1699 Ok(n)
1700 }
1701
1702 pub fn table_memtable_len(&self, table: &str) -> Result<usize> {
1704 let handle = self.inner.table(table).map_err(KitError::from)?;
1705 let n = handle.lock().memtable_len();
1706 Ok(n)
1707 }
1708
1709 pub fn table_mutable_run_len(&self, table: &str) -> Result<usize> {
1711 let handle = self.inner.table(table).map_err(KitError::from)?;
1712 let n = handle.lock().mutable_run_len();
1713 Ok(n)
1714 }
1715
1716 pub fn table_page_cache_len(&self, table: &str) -> Result<usize> {
1718 let handle = self.inner.table(table).map_err(KitError::from)?;
1719 let n = handle.lock().page_cache_len();
1720 Ok(n)
1721 }
1722
1723 pub fn table_decoded_cache_len(&self, table: &str) -> Result<usize> {
1725 let handle = self.inner.table(table).map_err(KitError::from)?;
1726 let n = handle.lock().decoded_cache_len();
1727 Ok(n)
1728 }
1729
1730 pub fn sql(&self, statement: &str) -> Result<Vec<arrow::record_batch::RecordBatch>> {
1748 self.sql_with_options(statement, SqlOptions::default())
1749 }
1750
1751 fn sql_session(&self) -> Result<Arc<mongreldb_query::MongrelSession>> {
1752 if let Some(session) = self.session.read().as_ref() {
1753 return Ok(Arc::clone(session));
1754 }
1755 let session = Arc::new(
1756 mongreldb_query::MongrelSession::open(self.core_arc()).map_err(KitError::from)?,
1757 );
1758 let mut cached = self.session.write();
1759 Ok(Arc::clone(cached.get_or_insert(session)))
1760 }
1761
1762 #[doc(hidden)]
1763 pub fn set_sql_test_hook(&self, hook: Option<mongreldb_query::SqlTestHook>) -> Result<()> {
1764 self.sql_session()?.set_test_hook(hook);
1765 Ok(())
1766 }
1767
1768 pub fn start_sql(
1769 &self,
1770 statement: impl Into<String>,
1771 options: SqlOptions,
1772 ) -> Result<SqlQueryHandle> {
1773 let session = self.sql_session()?;
1774 let query = session
1775 .register_query(mongreldb_query::SqlQueryOptions {
1776 query_id: options.query_id,
1777 timeout: options.timeout,
1778 ..mongreldb_query::SqlQueryOptions::default()
1779 })
1780 .map_err(KitError::from)?;
1781 let query_id = query.id();
1782 let registration = mongreldb_query::RegisteredQueryGuard::new(query);
1783 let worker_session = Arc::clone(&session);
1784 let statement = statement.into();
1785 let worker = std::thread::Builder::new()
1786 .name(format!("mongreldb-kit-sql-{query_id}"))
1787 .spawn(move || {
1788 sql_runtime().block_on(
1789 worker_session
1790 .run_with_query_for_serialization(&statement, registration.into_query()),
1791 )
1792 })
1793 .map_err(|error| KitError::Storage(error.to_string()))?;
1794 Ok(SqlQueryHandle {
1795 query_id,
1796 session,
1797 worker: Some(worker),
1798 })
1799 }
1800
1801 pub fn sql_with_options(
1802 &self,
1803 statement: &str,
1804 options: SqlOptions,
1805 ) -> Result<Vec<arrow::record_batch::RecordBatch>> {
1806 self.start_sql(statement, options)?.wait()
1807 }
1808
1809 pub fn cancel_sql(&self, query_id: mongreldb_query::QueryId) -> mongreldb_query::CancelOutcome {
1810 self.session
1811 .read()
1812 .as_ref()
1813 .map_or(mongreldb_query::CancelOutcome::NotFound, |session| {
1814 session.cancel_query(query_id)
1815 })
1816 }
1817
1818 pub fn sql_query_status(
1819 &self,
1820 query_id: mongreldb_query::QueryId,
1821 ) -> Result<Option<mongreldb_query::QueryStatus>> {
1822 Ok(self.sql_session()?.query_registry().status(query_id))
1823 }
1824
1825 pub fn refresh_sql_session(&self) -> Result<()> {
1830 let session =
1831 mongreldb_query::MongrelSession::open(self.core_arc()).map_err(KitError::from)?;
1832 *self.session.write() = Some(Arc::new(session));
1833 Ok(())
1834 }
1835
1836 pub fn sql_arrow(&self, statement: &str) -> Result<Vec<u8>> {
1842 self.sql_arrow_with_options(statement, SqlOptions::default())
1843 }
1844
1845 pub fn sql_arrow_with_options(&self, statement: &str, options: SqlOptions) -> Result<Vec<u8>> {
1846 self.sql_serialized_with_options(statement, options, |output| {
1847 crate::arrow_util::batches_to_ipc_controlled(output.batches(), output.query())
1848 })
1849 }
1850
1851 pub fn sql_rows(&self, statement: &str) -> Result<Vec<serde_json::Map<String, Value>>> {
1855 self.sql_rows_with_options(statement, SqlOptions::default())
1856 }
1857
1858 pub fn sql_rows_with_options(
1859 &self,
1860 statement: &str,
1861 options: SqlOptions,
1862 ) -> Result<Vec<serde_json::Map<String, Value>>> {
1863 self.sql_serialized_with_options(statement, options, |output| {
1864 crate::arrow_util::batches_to_rows_controlled(output.batches(), output.query())
1865 })
1866 }
1867
1868 fn sql_serialized_with_options<T>(
1869 &self,
1870 statement: &str,
1871 options: SqlOptions,
1872 serialize: impl FnOnce(&mongreldb_query::ManagedQueryBatches) -> Result<T>,
1873 ) -> Result<T> {
1874 let session = self.sql_session()?;
1875 let query = session
1876 .register_query(mongreldb_query::SqlQueryOptions {
1877 query_id: options.query_id,
1878 timeout: options.timeout,
1879 ..mongreldb_query::SqlQueryOptions::default()
1880 })
1881 .map_err(KitError::from)?;
1882 let query_id = query.id();
1883 let output = sql_runtime()
1884 .block_on(session.run_with_query_for_serialization(statement, query))
1885 .map_err(|error| {
1886 let status = session.query_registry().status(query_id);
1887 crate::error::query_error_with_status(error, status.as_ref())
1888 })?;
1889 session.fire_test_hook(mongreldb_query::SqlTestHookPoint::BeforeSerializationBatch);
1890 match serialize(&output) {
1891 Ok(value) => {
1892 complete_sql_output(output)?;
1893 Ok(value)
1894 }
1895 Err(error) => {
1896 fail_sql_output(output, &error);
1897 Err(error)
1898 }
1899 }
1900 }
1901
1902 pub(crate) fn lookup_row_id(&self, table: &str, key: &[u8]) -> Result<Option<RowId>> {
1906 let handle = self.inner.table(table).map_err(KitError::from)?;
1907 let mut guard = handle.lock();
1908 guard.ensure_indexes_complete()?;
1909 Ok(guard.lookup_pk(key))
1910 }
1911
1912 pub(crate) fn root(&self) -> &Path {
1913 &self.root
1914 }
1915
1916 pub(crate) fn visible_core_rows_at(
1920 &self,
1921 table_name: &str,
1922 snapshot: Snapshot,
1923 ) -> Result<Vec<CoreRow>> {
1924 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1925 let guard = handle.lock();
1926 guard.visible_rows(snapshot).map_err(KitError::from)
1927 }
1928
1929 pub(crate) fn query_core_rows_at(
1936 &self,
1937 table_name: &str,
1938 conditions: &[mongreldb_core::query::Condition],
1939 snapshot: Snapshot,
1940 ) -> Result<Vec<CoreRow>> {
1941 if conditions.is_empty() {
1942 return self.visible_core_rows_at(table_name, snapshot);
1943 }
1944 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1945 let mut guard = handle.lock();
1946 let q = conditions
1947 .iter()
1948 .cloned()
1949 .fold(mongreldb_core::query::Query::new(), |query, condition| {
1950 query.and(condition)
1951 });
1952 guard.query(&q).map_err(KitError::from)
1953 }
1954
1955 pub(crate) fn flush_table(&self, table_name: &str) -> Result<()> {
1963 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1964 handle.lock().flush().map_err(KitError::from)?;
1965 Ok(())
1966 }
1967
1968 pub(crate) fn count_core_rows_at(
1978 &self,
1979 table_name: &str,
1980 conditions: &[mongreldb_core::query::Condition],
1981 snapshot: Snapshot,
1982 ) -> Result<Option<u64>> {
1983 let handle = self.inner.table(table_name).map_err(KitError::from)?;
1984 let mut guard = handle.lock();
1985 if guard.snapshot().epoch != snapshot.epoch {
1986 return Ok(None); }
1988 guard
1989 .count_conditions(conditions, snapshot)
1990 .map_err(KitError::from)
1991 }
1992
1993 pub(crate) fn aggregate_core_at(
2002 &self,
2003 table_name: &str,
2004 column: Option<u16>,
2005 conditions: &[mongreldb_core::query::Condition],
2006 agg: NativeAgg,
2007 snapshot: Snapshot,
2008 ) -> Result<Option<NativeAggResult>> {
2009 let handle = self.inner.table(table_name).map_err(KitError::from)?;
2010 let guard = handle.lock();
2011 if guard.snapshot().epoch != snapshot.epoch {
2012 return Ok(None); }
2014 guard
2015 .aggregate_native(snapshot, column, conditions, agg)
2016 .map_err(KitError::from)
2017 }
2018
2019 pub(crate) fn count_distinct_core_at(
2028 &self,
2029 table_name: &str,
2030 column_id: u16,
2031 snapshot: Snapshot,
2032 ) -> Result<Option<u64>> {
2033 let handle = self.inner.table(table_name).map_err(KitError::from)?;
2034 let mut guard = handle.lock();
2035 if guard.snapshot().epoch != snapshot.epoch {
2036 return Ok(None); }
2038 guard
2039 .count_distinct_from_bitmap(column_id)
2040 .map_err(KitError::from)
2041 }
2042
2043 #[allow(dead_code)]
2045 pub(crate) fn get_core_row(&self, table_name: &str, row_id: u64) -> Result<Option<CoreRow>> {
2046 let handle = self.inner.table(table_name).map_err(KitError::from)?;
2047 let guard = handle.lock();
2048 let snapshot = guard.snapshot();
2049 Ok(guard.get(mongreldb_core::RowId(row_id), snapshot))
2050 }
2051}
2052
2053pub(crate) fn create_core_table(db: &CoreDatabase, name: &str, schema: CoreSchema) -> Result<()> {
2054 if db.table_id(name).is_ok() {
2055 return Ok(());
2056 }
2057 db.create_table(name, schema).map_err(KitError::from)?;
2058 Ok(())
2059}
2060
2061fn sql_runtime() -> &'static tokio::runtime::Runtime {
2065 use std::sync::OnceLock;
2066 static RT: OnceLock<tokio::runtime::Runtime> = OnceLock::new();
2067 RT.get_or_init(|| {
2068 tokio::runtime::Builder::new_multi_thread()
2069 .worker_threads(4)
2070 .enable_all()
2071 .build()
2072 .expect("failed to build kit SQL tokio runtime")
2073 })
2074}
2075
2076fn core_procedure(spec: &ProcedureSpec) -> Result<mongreldb_core::StoredProcedure> {
2077 let parsed: mongreldb_core::StoredProcedure =
2078 serde_json::from_value(spec.json.clone()).map_err(KitError::from)?;
2079 mongreldb_core::StoredProcedure::new(parsed.name, parsed.mode, parsed.params, parsed.body, 0)
2080 .map_err(KitError::from)
2081}
2082
2083fn core_trigger(spec: &TriggerSpec) -> Result<mongreldb_core::StoredTrigger> {
2084 let parsed: mongreldb_core::StoredTrigger =
2085 serde_json::from_value(spec.json.clone()).map_err(KitError::from)?;
2086 mongreldb_core::StoredTrigger::new(
2087 parsed.name,
2088 mongreldb_core::TriggerDefinition {
2089 target: parsed.target,
2090 timing: parsed.timing,
2091 event: parsed.event,
2092 update_of: parsed.update_of,
2093 target_columns: parsed.target_columns,
2094 when: parsed.when,
2095 program: parsed.program,
2096 },
2097 0,
2098 )
2099 .map_err(KitError::from)
2100}
2101
2102fn json_to_core_value(value: &Value) -> Result<CoreValue> {
2103 match value {
2104 Value::Null => Ok(CoreValue::Null),
2105 Value::Bool(value) => Ok(CoreValue::Bool(*value)),
2106 Value::Number(value) => {
2107 if let Some(value) = value.as_i64() {
2108 Ok(CoreValue::Int64(value))
2109 } else if let Some(value) = value.as_f64() {
2110 Ok(CoreValue::Float64(value))
2111 } else {
2112 Err(KitError::Validation("unsupported JSON number".into()))
2113 }
2114 }
2115 Value::String(value) => Ok(CoreValue::Bytes(value.as_bytes().to_vec())),
2116 Value::Array(_) | Value::Object(_) => Err(KitError::Validation(
2117 "procedure args only support scalar JSON values".into(),
2118 )),
2119 }
2120}
2121
2122pub(crate) fn internal_bytes(row: &CoreRow, col_id: u16) -> Option<String> {
2124 match row.columns.get(&col_id) {
2125 Some(CoreValue::Bytes(b)) => String::from_utf8(b.clone()).ok(),
2126 _ => None,
2127 }
2128}
2129
2130fn reap_rotated_wal_segments(db: &CoreDatabase) {
2145 let _ = db.gc();
2146}
2147
2148pub(crate) fn load_schema(path: &Path) -> Result<KitSchema> {
2149 let file = path.join(SCHEMA_FILE);
2150 let json = std::fs::read_to_string(&file)
2151 .map_err(|e| KitError::Migration(format!("cannot read schema file: {e}")))?;
2152 let schema: KitSchema = serde_json::from_str(&json)?;
2153 Ok(schema)
2154}
2155
2156pub(crate) fn store_schema(path: &Path, schema: &KitSchema) -> Result<()> {
2157 let file = path.join(SCHEMA_FILE);
2158 let json = serde_json::to_string_pretty(schema)?;
2159 std::fs::write(&file, json)?;
2160 Ok(())
2161}
2162
2163pub(crate) fn persist_schema(db: &Database, schema: &KitSchema) -> Result<()> {
2165 store_schema(&db.root, schema)
2166}
2167
2168#[cfg(test)]
2169mod tests {
2170 use super::open_core_with_retry;
2171
2172 fn lock_error() -> mongreldb_core::MongrelError {
2173 mongreldb_core::MongrelError::DatabaseLocked {
2174 path: "/tmp/db".into(),
2175 message: "would block".into(),
2176 }
2177 }
2178
2179 #[test]
2180 fn open_retry_waits_for_lock_contention_only() {
2181 let mut calls = 0;
2182 let value = open_core_with_retry(50, || {
2183 calls += 1;
2184 if calls < 3 {
2185 Err(lock_error())
2186 } else {
2187 Ok(7)
2188 }
2189 })
2190 .unwrap();
2191 assert_eq!(value, 7);
2192 assert_eq!(calls, 3);
2193
2194 let mut non_lock_calls = 0;
2195 let err: mongreldb_core::Result<()> = open_core_with_retry(50, || {
2196 non_lock_calls += 1;
2197 Err(mongreldb_core::MongrelError::Other("nope".into()))
2198 });
2199 let err = err.unwrap_err();
2200 assert_eq!(non_lock_calls, 1);
2201 assert!(matches!(err, mongreldb_core::MongrelError::Other(_)));
2202 }
2203}