celerix_store/engine/
persistence.rs

1use std::collections::HashMap;
2use std::fs;
3use std::path::{Path, PathBuf};
4use crate::{Result, Error};
5use log::warn;
6
7pub struct Persistence {
8    data_dir: PathBuf,
9}
10
11impl Persistence {
12    pub fn new<P: AsRef<Path>>(dir: P) -> Result<Self> {
13        let dir = dir.as_ref().to_path_buf();
14        if !dir.exists() {
15            fs::create_dir_all(&dir)?;
16        }
17        Ok(Self { data_dir: dir })
18    }
19
20    pub fn save_persona(&self, persona_id: &str, data: &HashMap<String, HashMap<String, serde_json::Value>>) -> Result<()> {
21        let file_path = self.data_dir.join(format!("{}.json", persona_id));
22        let temp_path = file_path.with_extension("json.tmp");
23
24        let bytes = serde_json::to_vec_pretty(data)?;
25        
26        fs::write(&temp_path, bytes)?;
27        fs::rename(&temp_path, &file_path)?;
28
29        Ok(())
30    }
31
32    pub fn load_all(&self) -> Result<HashMap<String, HashMap<String, HashMap<String, serde_json::Value>>>> {
33        let mut all_data = HashMap::new();
34
35        if !self.data_dir.exists() {
36            return Ok(all_data);
37        }
38
39        for entry in fs::read_dir(&self.data_dir)? {
40            let entry = entry?;
41            let path = entry.path();
42            
43            if path.extension().and_then(|s| s.to_str()) == Some("json") {
44                let persona_id = path.file_stem()
45                    .and_then(|s| s.to_str())
46                    .ok_or_else(|| Error::Internal("Invalid filename".to_string()))?
47                    .to_string();
48
49                let content = match fs::read(&path) {
50                    Ok(c) => c,
51                    Err(e) => {
52                        warn!("Could not read persona file {:?}: {}", path, e);
53                        continue;
54                    }
55                };
56
57                let persona_data: HashMap<String, HashMap<String, serde_json::Value>> = match serde_json::from_slice(&content) {
58                    Ok(d) => d,
59                    Err(e) => {
60                        warn!("Could not unmarshal persona data from {:?}: {}", path, e);
61                        continue;
62                    }
63                };
64
65                all_data.insert(persona_id, persona_data);
66            }
67        }
68
69        Ok(all_data)
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76    use tempfile::tempdir;
77    use serde_json::json;
78
79    #[test]
80    fn test_save_and_load_all() {
81        let dir = tempdir().unwrap();
82        let persistence = Persistence::new(dir.path()).unwrap();
83
84        let mut data = HashMap::new();
85        let mut app_data = HashMap::new();
86        app_data.insert("key1".to_string(), json!("value1"));
87        data.insert("app1".to_string(), app_data);
88
89        persistence.save_persona("p1", &data).unwrap();
90
91        let loaded = persistence.load_all().unwrap();
92        assert_eq!(loaded.len(), 1);
93        assert_eq!(loaded.get("p1").unwrap().get("app1").unwrap().get("key1").unwrap(), &json!("value1"));
94    }
95
96    #[test]
97    fn test_atomic_rename() {
98        let dir = tempdir().unwrap();
99        let persistence = Persistence::new(dir.path()).unwrap();
100
101        let mut data = HashMap::new();
102        let mut app_data = HashMap::new();
103        app_data.insert("key1".to_string(), json!("value1"));
104        data.insert("app1".to_string(), app_data);
105
106        persistence.save_persona("p1", &data).unwrap();
107
108        let file_path = dir.path().join("p1.json");
109        assert!(file_path.exists());
110        
111        let temp_path = dir.path().join("p1.json.tmp");
112        assert!(!temp_path.exists());
113    }
114
115    #[test]
116    fn test_go_compatibility() {
117        // Mock the Go test data structure
118        let go_json = r#"{
119  "test_app": {
120    "key_0": 0,
121    "key_1": "string_val"
122  }
123}"#;
124        let dir = tempdir().unwrap();
125        let file_path = dir.path().join("go_persona.json");
126        fs::write(&file_path, go_json).unwrap();
127
128        let persistence = Persistence::new(dir.path()).unwrap();
129        let loaded = persistence.load_all().unwrap();
130        
131        let persona = loaded.get("go_persona").unwrap();
132        let app = persona.get("test_app").unwrap();
133        assert_eq!(app.get("key_0").unwrap(), &json!(0));
134        assert_eq!(app.get("key_1").unwrap(), &json!("string_val"));
135    }
136}