1use std::path::Path;
24
25use serde::Serialize;
26use sha2::{Digest, Sha256};
27
28use crate::memtable::Row;
29use crate::schema::Schema;
30use crate::storage_mode::StorageMode;
31use crate::{MongrelError, Result};
32
33#[derive(Debug, Clone)]
37pub struct ImportTablePlan {
38 pub table_id: u64,
40 pub name: String,
42 pub schema: Schema,
44 pub row_count: u64,
46 pub rows_sha256: [u8; 32],
49}
50
51#[derive(Debug, Clone)]
54pub struct ImportPlan {
55 pub database: String,
57 pub source_storage_mode: Option<StorageMode>,
60 pub snapshot_epoch: u64,
62 pub tables: Vec<ImportTablePlan>,
64 pub total_rows: u64,
66 pub schema_sha256: [u8; 32],
68 pub rows_sha256: [u8; 32],
70}
71
72pub fn hash_rows_canonical(rows: &[Row]) -> [u8; 32] {
80 let mut ordered: Vec<&Row> = rows.iter().collect();
81 ordered.sort_by_key(|row| row.row_id);
82 let mut hasher = Sha256::new();
83 for row in ordered {
84 hasher.update(row.row_id.0.to_le_bytes());
85 hasher.update(row.committed_epoch.0.to_le_bytes());
86 hasher.update([u8::from(row.deleted)]);
87 let mut columns: Vec<(&u16, &crate::memtable::Value)> = row.columns.iter().collect();
88 columns.sort_by_key(|(column_id, _)| *column_id);
89 hasher.update((columns.len() as u64).to_le_bytes());
90 for (column_id, value) in columns {
91 hasher.update(column_id.to_le_bytes());
92 let encoded =
93 bincode::serialize(value).expect("Value serialization is infallible for hashing");
94 hasher.update((encoded.len() as u64).to_le_bytes());
95 hasher.update(encoded);
96 }
97 }
98 hasher.finalize().into()
99}
100
101fn canonical_json<T: Serialize>(value: &T) -> Result<Vec<u8>> {
105 let mut value = serde_json::to_value(value)
106 .map_err(|error| MongrelError::Other(format!("canonical json: {error}")))?;
107 canonicalize(&mut value);
108 serde_json::to_vec(&value)
109 .map_err(|error| MongrelError::Other(format!("canonical json: {error}")))
110}
111
112fn canonicalize(value: &mut serde_json::Value) {
113 match value {
114 serde_json::Value::Object(map) => {
115 map.sort_keys();
116 for value in map.values_mut() {
117 canonicalize(value);
118 }
119 }
120 serde_json::Value::Array(items) => {
121 for item in items {
122 canonicalize(item);
123 }
124 }
125 _ => {}
126 }
127}
128
129pub fn cluster_import_prepare(source: impl AsRef<Path>, database: &str) -> Result<ImportPlan> {
137 if database.is_empty() {
138 return Err(MongrelError::InvalidArgument(
139 "cluster import requires a target database name".into(),
140 ));
141 }
142 let options = crate::OpenOptions::default().with_offline_validation(true);
143 let db = crate::Database::open_with_options(source.as_ref(), options)?;
144 let source_storage_mode = db.storage_mode()?;
145 let snapshot_epoch = db.visible_epoch();
146 let snapshot = db.snapshot_for_epoch(snapshot_epoch);
148
149 let mut tables = Vec::new();
150 for name in db.table_names() {
151 let table_id = db.table_id(&name)?;
152 let handle = db.table(&name)?;
153 let (schema, rows) = {
154 let table = handle.lock();
155 (table.schema().clone(), table.visible_rows(snapshot)?)
156 };
157 tables.push(ImportTablePlan {
158 table_id,
159 name,
160 schema,
161 row_count: rows.len() as u64,
162 rows_sha256: hash_rows_canonical(&rows),
163 });
164 }
165 drop(db);
166 tables.sort_by_key(|table| table.table_id);
167
168 let mut schema_hasher = Sha256::new();
169 let mut rows_hasher = Sha256::new();
170 let mut total_rows = 0_u64;
171 for table in &tables {
172 let encoded = canonical_json(&table.schema)?;
173 schema_hasher.update((encoded.len() as u64).to_le_bytes());
174 schema_hasher.update(encoded);
175 rows_hasher.update(table.rows_sha256);
176 total_rows = total_rows.saturating_add(table.row_count);
177 }
178 Ok(ImportPlan {
179 database: database.to_string(),
180 source_storage_mode,
181 snapshot_epoch: snapshot_epoch.0,
182 tables,
183 total_rows,
184 schema_sha256: schema_hasher.finalize().into(),
185 rows_sha256: rows_hasher.finalize().into(),
186 })
187}
188
189#[cfg(test)]
190mod tests {
191 use super::*;
192 use crate::memtable::Value;
193 use crate::rowid::RowId;
194
195 #[test]
196 fn canonical_row_hash_is_order_independent() {
197 let mut a = Row::new(RowId(1), crate::epoch::Epoch(3));
198 a.columns.insert(1, Value::Int64(10));
199 a.columns.insert(2, Value::Bytes(b"x".to_vec()));
200 a.columns.insert(7, Value::Null);
201 let mut b = Row::new(RowId(2), crate::epoch::Epoch(3));
202 b.columns.insert(1, Value::Bool(true));
203
204 let forward = hash_rows_canonical(&[a.clone(), b.clone()]);
205 let backward = hash_rows_canonical(&[b, a.clone()]);
206 assert_eq!(forward, backward);
207
208 let mut changed = Row::new(RowId(2), crate::epoch::Epoch(3));
209 changed.columns.insert(1, Value::Bool(false));
210 assert_ne!(forward, hash_rows_canonical(&[a, changed]));
211 }
212
213 #[test]
214 fn canonical_json_sorts_object_keys() {
215 let first = canonical_json(&serde_json::json!({"b": 1, "a": {"d": 2, "c": 3}})).unwrap();
216 let second = canonical_json(&serde_json::json!({"a": {"c": 3, "d": 2}, "b": 1})).unwrap();
217 assert_eq!(first, second);
218 }
219
220 #[test]
221 fn import_plan_counts_and_hashes_match_source() {
222 use crate::memtable::Value;
223 use crate::schema::{ColumnDef, ColumnFlags, Schema, TypeId};
224
225 let dir = tempfile::tempdir().unwrap();
226 let db = crate::Database::create(dir.path()).unwrap();
227 let schema = Schema {
228 columns: vec![ColumnDef {
229 id: 1,
230 name: "id".into(),
231 ty: TypeId::Int64,
232 flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
233 default_value: None,
234 embedding_source: None,
235 }],
236 ..Schema::default()
237 };
238 db.create_table("items", schema).unwrap();
239 let mut txn = db.begin();
240 for value in [10_i64, 20, 30] {
241 txn.put("items", vec![(1, Value::Int64(value))]).unwrap();
242 }
243 txn.commit().unwrap();
244 let snapshot_epoch = db.visible_epoch();
245 let handle = db.table("items").unwrap();
246 let expected_rows = handle
247 .lock()
248 .visible_rows(db.snapshot_for_epoch(snapshot_epoch))
249 .unwrap();
250 drop(handle);
251 drop(db);
252
253 let plan = cluster_import_prepare(dir.path(), "app").unwrap();
254 assert_eq!(plan.database, "app");
255 assert_eq!(plan.source_storage_mode, Some(StorageMode::Standalone));
256 assert_eq!(plan.snapshot_epoch, snapshot_epoch.0);
257 assert_eq!(plan.tables.len(), 1);
258 let table = &plan.tables[0];
259 assert_eq!(table.name, "items");
260 assert_eq!(table.row_count, 3);
261 assert_eq!(table.rows_sha256, hash_rows_canonical(&expected_rows));
262 assert_eq!(plan.total_rows, 3);
263 assert_eq!(plan.schema_sha256.len(), 32);
264
265 let second = cluster_import_prepare(dir.path(), "app").unwrap();
268 assert_eq!(plan.schema_sha256, second.schema_sha256);
269 assert_eq!(plan.rows_sha256, second.rows_sha256);
270 assert_eq!(plan.total_rows, second.total_rows);
271
272 let db = crate::Database::open(dir.path()).unwrap();
274 let handle = db.table("items").unwrap();
275 let rows = handle
276 .lock()
277 .visible_rows(db.snapshot_for_epoch(snapshot_epoch))
278 .unwrap();
279 assert_eq!(rows.len(), 3);
280 assert_eq!(plan.database, "app");
281 }
282
283 #[test]
284 fn import_plan_rejects_empty_database_name() {
285 let dir = tempfile::tempdir().unwrap();
286 let db = crate::Database::create(dir.path()).unwrap();
287 drop(db);
288 assert!(cluster_import_prepare(dir.path(), "").is_err());
289 }
290}