Skip to main content

mongreldb_kit/
db.rs

1//! Database handle for `mongreldb-kit`.
2
3use 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};
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 serde_json::Value;
16
17use std::collections::HashMap;
18use std::path::{Path, PathBuf};
19
20const SCHEMA_FILE: &str = "kit_schema.json";
21
22/// A named default-value provider registered by the application.
23pub type DefaultProvider = Box<dyn Fn() -> Value + Send + Sync>;
24
25/// The result of [`Database::explain`]: a static description of a predicate's
26/// index push-down, without running the query.
27#[derive(Debug, Clone)]
28pub struct ExplainPlan {
29    /// Whether at least one native index condition pushes down (vs. a full scan).
30    pub index_accelerated: bool,
31    /// Whether the push-down is exact — the whole predicate translated, so no
32    /// Rust-side residual re-filtering is needed.
33    pub exact: bool,
34    /// The kind of each pushed condition (e.g. `BitmapEq`, `RangeInt`, `Ann`).
35    pub pushed_conditions: Vec<String>,
36}
37
38/// A row paired with its Jaccard set-similarity to a query set (`0.0..=1.0`).
39#[derive(Debug, Clone)]
40pub struct SimilarRow {
41    pub row: crate::schema::Row,
42    pub similarity: f64,
43}
44
45/// Collect the string members of a set-valued column cell. Accepts either a
46/// JSON array value or a JSON string holding an array (how the Kit stores
47/// `json`/`text` set columns); anything else yields the empty set.
48fn parse_string_set(value: Option<&Value>) -> std::collections::HashSet<String> {
49    let arr = match value {
50        Some(Value::Array(a)) => Some(a.clone()),
51        Some(Value::String(s)) => serde_json::from_str::<Value>(s)
52            .ok()
53            .and_then(|v| v.as_array().cloned()),
54        _ => None,
55    };
56    arr.into_iter()
57        .flatten()
58        .filter_map(|v| match v {
59            Value::String(s) => Some(s),
60            Value::Number(n) => Some(n.to_string()),
61            Value::Bool(b) => Some(b.to_string()),
62            _ => None,
63        })
64        .collect()
65}
66
67/// Which aggregate to maintain incrementally.
68#[derive(Debug, Clone, Copy, PartialEq, Eq)]
69pub enum IncrementalAggKind {
70    Count,
71    Sum,
72    Min,
73    Max,
74    Avg,
75}
76
77/// The result of [`Database::incremental_aggregate`].
78#[derive(Debug, Clone)]
79pub struct IncrementalAggregate {
80    /// The exact aggregate value at the current epoch: a JSON number, or `null`
81    /// when no rows matched (`COUNT` returns `0`, not null).
82    pub value: Value,
83    /// `true` when produced by merging only newly-committed rows (the fast
84    /// path); `false` when a full recompute was required (cold cache, a delete,
85    /// pending writes, or the same epoch as the cached state).
86    pub incremental: bool,
87    /// Rows processed in the delta pass (`0` for a full recompute).
88    pub delta_rows: u64,
89}
90
91/// Stable per-`(table, column, agg, filter)` cache key for the engine's
92/// incremental-aggregate cache. Deterministic within a process (fixed-seed
93/// hasher); the cache itself is per-`Db`, so cross-process stability is moot.
94fn incremental_cache_key(
95    table_id: u32,
96    column: Option<u16>,
97    agg: IncrementalAggKind,
98    conditions: &[mongreldb_core::query::Condition],
99) -> u64 {
100    use std::hash::{Hash, Hasher};
101    let mut h = std::collections::hash_map::DefaultHasher::new();
102    table_id.hash(&mut h);
103    column.hash(&mut h);
104    (agg as u8).hash(&mut h);
105    // `Condition` has no `Hash`; its `Debug` form is stable and unique enough.
106    format!("{conditions:?}").hash(&mut h);
107    h.finish()
108}
109
110/// Finalize a mergeable [`AggState`] to a JSON scalar, preserving integer-ness
111/// for `COUNT`/`MIN`/`MAX`/int `SUM` and using a float for averages / float
112/// columns. `null` when there were no matching inputs.
113fn agg_state_value(s: &AggState) -> Value {
114    let num_f64 = |x: f64| {
115        serde_json::Number::from_f64(x)
116            .map(Value::Number)
117            .unwrap_or(Value::Null)
118    };
119    match s {
120        AggState::Count(n) => Value::from(*n),
121        AggState::SumI { sum, .. } => i64::try_from(*sum)
122            .map(Value::from)
123            .unwrap_or_else(|_| num_f64(*sum as f64)),
124        AggState::SumF { sum, .. } => num_f64(*sum),
125        AggState::AvgI { sum, count } if *count > 0 => num_f64(*sum as f64 / *count as f64),
126        AggState::AvgF { sum, count } if *count > 0 => num_f64(*sum / *count as f64),
127        AggState::AvgI { .. } | AggState::AvgF { .. } => Value::Null,
128        AggState::MinI(n) | AggState::MaxI(n) => Value::from(*n),
129        AggState::MinF(f) | AggState::MaxF(f) => num_f64(*f),
130        AggState::Empty => Value::Null,
131    }
132}
133
134/// Which approximate aggregate to estimate from the reservoir sample.
135#[derive(Debug, Clone, Copy, PartialEq, Eq)]
136pub enum ApproxAggKind {
137    Count,
138    Sum,
139    Avg,
140}
141
142/// A reservoir-sampled approximate aggregate with a normal-theory confidence
143/// interval. `ci_low`/`ci_high` bracket `point` at the requested z-score; the
144/// interval collapses to zero width when the sample covers the whole table.
145#[derive(Debug, Clone)]
146pub struct ApproxAggregate {
147    pub point: f64,
148    pub ci_low: f64,
149    pub ci_high: f64,
150    pub n_population: u64,
151    pub n_sample_live: usize,
152    pub n_passing: usize,
153}
154
155/// Short kind label for a core `Condition` (the variant name), decoupled from
156/// the enum shape via its `Debug` form.
157fn condition_label(c: &mongreldb_core::query::Condition) -> String {
158    let dbg = format!("{c:?}");
159    dbg.split(['(', '{', ' ']).next().unwrap_or("").to_string()
160}
161
162/// A kit database handle.
163///
164/// Wraps a MongrelDB core database and a kit schema, providing table metadata
165/// and transaction creation.
166pub struct Database {
167    pub(crate) inner: CoreDatabase,
168    pub(crate) schema: KitSchema,
169    pub(crate) root: PathBuf,
170    /// Application-registered named default providers (`DefaultKind::CustomName`).
171    pub(crate) default_providers: HashMap<String, DefaultProvider>,
172}
173
174impl Database {
175    /// Open an existing kit database.
176    pub fn open(path: &Path) -> Result<Self> {
177        let inner = CoreDatabase::open(path)?;
178        let schema = load_schema(path)?;
179        // Ensure reserved tables exist for databases created by older versions.
180        ensure_internal_tables(&inner)?;
181        Ok(Self {
182            inner,
183            schema,
184            root: path.to_path_buf(),
185            default_providers: HashMap::new(),
186        })
187    }
188
189    /// Open an existing page-encrypted kit database with its passphrase.
190    pub fn open_encrypted(path: &Path, passphrase: &str) -> Result<Self> {
191        let inner = CoreDatabase::open_encrypted(path, passphrase)?;
192        let schema = load_schema(path)?;
193        ensure_internal_tables(&inner)?;
194        Ok(Self {
195            inner,
196            schema,
197            root: path.to_path_buf(),
198            default_providers: HashMap::new(),
199        })
200    }
201
202    /// Create a fresh page-encrypted kit database (AES-256-GCM; the passphrase
203    /// derives the key hierarchy). Columns flagged `encrypted` /
204    /// `encrypted_indexable` are encrypted at rest.
205    pub fn create_encrypted(path: &Path, schema: KitSchema, passphrase: &str) -> Result<Self> {
206        std::fs::create_dir_all(path)?;
207        let inner = CoreDatabase::create_encrypted(path, passphrase)?;
208        ensure_internal_tables(&inner)?;
209        store_schema(path, &schema)?;
210        for table in &schema.tables {
211            create_core_table(&inner, &table.name, to_core_schema(table))?;
212        }
213        Ok(Self {
214            inner,
215            schema,
216            root: path.to_path_buf(),
217            default_providers: HashMap::new(),
218        })
219    }
220
221    /// Create a fresh kit database with the given schema.
222    pub fn create(path: &Path, schema: KitSchema) -> Result<Self> {
223        std::fs::create_dir_all(path)?;
224        let inner = CoreDatabase::create(path)?;
225
226        // Create the reserved kit tables first so we can record migrations,
227        // reserve unique keys, and touch row guards.
228        ensure_internal_tables(&inner)?;
229
230        // Persist the kit schema to a sidecar file (core tables cannot update
231        // a specific row by id, so a file is the pragmatic stable store).
232        store_schema(path, &schema)?;
233
234        // Create core tables for every user table.
235        for table in &schema.tables {
236            create_core_table(&inner, &table.name, to_core_schema(table))?;
237        }
238
239        Ok(Self {
240            inner,
241            schema,
242            root: path.to_path_buf(),
243            default_providers: HashMap::new(),
244        })
245    }
246
247    /// Register a named default provider used by `DefaultKind::CustomName`
248    /// columns. Returns the database for chaining.
249    pub fn register_default(
250        &mut self,
251        name: impl Into<String>,
252        provider: impl Fn() -> Value + Send + Sync + 'static,
253    ) {
254        self.default_providers
255            .insert(name.into(), Box::new(provider));
256    }
257
258    /// The raw, unguarded MongrelDB core database. This is the Rust analogue of
259    /// the TypeScript kit's `nativeDb` escape hatch: writes made directly
260    /// against it bypass all kit constraints.
261    pub fn raw(&self) -> &CoreDatabase {
262        &self.inner
263    }
264
265    /// Application table names, excluding the reserved `__kit_*` tables.
266    pub fn table_names(&self) -> Vec<String> {
267        self.schema
268            .tables
269            .iter()
270            .map(|t| t.name.clone())
271            .filter(|n| !n.starts_with("__kit_"))
272            .collect()
273    }
274
275    /// Allocate `count` values from the named sequence, returning the first
276    /// allocated value. A fresh sequence starts at `1` (SQL AUTO_INCREMENT
277    /// semantics). The allocation
278    /// runs in its own committed transaction and retries on write-write
279    /// conflicts.
280    pub fn allocate_sequence(&self, name: &str, count: i64) -> Result<i64> {
281        use crate::internal::cols;
282        let mut attempt = 0;
283        loop {
284            let mut txn = self.inner.begin();
285            let snapshot = txn.read_snapshot();
286            let existing = self
287                .visible_core_rows_at(crate::internal::SEQUENCES, snapshot)?
288                .into_iter()
289                .find(|r| internal_bytes(r, cols::SEQ_NAME) == Some(name.to_string()));
290
291            let now = crate::internal::iso_now();
292            // Sequences are 1-based, matching SQL AUTO_INCREMENT / SERIAL. A
293            // starting value of 0 is unsafe: applications use 0 as the "unset"
294            // sentinel for nullable foreign keys.
295            let (start, next, old_row_id) = match &existing {
296                Some(row) => {
297                    let current = match row.columns.get(&cols::SEQ_NEXT) {
298                        Some(CoreValue::Int64(i)) => *i,
299                        _ => 1,
300                    };
301                    (current, current + count, Some(row.row_id))
302                }
303                None => (1, 1 + count, None),
304            };
305
306            if let Some(rid) = old_row_id {
307                txn.delete(crate::internal::SEQUENCES, rid)
308                    .map_err(KitError::from)?;
309            }
310            txn.put(
311                crate::internal::SEQUENCES,
312                vec![
313                    (cols::SEQ_NAME, CoreValue::Bytes(name.as_bytes().to_vec())),
314                    (cols::SEQ_NEXT, CoreValue::Int64(next)),
315                    (cols::SEQ_UPDATED, CoreValue::Bytes(now.into_bytes())),
316                ],
317            )
318            .map_err(KitError::from)?;
319            match txn.commit() {
320                Ok(_) => return Ok(start),
321                Err(mongreldb_core::MongrelError::Conflict(_)) if attempt < 10_000 => {
322                    attempt += 1;
323                    std::thread::yield_now();
324                    continue;
325                }
326                Err(e) => return Err(KitError::from(e)),
327            }
328        }
329    }
330
331    /// Run `f` inside a kit transaction, committing on success and retrying on
332    /// retryable write-write conflicts up to `max_retries` times.
333    pub fn transaction<T, F>(&self, max_retries: usize, mut f: F) -> Result<T>
334    where
335        F: FnMut(&mut crate::txn::Transaction<'_>) -> Result<T>,
336    {
337        let mut attempt = 0;
338        loop {
339            let mut txn = self.begin()?;
340            match f(&mut txn) {
341                Ok(value) => match txn.commit() {
342                    Ok(()) => return Ok(value),
343                    Err(KitError::Conflict(_)) if attempt < max_retries => {
344                        attempt += 1;
345                        continue;
346                    }
347                    Err(e) => return Err(e),
348                },
349                Err(KitError::Conflict(_)) if attempt < max_retries => {
350                    txn.rollback();
351                    attempt += 1;
352                    continue;
353                }
354                Err(e) => {
355                    txn.rollback();
356                    return Err(e);
357                }
358            }
359        }
360    }
361
362    /// Look up a table definition by name.
363    pub fn table(&self, name: &str) -> Option<&KitTable> {
364        self.schema.table(name)
365    }
366
367    /// Return the currently loaded schema.
368    pub fn schema(&self) -> &KitSchema {
369        &self.schema
370    }
371
372    /// Begin a new kit transaction.
373    pub fn begin(&self) -> Result<crate::txn::Transaction<'_>> {
374        let core_txn = self.inner.begin();
375        Ok(crate::txn::Transaction::new(self, core_txn))
376    }
377
378    /// Replace the in-memory schema, usually after a migration.
379    pub fn set_schema(&mut self, schema: KitSchema) {
380        self.schema = schema;
381    }
382
383    /// Verify that the sidecar schema file and all reserved `__kit_*` tables
384    /// are present.
385    pub fn check_internal_tables(&self) -> Result<()> {
386        let schema_file = self.root.join(SCHEMA_FILE);
387        if !schema_file.exists() {
388            return Err(KitError::Integrity(format!(
389                "schema file {} is missing",
390                schema_file.display()
391            )));
392        }
393        for (name, _) in internal_tables_core() {
394            if self.inner.table_id(name).is_err() {
395                return Err(KitError::Integrity(format!(
396                    "internal table {name} is missing"
397                )));
398            }
399        }
400        Ok(())
401    }
402
403    /// Reclaim orphaned runs and stale WAL/shadow files; returns the count
404    /// removed. Safe to run on a live database.
405    pub fn gc(&self) -> Result<usize> {
406        self.inner.gc().map_err(KitError::from)
407    }
408
409    /// Verify run footer checksums; returns any integrity issues as JSON objects
410    /// (`table_id`, `table_name`, `severity`, `description`). Empty ⇒ healthy.
411    pub fn check(&self) -> Vec<serde_json::Value> {
412        self.inner
413            .check()
414            .into_iter()
415            .map(|i| {
416                serde_json::json!({
417                    "table_id": i.table_id,
418                    "table_name": i.table_name,
419                    "severity": i.severity,
420                    "description": i.description,
421                })
422            })
423            .collect()
424    }
425
426    /// Drop corrupt runs; returns the ids of the runs that were dropped.
427    pub fn doctor(&self) -> Result<Vec<u64>> {
428        self.inner.doctor().map_err(KitError::from)
429    }
430
431    /// The current visible commit epoch — a monotonically increasing version
432    /// stamp. A committed write bumps it; a snapshot at this epoch sees all
433    /// currently-committed data.
434    pub fn snapshot_epoch(&self) -> u64 {
435        self.inner.snapshot().0.epoch.0
436    }
437
438    /// Export every visible row of `table` as a TSV document (header row of
439    /// column names, tab-separated cells, `NULL` = empty field). See
440    /// [`crate::tsv`] for the escaping rules.
441    pub fn export_tsv(&self, table: &str) -> Result<String> {
442        let t = self
443            .schema
444            .tables
445            .iter()
446            .find(|t| t.name == table)
447            .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?
448            .clone();
449        let tx = self.begin()?;
450        let rows = tx.all_rows(table)?;
451        Ok(crate::tsv::rows_to_tsv(&t, &rows))
452    }
453
454    /// Import a TSV document into `table` (one committed transaction). Each row
455    /// passes through defaults, validation, and constraint checks like a normal
456    /// insert. Returns the number of rows inserted.
457    pub fn import_tsv(&self, table: &str, text: &str) -> Result<usize> {
458        let t = self
459            .schema
460            .tables
461            .iter()
462            .find(|t| t.name == table)
463            .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?
464            .clone();
465        let rows = crate::tsv::tsv_to_rows(&t, text)?;
466        let n = rows.len();
467        self.transaction(1, |tx| {
468            tx.insert_many(table, rows.clone())?;
469            Ok(())
470        })?;
471        Ok(n)
472    }
473
474    /// Describe how `predicate` would be executed against `table`: which native
475    /// index conditions push down, whether the push-down is exact (no residual
476    /// re-filtering), and whether any index acceleration applies at all. A pure
477    /// diagnostic — it plans but does not run the query.
478    pub fn explain(
479        &self,
480        table: &str,
481        predicate: &mongreldb_kit_core::query::Expr,
482    ) -> Result<ExplainPlan> {
483        let t = self
484            .schema
485            .tables
486            .iter()
487            .find(|t| t.name == table)
488            .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
489        Ok(match crate::pushdown::translate_predicate(t, predicate) {
490            Some(p) => ExplainPlan {
491                index_accelerated: p.can_push(),
492                exact: p.fully_translated,
493                pushed_conditions: p.conditions.iter().map(condition_label).collect(),
494            },
495            None => ExplainPlan {
496                index_accelerated: false,
497                exact: false,
498                pushed_conditions: Vec::new(),
499            },
500        })
501    }
502
503    /// Read every row of `table` visible at commit `epoch` — a point-in-time
504    /// (MVCC time-travel) read. `epoch` must not exceed the current snapshot
505    /// epoch. Rows reclaimed by GC/compaction for retired snapshots may no
506    /// longer be reconstructable; this reads whatever the engine still retains
507    /// at that epoch.
508    pub fn rows_at_epoch(&self, table: &str, epoch: u64) -> Result<Vec<crate::schema::Row>> {
509        let t = self
510            .schema
511            .tables
512            .iter()
513            .find(|t| t.name == table)
514            .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
515        let current = self.snapshot_epoch();
516        if epoch > current {
517            return Err(KitError::Validation(format!(
518                "epoch {epoch} is in the future (current committed epoch is {current})"
519            )));
520        }
521        let snap = Snapshot::at(mongreldb_core::epoch::Epoch(epoch));
522        let rows = self.visible_core_rows_at(table, snap)?;
523        rows.iter()
524            .map(|r| crate::schema::core_row_to_json(r, t))
525            .collect()
526    }
527
528    /// Estimate an aggregate over `table` from the engine's reservoir sample,
529    /// returning a point estimate and a `z`-score confidence interval (e.g.
530    /// `z = 1.96` for ~95%). `column` is required for `Sum`/`Avg` and ignored
531    /// for `Count`. Returns `None` when the reservoir is empty (no sampled rows
532    /// yet). Fast and O(sample) — trades exactness for speed on large tables.
533    pub fn approx_aggregate(
534        &self,
535        table: &str,
536        column: Option<&str>,
537        agg: ApproxAggKind,
538        z: f64,
539    ) -> Result<Option<ApproxAggregate>> {
540        let t = self
541            .schema
542            .tables
543            .iter()
544            .find(|t| t.name == table)
545            .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
546        if matches!(agg, ApproxAggKind::Sum | ApproxAggKind::Avg) && column.is_none() {
547            return Err(KitError::Validation(
548                "approx sum/avg requires a column".into(),
549            ));
550        }
551        let cid = match column {
552            Some(name) => Some(
553                t.columns
554                    .iter()
555                    .find(|c| c.name == name)
556                    .ok_or_else(|| KitError::Validation(format!("unknown column '{name}'")))?
557                    .id as u16,
558            ),
559            None => None,
560        };
561        let core_agg = match agg {
562            ApproxAggKind::Count => ApproxAgg::Count,
563            ApproxAggKind::Sum => ApproxAgg::Sum,
564            ApproxAggKind::Avg => ApproxAgg::Avg,
565        };
566        let handle = self.inner.table(table).map_err(KitError::from)?;
567        let guard = handle.lock();
568        let res = guard
569            .approx_aggregate(&[], cid, core_agg, z)
570            .map_err(KitError::from)?;
571        Ok(res.map(|r| ApproxAggregate {
572            point: r.point,
573            ci_low: r.ci_low,
574            ci_high: r.ci_high,
575            n_population: r.n_population,
576            n_sample_live: r.n_sample_live,
577            n_passing: r.n_passing,
578        }))
579    }
580
581    /// Stream `table` in row batches without materializing the whole table at
582    /// once. `f` receives successive chunks of at most `batch_size` value-maps,
583    /// all from one snapshot. Backed by the engine's native scan cursor when the
584    /// table has a sorted run; for an overlay-only table (no run yet) it falls
585    /// back to a single in-memory pass, still chunked to `batch_size`.
586    pub fn scan_batched<F>(&self, table: &str, batch_size: usize, mut f: F) -> Result<()>
587    where
588        F: FnMut(&[serde_json::Map<String, Value>]) -> Result<()>,
589    {
590        let kit_t = self
591            .schema
592            .tables
593            .iter()
594            .find(|t| t.name == table)
595            .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
596        let batch_size = batch_size.max(1);
597        // Keep the pin guard alive for the whole scan so GC can't reclaim the
598        // snapshot's versions mid-stream.
599        let (snapshot, _pin) = self.inner.snapshot();
600        let handle = self.inner.table(table).map_err(KitError::from)?;
601        let guard = handle.lock();
602
603        // Projection + per-column (name, kit type), index-aligned, in core order.
604        let mut projection: Vec<(u16, mongreldb_core::schema::TypeId)> = Vec::new();
605        let mut meta: Vec<(String, mongreldb_kit_core::schema::ColumnType)> = Vec::new();
606        for c in &guard.schema().columns {
607            if let Some(kc) = kit_t.columns.iter().find(|kc| kc.id as u16 == c.id) {
608                projection.push((c.id, c.ty));
609                meta.push((kc.name.clone(), kc.storage_type));
610            }
611        }
612
613        match guard
614            .scan_cursor(snapshot, projection, &[])
615            .map_err(KitError::from)?
616        {
617            Some(mut cursor) => {
618                let mut buf: Vec<serde_json::Map<String, Value>> = Vec::with_capacity(batch_size);
619                while let Some(batch) = cursor.next_batch().map_err(KitError::from)? {
620                    let nrows = batch.first().map(|c| c.len()).unwrap_or(0);
621                    for j in 0..nrows {
622                        let mut m = serde_json::Map::new();
623                        for (ci, (name, ty)) in meta.iter().enumerate() {
624                            let cv = batch
625                                .get(ci)
626                                .and_then(|col| col.value_at(j))
627                                .unwrap_or(CoreValue::Null);
628                            m.insert(name.clone(), crate::schema::core_to_json(&cv, *ty)?);
629                        }
630                        buf.push(m);
631                        if buf.len() >= batch_size {
632                            f(&buf)?;
633                            buf.clear();
634                        }
635                    }
636                }
637                if !buf.is_empty() {
638                    f(&buf)?;
639                }
640                Ok(())
641            }
642            None => {
643                drop(guard);
644                let rows = self.visible_core_rows_at(table, snapshot)?;
645                let maps: Vec<serde_json::Map<String, Value>> = rows
646                    .iter()
647                    .map(|r| crate::schema::core_row_to_json(r, kit_t).map(|row| row.values))
648                    .collect::<Result<Vec<_>>>()?;
649                for chunk in maps.chunks(batch_size) {
650                    f(chunk)?;
651                }
652                Ok(())
653            }
654        }
655    }
656
657    /// Rank rows of `table` by Jaccard set-similarity between `query` and the
658    /// string set stored (as a JSON array) in `column`, returning the top `k`
659    /// with similarity `> 0`, highest first — the dedup/join primitive.
660    ///
661    /// When `column` has a `MinHash` index, candidate rows come from the engine's
662    /// LSH index (sub-linear) and are then re-verified with exact Jaccard, so the
663    /// top-k is exact for the recalled candidates (LSH recall is high but < 100%).
664    /// Without an index it is an exact linear scan.
665    pub fn set_similarity(
666        &self,
667        table: &str,
668        column: &str,
669        query: &[String],
670        k: usize,
671    ) -> Result<Vec<SimilarRow>> {
672        let t = self
673            .schema
674            .tables
675            .iter()
676            .find(|t| t.name == table)
677            .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
678        let col = t.columns.iter().find(|c| c.name == column).ok_or_else(|| {
679            KitError::Validation(format!("unknown column '{column}' on table '{table}'"))
680        })?;
681        let query_set: std::collections::HashSet<String> = query.iter().cloned().collect();
682
683        let has_minhash = t.indexes.iter().any(|idx| {
684            idx.kind == KitIndexKind::MinHash && idx.columns.iter().any(|c| c == column)
685        });
686        let rows = if has_minhash {
687            // Sub-linear candidate generation via the engine MinHash/LSH index.
688            let query_hashes: Vec<u64> = query
689                .iter()
690                .map(|s| mongreldb_core::index::minhash_token_hash(s))
691                .collect();
692            // Generous candidate budget so exact top-k keeps high recall.
693            let cand_k = k.saturating_mul(8).max(k + 64);
694            let cond = mongreldb_core::query::Condition::MinHashSimilar {
695                column_id: col.id as u16,
696                query: query_hashes,
697                k: cand_k,
698            };
699            let (snapshot, _pin) = self.inner.snapshot();
700            let core_rows = self.query_core_rows_at(table, &[cond], snapshot)?;
701            core_rows
702                .iter()
703                .map(|r| crate::schema::core_row_to_json(r, t))
704                .collect::<Result<Vec<_>>>()?
705        } else {
706            let tx = self.begin()?;
707            tx.all_rows(table)?
708        };
709
710        let mut scored: Vec<SimilarRow> = Vec::new();
711        for row in rows {
712            let set = parse_string_set(row.values.get(column));
713            let inter = set.iter().filter(|x| query_set.contains(*x)).count();
714            let union = set.len() + query_set.len() - inter;
715            let sim = if union == 0 {
716                0.0
717            } else {
718                inter as f64 / union as f64
719            };
720            if sim > 0.0 {
721                scored.push(SimilarRow {
722                    row,
723                    similarity: sim,
724                });
725            }
726        }
727        scored.sort_by(|a, b| {
728            b.similarity
729                .partial_cmp(&a.similarity)
730                .unwrap_or(std::cmp::Ordering::Equal)
731        });
732        scored.truncate(k);
733        Ok(scored)
734    }
735
736    /// Flush every table's in-memory writes to durable sorted runs. Besides
737    /// durability, this empties the memtable, which is what enables the engine's
738    /// incremental-aggregate fast path (see [`Self::incremental_aggregate`]).
739    pub fn flush(&self) -> Result<()> {
740        for name in self.inner.table_names() {
741            let handle = self.inner.table(&name).map_err(KitError::from)?;
742            let mut guard = handle.lock();
743            guard.flush().map_err(KitError::from)?;
744        }
745        Ok(())
746    }
747
748    /// Maintain and read an aggregate over `table` that updates by merging only
749    /// newly-committed rows instead of rescanning. `column` is required for
750    /// `Sum`/`Min`/`Max`/`Avg` and ignored for `Count`. An optional `filter`
751    /// restricts the aggregate; it must translate **exactly** to index
752    /// conditions (no residual), otherwise this errors — an inexact filter would
753    /// silently aggregate the wrong rows.
754    ///
755    /// The engine keeps a per-`(table, column, agg, filter)` cached state and,
756    /// on a warm cache with an advanced epoch and no deletes/pending writes,
757    /// folds in just the delta. The returned `value` is always exact; the
758    /// `incremental` flag reports whether the fast path was taken.
759    pub fn incremental_aggregate(
760        &self,
761        table: &str,
762        column: Option<&str>,
763        agg: IncrementalAggKind,
764        filter: Option<&mongreldb_kit_core::query::Expr>,
765    ) -> Result<IncrementalAggregate> {
766        let t = self
767            .schema
768            .tables
769            .iter()
770            .find(|t| t.name == table)
771            .ok_or_else(|| KitError::Validation(format!("unknown table '{table}'")))?;
772        if !matches!(agg, IncrementalAggKind::Count) && column.is_none() {
773            return Err(KitError::Validation(
774                "sum/min/max/avg incremental aggregate requires a column".into(),
775            ));
776        }
777        let cid = match column {
778            Some(name) => Some(
779                t.columns
780                    .iter()
781                    .find(|c| c.name == name)
782                    .ok_or_else(|| KitError::Validation(format!("unknown column '{name}'")))?
783                    .id as u16,
784            ),
785            None => None,
786        };
787        let conditions = match filter {
788            Some(expr) => {
789                let plan = crate::pushdown::translate_predicate(t, expr).ok_or_else(|| {
790                    KitError::Validation(
791                        "filter is not index-translatable for an incremental aggregate".into(),
792                    )
793                })?;
794                if !plan.fully_translated {
795                    return Err(KitError::Validation(
796                        "filter has a residual that an incremental aggregate cannot apply exactly"
797                            .into(),
798                    ));
799                }
800                plan.conditions
801            }
802            None => Vec::new(),
803        };
804        let core_agg = match agg {
805            IncrementalAggKind::Count => NativeAgg::Count,
806            IncrementalAggKind::Sum => NativeAgg::Sum,
807            IncrementalAggKind::Min => NativeAgg::Min,
808            IncrementalAggKind::Max => NativeAgg::Max,
809            IncrementalAggKind::Avg => NativeAgg::Avg,
810        };
811        let cache_key = incremental_cache_key(t.id, cid, agg, &conditions);
812        let handle = self.inner.table(table).map_err(KitError::from)?;
813        let mut guard = handle.lock();
814        let res = guard
815            .aggregate_incremental(cache_key, &conditions, cid, core_agg)
816            .map_err(KitError::from)?;
817        Ok(IncrementalAggregate {
818            value: agg_state_value(&res.state),
819            incremental: res.incremental,
820            delta_rows: res.delta_rows,
821        })
822    }
823
824    /// Return the migrations already recorded in `__kit_schema_migrations`.
825    pub fn applied_migrations(&self) -> Result<Vec<mongreldb_kit_core::migrations::Migration>> {
826        crate::migrate::load_applied_migrations(&self.inner)
827    }
828
829    pub(crate) fn core_db(&self) -> &CoreDatabase {
830        &self.inner
831    }
832
833    pub(crate) fn root(&self) -> &Path {
834        &self.root
835    }
836
837    /// All visible core rows for a table at a specific read snapshot. Used so
838    /// kit transactions read at their own snapshot (repeatable reads) rather
839    /// than the latest committed state.
840    pub(crate) fn visible_core_rows_at(
841        &self,
842        table_name: &str,
843        snapshot: Snapshot,
844    ) -> Result<Vec<CoreRow>> {
845        let handle = self.inner.table(table_name).map_err(KitError::from)?;
846        let guard = handle.lock();
847        guard.visible_rows(snapshot).map_err(KitError::from)
848    }
849
850    /// Query visible core rows with native `Condition`s at a specific read
851    /// snapshot (Kit Priority 1 pushdown). Resolves `conditions` via core's
852    /// indexes (HOT / bitmap / range) and returns only the matching rows —
853    /// avoiding the full scan that `visible_core_rows_at` does. Returns the
854    /// empty vec when no conditions match, and falls back to
855    /// `visible_core_rows_at` when `conditions` is empty (unfiltered).
856    pub(crate) fn query_core_rows_at(
857        &self,
858        table_name: &str,
859        conditions: &[mongreldb_core::query::Condition],
860        snapshot: Snapshot,
861    ) -> Result<Vec<CoreRow>> {
862        if conditions.is_empty() {
863            return self.visible_core_rows_at(table_name, snapshot);
864        }
865        let handle = self.inner.table(table_name).map_err(KitError::from)?;
866        let mut guard = handle.lock();
867        let q = mongreldb_core::query::Query {
868            conditions: conditions.to_vec(),
869        };
870        guard.query(&q).map_err(KitError::from)
871    }
872
873    /// Count visible rows matching `conditions` without materializing them
874    /// (Kit Priority 7 pushdown). Returns `None` when the conditions cannot be
875    /// served by indexes, or when `snapshot` is not the latest committed epoch
876    /// (caller falls back to a snapshot-correct row scan).
877    ///
878    /// `count_conditions` counts the engine's latest committed index state, not
879    /// a snapshot-filtered scan, so it only matches a repeatable-read row count
880    /// when the read snapshot IS the latest epoch. We hold the table lock while
881    /// comparing, so no commit can interleave between the check and the count.
882    pub(crate) fn count_core_rows_at(
883        &self,
884        table_name: &str,
885        conditions: &[mongreldb_core::query::Condition],
886        snapshot: Snapshot,
887    ) -> Result<Option<u64>> {
888        let handle = self.inner.table(table_name).map_err(KitError::from)?;
889        let mut guard = handle.lock();
890        if guard.snapshot().epoch != snapshot.epoch {
891            return Ok(None); // stale read snapshot ⇒ caller scans
892        }
893        guard
894            .count_conditions(conditions, snapshot)
895            .map_err(KitError::from)
896    }
897
898    /// Compute `SUM`/`MIN`/`MAX`/`AVG`/`COUNT(col)` over a column without
899    /// materializing rows (Kit Priority 7), via the engine's native aggregate.
900    /// `column` is the engine column id. Returns `None` when the engine fast
901    /// path does not apply (multi-run / non-empty overlay / non-numeric column),
902    /// or when `snapshot` is not the latest committed epoch — the same
903    /// guarantee as [`count_core_rows_at`](Self::count_core_rows_at): the engine
904    /// aggregate matches a snapshot-consistent row scan only at the latest epoch,
905    /// and we compare under the table lock so no commit can interleave.
906    pub(crate) fn aggregate_core_at(
907        &self,
908        table_name: &str,
909        column: Option<u16>,
910        conditions: &[mongreldb_core::query::Condition],
911        agg: NativeAgg,
912        snapshot: Snapshot,
913    ) -> Result<Option<NativeAggResult>> {
914        let handle = self.inner.table(table_name).map_err(KitError::from)?;
915        let guard = handle.lock();
916        if guard.snapshot().epoch != snapshot.epoch {
917            return Ok(None); // stale read snapshot ⇒ caller scans
918        }
919        guard
920            .aggregate_native(snapshot, column, conditions, agg)
921            .map_err(KitError::from)
922    }
923
924    /// `COUNT(DISTINCT col)` from the bitmap index's partition cardinality (Kit
925    /// Priority 7) — the number of distinct indexed values, no scan. Returns
926    /// `None` without a bitmap index on the column, when the table is not
927    /// insert-only, or when `snapshot` is not the latest committed epoch. The
928    /// engine method reads the latest committed index state (no snapshot
929    /// parameter), so — as with [`count_core_rows_at`](Self::count_core_rows_at)
930    /// — it only matches a repeatable-read scan at the latest epoch; we compare
931    /// under the table lock so no commit can interleave.
932    pub(crate) fn count_distinct_core_at(
933        &self,
934        table_name: &str,
935        column_id: u16,
936        snapshot: Snapshot,
937    ) -> Result<Option<u64>> {
938        let handle = self.inner.table(table_name).map_err(KitError::from)?;
939        let guard = handle.lock();
940        if guard.snapshot().epoch != snapshot.epoch {
941            return Ok(None); // stale read snapshot ⇒ caller scans
942        }
943        guard
944            .count_distinct_from_bitmap(column_id)
945            .map_err(KitError::from)
946    }
947
948    /// Materialize a single row by storage row id.
949    #[allow(dead_code)]
950    pub(crate) fn get_core_row(&self, table_name: &str, row_id: u64) -> Result<Option<CoreRow>> {
951        let handle = self.inner.table(table_name).map_err(KitError::from)?;
952        let guard = handle.lock();
953        let snapshot = guard.snapshot();
954        Ok(guard.get(mongreldb_core::RowId(row_id), snapshot))
955    }
956}
957
958pub(crate) fn create_core_table(db: &CoreDatabase, name: &str, schema: CoreSchema) -> Result<()> {
959    if db.table_id(name).is_ok() {
960        return Ok(());
961    }
962    db.create_table(name, schema).map_err(KitError::from)?;
963    Ok(())
964}
965
966/// Read a `Bytes` column from an internal-table core row as a UTF-8 string.
967pub(crate) fn internal_bytes(row: &CoreRow, col_id: u16) -> Option<String> {
968    match row.columns.get(&col_id) {
969        Some(CoreValue::Bytes(b)) => String::from_utf8(b.clone()).ok(),
970        _ => None,
971    }
972}
973
974pub(crate) fn load_schema(path: &Path) -> Result<KitSchema> {
975    let file = path.join(SCHEMA_FILE);
976    let json = std::fs::read_to_string(&file)
977        .map_err(|e| KitError::Migration(format!("cannot read schema file: {e}")))?;
978    let schema: KitSchema = serde_json::from_str(&json)?;
979    Ok(schema)
980}
981
982pub(crate) fn store_schema(path: &Path, schema: &KitSchema) -> Result<()> {
983    let file = path.join(SCHEMA_FILE);
984    let json = serde_json::to_string_pretty(schema)?;
985    std::fs::write(&file, json)?;
986    Ok(())
987}
988
989/// Persist a kit schema into the database. Used after migrations.
990pub(crate) fn persist_schema(db: &Database, schema: &KitSchema) -> Result<()> {
991    store_schema(&db.root, schema)
992}