use crate::error::{KitError, Result};
use crate::internal::{ensure_internal_tables, internal_tables_core};
use crate::schema::to_core_schema;
use mongreldb_core::epoch::Snapshot;
use mongreldb_core::memtable::Row as CoreRow;
use mongreldb_core::memtable::Value as CoreValue;
use mongreldb_core::schema::Schema as CoreSchema;
use mongreldb_core::Database as CoreDatabase;
use mongreldb_core::{AggState, ApproxAgg, NativeAgg, NativeAggResult, RowId};
use mongreldb_kit_core::schema::IndexKind as KitIndexKind;
use mongreldb_kit_core::schema::Schema as KitSchema;
use mongreldb_kit_core::schema::Table as KitTable;
use mongreldb_kit_core::{ProcedureSpec, TriggerSpec, ViewSpec};
use serde_json::Value;
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
const SCHEMA_FILE: &str = "kit_schema.json";
pub type DefaultProvider = Box<dyn Fn() -> Value + Send + Sync>;
#[derive(Debug, Clone)]
pub struct ExplainPlan {
pub index_accelerated: bool,
pub exact: bool,
pub pushed_conditions: Vec<String>,
}
#[derive(Debug, Clone)]
pub struct SimilarRow {
pub row: crate::schema::Row,
pub similarity: f64,
}
fn parse_string_set(value: Option<&Value>) -> std::collections::HashSet<String> {
let arr = match value {
Some(Value::Array(a)) => Some(a.clone()),
Some(Value::String(s)) => serde_json::from_str::<Value>(s)
.ok()
.and_then(|v| v.as_array().cloned()),
_ => None,
};
arr.into_iter()
.flatten()
.filter_map(|v| match v {
Value::String(s) => Some(s),
Value::Number(n) => Some(n.to_string()),
Value::Bool(b) => Some(b.to_string()),
_ => None,
})
.collect()
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum IncrementalAggKind {
Count,
Sum,
Min,
Max,
Avg,
}
#[derive(Debug, Clone)]
pub struct IncrementalAggregate {
pub value: Value,
pub incremental: bool,
pub delta_rows: u64,
}
fn incremental_cache_key(
table_id: u32,
column: Option<u16>,
agg: IncrementalAggKind,
conditions: &[mongreldb_core::query::Condition],
) -> u64 {
use std::hash::{Hash, Hasher};
let mut h = std::collections::hash_map::DefaultHasher::new();
table_id.hash(&mut h);
column.hash(&mut h);
(agg as u8).hash(&mut h);
format!("{conditions:?}").hash(&mut h);
h.finish()
}
fn agg_state_value(s: &AggState) -> Value {
let num_f64 = |x: f64| {
serde_json::Number::from_f64(x)
.map(Value::Number)
.unwrap_or(Value::Null)
};
match s {
AggState::Count(n) => Value::from(*n),
AggState::SumI { sum, .. } => i64::try_from(*sum)
.map(Value::from)
.unwrap_or_else(|_| num_f64(*sum as f64)),
AggState::SumF { sum, .. } => num_f64(*sum),
AggState::AvgI { sum, count } if *count > 0 => num_f64(*sum as f64 / *count as f64),
AggState::AvgF { sum, count } if *count > 0 => num_f64(*sum / *count as f64),
AggState::AvgI { .. } | AggState::AvgF { .. } => Value::Null,
AggState::MinI(n) | AggState::MaxI(n) => Value::from(*n),
AggState::MinF(f) | AggState::MaxF(f) => num_f64(*f),
AggState::Empty => Value::Null,
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ApproxAggKind {
Count,
Sum,
Avg,
}
#[derive(Debug, Clone)]
pub struct ApproxAggregate {
pub point: f64,
pub ci_low: f64,
pub ci_high: f64,
pub n_population: u64,
pub n_sample_live: usize,
pub n_passing: usize,
}
fn condition_label(c: &mongreldb_core::query::Condition) -> String {
let dbg = format!("{c:?}");
dbg.split(['(', '{', ' ']).next().unwrap_or("").to_string()
}
pub struct Database {
pub(crate) inner: Arc<CoreDatabase>,
pub(crate) schema: KitSchema,
pub(crate) root: PathBuf,
pub(crate) default_providers: HashMap<String, DefaultProvider>,
pub(crate) session: parking_lot::Mutex<Option<mongreldb_query::MongrelSession>>,
}
impl Database {
pub fn open(path: &Path) -> Result<Self> {
let inner = Arc::new(CoreDatabase::open(path)?);
let schema = load_schema(path)?;
ensure_internal_tables(&inner)?;
reap_rotated_wal_segments(&inner);
Ok(Self {
inner,
schema,
root: path.to_path_buf(),
default_providers: HashMap::new(),
session: parking_lot::Mutex::new(None),
})
}
pub fn open_encrypted(path: &Path, passphrase: &str) -> Result<Self> {
let inner = Arc::new(CoreDatabase::open_encrypted(path, passphrase)?);
let schema = load_schema(path)?;
ensure_internal_tables(&inner)?;
reap_rotated_wal_segments(&inner);
Ok(Self {
inner,
schema,
root: path.to_path_buf(),
default_providers: HashMap::new(),
session: parking_lot::Mutex::new(None),
})
}
pub fn create_encrypted(path: &Path, schema: KitSchema, passphrase: &str) -> Result<Self> {
std::fs::create_dir_all(path)?;
let inner = Arc::new(CoreDatabase::create_encrypted(path, passphrase)?);
ensure_internal_tables(&inner)?;
store_schema(path, &schema)?;
for table in &schema.tables {
create_core_table(&inner, &table.name, to_core_schema(table))?;
}
Ok(Self {
inner,
schema,
root: path.to_path_buf(),
default_providers: HashMap::new(),
session: parking_lot::Mutex::new(None),
})
}
pub fn create(path: &Path, schema: KitSchema) -> Result<Self> {
std::fs::create_dir_all(path)?;
let inner = Arc::new(CoreDatabase::create(path)?);
ensure_internal_tables(&inner)?;
store_schema(path, &schema)?;
for table in &schema.tables {
create_core_table(&inner, &table.name, to_core_schema(table))?;
}
Ok(Self {
inner,
schema,
root: path.to_path_buf(),
default_providers: HashMap::new(),
session: parking_lot::Mutex::new(None),
})
}
pub fn open_with_credentials(path: &Path, username: &str, password: &str) -> Result<Self> {
let inner = Arc::new(CoreDatabase::open_with_credentials(
path, username, password,
)?);
let schema = load_schema(path)?;
ensure_internal_tables(&inner)?;
reap_rotated_wal_segments(&inner);
Ok(Self {
inner,
schema,
root: path.to_path_buf(),
default_providers: HashMap::new(),
session: parking_lot::Mutex::new(None),
})
}
pub fn create_with_credentials(
path: &Path,
schema: KitSchema,
admin_username: &str,
admin_password: &str,
) -> Result<Self> {
std::fs::create_dir_all(path)?;
let inner = Arc::new(CoreDatabase::create_with_credentials(
path,
admin_username,
admin_password,
)?);
ensure_internal_tables(&inner)?;
store_schema(path, &schema)?;
for table in &schema.tables {
create_core_table(&inner, &table.name, to_core_schema(table))?;
}
Ok(Self {
inner,
schema,
root: path.to_path_buf(),
default_providers: HashMap::new(),
session: parking_lot::Mutex::new(None),
})
}
pub fn open_encrypted_with_credentials(
path: &Path,
passphrase: &str,
username: &str,
password: &str,
) -> Result<Self> {
let inner = Arc::new(CoreDatabase::open_encrypted_with_credentials(
path, passphrase, username, password,
)?);
let schema = load_schema(path)?;
ensure_internal_tables(&inner)?;
reap_rotated_wal_segments(&inner);
Ok(Self {
inner,
schema,
root: path.to_path_buf(),
default_providers: HashMap::new(),
session: parking_lot::Mutex::new(None),
})
}
pub fn create_encrypted_with_credentials(
path: &Path,
schema: KitSchema,
passphrase: &str,
admin_username: &str,
admin_password: &str,
) -> Result<Self> {
std::fs::create_dir_all(path)?;
let inner = Arc::new(CoreDatabase::create_encrypted_with_credentials(
path,
passphrase,
admin_username,
admin_password,
)?);
ensure_internal_tables(&inner)?;
store_schema(path, &schema)?;
for table in &schema.tables {
create_core_table(&inner, &table.name, to_core_schema(table))?;
}
Ok(Self {
inner,
schema,
root: path.to_path_buf(),
default_providers: HashMap::new(),
session: parking_lot::Mutex::new(None),
})
}
pub fn enable_auth(&self, admin_username: &str, admin_password: &str) -> Result<()> {
self.inner
.enable_auth(admin_username, admin_password)
.map_err(KitError::from)
}
pub fn require_auth_enabled(&self) -> bool {
self.inner.require_auth_enabled()
}
pub fn refresh_principal(&self) -> Result<()> {
self.inner.refresh_principal().map_err(KitError::from)
}
pub fn register_default(
&mut self,
name: impl Into<String>,
provider: impl Fn() -> Value + Send + Sync + 'static,
) {
self.default_providers
.insert(name.into(), Box::new(provider));
}
pub fn raw(&self) -> &CoreDatabase {
&self.inner
}
pub fn table_names(&self) -> Vec<String> {
self.schema
.tables
.iter()
.map(|t| t.name.clone())
.filter(|n| !n.starts_with("__kit_"))
.collect()
}
pub fn create_procedure(
&self,
spec: &ProcedureSpec,
) -> Result<mongreldb_core::StoredProcedure> {
let procedure = core_procedure(spec)?;
self.inner
.create_procedure(procedure)
.map_err(KitError::from)
}
pub fn replace_procedure(
&self,
spec: &ProcedureSpec,
) -> Result<mongreldb_core::StoredProcedure> {
let procedure = core_procedure(spec)?;
self.inner
.create_or_replace_procedure(procedure)
.map_err(KitError::from)
}
pub fn drop_procedure(&self, name: &str) -> Result<()> {
self.inner.drop_procedure(name).map_err(KitError::from)
}
pub fn call_procedure(
&self,
name: &str,
args: serde_json::Map<String, Value>,
) -> Result<mongreldb_core::ProcedureCallResult> {
let args = args
.iter()
.map(|(key, value)| Ok((key.clone(), json_to_core_value(value)?)))
.collect::<Result<HashMap<_, _>>>()?;
self.inner
.call_procedure(name, args)
.map_err(KitError::from)
}
pub fn create_trigger(&self, spec: &TriggerSpec) -> Result<mongreldb_core::StoredTrigger> {
let trigger = core_trigger(spec)?;
self.inner.create_trigger(trigger).map_err(KitError::from)
}
pub fn replace_trigger(&self, spec: &TriggerSpec) -> Result<mongreldb_core::StoredTrigger> {
let trigger = core_trigger(spec)?;
self.inner
.create_or_replace_trigger(trigger)
.map_err(KitError::from)
}
pub fn drop_trigger(&self, name: &str) -> Result<()> {
self.inner.drop_trigger(name).map_err(KitError::from)
}
pub fn triggers(&self) -> Vec<mongreldb_core::StoredTrigger> {
self.inner.triggers()
}
pub fn trigger(&self, name: &str) -> Option<mongreldb_core::StoredTrigger> {
self.inner.trigger(name)
}
pub fn allocate_sequence(&self, name: &str, count: i64) -> Result<i64> {
use crate::internal::cols;
let mut attempt = 0;
loop {
let mut txn = self.inner.begin();
let snapshot = txn.read_snapshot();
let existing = self
.visible_core_rows_at(crate::internal::SEQUENCES, snapshot)?
.into_iter()
.find(|r| internal_bytes(r, cols::SEQ_NAME) == Some(name.to_string()));
let now = crate::internal::iso_now();
let (start, next, old_row_id) = match &existing {
Some(row) => {
let current = match row.columns.get(&cols::SEQ_NEXT) {
Some(CoreValue::Int64(i)) => *i,
_ => 1,
};
(current, current + count, Some(row.row_id))
}
None => (1, 1 + count, None),
};
if let Some(rid) = old_row_id {
txn.delete(crate::internal::SEQUENCES, rid)
.map_err(KitError::from)?;
}
txn.put(
crate::internal::SEQUENCES,
vec![
(cols::SEQ_NAME, CoreValue::Bytes(name.as_bytes().to_vec())),
(cols::SEQ_NEXT, CoreValue::Int64(next)),
(cols::SEQ_UPDATED, CoreValue::Bytes(now.into_bytes())),
],
)
.map_err(KitError::from)?;
match txn.commit() {
Ok(_) => return Ok(start),
Err(mongreldb_core::MongrelError::Conflict(_)) if attempt < 10_000 => {
attempt += 1;
std::thread::yield_now();
continue;
}
Err(e) => return Err(KitError::from(e)),
}
}
}
pub fn transaction<T, F>(&self, max_retries: usize, mut f: F) -> Result<T>
where
F: FnMut(&mut crate::txn::Transaction<'_>) -> Result<T>,
{
let mut attempt = 0;
loop {
let mut txn = self.begin()?;
match f(&mut txn) {
Ok(value) => match txn.commit() {
Ok(()) => return Ok(value),
Err(KitError::Conflict(_)) if attempt < max_retries => {
attempt += 1;
continue;
}
Err(e) => return Err(e),
},
Err(KitError::Conflict(_)) if attempt < max_retries => {
txn.rollback();
attempt += 1;
continue;
}
Err(e) => {
txn.rollback();
return Err(e);
}
}
}
}
pub fn table(&self, name: &str) -> Option<&KitTable> {
self.schema.table(name)
}
pub fn schema(&self) -> &KitSchema {
&self.schema
}
pub fn begin(&self) -> Result<crate::txn::Transaction<'_>> {
let core_txn = self.inner.begin();
Ok(crate::txn::Transaction::new(self, core_txn))
}
pub fn set_schema(&mut self, schema: KitSchema) {
self.schema = schema;
}
pub fn check_internal_tables(&self) -> Result<()> {
let schema_file = self.root.join(SCHEMA_FILE);
if !schema_file.exists() {
return Err(KitError::Integrity(format!(
"schema file {} is missing",
schema_file.display()
)));
}
for (name, _) in internal_tables_core() {
if self.inner.table_id(name).is_err() {
return Err(KitError::Integrity(format!(
"internal table {name} is missing"
)));
}
}
Ok(())
}
pub fn gc(&self) -> Result<usize> {
self.inner.gc().map_err(KitError::from)
}
pub fn check(&self) -> Vec<serde_json::Value> {
self.inner
.check()
.into_iter()
.map(|i| {
serde_json::json!({
"table_id": i.table_id,
"table_name": i.table_name,
"severity": i.severity,
"description": i.description,
})
})
.collect()
}
pub fn doctor(&self) -> Result<Vec<u64>> {
self.inner.doctor().map_err(KitError::from)
}
pub fn snapshot_epoch(&self) -> u64 {
self.inner.snapshot().0.epoch.0
}
pub fn export_tsv(&self, table: &str) -> Result<String> {
let t = self
.schema
.tables
.iter()
.find(|t| t.name == table)
.ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?
.clone();
let tx = self.begin()?;
let rows = tx.all_rows(table)?;
Ok(crate::tsv::rows_to_tsv(&t, &rows))
}
pub fn import_tsv(&self, table: &str, text: &str) -> Result<usize> {
let t = self
.schema
.tables
.iter()
.find(|t| t.name == table)
.ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?
.clone();
let rows = crate::tsv::tsv_to_rows(&t, text)?;
let n = rows.len();
self.transaction(1, |tx| {
tx.insert_many(table, rows.clone())?;
Ok(())
})?;
Ok(n)
}
pub fn explain(
&self,
table: &str,
predicate: &mongreldb_kit_core::query::Expr,
) -> Result<ExplainPlan> {
let t = self
.schema
.tables
.iter()
.find(|t| t.name == table)
.ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
Ok(match crate::pushdown::translate_predicate(t, predicate) {
Some(p) => ExplainPlan {
index_accelerated: p.can_push(),
exact: p.fully_translated,
pushed_conditions: p.conditions.iter().map(condition_label).collect(),
},
None => ExplainPlan {
index_accelerated: false,
exact: false,
pushed_conditions: Vec::new(),
},
})
}
pub fn rows_at_epoch(&self, table: &str, epoch: u64) -> Result<Vec<crate::schema::Row>> {
let t = self
.schema
.tables
.iter()
.find(|t| t.name == table)
.ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
let current = self.snapshot_epoch();
if epoch > current {
return Err(KitError::Validation(format!(
"epoch {epoch} is in the future (current committed epoch is {current})"
)));
}
let snap = Snapshot::at(mongreldb_core::epoch::Epoch(epoch));
let rows = self.visible_core_rows_at(table, snap)?;
rows.iter()
.map(|r| crate::schema::core_row_to_json(r, t))
.collect()
}
pub fn approx_aggregate(
&self,
table: &str,
column: Option<&str>,
agg: ApproxAggKind,
z: f64,
) -> Result<Option<ApproxAggregate>> {
let t = self
.schema
.tables
.iter()
.find(|t| t.name == table)
.ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
if matches!(agg, ApproxAggKind::Sum | ApproxAggKind::Avg) && column.is_none() {
return Err(KitError::Validation(
"approx sum/avg requires a column".into(),
));
}
let cid = match column {
Some(name) => Some(
t.columns
.iter()
.find(|c| c.name == name)
.ok_or_else(|| KitError::Validation(format!("unknown column '{name}'")))?
.id as u16,
),
None => None,
};
let core_agg = match agg {
ApproxAggKind::Count => ApproxAgg::Count,
ApproxAggKind::Sum => ApproxAgg::Sum,
ApproxAggKind::Avg => ApproxAgg::Avg,
};
let handle = self.inner.table(table).map_err(KitError::from)?;
let mut guard = handle.lock();
let res = guard
.approx_aggregate(&[], cid, core_agg, z)
.map_err(KitError::from)?;
Ok(res.map(|r| ApproxAggregate {
point: r.point,
ci_low: r.ci_low,
ci_high: r.ci_high,
n_population: r.n_population,
n_sample_live: r.n_sample_live,
n_passing: r.n_passing,
}))
}
pub fn scan_batched<F>(&self, table: &str, batch_size: usize, mut f: F) -> Result<()>
where
F: FnMut(&[serde_json::Map<String, Value>]) -> Result<()>,
{
let kit_t = self
.schema
.tables
.iter()
.find(|t| t.name == table)
.ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
let batch_size = batch_size.max(1);
let (snapshot, _pin) = self.inner.snapshot();
let handle = self.inner.table(table).map_err(KitError::from)?;
let guard = handle.lock();
let mut projection: Vec<(u16, mongreldb_core::schema::TypeId)> = Vec::new();
let mut meta: Vec<(String, mongreldb_kit_core::schema::ColumnType)> = Vec::new();
for c in &guard.schema().columns {
if let Some(kc) = kit_t.columns.iter().find(|kc| kc.id as u16 == c.id) {
projection.push((c.id, c.ty));
meta.push((kc.name.clone(), kc.storage_type));
}
}
match guard
.scan_cursor(snapshot, projection, &[])
.map_err(KitError::from)?
{
Some(mut cursor) => {
let mut buf: Vec<serde_json::Map<String, Value>> = Vec::with_capacity(batch_size);
while let Some(batch) = cursor.next_batch().map_err(KitError::from)? {
let nrows = batch.first().map(|c| c.len()).unwrap_or(0);
for j in 0..nrows {
let mut m = serde_json::Map::new();
for (ci, (name, ty)) in meta.iter().enumerate() {
let cv = batch
.get(ci)
.and_then(|col| col.value_at(j))
.unwrap_or(CoreValue::Null);
m.insert(name.clone(), crate::schema::core_to_json(&cv, *ty)?);
}
buf.push(m);
if buf.len() >= batch_size {
f(&buf)?;
buf.clear();
}
}
}
if !buf.is_empty() {
f(&buf)?;
}
Ok(())
}
None => {
drop(guard);
let rows = self.visible_core_rows_at(table, snapshot)?;
let maps: Vec<serde_json::Map<String, Value>> = rows
.iter()
.map(|r| crate::schema::core_row_to_json(r, kit_t).map(|row| row.values))
.collect::<Result<Vec<_>>>()?;
for chunk in maps.chunks(batch_size) {
f(chunk)?;
}
Ok(())
}
}
}
pub fn set_similarity(
&self,
table: &str,
column: &str,
query: &[String],
k: usize,
) -> Result<Vec<SimilarRow>> {
let t = self
.schema
.tables
.iter()
.find(|t| t.name == table)
.ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
let col = t.columns.iter().find(|c| c.name == column).ok_or_else(|| {
KitError::Validation(format!("unknown column '{column}' on table '{table}'"))
})?;
let query_set: std::collections::HashSet<String> = query.iter().cloned().collect();
let has_minhash = t.indexes.iter().any(|idx| {
idx.kind == KitIndexKind::MinHash && idx.columns.iter().any(|c| c == column)
});
let rows = if has_minhash {
let query_hashes: Vec<u64> = query
.iter()
.map(|s| mongreldb_core::index::minhash_token_hash(s))
.collect();
let cand_k = k.saturating_mul(8).max(k + 64);
let cond = mongreldb_core::query::Condition::MinHashSimilar {
column_id: col.id as u16,
query: query_hashes,
k: cand_k,
};
let (snapshot, _pin) = self.inner.snapshot();
let core_rows = self.query_core_rows_at(table, &[cond], snapshot)?;
core_rows
.iter()
.map(|r| crate::schema::core_row_to_json(r, t))
.collect::<Result<Vec<_>>>()?
} else {
let tx = self.begin()?;
tx.all_rows(table)?
};
let mut scored: Vec<SimilarRow> = Vec::new();
for row in rows {
let set = parse_string_set(row.values.get(column));
let inter = set.iter().filter(|x| query_set.contains(*x)).count();
let union = set.len() + query_set.len() - inter;
let sim = if union == 0 {
0.0
} else {
inter as f64 / union as f64
};
if sim > 0.0 {
scored.push(SimilarRow {
row,
similarity: sim,
});
}
}
scored.sort_by(|a, b| {
b.similarity
.partial_cmp(&a.similarity)
.unwrap_or(std::cmp::Ordering::Equal)
});
scored.truncate(k);
Ok(scored)
}
pub fn flush(&self) -> Result<()> {
for name in self.inner.table_names() {
let handle = self.inner.table(&name).map_err(KitError::from)?;
let mut guard = handle.lock();
guard.flush().map_err(KitError::from)?;
}
Ok(())
}
pub fn incremental_aggregate(
&self,
table: &str,
column: Option<&str>,
agg: IncrementalAggKind,
filter: Option<&mongreldb_kit_core::query::Expr>,
) -> Result<IncrementalAggregate> {
let t = self
.schema
.tables
.iter()
.find(|t| t.name == table)
.ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
if !matches!(agg, IncrementalAggKind::Count) && column.is_none() {
return Err(KitError::Validation(
"sum/min/max/avg incremental aggregate requires a column".into(),
));
}
let cid = match column {
Some(name) => Some(
t.columns
.iter()
.find(|c| c.name == name)
.ok_or_else(|| KitError::Validation(format!("unknown column '{name}'")))?
.id as u16,
),
None => None,
};
let conditions = match filter {
Some(expr) => {
let plan = crate::pushdown::translate_predicate(t, expr).ok_or_else(|| {
KitError::Validation(
"filter is not index-translatable for an incremental aggregate".into(),
)
})?;
if !plan.fully_translated {
return Err(KitError::Validation(
"filter has a residual that an incremental aggregate cannot apply exactly"
.into(),
));
}
plan.conditions
}
None => Vec::new(),
};
let core_agg = match agg {
IncrementalAggKind::Count => NativeAgg::Count,
IncrementalAggKind::Sum => NativeAgg::Sum,
IncrementalAggKind::Min => NativeAgg::Min,
IncrementalAggKind::Max => NativeAgg::Max,
IncrementalAggKind::Avg => NativeAgg::Avg,
};
let cache_key = incremental_cache_key(t.id, cid, agg, &conditions);
let handle = self.inner.table(table).map_err(KitError::from)?;
let mut guard = handle.lock();
let res = guard
.aggregate_incremental(cache_key, &conditions, cid, core_agg)
.map_err(KitError::from)?;
Ok(IncrementalAggregate {
value: agg_state_value(&res.state),
incremental: res.incremental,
delta_rows: res.delta_rows,
})
}
pub fn applied_migrations(&self) -> Result<Vec<mongreldb_kit_core::migrations::Migration>> {
crate::migrate::load_applied_migrations(&self.inner)
}
pub(crate) fn core_db(&self) -> &CoreDatabase {
&self.inner
}
pub(crate) fn core_arc(&self) -> Arc<CoreDatabase> {
Arc::clone(&self.inner)
}
pub fn close(&self) -> Result<()> {
self.inner.close().map_err(KitError::from)
}
pub fn compact_all(&self) -> Result<(usize, usize)> {
self.inner.compact().map_err(KitError::from)
}
pub fn compact_table(&self, name: &str) -> Result<bool> {
self.inner.compact_table(name).map_err(KitError::from)
}
pub fn rename_table(&mut self, from: &str, to: &str) -> Result<()> {
if from.starts_with("__kit_") || to.starts_with("__kit_") {
return Err(KitError::Validation(
"rename_table: names beginning with '__kit_' are reserved for internal tables"
.into(),
));
}
self.inner.rename_table(from, to).map_err(KitError::from)?;
if !self.schema.rename_table(from, to) {
return Err(KitError::Integrity(format!(
"rename_table: kit schema has no table '{from}' (or '{to}' already exists)"
)));
}
for table in &mut self.schema.tables {
for fk in &mut table.foreign_keys {
if fk.references_table == from {
fk.references_table = to.to_string();
}
}
}
store_schema(&self.root, &self.schema)?;
Ok(())
}
pub fn analyze(&self) -> Result<()> {
for name in self.inner.table_names() {
let handle = self.inner.table(&name).map_err(KitError::from)?;
handle.lock().ensure_indexes_complete()?;
}
Ok(())
}
pub fn vacuum(&self) -> Result<usize> {
self.inner.compact().map_err(KitError::from)?;
self.inner.gc().map_err(KitError::from)
}
pub fn create_view(&self, spec: &ViewSpec) -> Result<()> {
self.sql(&spec.create_sql())?;
Ok(())
}
pub fn drop_view(&self, name: &str) -> Result<()> {
self.sql(&format!("DROP VIEW IF EXISTS {name}"))?;
Ok(())
}
pub fn reserve_auto_inc(&self, table: &str) -> Result<Option<i64>> {
let handle = self.inner.table(table).map_err(KitError::from)?;
let mut guard = handle.lock();
guard.reserve_auto_inc().map_err(KitError::from)
}
pub fn create_user(&self, username: &str, password: &str) -> Result<()> {
self.inner
.create_user(username, password)
.map_err(KitError::from)?;
Ok(())
}
pub fn drop_user(&self, username: &str) -> Result<()> {
self.inner.drop_user(username).map_err(KitError::from)
}
pub fn alter_user_password(&self, username: &str, new_password: &str) -> Result<()> {
self.inner
.alter_user_password(username, new_password)
.map_err(KitError::from)
}
pub fn verify_user(
&self,
username: &str,
password: &str,
) -> Result<Option<mongreldb_core::auth::UserEntry>> {
self.inner
.verify_user(username, password)
.map_err(KitError::from)
}
pub fn set_user_admin(&self, username: &str, is_admin: bool) -> Result<()> {
self.inner
.set_user_admin(username, is_admin)
.map_err(KitError::from)
}
pub fn users(&self) -> Vec<String> {
self.inner.users().into_iter().map(|u| u.username).collect()
}
pub fn create_role(&self, name: &str) -> Result<()> {
self.inner.create_role(name).map_err(KitError::from)?;
Ok(())
}
pub fn drop_role(&self, name: &str) -> Result<()> {
self.inner.drop_role(name).map_err(KitError::from)
}
pub fn roles(&self) -> Vec<String> {
self.inner.roles().into_iter().map(|r| r.name).collect()
}
pub fn grant_role(&self, username: &str, role_name: &str) -> Result<()> {
self.inner
.grant_role(username, role_name)
.map_err(KitError::from)
}
pub fn revoke_role(&self, username: &str, role_name: &str) -> Result<()> {
self.inner
.revoke_role(username, role_name)
.map_err(KitError::from)
}
pub fn grant_permission(
&self,
role_name: &str,
permission: mongreldb_core::auth::Permission,
) -> Result<()> {
self.inner
.grant_permission(role_name, permission)
.map_err(KitError::from)
}
pub fn revoke_permission(
&self,
role_name: &str,
permission: mongreldb_core::auth::Permission,
) -> Result<()> {
self.inner
.revoke_permission(role_name, permission)
.map_err(KitError::from)
}
pub fn set_spill_threshold(&self, bytes: u64) {
self.inner.set_spill_threshold(bytes);
}
pub fn set_recursive_triggers(&self, enabled: bool) {
self.inner.set_recursive_triggers(enabled);
}
pub fn trigger_config(&self) -> mongreldb_core::TriggerConfig {
self.inner.trigger_config()
}
pub fn set_trigger_config(&self, config: mongreldb_core::TriggerConfig) -> Result<()> {
self.inner
.set_trigger_config(config)
.map_err(KitError::from)
}
pub fn set_table_compaction_zstd_level(&self, table: &str, level: i32) -> Result<()> {
let handle = self.inner.table(table).map_err(KitError::from)?;
handle.lock().set_compaction_zstd_level(level);
Ok(())
}
pub fn set_table_result_cache_max_bytes(&self, table: &str, max_bytes: u64) -> Result<()> {
let handle = self.inner.table(table).map_err(KitError::from)?;
handle.lock().set_result_cache_max_bytes(max_bytes);
Ok(())
}
pub fn set_table_mutable_run_spill_bytes(&self, table: &str, bytes: u64) -> Result<()> {
let handle = self.inner.table(table).map_err(KitError::from)?;
handle.lock().set_mutable_run_spill_bytes(bytes);
Ok(())
}
pub fn set_table_sync_byte_threshold(&self, table: &str, threshold: u64) -> Result<()> {
let handle = self.inner.table(table).map_err(KitError::from)?;
handle.lock().set_sync_byte_threshold(threshold);
Ok(())
}
pub fn set_table_index_build_policy(
&self,
table: &str,
policy: mongreldb_core::IndexBuildPolicy,
) -> Result<()> {
let handle = self.inner.table(table).map_err(KitError::from)?;
handle.lock().set_index_build_policy(policy);
Ok(())
}
pub fn table_page_cache_stats(&self, table: &str) -> Result<mongreldb_core::cache::CacheStats> {
let handle = self.inner.table(table).map_err(KitError::from)?;
let stats = handle.lock().page_cache_stats();
Ok(stats)
}
pub fn table_run_count(&self, table: &str) -> Result<usize> {
let handle = self.inner.table(table).map_err(KitError::from)?;
let n = handle.lock().run_count();
Ok(n)
}
pub fn table_memtable_len(&self, table: &str) -> Result<usize> {
let handle = self.inner.table(table).map_err(KitError::from)?;
let n = handle.lock().memtable_len();
Ok(n)
}
pub fn table_mutable_run_len(&self, table: &str) -> Result<usize> {
let handle = self.inner.table(table).map_err(KitError::from)?;
let n = handle.lock().mutable_run_len();
Ok(n)
}
pub fn table_page_cache_len(&self, table: &str) -> Result<usize> {
let handle = self.inner.table(table).map_err(KitError::from)?;
let n = handle.lock().page_cache_len();
Ok(n)
}
pub fn table_decoded_cache_len(&self, table: &str) -> Result<usize> {
let handle = self.inner.table(table).map_err(KitError::from)?;
let n = handle.lock().decoded_cache_len();
Ok(n)
}
pub fn sql(&self, statement: &str) -> Result<Vec<arrow::record_batch::RecordBatch>> {
let session = match self.session.lock().take() {
Some(s) => s,
None => {
mongreldb_query::MongrelSession::open(self.core_arc()).map_err(KitError::from)?
}
};
let runtime = sql_runtime();
let result = runtime
.block_on(session.run(statement))
.map_err(KitError::from);
*self.session.lock() = Some(session);
result
}
pub fn refresh_sql_session(&self) -> Result<()> {
let session =
mongreldb_query::MongrelSession::open(self.core_arc()).map_err(KitError::from)?;
*self.session.lock() = Some(session);
Ok(())
}
pub fn sql_arrow(&self, statement: &str) -> Result<Vec<u8>> {
let batches = self.sql(statement)?;
crate::arrow_util::batches_to_ipc(&batches)
}
pub fn sql_rows(&self, statement: &str) -> Result<Vec<serde_json::Map<String, Value>>> {
let batches = self.sql(statement)?;
crate::arrow_util::batches_to_rows(&batches)
}
pub(crate) fn lookup_row_id(&self, table: &str, key: &[u8]) -> Result<Option<RowId>> {
let handle = self.inner.table(table).map_err(KitError::from)?;
let mut guard = handle.lock();
guard.ensure_indexes_complete()?;
Ok(guard.lookup_pk(key))
}
pub(crate) fn root(&self) -> &Path {
&self.root
}
pub(crate) fn visible_core_rows_at(
&self,
table_name: &str,
snapshot: Snapshot,
) -> Result<Vec<CoreRow>> {
let handle = self.inner.table(table_name).map_err(KitError::from)?;
let guard = handle.lock();
guard.visible_rows(snapshot).map_err(KitError::from)
}
pub(crate) fn query_core_rows_at(
&self,
table_name: &str,
conditions: &[mongreldb_core::query::Condition],
snapshot: Snapshot,
) -> Result<Vec<CoreRow>> {
if conditions.is_empty() {
return self.visible_core_rows_at(table_name, snapshot);
}
let handle = self.inner.table(table_name).map_err(KitError::from)?;
let mut guard = handle.lock();
let q = mongreldb_core::query::Query {
conditions: conditions.to_vec(),
};
guard.query(&q).map_err(KitError::from)
}
pub(crate) fn flush_table(&self, table_name: &str) -> Result<()> {
let handle = self.inner.table(table_name).map_err(KitError::from)?;
handle.lock().flush().map_err(KitError::from)?;
Ok(())
}
pub(crate) fn count_core_rows_at(
&self,
table_name: &str,
conditions: &[mongreldb_core::query::Condition],
snapshot: Snapshot,
) -> Result<Option<u64>> {
let handle = self.inner.table(table_name).map_err(KitError::from)?;
let mut guard = handle.lock();
if guard.snapshot().epoch != snapshot.epoch {
return Ok(None); }
guard
.count_conditions(conditions, snapshot)
.map_err(KitError::from)
}
pub(crate) fn aggregate_core_at(
&self,
table_name: &str,
column: Option<u16>,
conditions: &[mongreldb_core::query::Condition],
agg: NativeAgg,
snapshot: Snapshot,
) -> Result<Option<NativeAggResult>> {
let handle = self.inner.table(table_name).map_err(KitError::from)?;
let guard = handle.lock();
if guard.snapshot().epoch != snapshot.epoch {
return Ok(None); }
guard
.aggregate_native(snapshot, column, conditions, agg)
.map_err(KitError::from)
}
pub(crate) fn count_distinct_core_at(
&self,
table_name: &str,
column_id: u16,
snapshot: Snapshot,
) -> Result<Option<u64>> {
let handle = self.inner.table(table_name).map_err(KitError::from)?;
let mut guard = handle.lock();
if guard.snapshot().epoch != snapshot.epoch {
return Ok(None); }
guard
.count_distinct_from_bitmap(column_id)
.map_err(KitError::from)
}
#[allow(dead_code)]
pub(crate) fn get_core_row(&self, table_name: &str, row_id: u64) -> Result<Option<CoreRow>> {
let handle = self.inner.table(table_name).map_err(KitError::from)?;
let guard = handle.lock();
let snapshot = guard.snapshot();
Ok(guard.get(mongreldb_core::RowId(row_id), snapshot))
}
}
pub(crate) fn create_core_table(db: &CoreDatabase, name: &str, schema: CoreSchema) -> Result<()> {
if db.table_id(name).is_ok() {
return Ok(());
}
db.create_table(name, schema).map_err(KitError::from)?;
Ok(())
}
fn sql_runtime() -> &'static tokio::runtime::Runtime {
use std::sync::OnceLock;
static RT: OnceLock<tokio::runtime::Runtime> = OnceLock::new();
RT.get_or_init(|| {
tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("failed to build kit SQL tokio runtime")
})
}
fn core_procedure(spec: &ProcedureSpec) -> Result<mongreldb_core::StoredProcedure> {
let parsed: mongreldb_core::StoredProcedure =
serde_json::from_value(spec.json.clone()).map_err(KitError::from)?;
mongreldb_core::StoredProcedure::new(parsed.name, parsed.mode, parsed.params, parsed.body, 0)
.map_err(KitError::from)
}
fn core_trigger(spec: &TriggerSpec) -> Result<mongreldb_core::StoredTrigger> {
let parsed: mongreldb_core::StoredTrigger =
serde_json::from_value(spec.json.clone()).map_err(KitError::from)?;
mongreldb_core::StoredTrigger::new(
parsed.name,
mongreldb_core::TriggerDefinition {
target: parsed.target,
timing: parsed.timing,
event: parsed.event,
update_of: parsed.update_of,
target_columns: parsed.target_columns,
when: parsed.when,
program: parsed.program,
},
0,
)
.map_err(KitError::from)
}
fn json_to_core_value(value: &Value) -> Result<CoreValue> {
match value {
Value::Null => Ok(CoreValue::Null),
Value::Bool(value) => Ok(CoreValue::Bool(*value)),
Value::Number(value) => {
if let Some(value) = value.as_i64() {
Ok(CoreValue::Int64(value))
} else if let Some(value) = value.as_f64() {
Ok(CoreValue::Float64(value))
} else {
Err(KitError::Validation("unsupported JSON number".into()))
}
}
Value::String(value) => Ok(CoreValue::Bytes(value.as_bytes().to_vec())),
Value::Array(_) | Value::Object(_) => Err(KitError::Validation(
"procedure args only support scalar JSON values".into(),
)),
}
}
pub(crate) fn internal_bytes(row: &CoreRow, col_id: u16) -> Option<String> {
match row.columns.get(&col_id) {
Some(CoreValue::Bytes(b)) => String::from_utf8(b.clone()).ok(),
_ => None,
}
}
fn reap_rotated_wal_segments(db: &CoreDatabase) {
let _ = db.gc();
}
pub(crate) fn load_schema(path: &Path) -> Result<KitSchema> {
let file = path.join(SCHEMA_FILE);
let json = std::fs::read_to_string(&file)
.map_err(|e| KitError::Migration(format!("cannot read schema file: {e}")))?;
let schema: KitSchema = serde_json::from_str(&json)?;
Ok(schema)
}
pub(crate) fn store_schema(path: &Path, schema: &KitSchema) -> Result<()> {
let file = path.join(SCHEMA_FILE);
let json = serde_json::to_string_pretty(schema)?;
std::fs::write(&file, json)?;
Ok(())
}
pub(crate) fn persist_schema(db: &Database, schema: &KitSchema) -> Result<()> {
store_schema(&db.root, schema)
}