Merge branch 'AlboCode-feature/cache' into develop

core/cat/auth/connection.py

@@ -3,7 +3,7 @@
 #  to have a standard auth interface.
 
 from abc import ABC, abstractmethod
-from typing import Tuple
+from typing import Tuple, AsyncGenerator
 import asyncio
 from urllib.parse import urlencode
 
@@ -23,6 +23,7 @@
 from cat.looking_glass.stray_cat import StrayCat
 from cat.log import log
 
+
 class ConnectionAuth(ABC):
 
     def __init__(
@@ -36,7 +37,7 @@ def __init__(
     async def __call__(
         self,
         connection: HTTPConnection # Request | WebSocket,
-    ) -> StrayCat:
+    ) -> AsyncGenerator[StrayCat, None]:
 
         # get protocol from Starlette request
         protocol = connection.scope.get('type')
@@ -53,7 +54,12 @@ async def __call__(
                 protocol, credential, self.resource, self.permission, user_id=user_id
             )
             if user:
-                return await self.get_user_stray(user, connection)
+                stray = await self.get_user_stray(user, connection)
+                yield stray
+
+                stray.update_working_memory_cache()
+                del stray
+                return
 
         # if no stray was obtained, raise exception
         self.not_allowed(connection)
@@ -103,16 +109,14 @@ def extract_credentials(self, connection: Request) -> Tuple[str, str] | None:
 
 
     async def get_user_stray(self, user: AuthUserInfo, connection: Request) -> StrayCat:
-        strays = connection.app.state.strays
         event_loop = connection.app.state.event_loop
 
-        if user.id not in strays.keys():
-            strays[user.id] = StrayCat(
-                    # TODOV2: user_id should be the user.id
-                user_id=user.name, user_data=user, main_loop=event_loop
-            )
-        return strays[user.id]
-
+        return StrayCat(
+                # TODOV2: user_id should be the user.id
+            user_id=user.name, user_data=user, main_loop=event_loop
+        )
+
+
     def not_allowed(self, connection: Request):
         raise HTTPException(status_code=403, detail={"error": "Invalid Credentials"})
 
@@ -136,29 +140,13 @@ def extract_credentials(self, connection: WebSocket) -> Tuple[str, str] | None:
 
 
     async def get_user_stray(self, user: AuthUserInfo, connection: WebSocket) -> StrayCat:
-        strays = connection.app.state.strays
-
-        if user.id in strays.keys():
-            stray = strays[user.id]
-            await stray.close_connection()
+        return StrayCat(
+            ws=connection,
+            user_id=user.name, # TODOV2: user_id should be the user.id
+            user_data=user,
+            main_loop=asyncio.get_running_loop(),
+        )
 
-            # Set new ws connection
-            stray.reset_connection(connection)
-            log.info(
-                f"New websocket connection for user '{user.id}', the old one has been closed."
-            )
-            return stray
-
-        else:
-            stray = StrayCat(
-                ws=connection,
-                user_id=user.name, # TODOV2: user_id should be the user.id
-                user_data=user,
-                main_loop=asyncio.get_running_loop(),
-            )
-            strays[user.id] = stray
-            return stray
-
     def not_allowed(self, connection: WebSocket):
         raise WebSocketException(code=1004, reason="Invalid Credentials")
 

core/cat/cache/base_cache.py

@@ -0,0 +1,22 @@
+
+from abc import ABC, abstractmethod
+
+class BaseCache(ABC):
+
+    @abstractmethod
+    def insert(self, cache_item):
+        pass
+
+    @abstractmethod
+    def get_item(self, key):
+        pass
+
+    @abstractmethod
+    def get_value(self, key):
+        pass
+
+    @abstractmethod
+    def delete(self, key):
+        pass
+
+

core/cat/cache/cache_item.py

@@ -0,0 +1,19 @@
+import time
+
+
+class CacheItem:
+
+    def __init__(self, key, value, ttl):
+        self.key = key
+        self.value = value
+        self.ttl = ttl
+        self.created_at = time.time()
+
+    def is_expired(self):
+        if self.ttl == -1 or self.ttl is None:
+            return False
+
+        return (self.created_at + self.ttl) < time.time()
+
+    def __repr__(self):
+        return f'CacheItem(key={self.key}, value={self.value}, ttl={self.ttl}, created_at={self.created_at})'

core/cat/cache/cache_manager.py

@@ -0,0 +1,22 @@
+from cat.env import get_env
+
+from cat.utils import singleton
+
+
+@singleton
+class CacheManager:
+    """Class to instantiate different cache types."""
+
+    def __init__(self):
+
+        self.cache_type = get_env("CCAT_CACHE_TYPE")
+
+        if self.cache_type == "in_memory":
+            from cat.cache.in_memory_cache import InMemoryCache
+            self.cache = InMemoryCache()
+        elif self.cache_type == "file_system":
+            cache_dir = get_env("CCAT_CACHE_DIR")
+            from cat.cache.file_system_cache import FileSystemCache
+            self.cache = FileSystemCache(cache_dir)
+        else:
+            raise ValueError(f"Cache type {self.cache_type} not supported")

core/cat/cache/file_system_cache.py

@@ -0,0 +1,99 @@
+import os
+import pickle
+from cat.cache.base_cache import BaseCache
+from cat.utils import singleton
+
+
+@singleton
+class FileSystemCache(BaseCache):
+    """Cache implementation using the file system.
+
+    Attributes
+    ----------
+    cache_dir : str
+        Directory to store the cache.
+
+    """
+
+    def __init__(self, cache_dir):
+        self.cache_dir = cache_dir
+        if not os.path.exists(self.cache_dir):
+            os.makedirs(self.cache_dir)
+
+    def _get_file_path(self, key):
+        return os.path.join(self.cache_dir, f"{key}.cache")
+
+    def insert(self, cache_item):
+        """Insert a key-value pair in the cache.
+
+        Parameters
+        ----------
+        cache_item : CacheItem
+            Cache item to store.
+
+        """
+
+        with open(self._get_file_path(cache_item.key), "wb") as f:
+            pickle.dump(cache_item, f)
+
+    def get_item(self, key):
+        """Get the value stored in the cache.
+
+        Parameters
+        ----------
+        key : str
+            Key to retrieve the value.
+
+        Returns
+        -------
+        any
+            Value stored in the cache.
+
+        """
+        file_path = self._get_file_path(key)
+        if not os.path.exists(file_path):
+            return None
+
+        with open(file_path, "rb") as f:
+            cache_item = pickle.load(f)
+
+        if cache_item.is_expired():
+            os.remove(file_path)
+            return None
+
+        return cache_item
+
+    def get_value(self, key):
+        """Get the value stored in the cache.
+
+        Parameters
+        ----------
+        key : str
+            Key to retrieve the value.
+
+        Returns
+        -------
+        any
+            Value stored in the cache.
+
+        """
+
+        cache_item = self.get_item(key)
+        if cache_item:
+            return cache_item.value
+        return None
+
+    def delete(self, key):
+        """Delete a key-value pair from the cache.
+
+        Parameters
+        ----------
+        key : str
+            Key to delete the value.
+
+        """
+        file_path = self._get_file_path(key)
+        if os.path.exists(file_path):
+            os.remove(file_path)
+
+

core/cat/cache/in_memory_cache.py

@@ -0,0 +1,85 @@
+from cat.cache.base_cache import BaseCache
+from cat.cache.cache_item import CacheItem
+
+from cat.utils import singleton
+
+
+@singleton
+class InMemoryCache(BaseCache):
+    """Cache implementation using a python dictionary.
+
+    Attributes
+    ----------
+    cache : dict
+        Dictionary to store the cache.
+
+    """
+
+    def __init__(self):
+        self.cache = {}
+
+    def insert(self, cache_item):
+        """Insert a key-value pair in the cache.
+
+        Parameters
+        ----------
+        cache_item : CacheItem
+            Cache item to store.
+
+        """
+        self.cache[cache_item.key] = cache_item
+
+    def get_item(self, key) -> CacheItem:
+        """Get the value stored in the cache.
+
+        Parameters
+        ----------
+        key : str
+            Key to retrieve the value.
+
+        Returns
+        -------
+        any
+            Value stored in the cache.
+
+        """
+        item = self.cache.get(key)
+
+        if item and item.is_expired():
+            del self.cache[key]
+            return None
+
+        return item
+
+    def get_value(self, key):
+        """Get the value stored in the cache.
+
+        Parameters
+        ----------
+        key : str
+            Key to retrieve the value.
+
+        Returns
+        -------
+        any
+            Value stored in the cache.
+
+        """
+
+
+        item = self.get_item(key)
+        if item:
+            return item.value
+        return None
+
+    def delete(self, key):
+        """Delete a key-value pair from the cache.
+
+        Parameters
+        ----------
+        key : str
+            Key to delete the value.
+
+        """
+        if key in self.cache:
+            del self.cache[key]

core/cat/env.py

@@ -21,7 +21,9 @@ def get_supported_env_variables():
         "CCAT_JWT_EXPIRE_MINUTES": str(60 * 24),  # JWT expires after 1 day
         "CCAT_HTTPS_PROXY_MODE": False,
         "CCAT_CORS_FORWARDED_ALLOW_IPS": "*",
-        "CCAT_CORS_ENABLED": "true"
+        "CCAT_CORS_ENABLED": "true",
+        "CCAT_CACHE_TYPE": "file_system",
+        "CCAT_CACHE_DIR": "/tmp",
     }