Changes after testing

Author: ptiurin

src/firebolt/async_db/connection.py

@@ -299,8 +299,7 @@ async def connect(
     if not auth:
         raise ConfigurationError("auth is required to connect.")
 
-    if account_name:
-        auth._account_name = account_name
+    auth.account = account_name
 
     api_endpoint = fix_url_schema(api_endpoint)
     # Type checks

src/firebolt/client/auth/base.py

@@ -2,7 +2,7 @@
 from abc import abstractmethod
 from enum import IntEnum
 from time import time
-from typing import AsyncGenerator, Generator, Optional
+from typing import AsyncGenerator, Generator, Optional, Tuple
 
 from anyio import Lock
 from httpx import Auth as HttpxAuth
@@ -57,6 +57,17 @@ def __init__(self, use_token_cache: bool = True):
         self._expires: Optional[int] = None
         self._lock = Lock()
 
+    @property
+    def account(self) -> Optional[str]:
+        return self._account_name
+
+    @account.setter
+    def account(self, value: str) -> None:
+        self._account_name = value
+        # Now we have all the elements to fetch the cached token
+        if not self._token:
+            self._token, self._expires = self._get_cached_token()
+
     def copy(self) -> "Auth":
         """Make another auth object with same credentials.
 
@@ -109,7 +120,7 @@ def expired(self) -> bool:
         """
         return self._expires is not None and self._expires <= int(time())
 
-    def _get_cached_token(self) -> Optional[str]:
+    def _get_cached_token(self) -> Tuple[Optional[str], Optional[int]]:
         """If caching is enabled, get token from cache.
 
         If caching is disabled, None is returned.
@@ -118,17 +129,17 @@ def _get_cached_token(self) -> Optional[str]:
             Optional[str]: Token if any, and if caching is enabled; None otherwise
         """
         if not self._use_token_cache:
-            return None
+            return (None, None)
 
         cache_key = SecureCacheKey(
             [self.principal, self.secret, self._account_name], self.secret
         )
         connection_info = _firebolt_cache.get(cache_key)
 
         if connection_info and connection_info.token:
-            return connection_info.token
+            return (connection_info.token, connection_info.expiry_time)
 
-        return None
+        return (None, None)
 
     def _cache_token(self) -> None:
         """If caching is enabled, cache token."""

src/firebolt/db/connection.py

@@ -68,8 +68,7 @@ def connect(
     if not auth:
         raise ConfigurationError("auth is required to connect.")
 
-    if account_name:
-        auth._account_name = account_name
+    auth.account = account_name
 
     api_endpoint = fix_url_schema(api_endpoint)
     # Type checks

src/firebolt/utils/cache.py

@@ -71,14 +71,24 @@ def __post_init__(self) -> None:
         """
         if self.system_engine and isinstance(self.system_engine, dict):
             self.system_engine = EngineInfo(**self.system_engine)
-        self.databases = {
-            k: DatabaseInfo(**v)
-            for k, v in self.databases.items()
-            if isinstance(v, dict)
-        }
-        self.engines = {
-            k: EngineInfo(**v) for k, v in self.engines.items() if isinstance(v, dict)
-        }
+
+        # Convert dict values to dataclasses, keep existing dataclass objects
+        new_databases = {}
+        for k, db in self.databases.items():
+            if isinstance(db, dict):
+                new_databases[k] = DatabaseInfo(**db)
+            else:
+                new_databases[k] = db
+        self.databases = new_databases
+
+        # Convert dict values to dataclasses, keep existing dataclass objects
+        new_engines = {}
+        for k, engine in self.engines.items():
+            if isinstance(engine, dict):
+                new_engines[k] = EngineInfo(**engine)
+            else:
+                new_engines[k] = engine
+        self.engines = new_engines
 
 
 def noop_if_disabled(func: Callable) -> Callable: