Merge pull request #7 from AIBlockOfficial/restructure_data_entry

Restructure data entry for all routes
AIBlockOfficial · Jun 26, 2024 · 14c8087 · 14c8087
2 parents fab2aac + 150089d
commit 14c8087
Show file tree

Hide file tree

Showing 5 changed files with 138 additions and 41 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "valence_core"
-version = "0.1.5"
+version = "0.1.6"
 edition = "2021"
 license = "MIT"
 keywords = ["blockchain", "L2", "peer-to-peer", "P2P"]

diff --git a/src/api/utils.rs b/src/api/utils.rs
@@ -47,7 +47,7 @@ pub fn get_cors() -> warp::cors::Builder {
 }
 
 /// Easy and simple DELETE CORS
-pub fn delete_cors() -> warp::cors::Builder {
+pub fn del_cors() -> warp::cors::Builder {
     cors_builder(vec!["DELETE", "OPTIONS"])
 }
 

diff --git a/src/db/handler.rs b/src/db/handler.rs
@@ -1,5 +1,6 @@
 use async_trait::async_trait;
 use serde::{de::DeserializeOwned, Serialize};
+use std::collections::HashMap;
 
 /// Trait for a key-value data store connection
 #[async_trait]
@@ -18,10 +19,12 @@ pub trait KvStoreConnection {
     /// ### Arguments
     ///
     /// * `key` - Key of the data entry to set
+    /// * `value_id` - ID of the value to set
     /// * `value` - Value of the data entry to set
     async fn set_data<T: Serialize + Send + DeserializeOwned>(
         &mut self,
         key: &str,
+        value_id: &str,
         value: T,
     ) -> Result<(), Box<dyn std::error::Error + Send + Sync>>;
 
@@ -30,11 +33,13 @@ pub trait KvStoreConnection {
     /// ### Arguments
     ///
     /// * `key` - Key of the data entry to set
+    /// * `value_id` - ID of the value to set
     /// * `value` - Value of the data entry to set
     /// * `seconds` - Number of seconds to expire the data entry in
     async fn set_data_with_expiry<T: Serialize + DeserializeOwned + Send>(
         &mut self,
         key: &str,
+        value_id: &str,
         value: T,
         seconds: usize,
     ) -> Result<(), Box<dyn std::error::Error + Send + Sync>>;
@@ -44,20 +49,24 @@ pub trait KvStoreConnection {
     /// ### Arguments
     ///
     /// * `key` - Key of the data entry to delete
-    async fn delete_data(
+    /// * `value_id` - ID of the value to delete. If not provided, all values for the key are deleted
+    async fn del_data(
         &mut self,
         key: &str,
+        value_id: Option<&str>,
     ) -> Result<(), Box<dyn std::error::Error + Send + Sync>>;
 
     /// Gets a data entry from the cache
     ///
     /// ### Arguments
     ///
     /// * `key` - Key of the data entry to get
-    async fn get_data<T: DeserializeOwned>(
+    /// * `value_id` - ID of the value to get. If not provided, all values for the key are retrieved
+    async fn get_data<T: Clone + DeserializeOwned>(
         &mut self,
         key: &str,
-    ) -> Result<Option<Vec<T>>, Box<dyn std::error::Error + Send + Sync>>;
+        value_id: Option<&str>,
+    ) -> Result<Option<HashMap<String, T>>, Box<dyn std::error::Error + Send + Sync>>;
 }
 
 #[async_trait]

diff --git a/src/db/mongo_db.rs b/src/db/mongo_db.rs
@@ -2,6 +2,7 @@ use async_trait::async_trait;
 use mongodb::bson::{doc, DateTime, Document};
 use mongodb::{options::ClientOptions, Client};
 use serde::{de::DeserializeOwned, Serialize};
+use std::collections::HashMap;
 use tracing::{event, span, trace, warn, Level};
 
 use super::handler::KvStoreConnection;
@@ -75,6 +76,7 @@ impl KvStoreConnection for MongoDbConn {
     async fn set_data<T: Serialize + std::marker::Send + DeserializeOwned>(
         &mut self,
         key: &str,
+        value_id: &str,
         value: T,
     ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
         // Tracing
@@ -90,22 +92,22 @@ impl KvStoreConnection for MongoDbConn {
         let filter = doc! { "_id": key };
         let existing_doc = collection.find_one(filter.clone(), None).await?;
 
-        let mut vec: Vec<T> = if let Some(doc) = existing_doc {
+        let mut mapping: HashMap<String, T> = if let Some(doc) = existing_doc {
             if doc.contains_key("data") {
                 // Deserialize the existing data
                 mongodb::bson::from_bson(doc.get("data").unwrap().clone())?
             } else {
-                Vec::new()
+                HashMap::new()
             }
         } else {
-            Vec::new()
+            HashMap::new()
         };
 
         // Append the new data to the vec
-        vec.push(value);
+        mapping.insert(value_id.to_string(), value);
 
         // Serialize the vec back to a BSON array
-        let serialized_vec = mongodb::bson::to_bson(&vec)?;
+        let serialized_vec = mongodb::bson::to_bson(&mapping)?;
 
         // Create or update the document
         let update = doc! {
@@ -135,6 +137,7 @@ impl KvStoreConnection for MongoDbConn {
     async fn set_data_with_expiry<T: Serialize + std::marker::Send + DeserializeOwned>(
         &mut self,
         key: &str,
+        value_id: &str,
         value: T,
         seconds: usize,
     ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
@@ -151,18 +154,18 @@ impl KvStoreConnection for MongoDbConn {
         let filter = doc! { "_id": key };
         let existing_doc = collection.find_one(filter.clone(), None).await?;
 
-        let mut vec: Vec<T> = if let Some(doc) = existing_doc {
+        let mut mapping: HashMap<String, T> = if let Some(doc) = existing_doc {
             // Deserialize the existing data
             mongodb::bson::from_bson(doc.get("data").unwrap().clone())?
         } else {
-            Vec::new()
+            HashMap::new()
         };
 
         // Append the new data to the vec
-        vec.push(value);
+        mapping.insert(value_id.to_string(), value);
 
         // Serialize the vec back to a BSON array
-        let serialized_vec = mongodb::bson::to_bson(&vec)?;
+        let serialized_vec = mongodb::bson::to_bson(&mapping)?;
 
         // Calculate the expiry time
         let expiry_time = (seconds * 1000) as i64;
@@ -190,36 +193,68 @@ impl KvStoreConnection for MongoDbConn {
         Ok(())
     }
 
-    async fn delete_data(
+    async fn del_data(
         &mut self,
         key: &str,
+        value_id: Option<&str>,
     ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
         // Tracing
-        let span = span!(Level::TRACE, "MongoDbConn::delete_data");
+        let span = span!(Level::TRACE, "MongoDbConn::del_data");
         let _enter = span.enter();
 
         let collection = self
             .client
             .database(&self.index.db_name)
             .collection::<Document>(&self.index.coll_name);
 
+        // Build the filter based on the key
         let filter = doc! { "_id": key };
-        match collection.delete_one(filter, None).await {
-            Ok(_) => (),
-            Err(e) => {
-                event!(Level::ERROR, "Failed to delete data with error: {e}");
-            }
-        };
 
-        trace!("Data deleted successfully");
+        // If value_id is provided, we need to fetch the document and update it
+        if let Some(value_id) = value_id {
+            let update = doc! {
+                "$unset": {
+                    &format!("data.{}", value_id): ""
+                }
+            };
+
+            match collection.find_one_and_update(filter, update, None).await {
+                Ok(result) => {
+                    if let Some(_) = result {
+                        // Document was found and updated, log success or handle as needed
+                        trace!("Data updated successfully");
+                    } else {
+                        // Document not found
+                        event!(Level::ERROR, "Document not found for key: {}", key);
+                    }
+                }
+                Err(e) => {
+                    // Handle error from MongoDB
+                    event!(Level::ERROR, "Failed to update data with error: {:?}", e);
+                    return Err(Box::new(e));
+                }
+            }
+        } else {
+            // value_id is None, so delete the entire document
+            match collection.delete_one(filter.clone(), None).await {
+                Ok(_) => {
+                    trace!("Data deleted successfully");
+                }
+                Err(e) => {
+                    event!(Level::ERROR, "Failed to delete data with error: {:?}", e);
+                    return Err(Box::new(e));
+                }
+            };
+        }
 
         Ok(())
     }
 
-    async fn get_data<T: DeserializeOwned>(
+    async fn get_data<T: Clone + DeserializeOwned>(
         &mut self,
         key: &str,
-    ) -> Result<Option<Vec<T>>, Box<dyn std::error::Error + Send + Sync>> {
+        value_id: Option<&str>,
+    ) -> Result<Option<HashMap<String, T>>, Box<dyn std::error::Error + Send + Sync>> {
         // Tracing
         let span = span!(Level::TRACE, "MongoDbConn::get_data");
         let _enter = span.enter();
@@ -241,8 +276,22 @@ impl KvStoreConnection for MongoDbConn {
 
         if let Some(doc) = doc_find {
             // Deserialize the existing data
-            let vec: Vec<T> = mongodb::bson::from_bson(doc.get("data").unwrap().clone())?;
-            return Ok(Some(vec));
+            let mapping: HashMap<String, T> =
+                mongodb::bson::from_bson(doc.get("data").unwrap().clone())?;
+
+            if let Some(id) = value_id {
+                // If value_id is provided, return only the value with the given ID
+                if let Some(value) = mapping.get(id) {
+                    let mut result: HashMap<String, T> = HashMap::new();
+                    result.insert(id.to_string(), value.clone());
+                    return Ok(Some(result));
+                } else {
+                    // Value with the given ID not found
+                    event!(Level::ERROR, "Value with ID {id} not found for key {key}");
+                    return Ok(None);
+                }
+            }
+            return Ok(Some(mapping));
         }
 
         warn!("Data unsuccessfully deserialized");

diff --git a/src/db/redis_cache.rs b/src/db/redis_cache.rs
@@ -1,7 +1,10 @@
+use std::collections::HashMap;
+
 use crate::db::handler::{CacheHandler, KvStoreConnection};
 use async_trait::async_trait;
 use redis::{aio::ConnectionManager, AsyncCommands};
 use serde::{de::DeserializeOwned, Serialize};
+use tracing::{event, span, Level};
 
 #[derive(Clone)]
 pub struct RedisCacheConn {
@@ -34,22 +37,23 @@ impl KvStoreConnection for RedisCacheConn {
     async fn set_data<T: Serialize + DeserializeOwned + Send>(
         &mut self,
         key: &str,
+        value_id: &str,
         value: T,
     ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
         let exists: bool = self.connection.exists(key).await?;
 
-        let mut vec: Vec<T> = if exists {
+        let mut mapping: HashMap<String, T> = if exists {
             // Get the existing data
             let data: String = self.connection.get(key).await?;
             serde_json::from_str(&data)?
         } else {
-            Vec::new()
+            HashMap::new()
         };
 
         // Append the new data to the vec
-        vec.push(value);
+        mapping.insert(value_id.to_string(), value);
 
-        let serialized = serde_json::to_string(&vec)?;
+        let serialized = serde_json::to_string(&mapping)?;
         self.connection.set(key, serialized).await?;
 
         Ok(())
@@ -58,25 +62,26 @@ impl KvStoreConnection for RedisCacheConn {
     async fn set_data_with_expiry<T: Serialize + DeserializeOwned + Send>(
         &mut self,
         key: &str,
+        value_id: &str,
         value: T,
         seconds: usize,
     ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
         // Check if the key exists
         let exists: bool = self.connection.exists(key).await?;
 
-        let mut vec: Vec<T> = if exists {
+        let mut mapping: HashMap<String, T> = if exists {
             // Get the existing data
             let data: String = self.connection.get(key).await?;
             serde_json::from_str(&data)?
         } else {
-            Vec::new()
+            HashMap::new()
         };
 
-        // Append the new data to the vec
-        vec.push(value);
+        // Append the new data to the hashmap
+        mapping.insert(value_id.to_string(), value);
 
         // Serialize the vec back to a string
-        let serialized = serde_json::to_string(&vec)?;
+        let serialized = serde_json::to_string(&mapping)?;
 
         // Set the data back to Redis
         self.connection.set(key, serialized).await?;
@@ -87,26 +92,60 @@ impl KvStoreConnection for RedisCacheConn {
         Ok(())
     }
 
-    async fn delete_data(
+    async fn del_data(
         &mut self,
         key: &str,
+        value_id: Option<&str>,
     ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+        if let Some(value_id) = value_id {
+            let exists: bool = self.connection.exists(key).await?;
+
+            if exists {
+                let mut mapping: HashMap<String, String> = self.get_data(key, None).await?.unwrap();
+                mapping.remove(value_id);
+                let serialized = serde_json::to_string(&mapping)?;
+                self.connection.set(key, serialized).await?;
+            }
+            return Ok(());
+        }
+
         let _: () = self.connection.del(key).await?;
         Ok(())
     }
 
-    async fn get_data<T: DeserializeOwned>(
+    async fn get_data<T: Clone + DeserializeOwned>(
         &mut self,
         key: &str,
-    ) -> Result<Option<Vec<T>>, Box<dyn std::error::Error + Send + Sync>> {
+        value_id: Option<&str>,
+    ) -> Result<Option<HashMap<String, T>>, Box<dyn std::error::Error + Send + Sync>> {
+        let span = span!(Level::TRACE, "MongoDbConn::get_data");
+        let _enter = span.enter();
+
         // Check if the key exists
         let exists: bool = self.connection.exists(key).await?;
 
         if exists {
             // Get the existing data
             let data: String = self.connection.get(key).await?;
-            let vec: Vec<T> = serde_json::from_str(&data)?;
-            return Ok(Some(vec));
+            let mapping: HashMap<String, T> = serde_json::from_str(&data)?;
+
+            if let Some(value_id) = value_id {
+                let value = mapping.get(value_id);
+                if let Some(value) = value {
+                    let mut new_mapping: HashMap<String, T> = HashMap::new();
+                    new_mapping.insert(value_id.to_string(), value.clone());
+                    return Ok(Some(new_mapping));
+                } else {
+                    // Value with the given ID not found
+                    event!(
+                        Level::ERROR,
+                        "Value with ID {value_id} not found for key {key}"
+                    );
+                    return Ok(None);
+                }
+            }
+
+            return Ok(Some(mapping));
         }
 
         Ok(None)