[change] Bound RoleCache to AtlanClient, moved class var to TLS

Author: Aryamanz29

pyatlan/cache/user_cache.py

@@ -1,35 +1,49 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright 2022 Atlan Pte. Ltd.
-from threading import Lock
-from typing import Dict, Iterable, Optional
+from __future__ import annotations
+
+from threading import Lock, local
+from typing import TYPE_CHECKING, Dict, Iterable, Optional
 
-from pyatlan.client.token import TokenClient
-from pyatlan.client.user import UserClient
 from pyatlan.errors import ErrorCode
 from pyatlan.model.constants import SERVICE_ACCOUNT_
 
+if TYPE_CHECKING:
+    from pyatlan.client.atlan import AtlanClient
+
 lock = Lock()
+user_cache_tls = local()  # Thread-local storage (TLS)
 
 
 class UserCache:
     """
     Lazily-loaded cache for translating Atlan-internal users into their various IDs.
     """
 
-    caches: Dict[int, "UserCache"] = {}
+    def __init__(self, client: AtlanClient):
+        self.client: AtlanClient = client
+        self.map_id_to_name: Dict[str, str] = {}
+        self.map_name_to_id: Dict[str, str] = {}
+        self.map_email_to_id: Dict[str, str] = {}
+        self.lock: Lock = Lock()
 
     @classmethod
-    def get_cache(cls) -> "UserCache":
+    def get_cache(cls, client: Optional[AtlanClient] = None) -> UserCache:
         from pyatlan.client.atlan import AtlanClient
 
         with lock:
-            client = AtlanClient.get_default_client()
+            client = client or AtlanClient.get_default_client()
             cache_key = client.cache_key
-            if cache_key not in cls.caches:
-                cls.caches[cache_key] = UserCache(
-                    user_client=client.user, token_client=client.token
-                )
-            return cls.caches[cache_key]
+
+            if not hasattr(user_cache_tls, "caches"):
+                user_cache_tls.caches = {}
+
+            if cache_key not in user_cache_tls.caches:
+                cache_instance = UserCache(client=client)
+                cache_instance._refresh_cache()  # Refresh on new cache instance
+                user_cache_tls.caches[cache_key] = cache_instance
+
+            return user_cache_tls.caches[cache_key]
 
     @classmethod
     def get_id_for_name(cls, name: str) -> Optional[str]:
@@ -70,28 +84,21 @@ def validate_names(cls, names: Iterable[str]):
         """
         return cls.get_cache()._validate_names(names)
 
-    def __init__(self, user_client: UserClient, token_client: TokenClient):
-        self.user_client: UserClient = user_client
-        self.token_client: TokenClient = token_client
-        self.map_id_to_name: Dict[str, str] = {}
-        self.map_name_to_id: Dict[str, str] = {}
-        self.map_email_to_id: Dict[str, str] = {}
-        self.lock: Lock = Lock()
-
     def _refresh_cache(self) -> None:
         with self.lock:
-            users = self.user_client.get_all()
-            if users is not None:
-                self.map_id_to_name = {}
-                self.map_name_to_id = {}
-                self.map_email_to_id = {}
-                for user in users:
-                    user_id = str(user.id)
-                    username = str(user.username)
-                    user_email = str(user.email)
-                    self.map_id_to_name[user_id] = username
-                    self.map_name_to_id[username] = user_id
-                    self.map_email_to_id[user_email] = user_id
+            users = self.client.user.get_all()
+            if not users:
+                return
+            self.map_id_to_name = {}
+            self.map_name_to_id = {}
+            self.map_email_to_id = {}
+            for user in users:
+                user_id = str(user.id)
+                username = str(user.username)
+                user_email = str(user.email)
+                self.map_id_to_name[user_id] = username
+                self.map_name_to_id[username] = user_id
+                self.map_email_to_id[user_email] = user_id
 
     def _get_id_for_name(self, name: str) -> Optional[str]:
         """
@@ -105,7 +112,7 @@ def _get_id_for_name(self, name: str) -> Optional[str]:
         # If we are translating an API token,
         # short-circuit any further cache refresh
         if name.startswith(SERVICE_ACCOUNT_):
-            token = self.token_client.get_by_id(client_id=name)
+            token = self.client.token.get_by_id(client_id=name)
             if token and token.guid:
                 self.map_name_to_id[name] = token.guid
                 return token.guid
@@ -138,7 +145,7 @@ def _get_name_for_id(self, idstr: str) -> Optional[str]:
         if username := self.map_id_to_name.get(idstr):
             return username
         # If the username isn't found, check if it is an API token
-        token = self.token_client.get_by_guid(guid=idstr)
+        token = self.client.token.get_by_guid(guid=idstr)
         if token and token.client_id:
             return token.username
         else:
@@ -152,7 +159,7 @@ def _validate_names(self, names: Iterable[str]):
         :param names: a collection of usernames to be checked
         """
         for username in names:
-            if not self.get_id_for_name(username) and not self.token_client.get_by_id(
+            if not self.get_id_for_name(username) and not self.client.token.get_by_id(
                 username
             ):
                 raise ValueError(

pyatlan/client/atlan.py

@@ -46,6 +46,7 @@
 from pyatlan.cache.enum_cache import EnumCache
 from pyatlan.cache.group_cache import GroupCache
 from pyatlan.cache.role_cache import RoleCache
+from pyatlan.cache.user_cache import UserCache
 from pyatlan.client.admin import AdminClient
 from pyatlan.client.asset import A, AssetClient, IndexSearchResults, LineageListResults
 from pyatlan.client.audit import AuditClient
@@ -179,6 +180,7 @@ class AtlanClient(BaseSettings):
     _enum_cache: Optional[EnumCache] = PrivateAttr(default=None)
     _group_cache: Optional[GroupCache] = PrivateAttr(default=None)
     _role_cache: Optional[RoleCache] = PrivateAttr(default=None)
+    _user_cache: Optional[UserCache] = PrivateAttr(default=None)
     _custom_metadata_cache: Optional[CustomMetadataCache] = PrivateAttr(default=None)
 
     class Config:
@@ -353,6 +355,12 @@ def role_cache(self) -> RoleCache:
             self._role_cache = RoleCache.get_cache(client=self)
         return self._role_cache
 
+    @property
+    def user_cache(self) -> UserCache:
+        if self._user_cache is None:
+            self._user_cache = UserCache.get_cache(client=self)
+        return self._user_cache
+
     @property
     def custom_metadata_cache(self) -> CustomMetadataCache:
         if self._custom_metadata_cache is None: