markbase/filetree-rocksdb/src/migrate.rs

use anyhow::{Context, Result};
use filetree_rocksdb::FileTreeRocksDB;
use rusqlite::Connection;
use std::time::Instant;

fn main() -> Result<()> {
    println!("=== SQLite → RocksDB Migration Test ===\n");

    let sqlite_path = "data/users/warren.sqlite";
    let rocksdb_path = "data/users_rocksdb/warren.rocksdb";

    println!("Step 1: Open SQLite database...");
    let conn =
        Connection::open(sqlite_path).with_context(|| format!("Failed to open {}", sqlite_path))?;

    let node_count: i64 =
        conn.query_row("SELECT COUNT(*) FROM file_nodes", [], |row| row.get(0))?;
    println!("  ✓ SQLite nodes count: {}", node_count);

    println!("\nStep 2: Read all nodes from SQLite...");
    let start = Instant::now();

    let mut stmt = conn.prepare(
        "SELECT node_id, label, aliases_json, file_uuid, sha256, parent_id, children_json,
                node_type, icon, color, bg_color, file_size, registered_at,
                created_at, updated_at, sort_order
         FROM file_nodes",
    )?;

    let nodes: Vec<filetree_rocksdb::FileNode> = stmt
        .query_map([], |row| {
            let children_json: String = row.get(6)?;
            let children: Vec<String> = serde_json::from_str(&children_json).unwrap_or_default();
            let node_type_str: String = row.get(7)?;

            Ok(filetree_rocksdb::FileNode {
                node_id: row.get(0)?,
                label: row.get(1)?,
                aliases: filetree_rocksdb::Aliases::from_json(&row.get::<_, String>(2)?),
                file_uuid: row.get(3)?,
                sha256: row.get(4)?,
                parent_id: row.get(5)?,
                children,
                node_type: std::str::FromStr::from_str(&node_type_str)
                    .unwrap_or(filetree_rocksdb::NodeType::Folder),
                icon: row.get(8)?,
                color: row.get(9)?,
                bg_color: row.get(10)?,
                file_size: row.get(11)?,
                registered_at: row.get(12)?,
                created_at: row.get(13)?,
                updated_at: row.get(14)?,
                sort_order: row.get(15)?,
            })
        })?
        .collect::<Result<Vec<_>, _>>()?;

    let read_time = start.elapsed();
    let read_throughput = nodes.len() as f64 / read_time.as_secs_f64();
    println!("  ✓ Read time: {:?}", read_time);
    println!("  ✓ Nodes read: {}", nodes.len());
    println!("  ✓ Throughput: {:.2} nodes/sec", read_throughput);

    println!("\nStep 3: Initialize RocksDB database...");
    let start = Instant::now();
    let rocksdb_tree = FileTreeRocksDB::init_user_db("warren")?;
    let init_time = start.elapsed();
    println!("  ✓ Init time: {:?}", init_time);

    println!("\nStep 4: Import nodes to RocksDB (batch insert)...");
    let start = Instant::now();

    rocksdb_tree.insert_node_batch(&nodes)?;

    let import_time = start.elapsed();
    let import_throughput = nodes.len() as f64 / import_time.as_secs_f64();
    println!("  ✓ Import time: {:?}", import_time);
    println!("  ✓ Throughput: {:.2} nodes/sec", import_throughput);

    println!("\nStep 5: Verify import...");
    let rocksdb_count = rocksdb_tree.count_nodes()?;
    println!("  ✓ RocksDB nodes count: {}", rocksdb_count);
    println!("  ✓ Match: {}", rocksdb_count == nodes.len());

    println!("\nStep 6: Query test (1000 random nodes)...");
    let test_nodes = &nodes[..1000.min(nodes.len())];
    let start = Instant::now();

    for node in test_nodes {
        let _ = rocksdb_tree.get_node(&node.node_id)?;
    }

    let query_time = start.elapsed();
    let query_latency = query_time.as_nanos() as f64 / test_nodes.len() as f64;
    println!("  ✓ Query time: {:?}", query_time);
    println!("  ✓ Average latency: {:.2} ns", query_latency);

    println!("\nStep 7: Database size comparison...");
    let sqlite_size = std::fs::metadata(sqlite_path)?.len();
    let rocksdb_size = get_directory_size(rocksdb_path)?;
    println!(
        "  ✓ SQLite size: {} bytes ({:.2} MB)",
        sqlite_size,
        sqlite_size as f64 / 1024.0 / 1024.0
    );
    println!(
        "  ✓ RocksDB size: {} bytes ({:.2} MB)",
        rocksdb_size,
        rocksdb_size as f64 / 1024.0 / 1024.0
    );
    println!(
        "  ✓ Size ratio: {:.2}x",
        rocksdb_size as f64 / sqlite_size as f64
    );

    println!("\n=== Migration Summary ===");
    println!("SQLite nodes: {}", node_count);
    println!("Imported nodes: {}", nodes.len());
    println!("Import throughput: {:.2} nodes/sec", import_throughput);
    println!("Query latency: {:.2} ns", query_latency);
    println!(
        "Size ratio: {:.2}x",
        rocksdb_size as f64 / sqlite_size as f64
    );

    println!("\nStep 8: Cleanup...");
    std::fs::remove_dir_all(rocksdb_path)?;
    println!("  ✓ Test database removed");

    println!("\n✅ Migration test completed successfully!");

    Ok(())
}

fn get_directory_size(path: &str) -> Result<u64> {
    let mut total_size = 0;
    for entry in std::fs::read_dir(path)? {
        let entry = entry?;
        let metadata = entry.metadata()?;
        if metadata.is_file() {
            total_size += metadata.len();
        } else if metadata.is_dir() {
            total_size += get_directory_size(entry.path().to_str().unwrap())?;
        }
    }
    Ok(total_size)
}