From 3c084a583e6b9c3b279718c90738a39bdbc1d7d6 Mon Sep 17 00:00:00 2001
From: PascalDR <akira.zeta@gmail.com>
Date: Wed, 20 Dec 2023 13:29:28 +0100
Subject: [PATCH] feat: documentation and refactoring

---
 pyeudiw/storage/mongo_cache.py | 61 +++++++++++++++++++++++++---------
 1 file changed, 46 insertions(+), 15 deletions(-)

diff --git a/pyeudiw/storage/mongo_cache.py b/pyeudiw/storage/mongo_cache.py
index f44dc859..12a877f8 100644
--- a/pyeudiw/storage/mongo_cache.py
+++ b/pyeudiw/storage/mongo_cache.py
@@ -4,32 +4,40 @@
 import pymongo
 
 from pyeudiw.storage.base_cache import BaseCache, RetrieveStatus
+from pymongo.collection import Collection
+from pymongo.mongo_client import MongoClient
+from pymongo.database import Database
 
 
 class MongoCache(BaseCache):
+    """
+    MongoDB cache implementation.
+    """
+
     def __init__(self, conf: dict, url: str, connection_params: dict = None) -> None:
+        """
+        Create a MongoCache istance.
+
+        :param conf: the configuration of the cache.
+        :type conf: dict
+        :param url: the url of the MongoDB server.
+        :type url: str
+        :param connection_params: the connection parameters.
+        :type connection_params: dict, optional
+        """
         super().__init__()
 
         self.storage_conf = conf
         self.url = url
         self.connection_params = connection_params
 
-        self.client = None
-        self.db = None
-
-    def _connect(self):
-        if not self.client or not self.client.server_info():
-            self.client = pymongo.MongoClient(
-                self.url, **self.connection_params)
-            self.db = getattr(self.client, self.storage_conf["db_name"])
-            self.collection = getattr(self.db, "cache_storage")
+        self.client: MongoClient = None
+        self.db: Database = None
+        self.collection: Collection = None
 
-    def _gen_cache_object(self, object_name: str, data: str):
-        return {
-            "object_name": object_name,
-            "data": data,
-            "creation_date": datetime.now().isoformat()
-        }
+    def close(self) -> None:
+        self._connect()
+        self.client.close()
 
     def try_retrieve(self, object_name: str, on_not_found: Callable[[], str]) -> tuple[dict, RetrieveStatus]:
         self._connect()
@@ -72,3 +80,26 @@ def set(self, data: dict) -> dict:
         self._connect()
 
         return self.collection.insert_one(data)
+    
+    def _connect(self) -> None:
+        if not self.client or not self.client.server_info():
+            self.client = pymongo.MongoClient(
+                self.url, **self.connection_params)
+            self.db = getattr(self.client, self.storage_conf["db_name"])
+            self.collection = getattr(self.db, "cache_storage")
+
+    def _gen_cache_object(self, object_name: str, data: str) -> dict:
+        """
+        Helper function to generate a cache object.
+
+        :param object_name: the name of the object.
+        :type object_name: str
+        :param data: the data to store.
+        :type data: str
+        """
+
+        return {
+            "object_name": object_name,
+            "data": data,
+            "creation_date": datetime.now().isoformat()
+        }