Added critical logging

pyeudiw/storage/base_cache.py

@@ -1,9 +1,16 @@
+from enum import Enum
 from typing import Callable
 
+class RetrieveStatus(Enum):
+    RETRIEVED = 0
+    ADDED = 1
 
 class BaseCache():
-    def try_retrieve(self, object_name: str, on_not_found: Callable[[], str]) -> dict:
+    def try_retrieve(self, object_name: str, on_not_found: Callable[[], str]) -> tuple[dict, RetrieveStatus]:
         raise NotImplementedError()
 
     def overwrite(self, object_name: str, value_gen_fn: Callable[[], str]) -> dict:
         raise NotImplementedError()
+
+    def set(self, data: dict) -> dict:
+        raise NotImplementedError()

pyeudiw/storage/base_storage.py

@@ -2,8 +2,8 @@ class BaseStorage(object):
     def init_session(self, document_id: str, dpop_proof: dict, attestation: dict):
         NotImplementedError()
 
-    def update_request_object(self, document_id: str, request_object: dict):
+    def update_request_object(self, document_id: str, nonce: str, state: str | None, request_object: dict):
         NotImplementedError()
 
-    def update_response_object(self, nonce: str, state: str, response_object: dict):
+    def update_response_object(self, nonce: str, state: str | None, response_object: dict):
         NotImplementedError()

pyeudiw/storage/db_engine.py

@@ -1,23 +1,27 @@
 import uuid
+import logging
 import importlib
-from pyeudiw.storage.base_cache import BaseCache
+from typing import Callable
+from pyeudiw.storage.base_cache import BaseCache, RetrieveStatus
 from pyeudiw.storage.base_storage import BaseStorage
 
+logger = logging.getLogger("openid4vp.storage.db")
+
 class DBEngine():
     def __init__(self, config: dict):
-        self.cache = []
+        self.caches = []
         self.storages = []
 
         for db_name, db_conf in config.items():
             storage_instance, cache_instance = self._handle_instance(db_conf)
 
             if storage_instance:
-                self.storages.append(storage_instance)
+                self.storages.append((db_name, storage_instance))
 
             if cache_instance:
-                self.cache.append(cache_instance)
+                self.caches.append((db_name, cache_instance))
 
-    def _handle_instance(instance: dict) -> dict[BaseStorage | None, BaseCache | None]:
+    def _handle_instance(self, instance: dict) -> dict[BaseStorage | None, BaseCache | None]:
         cache_conf = instance.get("cache", None)
         storage_conf = instance.get("storage", None)
 
