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
7#[allow(unused_imports)]
8use crate::engine::MemStore;
9
10/// Handles disk I/O for the [`MemStore`].
11/// 
12/// Persistence uses an atomic "write-then-rename" strategy to ensure data integrity.
13/// Each persona is stored in its own `.json` file.
14pub struct Persistence {
15    data_dir: PathBuf,
16}
17
18impl Persistence {
19    /// Initializes a new `Persistence` handler in the specified directory.
20    /// 
21    /// If the directory does not exist, it will be created.
22    pub fn new<P: AsRef<Path>>(dir: P) -> Result<Self> {
23        let dir = dir.as_ref().to_path_buf();
24        if !dir.exists() {
25            fs::create_dir_all(&dir)?;
26        }
27        Ok(Self { data_dir: dir })
28    }
29
30    /// Writes a single persona's data to a JSON file atomically.
31    /// 
32    /// This method writes to a temporary file first and then renames it to the
33    /// final destination, preventing file corruption during power failures.
34    pub fn save_persona(&self, persona_id: &str, data: &HashMap<String, HashMap<String, serde_json::Value>>) -> Result<()> {
35        let file_path = self.data_dir.join(format!("{}.json", persona_id));
36        let temp_path = file_path.with_extension("json.tmp");
37
38        let bytes = serde_json::to_vec_pretty(data)?;
39        
40        fs::write(&temp_path, bytes)?;
41        fs::rename(&temp_path, &file_path)?;
42
43        Ok(())
44    }
45
46    /// Loads all persona data found in the data directory.
47    /// 
48    /// Scans for all `.json` files in the `data_dir` and parses them into the
49    /// store's internal data structure.
50    pub fn load_all(&self) -> Result<HashMap<String, HashMap<String, HashMap<String, serde_json::Value>>>> {
51        let mut all_data = HashMap::new();
52
53        if !self.data_dir.exists() {
54            return Ok(all_data);
55        }
56
57        for entry in fs::read_dir(&self.data_dir)? {
58            let entry = entry?;
59            let path = entry.path();
60            
61            if path.extension().and_then(|s| s.to_str()) == Some("json") {
62                let persona_id = path.file_stem()
63                    .and_then(|s| s.to_str())
64                    .ok_or_else(|| Error::Internal("Invalid filename".to_string()))?
65                    .to_string();
66
67                let content = match fs::read(&path) {
68                    Ok(c) => c,
69                    Err(e) => {
70                        warn!("Could not read persona file {:?}: {}", path, e);
71                        continue;
72                    }
73                };
74
75                let persona_data: HashMap<String, HashMap<String, serde_json::Value>> = match serde_json::from_slice(&content) {
76                    Ok(d) => d,
77                    Err(e) => {
78                        warn!("Could not unmarshal persona data from {:?}: {}", path, e);
79                        continue;
80                    }
81                };
82
83                all_data.insert(persona_id, persona_data);
84            }
85        }
86
87        Ok(all_data)
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94    use tempfile::tempdir;
95    use serde_json::json;
96
97    #[test]
98    fn test_save_and_load_all() {
99        let dir = tempdir().unwrap();
100        let persistence = Persistence::new(dir.path()).unwrap();
101
102        let mut data = HashMap::new();
103        let mut app_data = HashMap::new();
104        app_data.insert("key1".to_string(), json!("value1"));
105        data.insert("app1".to_string(), app_data);
106
107        persistence.save_persona("p1", &data).unwrap();
108
109        let loaded = persistence.load_all().unwrap();
110        assert_eq!(loaded.len(), 1);
111        assert_eq!(loaded.get("p1").unwrap().get("app1").unwrap().get("key1").unwrap(), &json!("value1"));
112    }
113
114    #[test]
115    fn test_atomic_rename() {
116        let dir = tempdir().unwrap();
117        let persistence = Persistence::new(dir.path()).unwrap();
118
119        let mut data = HashMap::new();
120        let mut app_data = HashMap::new();
121        app_data.insert("key1".to_string(), json!("value1"));
122        data.insert("app1".to_string(), app_data);
123
124        persistence.save_persona("p1", &data).unwrap();
125
126        let file_path = dir.path().join("p1.json");
127        assert!(file_path.exists());
128        
129        let temp_path = dir.path().join("p1.json.tmp");
130        assert!(!temp_path.exists());
131    }
132
133    #[test]
134    fn test_go_compatibility() {
135        // Mock the Go test data structure
136        let go_json = r#"{
137  "test_app": {
138    "key_0": 0,
139    "key_1": "string_val"
140  }
141}"#;
142        let dir = tempdir().unwrap();
143        let file_path = dir.path().join("go_persona.json");
144        fs::write(&file_path, go_json).unwrap();
145
146        let persistence = Persistence::new(dir.path()).unwrap();
147        let loaded = persistence.load_all().unwrap();
148        
149        let persona = loaded.get("go_persona").unwrap();
150        let app = persona.get("test_app").unwrap();
151        assert_eq!(app.get("key_0").unwrap(), &json!(0));
152        assert_eq!(app.get("key_1").unwrap(), &json!("string_val"));
153    }
154}