feat: support merging file and index databases

Author: jyeshe

.gitignore

@@ -7,8 +7,9 @@ node_modules
 foo.js
 .env
 ao-cache
+tools/rocksdb-cli/target
 
 # terraform
 .terraform
 .terragrunt-cache
-terraform/*.plan
+terraform/*.plan

servers/su/Cargo.lock

@@ -3,14 +3,16 @@ use std::io;
 use su::domain::migrate_to_disk;
 use su::domain::migrate_to_local;
 use su::domain::sync_local_drives;
+use su::domain::merge_dbs;
 
 #[tokio::main]
 async fn main() -> io::Result<()> {
     let args: Vec<String> = env::args().collect();
 
     if args.len() < 2 {
-        eprintln!("Usage: {} <function_name>", args[0]);
-        eprintln!("Available functions: migrate_to_disk, migrate_to_local, sync_local_drives");
+        eprintln!("Usage: {} <function_name> [args...]", args[0]);
+        eprintln!("Available functions: migrate_to_disk, migrate_to_local, sync_local_drives, merge_dbs");
+        eprintln!("  merge_dbs: {} merge_dbs <src_file_db> <src_index_db> [dest_file_db] [dest_index_db]", args[0]);
         return Ok(());
     }
 
@@ -36,9 +38,22 @@ async fn main() -> io::Result<()> {
         "sync_local_drives" => {
             sync_local_drives(interval).await.unwrap();
         }
+        "merge_dbs" => {
+            if args.len() < 4 || args.len() > 6 {
+                eprintln!("Usage: {} merge_dbs <src_file_db> <src_index_db> [dest_file_db] [dest_index_db]", args[0]);
+                return Ok(());
+            }
+            let dest_file = if args.len() >= 5 { Some(args[4].as_str()) } else { None };
+            let dest_index = if args.len() >= 6 { Some(args[5].as_str()) } else { None };
+            
+            match merge_dbs(&args[2], &args[3], dest_file, dest_index) {
+                Ok(()) => println!("Database merge completed successfully"),
+                Err(e) => eprintln!("Error merging databases: {}", e),
+            }
+        }
         _ => {
             eprintln!("Invalid function name: {}", args[1]);
-            eprintln!("Available functions: migrate_to_disk, migrate_to_local, sync_local_drives");
+            eprintln!("Available functions: migrate_to_disk, migrate_to_local, sync_local_drives, merge_dbs");
         }
     }
 

servers/su/src/bin/cli.rs

@@ -0,0 +1,78 @@
+use rocksdb::WriteBatch;
+use std::sync::Arc;
+
+use crate::domain::config::AoConfig;
+use super::store::LocalStoreClient;
+
+pub fn merge_dbs(src_file_db_dir: &str, src_index_db_dir: &str, dest_file_db_dir: Option<&str>, dest_index_db_dir: Option<&str>) -> Result<(), Box<dyn std::error::Error>> {
+    // Open source databases using LocalStoreClient read-only method
+    let src_store = LocalStoreClient::new_read_only(
+        &src_file_db_dir.to_string(), 
+        &src_index_db_dir.to_string()
+    ).map_err(|e| format!("Failed to open source LocalStoreClient: {:?}", e))?;
+    
+    // Create destination store using provided paths or config
+    let dest_store = Arc::new(
+        if let (Some(file_dir), Some(index_dir)) = (dest_file_db_dir, dest_index_db_dir) {
+            LocalStoreClient::new(&file_dir.to_string(), &index_dir.to_string())
+                .map_err(|e| format!("Failed to create LocalStoreClient: {:?}", e))?
+        } else {
+            let config = AoConfig::new(None)?;
+            LocalStoreClient::new(&config.su_file_db_dir, &config.su_index_db_dir)
+                .map_err(|e| format!("Failed to create LocalStoreClient: {:?}", e))?
+        }
+    );
+    
+    // Merge file database
+    let mut file_count = 0;
+    
+    let file_iterator = src_store.file_db.iterator(rocksdb::IteratorMode::Start);
+    
+    for item in file_iterator {
+        let (key, value) = item?;
+        dest_store.file_db.put(key, value)?;
+        file_count += 1;
+        
+        if file_count % 100 == 0 {
+            println!("Processed file record {}", file_count);
+        }
+    }
+    
+    // Merge index database
+    let mut index_batch = WriteBatch::default();
+    let mut index_count = 0;
+    
+    let index_iterator = src_store.index_db.iterator(rocksdb::IteratorMode::Start);
+    
+    for item in index_iterator {
+        let (key, value) = item?;
+        index_batch.put(key, value);
+        index_count += 1;
+        
+        if index_count % 10000 == 0 {
+            dest_store.index_db.write(index_batch)?;
+            index_batch = WriteBatch::default();
+            println!("Committed index batch at record {}", index_count);
+        }
+    }
+    
+    if !index_batch.is_empty() {
+        dest_store.index_db.write(index_batch)?;
+        println!("Committed final index batch at record {}", index_count);
+    }
+    
+    println!("Successfully merged {} file records and {} index records", file_count, index_count);
+    println!("Source file DB: {}", src_file_db_dir);
+    println!("Source index DB: {}", src_index_db_dir);
+    
+    if let (Some(file_dir), Some(index_dir)) = (dest_file_db_dir, dest_index_db_dir) {
+        println!("Destination file DB: {}", file_dir);
+        println!("Destination index DB: {}", index_dir);
+    } else {
+        let config = AoConfig::new(None)?;
+        println!("Destination file DB: {}", config.su_file_db_dir);
+        println!("Destination index DB: {}", config.su_index_db_dir);
+    }
+    
+    Ok(())
+}

servers/su/src/domain/clients/local_store/import.rs

@@ -1,4 +1,5 @@
 pub mod migration;
+pub mod import;
 pub mod store;
 pub mod sync_local;
 pub mod tests;

servers/su/src/domain/clients/local_store/mod.rs

@@ -1,8 +1,10 @@
 #[cfg(test)]
 mod tests {
     use super::super::store::LocalStoreClient;
+    use crate::domain::clients::local_store::import::merge_dbs;
     use crate::domain::core::dal::{DataStore, Message, Process, StoreErrorType};
     use base64_url::decode;
+    use rocksdb::{DB, Options};
     use std::fs;
     use std::path::PathBuf;
 
@@ -427,4 +429,56 @@ mod tests {
         let messages = bundle_strings.iter().map(|b| decode(b).unwrap()).collect();
         (decode(process_string).unwrap(), messages)
     }
+
+    #[tokio::test]
+    async fn test_merge_dbs() {
+        let src_test_db = TestDb::new(100);
+        let dest_test_db = TestDb::new(101);
+
+        // Create source store and populate with keys 1, 2, 3
+        let src_store = LocalStoreClient::new(
+            &src_test_db.file_db_path(),
+            &src_test_db.index_db_path()
+        ).unwrap();
+
+        src_store.file_db.put(b"1", b"value1_src").unwrap();
+        src_store.file_db.put(b"2", b"value2_src").unwrap();
+        src_store.file_db.put(b"3", b"value3_src").unwrap();
+
+        // Create destination store and populate with keys 3, 4, 5        
+        let dest_store = LocalStoreClient::new(
+            &dest_test_db.file_db_path(),
+            &dest_test_db.index_db_path()
+        ).unwrap();
+
+        dest_store.file_db.put(b"3", b"value3_dest").unwrap();
+        dest_store.file_db.put(b"4", b"value4_dest").unwrap();
+        dest_store.file_db.put(b"5", b"value5_dest").unwrap();
+
+        // Close source and destination stores before merge
+        drop(src_store);
+        drop(dest_store);
+
+        // Run merge_dbs with explicit destination paths
+        merge_dbs(
+            &src_test_db.file_db_path(), 
+            &src_test_db.index_db_path(),
+            Some(&dest_test_db.file_db_path()),
+            Some(&dest_test_db.index_db_path())
+        ).expect("merge_dbs should succeed");
+
+        // Verify results
+        let mut verify_opts = Options::default();
+        verify_opts.create_if_missing(false);
+        verify_opts.set_enable_blob_files(true);
+        
+        let verify_db = DB::open(&verify_opts, &dest_test_db.file_db_path()).unwrap();
+        
+        // Check all keys exist
+        assert_eq!(verify_db.get(b"1").unwrap().unwrap(), b"value1_src");
+        assert_eq!(verify_db.get(b"2").unwrap().unwrap(), b"value2_src");
+        assert_eq!(verify_db.get(b"3").unwrap().unwrap(), b"value3_src"); // Should be overwritten from src
+        assert_eq!(verify_db.get(b"4").unwrap().unwrap(), b"value4_dest");
+        assert_eq!(verify_db.get(b"5").unwrap().unwrap(), b"value5_dest");
+    }
 }

servers/su/src/domain/clients/local_store/tests.rs

@@ -5,7 +5,7 @@ use tokio::task::spawn_blocking;
 
 use dashmap::DashMap;
 
-mod clients;
+pub mod clients;
 pub mod config;
 mod core;
 mod logger;
@@ -23,6 +23,7 @@ pub use core::flows;
 pub use core::router;
 pub use core::bytes;
 pub use flows::Deps;
+pub use local_store::import::merge_dbs;
 pub use local_store::migration::migrate_to_local;
 pub use local_store::sync_local::sync_local_drives;
 pub use store::migrate_to_disk;