fix pre commit

Author: ptiurin

src/firebolt/utils/cache.py

@@ -192,20 +192,34 @@ def __hash__(self) -> int:
         return hash(self.key)
 
 
-class FileBasedCache(UtilCache[ConnectionInfo]):
+class FileBasedCache:
     """
     File-based cache that persists to disk with encryption.
-    Extends UtilCache to provide persistent storage using encrypted files.
+    Uses composition to combine in-memory caching with persistent storage
+    using encrypted files.
     """
 
-    def __init__(self, cache_name: str = ""):
-        super().__init__(cache_name)
+    def __init__(self, memory_cache: UtilCache[ConnectionInfo], cache_name: str = ""):
+        self.memory_cache = memory_cache
         self._data_dir = user_data_dir(appname=APPNAME)  # TODO: change to new dir
         makedirs(self._data_dir, exist_ok=True)
+        # FileBasedCache has its own disabled state, independent of memory cache
+        cache_env_var = f"FIREBOLT_SDK_DISABLE_CACHE_${cache_name}"
+        self.disabled = os.getenv("FIREBOLT_SDK_DISABLE_CACHE", False) or os.getenv(
+            cache_env_var, False
+        )
+
+    def disable(self) -> None:
+        """Disable the file-based cache."""
+        self.disabled = True
+
+    def enable(self) -> None:
+        """Enable the file-based cache."""
+        self.disabled = False
 
     def _get_file_path(self, key: SecureCacheKey) -> str:
         """Get the file path for a cache key."""
-        cache_key = self.create_key(key)
+        cache_key = self.memory_cache.create_key(key)
         encrypted_filename = generate_encrypted_file_name(cache_key, key.encryption_key)
         return path.join(self._data_dir, encrypted_filename)
 
@@ -250,7 +264,7 @@ def get(self, key: SecureCacheKey) -> Optional[ConnectionInfo]:
             return None
 
         # First try memory cache
-        memory_result = super().get(key)
+        memory_result = self.memory_cache.get(key)
         if memory_result is not None:
             logger.debug("Cache hit in memory")
             return memory_result
@@ -265,7 +279,7 @@ def get(self, key: SecureCacheKey) -> Optional[ConnectionInfo]:
         data = ConnectionInfo(**raw_data)
 
         # Add to memory cache and return
-        super().set(key, data)
+        self.memory_cache.set(key, data)
         return data
 
     def set(self, key: SecureCacheKey, value: ConnectionInfo) -> None:
@@ -275,7 +289,7 @@ def set(self, key: SecureCacheKey, value: ConnectionInfo) -> None:
 
         logger.debug("Setting value in cache")
         # First set in memory
-        super().set(key, value)
+        self.memory_cache.set(key, value)
 
         file_path = self._get_file_path(key)
         encrypter = FernetEncrypter(generate_salt(), key.encryption_key)
@@ -289,7 +303,7 @@ def delete(self, key: SecureCacheKey) -> None:
             return
 
         # Delete from memory
-        super().delete(key)
+        self.memory_cache.delete(key)
 
         # Delete from disk
         file_path = self._get_file_path(key)
@@ -303,7 +317,9 @@ def delete(self, key: SecureCacheKey) -> None:
     def clear(self) -> None:
         # Clear memory only, as deleting every file is not safe
         logger.debug("Clearing memory cache")
-        super().clear()
+        self.memory_cache.clear()
 
 
-_firebolt_cache = FileBasedCache(cache_name="connection_info")
+_firebolt_cache = FileBasedCache(
+    UtilCache[ConnectionInfo](cache_name="memory_cache"), cache_name="file_cache"
+)

src/firebolt/utils/usage_tracker.py

@@ -8,7 +8,11 @@
 from typing import Dict, List, Optional, Tuple
 
 from firebolt import __version__
-from firebolt.utils.cache import ConnectionInfo, ReprCacheable, _firebolt_cache
+from firebolt.utils.cache import (
+    ConnectionInfo,
+    SecureCacheKey,
+    _firebolt_cache,
+)
 
 
 @dataclass
@@ -228,7 +232,7 @@ def get_user_agent_header(
 
 
 def get_cache_tracking_params(
-    cache_key: ReprCacheable, conn_id: str
+    cache_key: SecureCacheKey, conn_id: str
 ) -> List[Tuple[str, str]]:
     ua_parameters = []
     ua_parameters.append(("connId", conn_id))

tests/unit/utils/test_cache.py

@@ -355,25 +355,25 @@ def test_cache_disable_enable_behavior(
 
 def test_helper_functions():
     """Test the backward compatibility helper functions."""
-    from tests.unit.test_cache_helpers import cache_token, get_cached_token
     from firebolt.utils.cache import _firebolt_cache
-    
+    from tests.unit.test_cache_helpers import cache_token, get_cached_token
+
     _firebolt_cache.enable()
     _firebolt_cache.clear()
-    
+
     # Test caching and retrieving tokens
     principal = "test_user"
     secret = "test_secret"
     token = "test_token"
     account_name = "test_account"
-    
+
     # Cache token
     cache_token(principal, secret, token, 9999, account_name)
-    
+
     # Retrieve token
     cached_token = get_cached_token(principal, secret, account_name)
     assert cached_token == token
-    
+
     # Test with None account name
     cache_token(principal, secret, token, 9999, None)
     cached_token_none = get_cached_token(principal, secret, None)
@@ -385,31 +385,32 @@ def test_connection_info_post_init():
     # Test with dictionary inputs that should be converted to dataclasses
     engine_dict = {"url": "http://test.com", "params": {"key": "value"}}
     db_dict = {"name": "test_db"}
-    
+
     connection_info = ConnectionInfo(
         id="test",
         system_engine=engine_dict,
         databases={"db1": db_dict},
-        engines={"engine1": engine_dict}
+        engines={"engine1": engine_dict},
     )
-    
+
     # Should convert dicts to dataclasses
-    from firebolt.utils.cache import EngineInfo, DatabaseInfo
+    from firebolt.utils.cache import DatabaseInfo, EngineInfo
+
     assert isinstance(connection_info.system_engine, EngineInfo)
     assert isinstance(connection_info.databases["db1"], DatabaseInfo)
     assert isinstance(connection_info.engines["engine1"], EngineInfo)
-    
+
     # Test with already converted dataclass objects
     engine_obj = EngineInfo(url="http://test.com", params={"key": "value"})
     db_obj = DatabaseInfo(name="test_db")
-    
+
     connection_info2 = ConnectionInfo(
         id="test2",
         system_engine=engine_obj,
         databases={"db1": db_obj},
-        engines={"engine1": engine_obj}
+        engines={"engine1": engine_obj},
     )
-    
+
     # Should remain as dataclasses
     assert connection_info2.system_engine is engine_obj
     assert connection_info2.databases["db1"] is db_obj