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

Author: Aryamanz29

pyatlan/cache/role_cache.py

@@ -1,31 +1,48 @@
 # SPDX-License-Identifier: Apache-2.0
-# Copyright 2022 Atlan Pte. Ltd.
-from threading import Lock
-from typing import Dict, Iterable, Optional
+# Copyright 2025 Atlan Pte. Ltd.
+from __future__ import annotations
+
+from threading import Lock, local
+from typing import TYPE_CHECKING, Dict, Iterable, Optional
 
-from pyatlan.client.role import RoleClient
 from pyatlan.model.role import AtlanRole
 
+if TYPE_CHECKING:
+    from pyatlan.client.atlan import AtlanClient
+
 lock: Lock = Lock()
+role_cache_tls = local()  # Thread-local storage (TLS)
 
 
 class RoleCache:
     """
     Lazily-loaded cache for translating Atlan-internal roles into their various IDs.
     """
 
-    caches: Dict[int, "RoleCache"] = {}
+    def __init__(self, client: AtlanClient):
+        self.client: AtlanClient = client
+        self.cache_by_id: Dict[str, AtlanRole] = {}
+        self.map_id_to_name: Dict[str, str] = {}
+        self.map_name_to_id: Dict[str, str] = {}
+        self.lock: Lock = Lock()
 
     @classmethod
-    def get_cache(cls) -> "RoleCache":
+    def get_cache(cls, client: Optional[AtlanClient] = None) -> RoleCache:
         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] = RoleCache(role_client=client.role)
-            return cls.caches[cache_key]
+
+            if not hasattr(role_cache_tls, "caches"):
+                role_cache_tls.caches = {}
+
+            if cache_key not in role_cache_tls.caches:
+                cache_instance = RoleCache(client=client)
+                cache_instance._refresh_cache()  # Refresh on new cache instance
+                role_cache_tls.caches[cache_key] = cache_instance
+
+            return role_cache_tls.caches[cache_key]
 
     @classmethod
     def get_id_for_name(cls, name: str) -> Optional[str]:
@@ -56,28 +73,22 @@ def validate_idstrs(cls, idstrs: Iterable[str]):
         """
         return cls.get_cache()._validate_idstrs(idstrs=idstrs)
 
-    def __init__(self, role_client: RoleClient):
-        self.role_client: RoleClient = role_client
-        self.cache_by_id: Dict[str, AtlanRole] = {}
-        self.map_id_to_name: Dict[str, str] = {}
-        self.map_name_to_id: Dict[str, str] = {}
-        self.lock: Lock = Lock()
-
     def _refresh_cache(self) -> None:
         with self.lock:
-            response = self.role_client.get(
+            response = self.client.role.get(
                 limit=100, post_filter='{"name":{"$ilike":"$%"}}'
             )
-            if response is not None:
-                self.cache_by_id = {}
-                self.map_id_to_name = {}
-                self.map_name_to_id = {}
-                for role in response.records:
-                    role_id = role.id
-                    role_name = role.name
-                    self.cache_by_id[role_id] = role
-                    self.map_id_to_name[role_id] = role_name
-                    self.map_name_to_id[role_name] = role_id
+            if not response:
+                return
+            self.cache_by_id = {}
+            self.map_id_to_name = {}
+            self.map_name_to_id = {}
+            for role in response.records:
+                role_id = role.id
+                role_name = role.name
+                self.cache_by_id[role_id] = role
+                self.map_id_to_name[role_id] = role_name
+                self.map_name_to_id[role_name] = role_id
 
     def _get_id_for_name(self, name: str) -> Optional[str]:
         """

pyatlan/client/atlan.py

@@ -45,6 +45,7 @@
 from pyatlan.cache.custom_metadata_cache import CustomMetadataCache
 from pyatlan.cache.enum_cache import EnumCache
 from pyatlan.cache.group_cache import GroupCache
+from pyatlan.cache.role_cache import RoleCache
 from pyatlan.client.admin import AdminClient
 from pyatlan.client.asset import A, AssetClient, IndexSearchResults, LineageListResults
 from pyatlan.client.audit import AuditClient
@@ -177,6 +178,7 @@ class AtlanClient(BaseSettings):
     _atlan_tag_cache: Optional[AtlanTagCache] = PrivateAttr(default=None)
     _enum_cache: Optional[EnumCache] = PrivateAttr(default=None)
     _group_cache: Optional[GroupCache] = PrivateAttr(default=None)
+    _role_cache: Optional[RoleCache] = PrivateAttr(default=None)
     _custom_metadata_cache: Optional[CustomMetadataCache] = PrivateAttr(default=None)
 
     class Config:
@@ -345,6 +347,12 @@ def group_cache(self) -> GroupCache:
             self._group_cache = GroupCache.get_cache(client=self)
         return self._group_cache
 
+    @property
+    def role_cache(self) -> RoleCache:
+        if self._role_cache is None:
+            self._role_cache = RoleCache.get_cache(client=self)
+        return self._role_cache
+
     @property
     def custom_metadata_cache(self) -> CustomMetadataCache:
         if self._custom_metadata_cache is None: