fix

Signed-off-by: Pushkar Mishra <[email protected]>

fs-storage/examples/cli.rs

@@ -2,7 +2,6 @@ use anyhow::{Context, Result};
 use fs_storage::base_storage::BaseStorage;
 use fs_storage::file_storage::FileStorage;
 use serde_json::Value;
-use std::collections::BTreeMap;
 use std::env;
 use std::fs;
 use std::path::Path;
@@ -44,21 +43,22 @@ fn read_command(args: &[String], path: &str) -> Result<()> {
         vec![]
     };
 
-    let mut fs = FileStorage::new("cli".to_string(), Path::new(path));
-    let map: BTreeMap<String, String> =
-        fs.read_fs().context("Failed to read file")?;
+    let mut fs: FileStorage<String, String> =
+        FileStorage::new("cli".to_string(), Path::new(path));
 
+    let map = fs
+        .read_fs()
+        .expect("No Data is present on this path");
     if keys.is_empty() {
         for (key, value) in map {
             println!("{}: {}", key, value);
         }
-    } else {
-        for key in &keys {
-            if let Some(value) = map.get(key) {
-                println!("{}: {}", key, value);
-            } else {
-                eprintln!("Key '{}' not found", key);
-            }
+    }
+    for key in &keys {
+        if let Some(value) = fs.get(key) {
+            println!("{}: {}", key, value);
+        } else {
+            eprintln!("Key '{}' not found", key);
         }
     }
 
@@ -77,7 +77,8 @@ fn write_command(args: &[String], path: &str) -> Result<()> {
         .extension()
         .map_or(false, |ext| ext == "json");
 
-    let mut fs = FileStorage::new("cli".to_string(), Path::new(path));
+    let mut fs: FileStorage<String, String> =
+        FileStorage::new("cli".to_string(), Path::new(path));
     if content_json {
         let content =
             fs::read_to_string(content).context("Failed to read JSON file")?;
@@ -107,8 +108,5 @@ fn write_command(args: &[String], path: &str) -> Result<()> {
             }
         }
     }
-
-    fs.write_fs().context("Failed to write file")?;
-
     Ok(())
 }

fs-storage/src/base_storage.rs

@@ -5,7 +5,7 @@ use std::str::FromStr;
 
 use data_error::Result;
 
-pub trait BaseStorage<K, V>
+pub trait BaseStorage<K, V>: AsRef<BTreeMap<K, V>>
 where
     K: FromStr + Hash + Eq + Ord + Debug + Clone,
     V: Debug + Clone,
@@ -17,11 +17,6 @@ where
     /// Remove file at stored path
     fn erase(&self) -> Result<()>;
 
-    /// Get immutable BTreeMap
-    ///
-    /// This can be used to get most immutable BtreeMap related functions for free.
-    fn as_ref(&self) -> &BTreeMap<K, V>;
-
     /// Check if storage is updated
     ///
     /// This check can be used before reading the file.

fs-storage/src/file_storage.rs

@@ -24,17 +24,34 @@ pub struct FileStorage<K, V> {
 
 impl<K, V> FileStorage<K, V>
 where
-    K: serde::Serialize,
-    V: serde::Serialize,
+    K: FromStr
+        + Hash
+        + Eq
+        + Ord
+        + Debug
+        + Clone
+        + serde::Serialize
+        + serde::de::DeserializeOwned,
+    V: Debug + Clone + serde::Serialize + serde::de::DeserializeOwned,
 {
     /// Create a new file storage with a diagnostic label and file path
     pub fn new(label: String, path: &Path) -> Self {
-        Self {
+        let mut file_storage = Self {
             label,
             path: PathBuf::from(path),
             timestamp: SystemTime::now(),
             data: BTreeMap::new(),
-        }
+        };
+
+        // Load the data from the file
+        file_storage.data = match file_storage.read_fs() {
+            Ok(data) => data,
+            Err(err) => {
+                log::error!("Error reading storage file: {}", err);
+                BTreeMap::new()
+            }
+        };
+        file_storage
     }
 
     /// Verify the version stored in the file header
@@ -88,13 +105,17 @@ where
     fn set(&mut self, id: K, value: V) {
         self.data.insert(id, value);
         self.timestamp = std::time::SystemTime::now();
+        self.write_fs()
+            .expect("Failed to write data to disk");
     }
 
     fn remove(&mut self, id: &K) -> Result<()> {
         self.data.remove(id).ok_or_else(|| {
             ArklibError::Storage(self.label.clone(), "Key not found".to_owned())
         })?;
         self.timestamp = std::time::SystemTime::now();
+        self.write_fs()
+            .expect("Failed to remove data from disk");
         Ok(())
     }
 
@@ -104,10 +125,6 @@ where
         })
     }
 
-    fn as_ref(&self) -> &BTreeMap<K, V> {
-        &self.data
-    }
-
     fn is_storage_updated(&self) -> Result<bool> {
         let file_timestamp = fs::metadata(&self.path)?.modified()?;
         Ok(self.timestamp < file_timestamp)
@@ -177,6 +194,23 @@ where
     }
 }
 
+impl<K, V> AsRef<BTreeMap<K, V>> for FileStorage<K, V>
+where
+    K: FromStr
+        + Hash
+        + Eq
+        + Ord
+        + Debug
+        + Clone
+        + serde::Serialize
+        + serde::de::DeserializeOwned,
+    V: Debug + Clone + serde::Serialize + serde::de::DeserializeOwned,
+{
+    fn as_ref(&self) -> &BTreeMap<K, V> {
+        &self.data
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::collections::BTreeMap;
@@ -196,10 +230,6 @@ mod tests {
         file_storage.set("key1".to_string(), "value1".to_string());
         file_storage.set("key1".to_string(), "value2".to_string());
 
-        file_storage
-            .write_fs()
-            .expect("Failed to write data to disk");
-
         let data_read: BTreeMap<_, _> = file_storage
             .read_fs()
             .expect("Failed to read data from disk");
@@ -220,10 +250,6 @@ mod tests {
         file_storage.set("key1".to_string(), "value1".to_string());
         file_storage.set("key1".to_string(), "value2".to_string());
 
-        file_storage
-            .write_fs()
-            .expect("Failed to write data to disk");
-
         assert_eq!(storage_path.exists(), true);
 
         if let Err(err) = file_storage.erase() {