@@ -26,28 +30,98 @@ def _handle_instance(instance: dict) -> dict[BaseStorage | None, BaseCache | Non
             module = importlib.import_module(storage_conf["module"])
             instance_class = getattr(module, storage_conf["class"])
 
-            storage_instance = instance_class(storage_conf["config"])
+            storage_instance = instance_class(**storage_conf["init_params"])
 
         cache_instance = None
         if cache_conf:
             module = importlib.import_module(cache_conf["module"])
             instance_class = getattr(module, cache_conf["class"])
 
-            cache_instance = instance_class(cache_conf["config"])
+            cache_instance = instance_class(**cache_conf["init_params"])
 
         return storage_instance, cache_instance
 
-    def init_session(self, dpop_proof: dict, attestation: dict):
+    def init_session(self, dpop_proof: dict, attestation: dict) -> str:
         document_id = str(uuid.uuid4())
-        for storage in self.storages:
-            storage.init_session(dpop_proof, attestation)
+        for db_name, storage in self.storages:
+            try:
+                storage.init_session(document_id, dpop_proof, attestation)
+            except Exception as e:
+                logger.critical(f"Error {str(e)}")
+                logger.critical(f"Cannot write document with id {document_id} on {db_name}")
 
         return document_id
 
-    def update_request_object(self, document_id: str, request_object: dict):
-        for storage in self.storages:
-            storage.update_request_object(document_id, request_object)
+    def update_request_object(self, document_id: str, nonce: str, state: str | None, request_object: dict) -> tuple[str, str, int]:
+        replica_count = 0
+        for db_name, storage in self.storages:
+            try:
+                storage.update_request_object(document_id, nonce, state, request_object)        
+                replica_count += 1
+            except Exception as e:
+                logger.critical(f"Error {str(e)}")
+                logger.critical(f"Cannot update document with id {document_id} on {db_name}")
+
+        if replica_count == 0:
+            raise Exception(f"Cannot update document {document_id} on any instance")
+
+        return nonce, state, replica_count
+
+    def update_response_object(self, nonce: str, state: str, response_object: dict) -> int:
+        replica_count = 0
+        for db_name, storage in self.storages:
+            try:
+                storage.update_response_object(nonce, state, response_object)
+                replica_count += 1
+            except Exception as e:
+                logger.critical(f"Error {str(e)}")
+                logger.critical(f"Cannot update document with nonce {nonce} and state {state} on {db_name}")
+
+        if replica_count == 0:
+            raise Exception(f"Cannot update document with state {state} and nonce {nonce} on any instance")
+
+        return replica_count
+
+    def _cache_try_retrieve(self, object_name: str, on_not_found: Callable[[], str]) -> tuple[dict, RetrieveStatus, int]:
+        for i, cache in enumerate(self.caches):
+            try:
+                cache_object, status = cache.try_retrieve(object_name, on_not_found)
+                return cache_object, status, i
+            except Exception as e:
+                logger.critical("Cannot retrieve or write cache object with identifier {object_name} on database {db_name}")
+        raise ConnectionRefusedError("Cannot write cache object on any instance")
 
-    def update_response_object(self, nonce: str, state: str, response_object: dict):
-        for storage in self.storages:
-            storage.update_response_object(nonce, state, response_object)
+    def try_retrieve(self, object_name: str, on_not_found: Callable[[], str]) -> dict:
+        istances_len = len(self.caches)
+
+        # if no cache instance exist return the object
+        if istances_len == 0:
+            return on_not_found()
+
+        # if almost one cache instance exist try to retrieve
+        cache_object, status, idx = self._cache_try_retrieve(object_name, on_not_found)
+
+        # if the status is retrieved return the object
+        if status == RetrieveStatus.RETRIEVED:
+            return cache_object
+
+        # else try replicate the data on all the other istances
+        replica_instances = self.caches[:idx] + self.caches[idx + 1:] 
+
+        for cache_name, cache in replica_instances:
+            try:
+                cache.set(cache_object)
+            except Exception as e:
+                logger.critical("Cannot replicate cache object with identifier {object_name} on cache {cache_name}")
+
+        return cache_object
+
+    def overwrite(self, object_name: str, value_gen_fn: Callable[[], str]) -> dict:
+        for cache_name, cache in self.caches:
+            cache_object = None
+            try:
+                cache_object = cache.overwrite(object_name, value_gen_fn)
+            except Exception as e:
+                logger.critical("Cannot overwrite cache object with identifier {object_name} on cache {cache_name}")
+
+            return cache_object

pyeudiw/storage/mongo_cache.py

@@ -2,14 +2,13 @@
 from datetime import datetime
 from typing import Callable
 
-from .base_cache import BaseCache
-
+from pyeudiw.storage.base_cache import BaseCache, RetrieveStatus
 
 class MongoCache(BaseCache):
-    def __init__(self, storage_conf: dict, url: str, connection_params: dict = None) -> None:
+    def __init__(self, conf: dict, url: str, connection_params: dict = None) -> None:
         super().__init__()
 
-        self.storage_conf = storage_conf
+        self.storage_conf = conf
         self.url = url
         self.connection_params = connection_params
 
@@ -22,25 +21,29 @@ def _connect(self):
                 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):
+        creation_date = datetime.timestamp(datetime.now())
+        return {
+            "object_name": object_name,
+            "data": data,
+            "creation_date": creation_date
+        }
+
 
-    def try_retrieve(self, object_name: str, on_not_found: Callable[[], str]) -> dict:
+    def try_retrieve(self, object_name: str, on_not_found: Callable[[], str]) -> tuple[dict, RetrieveStatus]:
         self._connect()
 
         query = {"object_name": object_name}
 
         cache_object = self.collection.find_one(query)
 
         if cache_object is None:
-            creation_date = datetime.timestamp(datetime.now())
-            cache_object = {
-                "object_name": object_name,
-                "data": on_not_found(),
-                "creation_date": creation_date
-            }
-
+            cache_object = self._gen_cache_object(object_name, on_not_found())
             self.collection.insert_one(cache_object)
+            return cache_object, RetrieveStatus.ADDED
 
-        return cache_object
+        return cache_object, RetrieveStatus.RETRIEVED
 
     def overwrite(self, object_name: str, value_gen_fn: Callable[[], str]) -> dict:
         self._connect()
@@ -64,3 +67,8 @@ def overwrite(self, object_name: str, value_gen_fn: Callable[[], str]) -> dict:
         })
 
         return cache_object
+
+    def set(self, data: dict) -> dict:
+        self._connect()
+
+        return self.collection.insert_one(data)

pyeudiw/storage/mongo_storage.py

@@ -5,10 +5,10 @@
 
 
 class MongoStorage(BaseStorage):
-    def __init__(self, storage_conf: dict, url: str, connection_params: dict = None) -> None:
+    def __init__(self, conf: dict, url: str, connection_params: dict = None) -> None:
         super().__init__()
 
-        self.storage_conf = storage_conf
+        self.storage_conf = conf
         self.url = url
         self.connection_params = connection_params
 
@@ -33,7 +33,7 @@ def _retrieve_document_by_id(self, document_id: str) -> dict:
 
         return document
 
