Merge pull request #18 from AIBlockOfficial/sophie/add-docstrings-mon…

…go-redis

Add comprehensive docstrings to MongoDB and Redis modules

src/db/mongo_db.rs

@@ -14,6 +14,11 @@ pub struct MongoDbIndex {
     pub coll_name: String,
 }
 
+/// Manages connections to MongoDB, handling connection pooling and database operations.
+///
+/// This struct encapsulates a MongoDB client and index information, providing methods
+/// for performing CRUD operations. It uses the MongoDB driver's built-in connection pool
+/// for efficient database interactions.
 #[derive(Debug, Clone)]
 pub struct MongoDbConn {
     pub client: Client,
@@ -45,8 +50,22 @@ impl MongoDbConn {
     }
 }
 
+/// Implements the `KvStoreConnection` trait for MongoDB, enabling key-value operations.
+///
+/// Methods interact with MongoDB documents where each key corresponds to a document ID.
+/// Values are stored in a `data` map within the document, supporting nested key-value storage.
 #[async_trait]
 impl KvStoreConnection for MongoDbConn {
+    /// Initializes a connection to the MongoDB server using the provided URI.
+    ///
+    /// # Arguments
+    /// * `url` - MongoDB connection URI (e.g., `mongodb://user:pass@host:port/dbname?options`).
+    ///
+    /// # Connection Pooling
+    /// The client utilizes connection pooling as configured via the URI options.
+    ///
+    /// # Errors
+    /// Returns an error if the URI is invalid or if connecting to the server fails.
     async fn init(url: &str) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
         // Tracing
         let span = span!(Level::TRACE, "MongoDbConn::init");
@@ -74,6 +93,15 @@ impl KvStoreConnection for MongoDbConn {
         Ok(MongoDbConn { client, index })
     }
 
+    /// Inserts or updates a document in the specified collection.
+    ///
+    /// # Arguments
+    /// * `key` - The document's unique `_id` field value.
+    /// * `value_id` - Key within the document's `data` map to store the value.
+    /// * `value` - Data to store, serialized into BSON.
+    ///
+    /// The value is serialized using `bson::to_bson`, which may return an error
+    /// for unsupported types. Existing entries with the same `value_id` are overwritten.
     async fn set_data<T: Serialize + std::marker::Send + DeserializeOwned>(
         &mut self,
         key: &str,
@@ -136,6 +164,14 @@ impl KvStoreConnection for MongoDbConn {
         Ok(())
     }
 
+    /// Retrieves data from the specified collection's document.
+    ///
+    /// # Arguments
+    /// * `key` - The document's `_id` field value to query.
+    /// * `value_id` - Optional specific entry within the document's `data` map.
+    ///
+    /// Returns `None` if the document or entry doesn't exist. Data is deserialized
+    /// from BSON, which may fail if the stored format doesn't match `T`.
     async fn get_data<T: DeserializeOwned + Clone>(
         &mut self,
         key: &str,
@@ -183,6 +219,16 @@ impl KvStoreConnection for MongoDbConn {
         return Ok(None);
     }
 
+    /// Stores data with an expiration time using MongoDB's TTL feature.
+    ///
+    /// # Arguments
+    /// * `key` - Document's `_id` field value.
+    /// * `value_id` - Key within the document's `data` map.
+    /// * `value` - Data to store, serialized to BSON.
+    /// * `seconds` - TTL in seconds until document expiration.
+    ///
+    /// Requires a TTL index on the `expiry` field. The document is automatically
+    /// removed by MongoDB after `seconds` seconds.
     async fn set_data_with_expiry<T: Serialize + std::marker::Send + DeserializeOwned>(
         &mut self,
         key: &str,
@@ -242,6 +288,14 @@ impl KvStoreConnection for MongoDbConn {
         Ok(())
     }
 
+    /// Deletes data from the specified collection.
+    ///
+    /// # Arguments
+    /// * `key` - Document's `_id` field to delete.
+    /// * `value_id` - If provided, removes only this entry from the `data` map.
+    ///
+    /// When `value_id` is specified, only that entry is removed. If `None`,
+    /// the entire document is deleted.
     async fn del_data(
         &mut self,
         key: &str,
@@ -298,4 +352,4 @@ impl KvStoreConnection for MongoDbConn {
 
         Ok(())
     }
-}
+}

src/db/redis_cache.rs

@@ -6,13 +6,20 @@ use redis::{aio::ConnectionManager, AsyncCommands};
 use serde::{de::DeserializeOwned, Serialize};
 use tracing::{event, span, Level};
 
+/// A Redis-backed cache connection handling key-value storage and connection lifecycle.
+///
+/// Manages connections to Redis, providing methods to set, retrieve, and delete
+/// serialized data. Automatically handles connection management and expiration.
 #[derive(Clone)]
 pub struct RedisCacheConn {
     pub connection: ConnectionManager,
 }
 
 #[async_trait]
 impl CacheHandler for RedisCacheConn {
+    /// Implements cache expiration functionality for Redis keys.
+    ///
+    /// Provides atomic expiration commands through the Redis EXPIRE interface.
     async fn expire_entry(
         &mut self,
         key: &str,
@@ -23,6 +30,10 @@ impl CacheHandler for RedisCacheConn {
     }
 }
 
+/// Core connection handler implementing key-value store operations.
+///
+/// Manages Redis connections and implements CRUD operations with JSON serialization.
+/// Handles connection lifecycle through `init` method.
 #[async_trait]
 impl KvStoreConnection for RedisCacheConn {
     async fn init(url: &str) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
@@ -34,6 +45,19 @@ impl KvStoreConnection for RedisCacheConn {
         })
     }
 
+    /// Stores a value in the cache without setting expiration.
+    ///
+    /// # Arguments
+    /// * `key` - Redis key for the hashmap
+    /// * `value_id` - Field within the hashmap
+    /// * `value` - Serializable value to store
+    ///
+    /// # Serialization
+    /// Values are serialized to JSON using `serde_json`. Ensure `T` implements
+    /// `Serialize` and `DeserializeOwned`.
+    ///
+    /// # Errors
+    /// Returns errors on connection failures, serialization issues, or Redis operation failures.
     async fn set_data<T: Serialize + DeserializeOwned + Send>(
         &mut self,
         key: &str,
@@ -43,14 +67,12 @@ impl KvStoreConnection for RedisCacheConn {
         let exists: bool = self.connection.exists(key).await?;
 
         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 {
             HashMap::new()
         };
 
-        // Append the new data to the vec
         mapping.insert(value_id.to_string(), value);
 
         let serialized = serde_json::to_string(&mapping)?;
@@ -59,39 +81,50 @@ impl KvStoreConnection for RedisCacheConn {
         Ok(())
     }
 
+    /// Stores a value with expiration time (TTL).
+    ///
+    /// # Arguments
+    /// * `key` - Redis key for the hashmap
+    /// * `value_id` - Field within the hashmap
+    /// * `value` - Serializable value to store
+    /// * `seconds` - TTL in seconds for automatic expiration
+    ///
+    /// The TTL applies to the entire Redis key. Subsequent updates to the hashmap
+    /// will maintain the TTL unless explicitly modified.
     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 mapping: HashMap<String, T> = if exists {
-            // Get the existing data
             let data: String = self.connection.get(key).await?;
             serde_json::from_str(&data)?
         } else {
             HashMap::new()
         };
 
-        // 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(&mapping)?;
-
-        // Set the data back to Redis
         self.connection.set(key, serialized).await?;
-
-        // Set the expiry time for the key
         self.connection.expire(key, seconds).await?;
 
         Ok(())
     }
 
+    /// Deletes cache entries with configurable scope.
+    ///
+    /// When `value_id` is provided:
+    /// - Performs non-atomic read-modify-write operation on the hashmap
+    /// - May impact performance with large datasets due to full value retrieval
+    ///
+    /// When `value_id` is `None`:
+    /// - Deletes entire key atomically via Redis DEL command
+    /// - Optimal for bulk deletion operations
     async fn del_data(
         &mut self,
         key: &str,
@@ -113,6 +146,14 @@ impl KvStoreConnection for RedisCacheConn {
         Ok(())
     }
 
+    /// Retrieves data from the cache.
+    ///
+    /// # Arguments
+    /// * `key` - Redis key to retrieve
+    /// * `value_id` - Optional specific field to extract from hashmap
+    ///
+    /// Returns `None` if key doesn't exist or specific value_id not found.
+    /// Deserialization errors and connection issues will return error variants.
     async fn get_data<T: Clone + DeserializeOwned>(
         &mut self,
         key: &str,
@@ -121,11 +162,9 @@ impl KvStoreConnection for RedisCacheConn {
         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 mapping: HashMap<String, T> = serde_json::from_str(&data)?;
 
@@ -136,7 +175,6 @@ impl KvStoreConnection for RedisCacheConn {
                     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}"
@@ -150,4 +188,4 @@ impl KvStoreConnection for RedisCacheConn {
 
         Ok(None)
     }
-}
+}