use std::path::Path;
use serde::Serialize;
use sha2::{Digest, Sha256};
use crate::memtable::Row;
use crate::schema::Schema;
use crate::storage_mode::StorageMode;
use crate::{MongrelError, Result};
#[derive(Debug, Clone)]
pub struct ImportTablePlan {
pub table_id: u64,
pub name: String,
pub schema: Schema,
pub row_count: u64,
pub rows_sha256: [u8; 32],
}
#[derive(Debug, Clone)]
pub struct ImportPlan {
pub database: String,
pub source_storage_mode: Option<StorageMode>,
pub snapshot_epoch: u64,
pub tables: Vec<ImportTablePlan>,
pub total_rows: u64,
pub schema_sha256: [u8; 32],
pub rows_sha256: [u8; 32],
}
pub fn hash_rows_canonical(rows: &[Row]) -> [u8; 32] {
let mut ordered: Vec<&Row> = rows.iter().collect();
ordered.sort_by_key(|row| row.row_id);
let mut hasher = Sha256::new();
for row in ordered {
hasher.update(row.row_id.0.to_le_bytes());
hasher.update(row.committed_epoch.0.to_le_bytes());
hasher.update([u8::from(row.deleted)]);
let mut columns: Vec<(&u16, &crate::memtable::Value)> = row.columns.iter().collect();
columns.sort_by_key(|(column_id, _)| *column_id);
hasher.update((columns.len() as u64).to_le_bytes());
for (column_id, value) in columns {
hasher.update(column_id.to_le_bytes());
let encoded =
bincode::serialize(value).expect("Value serialization is infallible for hashing");
hasher.update((encoded.len() as u64).to_le_bytes());
hasher.update(encoded);
}
}
hasher.finalize().into()
}
fn canonical_json<T: Serialize>(value: &T) -> Result<Vec<u8>> {
let mut value = serde_json::to_value(value)
.map_err(|error| MongrelError::Other(format!("canonical json: {error}")))?;
canonicalize(&mut value);
serde_json::to_vec(&value)
.map_err(|error| MongrelError::Other(format!("canonical json: {error}")))
}
fn canonicalize(value: &mut serde_json::Value) {
match value {
serde_json::Value::Object(map) => {
let sorted: serde_json::Map<String, serde_json::Value> = map
.iter()
.map(|(key, value)| (key.clone(), value.clone()))
.collect();
*map = sorted;
for value in map.values_mut() {
canonicalize(value);
}
}
serde_json::Value::Array(items) => {
for item in items {
canonicalize(item);
}
}
_ => {}
}
}
pub fn cluster_import_prepare(source: impl AsRef<Path>, database: &str) -> Result<ImportPlan> {
if database.is_empty() {
return Err(MongrelError::InvalidArgument(
"cluster import requires a target database name".into(),
));
}
let options = crate::OpenOptions::default().with_offline_validation(true);
let db = crate::Database::open_with_options(source.as_ref(), options)?;
let source_storage_mode = db.storage_mode()?;
let snapshot_epoch = db.visible_epoch();
let snapshot = crate::epoch::Snapshot::at(snapshot_epoch);
let mut tables = Vec::new();
for name in db.table_names() {
let table_id = db.table_id(&name)?;
let handle = db.table(&name)?;
let (schema, rows) = {
let table = handle.lock();
(table.schema().clone(), table.visible_rows(snapshot)?)
};
tables.push(ImportTablePlan {
table_id,
name,
schema,
row_count: rows.len() as u64,
rows_sha256: hash_rows_canonical(&rows),
});
}
drop(db);
tables.sort_by_key(|table| table.table_id);
let mut schema_hasher = Sha256::new();
let mut rows_hasher = Sha256::new();
let mut total_rows = 0_u64;
for table in &tables {
let encoded = canonical_json(&table.schema)?;
schema_hasher.update((encoded.len() as u64).to_le_bytes());
schema_hasher.update(encoded);
rows_hasher.update(table.rows_sha256);
total_rows = total_rows.saturating_add(table.row_count);
}
Ok(ImportPlan {
database: database.to_string(),
source_storage_mode,
snapshot_epoch: snapshot_epoch.0,
tables,
total_rows,
schema_sha256: schema_hasher.finalize().into(),
rows_sha256: rows_hasher.finalize().into(),
})
}
#[cfg(test)]
mod tests {
use super::*;
use crate::memtable::Value;
use crate::rowid::RowId;
#[test]
fn canonical_row_hash_is_order_independent() {
let mut a = Row::new(RowId(1), crate::epoch::Epoch(3));
a.columns.insert(1, Value::Int64(10));
a.columns.insert(2, Value::Bytes(b"x".to_vec()));
a.columns.insert(7, Value::Null);
let mut b = Row::new(RowId(2), crate::epoch::Epoch(3));
b.columns.insert(1, Value::Bool(true));
let forward = hash_rows_canonical(&[a.clone(), b.clone()]);
let backward = hash_rows_canonical(&[b, a.clone()]);
assert_eq!(forward, backward);
let mut changed = Row::new(RowId(2), crate::epoch::Epoch(3));
changed.columns.insert(1, Value::Bool(false));
assert_ne!(forward, hash_rows_canonical(&[a, changed]));
}
#[test]
fn canonical_json_sorts_object_keys() {
let first = canonical_json(&serde_json::json!({"b": 1, "a": {"d": 2, "c": 3}})).unwrap();
let second = canonical_json(&serde_json::json!({"a": {"c": 3, "d": 2}, "b": 1})).unwrap();
assert_eq!(first, second);
}
#[test]
fn import_plan_counts_and_hashes_match_source() {
use crate::memtable::Value;
use crate::schema::{ColumnDef, ColumnFlags, Schema, TypeId};
let dir = tempfile::tempdir().unwrap();
let db = crate::Database::create(dir.path()).unwrap();
let schema = Schema {
columns: vec![ColumnDef {
id: 1,
name: "id".into(),
ty: TypeId::Int64,
flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
default_value: None,
embedding_source: None,
}],
..Schema::default()
};
db.create_table("items", schema).unwrap();
let mut txn = db.begin();
for value in [10_i64, 20, 30] {
txn.put("items", vec![(1, Value::Int64(value))]).unwrap();
}
txn.commit().unwrap();
let snapshot_epoch = db.visible_epoch();
let handle = db.table("items").unwrap();
let expected_rows = handle
.lock()
.visible_rows(crate::epoch::Snapshot::at(snapshot_epoch))
.unwrap();
drop(handle);
drop(db);
let plan = cluster_import_prepare(dir.path(), "app").unwrap();
assert_eq!(plan.database, "app");
assert_eq!(plan.source_storage_mode, Some(StorageMode::Standalone));
assert_eq!(plan.snapshot_epoch, snapshot_epoch.0);
assert_eq!(plan.tables.len(), 1);
let table = &plan.tables[0];
assert_eq!(table.name, "items");
assert_eq!(table.row_count, 3);
assert_eq!(table.rows_sha256, hash_rows_canonical(&expected_rows));
assert_eq!(plan.total_rows, 3);
assert_eq!(plan.schema_sha256.len(), 32);
let second = cluster_import_prepare(dir.path(), "app").unwrap();
assert_eq!(plan.schema_sha256, second.schema_sha256);
assert_eq!(plan.rows_sha256, second.rows_sha256);
assert_eq!(plan.total_rows, second.total_rows);
let db = crate::Database::open(dir.path()).unwrap();
let handle = db.table("items").unwrap();
let rows = handle
.lock()
.visible_rows(crate::epoch::Snapshot::at(snapshot_epoch))
.unwrap();
assert_eq!(rows.len(), 3);
assert_eq!(plan.database, "app");
}
#[test]
fn import_plan_rejects_empty_database_name() {
let dir = tempfile::tempdir().unwrap();
let db = crate::Database::create(dir.path()).unwrap();
drop(db);
assert!(cluster_import_prepare(dir.path(), "").is_err());
}
}