-    def _retrieve_document_by_nonce_state(self, nonce: str, state: str) -> dict:
+    def _retrieve_document_by_nonce_state(self, nonce: str | None, state: str) -> dict:
         self._connect()
 
         query = {"state": state, "nonce": nonce}
@@ -46,7 +46,7 @@ def _retrieve_document_by_nonce_state(self, nonce: str, state: str) -> dict:
 
         return document
 
-    def init_session(self, document_id: str, dpop_proof: dict, attestation: dict):
+    def init_session(self, document_id: str, dpop_proof: dict, attestation: dict) -> str:
         creation_date = datetime.timestamp(datetime.now())
 
         entity = {
@@ -63,11 +63,9 @@ def init_session(self, document_id: str, dpop_proof: dict, attestation: dict):
 
         return document_id
 
-    def update_request_object(self, document_id: str, request_object: dict):
-        nonce = request_object["nonce"]
-        state = request_object["state"]
-
-        self._connect()
+    def update_request_object(self, document_id: str, nonce: str, state: str, request_object: dict) -> tuple[str, str, dict]:        
+        self._retrieve_document_by_id(document_id)
+
         documentStatus = self.collection.update_one(
             {"document_id": document_id},
             {
@@ -78,7 +76,6 @@ def update_request_object(self, document_id: str, request_object: dict):
                 }
             }
         )
-
         return nonce, state, documentStatus
 
     def update_response_object(self, nonce: str, state: str, response_object: dict):

pyeudiw/tests/storage/test_db_engine.py

@@ -0,0 +1,82 @@
+import uuid
+import pytest
+
+from pyeudiw.storage.db_engine import DBEngine
+
+conf = {
+    "mongo_db": {
+        "cache": {
+            "module": "pyeudiw.storage.mongo_cache",
+            "class": "MongoCache",
+            "init_params": {
+                "url": "mongodb://localhost:27017/",
+                "conf": {
+                    "db_name": "eudiw"
+                },
+                "connection_params": {}
+            }
+        },
+        "storage": {
+            "module": "pyeudiw.storage.mongo_storage",
+            "class": "MongoStorage",
+            "init_params": {
+                "url": "mongodb://localhost:27017/",
+                "conf": {
+                    "db_name": "eudiw",
+                    "db_collection": "sessions"   
+                },
+                "connection_params": {}
+            }
+        }
+    }
+}
+
+class TestMongoDBEngine:
+    @pytest.fixture(autouse=True)
+    def create_engine_instance(self):
+        self.engine = DBEngine(conf)
+
+    @pytest.fixture(autouse=True)
+    def test_init_session(self):
+        dpop_proof = {"dpop": "dpop"}
+        attestation = {"attestation": "attestation"}
+
+        document_id = self.engine.init_session(dpop_proof, attestation)
+
+        assert document_id
+
+        self.document_id = document_id
+
+    @pytest.fixture(autouse=True)
+    def test_update_request_object(self):
+        self.nonce = str(uuid.uuid4())
+        self.state = str(uuid.uuid4())
+        request_object = {"request_object": "request_object"}
+
+        r_nonce, r_state, _ = self.engine.update_request_object(self.document_id, self.nonce, self.state, request_object)
+
+        assert self.nonce == r_nonce
+        assert self.state == r_state
+
+    def test_update_request_object_with_unexistent_id_object(self):
+        nonce = str(uuid.uuid4())
+        state = str(uuid.uuid4())
+        unx_document_id = str(uuid.uuid4())
+        request_object = {"request_object": "request_object"}
+
+        try:
+            self.engine.update_request_object(unx_document_id, nonce, state, request_object)
+        except:
+            return
+
+    def test_update_response_object(self):
+        response_object = {"response_object": "response_object"}
+        self.engine.update_response_object(self.nonce, self.state, response_object)
+
+    def test_update_response_object_unexistent_id_object(self):
+        response_object = {"response_object": "response_object"}
+
+        try:
+            replica_count = self.engine.update_response_object(str(uuid.uuid4()), str(uuid.uuid4()), response_object)
+        except:
+            return

pyeudiw/tests/storage/test_mongo_cache.py

@@ -17,7 +17,7 @@ def test_try_retrieve(self):
         object_name = str(uuid.uuid4())
         data = str(uuid.uuid4())
 
-        obj = self.cache.try_retrieve(object_name, lambda: data)
+        obj, _ = self.cache.try_retrieve(object_name, lambda: data)
 
         assert obj
         assert obj["object_name"] == object_name
@@ -34,7 +34,7 @@ def test_overwrite(self):
         object_name = str(uuid.uuid4())
         data = str(uuid.uuid4())
 
-        obj = self.cache.try_retrieve(object_name, lambda: data)
+        obj, _ = self.cache.try_retrieve(object_name, lambda: data)
 
         data_updated = str(uuid.uuid4())