use anyhow::{anyhow, Result};
use chrono::{DateTime, Utc};
use rocksdb::{
backup::{BackupEngine, BackupEngineOptions},
checkpoint::Checkpoint,
Env, DB,
};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::Arc;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BackupMetadata {
pub backup_id: u32,
pub created_at: DateTime<Utc>,
pub user_id: String,
pub backup_type: BackupType,
pub size_bytes: u64,
pub checksum: String,
pub memory_count: usize,
pub sequence_number: u64,
#[serde(default)]
pub secondary_stores: Vec<String>,
#[serde(default)]
pub secondary_size_bytes: u64,
}
pub struct SecondaryStoreRef<'a> {
pub name: &'a str,
pub db: &'a Arc<DB>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum BackupType {
Full,
Incremental,
}
pub struct ShodhBackupEngine {
backup_path: PathBuf,
}
impl ShodhBackupEngine {
pub fn new(backup_path: PathBuf) -> Result<Self> {
fs::create_dir_all(&backup_path)?;
Ok(Self { backup_path })
}
pub fn backup_path(&self) -> &Path {
&self.backup_path
}
pub fn create_backup(&self, db: &DB, user_id: &str) -> Result<BackupMetadata> {
let backup_dir = self.backup_path.join(user_id);
fs::create_dir_all(&backup_dir)?;
let backup_opts = BackupEngineOptions::new(&backup_dir)?;
let env = Env::new()?;
let mut backup_engine = BackupEngine::open(&backup_opts, &env)?;
let before_count = backup_engine.get_backup_info().len();
backup_engine.create_new_backup(db)?;
let backup_info = backup_engine.get_backup_info();
let latest_backup = backup_info
.last()
.ok_or_else(|| anyhow!("No backup created"))?;
let backup_id = latest_backup.backup_id;
let size_bytes = latest_backup.size;
let sequence_number = db.latest_sequence_number();
let memory_count = self.estimate_memory_count(db)?;
let checksum = self.calculate_backup_checksum(&backup_dir, backup_id)?;
let backup_type = if before_count == 0 {
BackupType::Full
} else {
BackupType::Incremental
};
let metadata = BackupMetadata {
backup_id,
created_at: Utc::now(),
user_id: user_id.to_string(),
backup_type,
size_bytes,
checksum,
memory_count,
sequence_number,
secondary_stores: Vec::new(),
secondary_size_bytes: 0,
};
self.save_metadata(&metadata)?;
tracing::info!(
backup_id = backup_id,
user_id = user_id,
size_mb = size_bytes / 1024 / 1024,
"Backup created successfully"
);
Ok(metadata)
}
pub fn create_comprehensive_backup(
&self,
db: &DB,
user_id: &str,
secondary_stores: &[SecondaryStoreRef<'_>],
) -> Result<BackupMetadata> {
self.create_comprehensive_backup_with_graph(db, user_id, secondary_stores, None)
}
pub fn create_comprehensive_backup_with_graph(
&self,
db: &DB,
user_id: &str,
secondary_stores: &[SecondaryStoreRef<'_>],
graph_db: Option<&DB>,
) -> Result<BackupMetadata> {
let mut metadata = self.create_backup(db, user_id)?;
let secondary_dir = self
.backup_path
.join(user_id)
.join(format!("secondary_{}", metadata.backup_id));
fs::create_dir_all(&secondary_dir)?;
if let Some(graph) = graph_db {
let graph_checkpoint_dir = secondary_dir.join("graph");
let checkpoint = Checkpoint::new(graph)
.map_err(|e| anyhow!("Failed to create checkpoint handle for graph DB: {}", e))?;
checkpoint
.create_checkpoint(&graph_checkpoint_dir)
.map_err(|e| {
let _ = fs::remove_dir_all(&graph_checkpoint_dir);
anyhow!("Failed to checkpoint graph DB: {}", e)
})?;
let graph_size = dir_size(&graph_checkpoint_dir).unwrap_or(0);
tracing::debug!(size_kb = graph_size / 1024, "Graph DB checkpointed");
}
let mut backed_up_stores = Vec::new();
let mut total_secondary_bytes: u64 = 0;
for store_ref in secondary_stores {
let store_checkpoint_dir = secondary_dir.join(store_ref.name);
if store_checkpoint_dir.exists() {
tracing::warn!(
store = store_ref.name,
"Checkpoint directory already exists, skipping"
);
continue;
}
let checkpoint = Checkpoint::new(store_ref.db).map_err(|e| {
anyhow!(
"Failed to create checkpoint handle for secondary store '{}': {}",
store_ref.name,
e
)
})?;
if let Err(e) = checkpoint.create_checkpoint(&store_checkpoint_dir) {
let _ = fs::remove_dir_all(&store_checkpoint_dir);
return Err(anyhow!(
"Failed to checkpoint secondary store '{}': {}",
store_ref.name,
e
));
}
let store_size = dir_size(&store_checkpoint_dir).unwrap_or(0);
total_secondary_bytes += store_size;
backed_up_stores.push(store_ref.name.to_string());
tracing::debug!(
store = store_ref.name,
size_kb = store_size / 1024,
"Secondary store checkpointed"
);
}
if graph_db.is_some() {
backed_up_stores.push("graph".to_string());
}
metadata.secondary_stores = backed_up_stores;
metadata.secondary_size_bytes = total_secondary_bytes;
let backup_dir = self.backup_path.join(user_id);
metadata.checksum = self.calculate_backup_checksum(&backup_dir, metadata.backup_id)?;
self.save_metadata(&metadata)?;
tracing::info!(
backup_id = metadata.backup_id,
user_id = user_id,
secondary_stores = metadata.secondary_stores.len(),
secondary_size_kb = total_secondary_bytes / 1024,
"Comprehensive backup created"
);
Ok(metadata)
}
pub fn restore_backup(
&self,
user_id: &str,
backup_id: Option<u32>,
restore_path: &Path,
) -> Result<()> {
let backup_dir = self.backup_path.join(user_id);
if !backup_dir.exists() {
return Err(anyhow!("No backups found for user: {user_id}"));
}
let backup_opts = BackupEngineOptions::new(&backup_dir)?;
let env = Env::new()?;
let mut backup_engine = BackupEngine::open(&backup_opts, &env)?;
match backup_id {
Some(id) => {
tracing::info!(backup_id = id, "Restoring from specific backup");
backup_engine.restore_from_backup(
restore_path,
restore_path,
&rocksdb::backup::RestoreOptions::default(),
id,
)?;
}
None => {
tracing::info!("Restoring from latest backup");
backup_engine.restore_from_latest_backup(
restore_path,
restore_path,
&rocksdb::backup::RestoreOptions::default(),
)?;
}
}
tracing::info!(
user_id = user_id,
restore_path = ?restore_path,
"Restore completed successfully"
);
Ok(())
}
pub fn list_backups(&self, user_id: &str) -> Result<Vec<BackupMetadata>> {
let backup_dir = self.backup_path.join(user_id);
if !backup_dir.exists() {
return Ok(Vec::new());
}
let backup_opts = BackupEngineOptions::new(&backup_dir)?;
let env = Env::new()?;
let backup_engine = BackupEngine::open(&backup_opts, &env)?;
let backup_info = backup_engine.get_backup_info();
let mut metadata_list = Vec::new();
for info in backup_info {
if let Ok(metadata) = self.load_metadata(user_id, info.backup_id) {
metadata_list.push(metadata);
}
}
Ok(metadata_list)
}
pub fn restore_comprehensive_backup(
&self,
user_id: &str,
backup_id: Option<u32>,
restore_path: &Path,
secondary_restore_paths: &[(&str, &Path)],
) -> Result<Vec<String>> {
self.restore_backup(user_id, backup_id, restore_path)?;
let resolved_backup_id = match backup_id {
Some(id) => id,
None => {
let backup_dir = self.backup_path.join(user_id);
let backup_opts = BackupEngineOptions::new(&backup_dir)?;
let env = Env::new()?;
let backup_engine = BackupEngine::open(&backup_opts, &env)?;
let info = backup_engine.get_backup_info();
info.last()
.map(|i| i.backup_id)
.ok_or_else(|| anyhow!("No backups available"))?
}
};
let secondary_dir = self
.backup_path
.join(user_id)
.join(format!("secondary_{resolved_backup_id}"));
let mut restored_stores = Vec::new();
if secondary_dir.exists() {
for (store_name, target_path) in secondary_restore_paths {
let checkpoint_dir = secondary_dir.join(store_name);
if !checkpoint_dir.exists() {
tracing::debug!(
store = *store_name,
"No checkpoint found in backup, skipping"
);
continue;
}
let mut tmp_os = target_path.as_os_str().to_os_string();
tmp_os.push(".restore_tmp");
let temp_path = PathBuf::from(tmp_os);
if temp_path.exists() {
fs::remove_dir_all(&temp_path).map_err(|e| {
anyhow!(
"Failed to clean up stale temp dir for {}: {}",
store_name,
e
)
})?;
}
if let Err(e) = copy_dir_recursive(&checkpoint_dir, &temp_path) {
let _ = fs::remove_dir_all(&temp_path);
tracing::warn!(
store = *store_name,
error = %e,
"Failed to copy checkpoint for restore, skipping (original data preserved)"
);
continue;
}
if target_path.exists() {
if let Err(e) = fs::remove_dir_all(target_path) {
let _ = fs::remove_dir_all(&temp_path);
return Err(anyhow!(
"Failed to remove existing {} directory at {:?}: {}",
store_name,
target_path,
e
));
}
}
if let Err(e) = fs::rename(&temp_path, target_path) {
if let Err(copy_err) = copy_dir_recursive(&temp_path, target_path) {
let _ = fs::remove_dir_all(&temp_path);
return Err(anyhow!(
"Failed to finalize restore for {}: rename={}, copy={}",
store_name,
e,
copy_err
));
}
let _ = fs::remove_dir_all(&temp_path);
}
restored_stores.push(store_name.to_string());
tracing::info!(
store = *store_name,
target = ?target_path,
"Secondary store restored from checkpoint"
);
}
}
tracing::info!(
user_id = user_id,
backup_id = resolved_backup_id,
restored_secondary = restored_stores.len(),
"Comprehensive restore completed"
);
Ok(restored_stores)
}
pub fn purge_old_backups(&self, user_id: &str, keep_count: usize) -> Result<usize> {
if keep_count == 0 {
return Err(anyhow!(
"keep_count must be >= 1 to prevent deleting all backups"
));
}
let backup_dir = self.backup_path.join(user_id);
if !backup_dir.exists() {
return Ok(0);
}
let backup_opts = BackupEngineOptions::new(&backup_dir)?;
let env = Env::new()?;
let mut backup_engine = BackupEngine::open(&backup_opts, &env)?;
let backup_info = backup_engine.get_backup_info();
let total_backups = backup_info.len();
if total_backups <= keep_count {
return Ok(0);
}
let to_delete = total_backups - keep_count;
let mut purge_ids: Vec<u32> = backup_info.iter().map(|b| b.backup_id).collect();
purge_ids.sort();
let purge_ids: Vec<u32> = purge_ids.into_iter().take(to_delete).collect();
backup_engine.purge_old_backups(keep_count)?;
for purged_id in &purge_ids {
let secondary_dir = backup_dir.join(format!("secondary_{purged_id}"));
if secondary_dir.exists() {
if let Err(e) = fs::remove_dir_all(&secondary_dir) {
tracing::warn!(
backup_id = purged_id,
error = %e,
"Failed to clean up secondary store checkpoint"
);
}
}
let metadata_path = backup_dir.join(format!("backup_{purged_id}.json"));
if let Err(e) = fs::remove_file(&metadata_path) {
tracing::warn!(
backup_id = purged_id,
error = %e,
"Failed to remove backup metadata file"
);
}
}
tracing::info!(
purged_count = to_delete,
kept_count = keep_count,
user_id = user_id,
"Purged old backups"
);
Ok(to_delete)
}
pub fn verify_backup(&self, user_id: &str, backup_id: u32) -> Result<bool> {
let metadata = self.load_metadata(user_id, backup_id)?;
let backup_dir = self.backup_path.join(user_id);
let current_checksum = self.calculate_backup_checksum(&backup_dir, backup_id)?;
Ok(current_checksum == metadata.checksum)
}
fn save_metadata(&self, metadata: &BackupMetadata) -> Result<()> {
let metadata_path = self
.backup_path
.join(&metadata.user_id)
.join(format!("backup_{}.json", metadata.backup_id));
let json = serde_json::to_string_pretty(metadata)?;
fs::write(metadata_path, json)?;
Ok(())
}
fn load_metadata(&self, user_id: &str, backup_id: u32) -> Result<BackupMetadata> {
let metadata_path = self
.backup_path
.join(user_id)
.join(format!("backup_{backup_id}.json"));
let json = fs::read_to_string(metadata_path)?;
let metadata = serde_json::from_str(&json)?;
Ok(metadata)
}
fn calculate_backup_checksum(&self, backup_dir: &Path, backup_id: u32) -> Result<String> {
let mut hasher = Sha256::new();
let backup_path = backup_dir.join(format!("private/{backup_id}"));
self.hash_directory_sorted(&backup_path, &mut hasher)?;
let secondary_path = backup_dir.join(format!("secondary_{backup_id}"));
self.hash_directory_sorted(&secondary_path, &mut hasher)?;
let result = hasher.finalize();
Ok(format!("{result:x}"))
}
fn hash_directory_sorted(&self, dir: &Path, hasher: &mut Sha256) -> Result<()> {
if !dir.exists() {
return Ok(());
}
let mut entries: Vec<_> = fs::read_dir(dir)?.filter_map(|e| e.ok()).collect();
entries.sort_by_key(|e| e.file_name());
for entry in entries {
let path = entry.path();
hasher.update(entry.file_name().to_string_lossy().as_bytes());
if path.is_dir() {
self.hash_directory_sorted(&path, hasher)?;
} else {
let file_contents = fs::read(&path)?;
hasher.update(&file_contents);
}
}
Ok(())
}
fn estimate_memory_count(&self, db: &DB) -> Result<usize> {
let mut count = 0;
let iter = db.iterator(rocksdb::IteratorMode::Start);
for _ in iter {
count += 1;
}
Ok(count)
}
}
pub fn copy_dir_recursive_pub(src: &Path, dst: &Path) -> Result<()> {
copy_dir_recursive(src, dst)
}
fn dir_size(path: &Path) -> Result<u64> {
let mut total = 0u64;
if path.is_dir() {
for entry in fs::read_dir(path)? {
let entry = entry?;
let entry_path = entry.path();
if entry_path.is_dir() {
total += dir_size(&entry_path)?;
} else {
total += entry.metadata()?.len();
}
}
}
Ok(total)
}
fn copy_dir_recursive(src: &Path, dst: &Path) -> Result<()> {
fs::create_dir_all(dst)?;
for entry in fs::read_dir(src)? {
let entry = entry?;
let src_path = entry.path();
let dst_path = dst.join(entry.file_name());
if src_path.is_dir() {
copy_dir_recursive(&src_path, &dst_path)?;
} else {
fs::copy(&src_path, &dst_path)?;
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use rocksdb::Options;
use serde_json::Value;
use tempfile::TempDir;
#[test]
fn test_backup_engine_creation() {
let temp_dir = TempDir::new().unwrap();
let backup_engine = ShodhBackupEngine::new(temp_dir.path().to_path_buf());
assert!(backup_engine.is_ok());
}
#[test]
fn test_backup_metadata_serialization() {
let metadata = BackupMetadata {
backup_id: 1,
created_at: Utc::now(),
user_id: "test_user".to_string(),
backup_type: BackupType::Full,
size_bytes: 1024,
checksum: "abc123".to_string(),
memory_count: 100,
sequence_number: 42,
secondary_stores: vec!["todo_items".to_string(), "prospective_tasks".to_string()],
secondary_size_bytes: 2048,
};
let json = serde_json::to_string(&metadata).unwrap();
let deserialized: BackupMetadata = serde_json::from_str(&json).unwrap();
assert_eq!(metadata.backup_id, deserialized.backup_id);
assert_eq!(metadata.user_id, deserialized.user_id);
}
#[test]
fn test_dir_size_counts_nested_files() {
let temp_dir = TempDir::new().unwrap();
let root = temp_dir.path();
let nested = root.join("nested");
fs::create_dir_all(&nested).unwrap();
fs::write(root.join("a.txt"), b"12345").unwrap();
fs::write(nested.join("b.txt"), b"1234567890").unwrap();
let size = dir_size(root).unwrap();
assert_eq!(size, 15);
}
#[test]
fn test_copy_dir_recursive_pub_copies_files() {
let temp_dir = TempDir::new().unwrap();
let src = temp_dir.path().join("src");
let dst = temp_dir.path().join("dst");
fs::create_dir_all(src.join("deep")).unwrap();
fs::write(src.join("file1.txt"), b"alpha").unwrap();
fs::write(src.join("deep").join("file2.txt"), b"beta").unwrap();
copy_dir_recursive_pub(&src, &dst).unwrap();
assert_eq!(fs::read(dst.join("file1.txt")).unwrap(), b"alpha");
assert_eq!(
fs::read(dst.join("deep").join("file2.txt")).unwrap(),
b"beta"
);
}
#[test]
fn test_list_backups_empty_when_user_missing() {
let temp_dir = TempDir::new().unwrap();
let engine = ShodhBackupEngine::new(temp_dir.path().to_path_buf()).unwrap();
let backups = engine.list_backups("missing-user").unwrap();
assert!(backups.is_empty());
}
#[test]
fn test_verify_backup_round_trip() {
let temp_dir = TempDir::new().unwrap();
let backup_root = temp_dir.path().join("backups");
let db_path = temp_dir.path().join("db");
let user_id = "user1";
let mut opts = Options::default();
opts.create_if_missing(true);
let db = DB::open(&opts, &db_path).unwrap();
db.put(b"k1", b"v1").unwrap();
db.put(b"k2", b"v2").unwrap();
let engine = ShodhBackupEngine::new(backup_root.clone()).unwrap();
let metadata = engine.create_backup(&db, user_id).unwrap();
let verified = engine.verify_backup(user_id, metadata.backup_id).unwrap();
assert!(verified);
}
#[test]
fn test_verify_backup_detects_checksum_mismatch() {
let temp_dir = TempDir::new().unwrap();
let backup_root = temp_dir.path().join("backups");
let db_path = temp_dir.path().join("db");
let user_id = "user2";
let mut opts = Options::default();
opts.create_if_missing(true);
let db = DB::open(&opts, &db_path).unwrap();
db.put(b"k", b"v").unwrap();
let engine = ShodhBackupEngine::new(backup_root.clone()).unwrap();
let metadata = engine.create_backup(&db, user_id).unwrap();
let metadata_path = backup_root
.join(user_id)
.join(format!("backup_{}.json", metadata.backup_id));
let json = fs::read_to_string(&metadata_path).unwrap();
let mut parsed: Value = serde_json::from_str(&json).unwrap();
parsed["checksum"] = Value::String("0000badchecksum".to_string());
fs::write(
&metadata_path,
serde_json::to_string_pretty(&parsed).unwrap(),
)
.unwrap();
let verified = engine.verify_backup(user_id, metadata.backup_id).unwrap();
assert!(!verified);
}
#[test]
fn test_purge_old_backups_validates_keep_count() {
let temp_dir = TempDir::new().unwrap();
let engine = ShodhBackupEngine::new(temp_dir.path().to_path_buf()).unwrap();
let err = engine.purge_old_backups("user", 0).unwrap_err();
assert!(err.to_string().contains("keep_count must be >= 1"));
}
#[test]
fn test_purge_old_backups_removes_old_entries() {
let temp_dir = TempDir::new().unwrap();
let backup_root = temp_dir.path().join("backups");
let db_path = temp_dir.path().join("db");
let user_id = "purge-user";
let mut opts = Options::default();
opts.create_if_missing(true);
let db = DB::open(&opts, &db_path).unwrap();
let engine = ShodhBackupEngine::new(backup_root.clone()).unwrap();
db.put(b"k1", b"v1").unwrap();
let first = engine.create_backup(&db, user_id).unwrap();
db.put(b"k2", b"v2").unwrap();
let second = engine.create_backup(&db, user_id).unwrap();
let purged = engine.purge_old_backups(user_id, 1).unwrap();
assert_eq!(purged, 1);
let remaining = engine.list_backups(user_id).unwrap();
assert_eq!(remaining.len(), 1);
assert_eq!(remaining[0].backup_id, second.backup_id);
let first_metadata_path = backup_root
.join(user_id)
.join(format!("backup_{}.json", first.backup_id));
assert!(!first_metadata_path.exists());
}
}