AkurAI-Notes / crates/notes-cli/src/cmd_repair.rs

Latest change 7f7b543549487ff8ff6802fb420feff1a197c507 - Checkpoint WIP + add AkurAI Build CI/CD pipeline by Ólafur Búi Ólafsson

use akurai_storage::BTree;
use std::path::Path;

pub fn run(source: &str, output: &str) -> Result<(), String> {
    if source == output {
        return Err("repair output must differ from the source database".into());
    }
    if Path::new(output).exists() {
        return Err(format!("repair output already exists: {output}"));
    }

    let mut source_db = BTree::open(source).map_err(|e| format!("open source: {e}"))?;
    let mut output_db = BTree::open(output).map_err(|e| format!("open output: {e}"))?;
    let mut records = 0_u64;
    let mut embeddings = 0_u64;

    for collection in ["notes", "notebooks", "tags"] {
        let prefix = format!("{collection}:").into_bytes();
        let mut end = prefix.clone();
        end.push(0xff);
        let mut last_id = 0_u64;

        for (key, value) in source_db
            .salvage_range(&prefix, &end)
            .map_err(|e| format!("salvage {collection}: {e}"))?
        {
            let suffix = &key[prefix.len()..];
            if suffix.len() != 8 || suffix == b"_counter" {
                continue;
            }
            let id = u64::from_be_bytes(suffix.try_into().unwrap());
            output_db
                .insert(&key, &value)
                .map_err(|e| format!("copy {collection} record {id}: {e}"))?;
            last_id = last_id.max(id);
            records += 1;
        }

        let counter_key = format!("{collection}:_counter");
        output_db
            .insert(counter_key.as_bytes(), last_id.to_string().as_bytes())
            .map_err(|e| format!("write {counter_key}: {e}"))?;
    }

    let prefix = b"note_emb:";
    let mut end = prefix.to_vec();
    end.push(0xff);
    for (key, value) in source_db
        .salvage_range(prefix, &end)
        .map_err(|e| format!("salvage embeddings: {e}"))?
    {
        output_db
            .insert(&key, &value)
            .map_err(|e| format!("copy embedding: {e}"))?;
        embeddings += 1;
    }

    output_db
        .commit()
        .map_err(|e| format!("commit repaired database: {e}"))?;
    println!("repaired database: {records} records, {embeddings} embeddings");
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn repairs_numeric_records_without_copying_sessions() {
        let base = std::env::temp_dir().join(format!(
            "notes-repair-{}-{}",
            std::process::id(),
            std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_nanos()
        ));
        let source = base.with_extension("source.db");
        let output = base.with_extension("output.db");
        let mut db = BTree::open(&source).unwrap();
        db.insert(b"notes:_counter", b"18446744073709551614")
            .unwrap();
        let mut note_key = b"notes:".to_vec();
        note_key.extend_from_slice(&4_u64.to_be_bytes());
        db.insert(&note_key, br#"{"id":4,"title":"Keep"}"#).unwrap();
        db.insert(b"note_emb:4", &[1, 2, 3]).unwrap();
        db.insert(b"_session:discard", b"private").unwrap();
        db.commit().unwrap();
        drop(db);

        run(source.to_str().unwrap(), output.to_str().unwrap()).unwrap();

        let mut repaired = BTree::open(&output).unwrap();
        assert!(repaired.get(&note_key).unwrap().is_some());
        assert_eq!(repaired.get(b"note_emb:4").unwrap(), Some(vec![1, 2, 3]));
        assert_eq!(
            repaired.get(b"notes:_counter").unwrap(),
            Some(b"4".to_vec())
        );
        assert!(repaired.get(b"_session:discard").unwrap().is_none());
    }
}