use mongreldb_core::{
schema::{ColumnDef, ColumnFlags, IndexDef, IndexKind, Schema, TypeId},
Condition, Query, RowId, Snapshot, Table, Value,
};
use tempfile::tempdir;
fn schema() -> Schema {
Schema {
schema_id: 1,
columns: vec![
ColumnDef {
id: 1,
name: "id".into(),
ty: TypeId::Int64,
flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
},
ColumnDef {
id: 2,
name: "label".into(),
ty: TypeId::Bytes,
flags: ColumnFlags::empty(),
},
ColumnDef {
id: 3,
name: "emb".into(),
ty: TypeId::Embedding { dim: 8 },
flags: ColumnFlags::empty().with(ColumnFlags::NULLABLE),
},
],
indexes: vec![
IndexDef {
name: "label_bitmap".into(),
column_id: 2,
kind: IndexKind::Bitmap,
},
IndexDef {
name: "label_fm".into(),
column_id: 2,
kind: IndexKind::FmIndex,
},
IndexDef {
name: "emb_ann".into(),
column_id: 3,
kind: IndexKind::Ann,
},
],
colocation: vec![],
constraints: Default::default(),
}
}
fn red_emb() -> Vec<f32> {
vec![1.0, 1.0, 1.0, 1.0, -1.0, -1.0, -1.0, -1.0]
}
fn blue_emb() -> Vec<f32> {
vec![-1.0, -1.0, -1.0, -1.0, 1.0, 1.0, 1.0, 1.0]
}
fn seed(db: &mut Table) -> Vec<RowId> {
let specs: [(i64, &[u8], Vec<f32>); 4] = [
(10, b"red", red_emb()),
(20, b"blue", blue_emb()),
(30, b"red", red_emb()),
(40, b"green", red_emb()),
];
let mut ids = Vec::new();
for (id, label, emb) in specs {
ids.push(
db.put(vec![
(1, Value::Int64(id)),
(2, Value::Bytes(label.to_vec())),
(3, Value::Embedding(emb)),
])
.unwrap(),
);
}
ids
}
#[test]
fn flush_then_read_across_run() {
let dir = tempdir().unwrap();
let mut db = Table::create(dir.path(), schema(), 1).unwrap();
let ids = seed(&mut db);
db.set_mutable_run_spill_bytes(1); db.flush().unwrap(); assert_eq!(db.run_count(), 1);
assert_eq!(db.memtable_len(), 0);
let snap = db.snapshot();
let row = db.get(ids[1], snap).unwrap();
assert!(matches!(row.columns.get(&2), Some(Value::Bytes(_))));
let rid = db.lookup_pk(&Value::Int64(20).encode_key()).unwrap();
assert_eq!(rid, ids[1]);
}
#[test]
fn count_conditions_returns_survivor_cardinality() {
let dir = tempdir().unwrap();
let mut db = Table::create(dir.path(), schema(), 1).unwrap();
seed(&mut db);
db.flush().unwrap();
let snap = db.snapshot();
let red = [Condition::BitmapEq {
column_id: 2,
value: b"red".to_vec(),
}];
assert_eq!(db.count_conditions(&red, snap).unwrap(), Some(2));
let pk = [Condition::Pk(Value::Int64(20).encode_key())];
assert_eq!(db.count_conditions(&pk, snap).unwrap(), Some(1));
}
#[test]
fn mvcc_snapshot_isolation_after_update() {
let dir = tempdir().unwrap();
let mut db = Table::create(dir.path(), schema(), 1).unwrap();
let id = db
.put(vec![
(1, Value::Int64(1)),
(2, Value::Bytes(b"v1".to_vec())),
])
.unwrap();
let e1 = db.commit().unwrap();
db.put(vec![
(1, Value::Int64(1)),
(2, Value::Bytes(b"v2".to_vec())),
])
.unwrap();
let _ = db
.put(vec![
(1, Value::Int64(99)),
(2, Value::Bytes(b"noise".to_vec())),
])
.unwrap(); let e2 = db.commit().unwrap();
assert!(matches!(
db.get(id, Snapshot::at(e1)).unwrap().columns.get(&2),
Some(Value::Bytes(b))
if b == b"v1"
));
assert_eq!(db.visible_rows(Snapshot::at(e1)).unwrap().len(), 1);
let _ = e2;
}
#[test]
fn conjunctive_query_intersects_row_id_space() {
let dir = tempdir().unwrap();
let mut db = Table::create(dir.path(), schema(), 1).unwrap();
let _ = seed(&mut db);
db.commit().unwrap();
let reds = db
.query(&Query::new().and(Condition::BitmapEq {
column_id: 2,
value: b"red".to_vec(),
}))
.unwrap();
assert_eq!(reds.len(), 2);
let blues = db
.query(&Query::new().and(Condition::FmContains {
column_id: 2,
pattern: b"lu".to_vec(),
}))
.unwrap();
assert_eq!(blues.len(), 1);
let nearest = db
.query(&Query::new().and(Condition::Ann {
column_id: 3,
query: red_emb(),
k: 10,
}))
.unwrap();
assert!(nearest.len() >= 2);
}
#[test]
fn reopen_recovers_state_from_wal_and_runs() {
let dir = tempdir().unwrap();
let path = dir.path().to_path_buf();
let committed_id = {
let mut db = Table::create(&path, schema(), 1).unwrap();
let ids = seed(&mut db);
db.flush().unwrap(); db.put(vec![
(1, Value::Int64(50)),
(2, Value::Bytes(b"yellow".to_vec())),
(3, Value::Embedding(red_emb())),
])
.unwrap();
db.commit().unwrap();
ids[0]
};
let mut db = Table::open(&path).unwrap();
let snap = db.snapshot();
assert!(db.get(committed_id, snap).is_some(), "flushed row survives");
let all = db.visible_rows(snap).unwrap();
assert_eq!(all.len(), 5, "4 flushed + 1 WAL-replayed");
let yellow = db
.query(&Query::new().and(Condition::BitmapEq {
column_id: 2,
value: b"yellow".to_vec(),
}))
.unwrap();
assert_eq!(yellow.len(), 1);
}
#[test]
fn compaction_preserves_visible_state() {
let dir = tempdir().unwrap();
let mut db = Table::create(dir.path(), schema(), 1).unwrap();
let _ = seed(&mut db);
db.set_mutable_run_spill_bytes(1); db.flush().unwrap();
db.put(vec![
(1, Value::Int64(60)),
(2, Value::Bytes(b"red".to_vec())),
(3, Value::Embedding(red_emb())),
])
.unwrap();
db.flush().unwrap();
assert_eq!(db.run_count(), 2);
let before = db.visible_rows(db.snapshot()).unwrap().len();
db.compact().unwrap();
assert_eq!(db.run_count(), 1);
let after = db.visible_rows(db.snapshot()).unwrap().len();
assert_eq!(before, after, "compaction must not change the visible set");
}
#[cfg(feature = "encryption")]
#[test]
fn encrypted_flush_and_read_round_trips() {
let dir = tempdir().unwrap();
let path = dir.path().to_path_buf();
{
let mut db =
Table::create_encrypted(&path, schema(), 1, "correct horse battery staple").unwrap();
db.set_mutable_run_spill_bytes(1); let id = db
.put(vec![
(1, Value::Int64(7)),
(2, Value::Bytes(b"secret".to_vec())),
])
.unwrap();
db.flush().unwrap();
let snap = db.snapshot();
let row = db.get(id, snap).unwrap();
assert!(matches!(row.columns.get(&2), Some(Value::Bytes(b)) if b == b"secret"));
}
let db = Table::open_encrypted(&path, "correct horse battery staple").unwrap();
let all = db.visible_rows(db.snapshot()).unwrap();
assert_eq!(all.len(), 1);
let run_bytes = std::fs::read(
std::fs::read_dir(path.join("_runs"))
.unwrap()
.next()
.unwrap()
.unwrap()
.path(),
)
.unwrap();
assert!(!std::str::from_utf8(&run_bytes)
.unwrap_or("")
.contains("secret"));
}