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

Author: Aryamanz29

pyatlan/cache/group_cache.py

@@ -1,11 +1,15 @@
 # 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 pyatlan.client.group import GroupClient
+from threading import Lock, local
+from typing import TYPE_CHECKING, Dict, Iterable, Optional
+
+if TYPE_CHECKING:
+    from pyatlan.client.atlan import AtlanClient
 
 lock: Lock = Lock()
+group_cache_tls = local()  # Thread-local storage (TLS)
 
 
 class GroupCache:
@@ -15,16 +19,30 @@ class GroupCache:
 
     caches: Dict[int, "GroupCache"] = {}
 
+    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_alias_to_id: Dict[str, str] = {}
+        self.lock: Lock = Lock()
+
     @classmethod
-    def get_cache(cls) -> "GroupCache":
+    def get_cache(cls, client: Optional[AtlanClient] = None) -> GroupCache:
         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] = GroupCache(group_client=client.group)
-            return cls.caches[cache_key]
+
+            if not hasattr(group_cache_tls, "caches"):
+                group_cache_tls.caches = {}
+
+            if cache_key not in group_cache_tls.caches:
+                cache_instance = GroupCache(client=client)
+                cache_instance._refresh_cache()  # Refresh on new cache instance
+                group_cache_tls.caches[cache_key] = cache_instance
+
+            return group_cache_tls.caches[cache_key]
 
     @classmethod
     def get_id_for_name(cls, name: str) -> Optional[str]:
@@ -65,27 +83,21 @@ def validate_aliases(cls, aliases: Iterable[str]):
         """
         return cls.get_cache()._validate_aliases(aliases)
 
-    def __init__(self, group_client: GroupClient):
-        self.group_client: GroupClient = group_client
-        self.map_id_to_name: Dict[str, str] = {}
-        self.map_name_to_id: Dict[str, str] = {}
-        self.map_alias_to_id: Dict[str, str] = {}
-        self.lock: Lock = Lock()
-
     def _refresh_cache(self) -> None:
         with self.lock:
-            groups = self.group_client.get_all()
-            if groups is not None:
-                self.map_id_to_name = {}
-                self.map_name_to_id = {}
-                self.map_alias_to_id = {}
-                for group in groups:
-                    group_id = str(group.id)
-                    group_name = str(group.name)
-                    group_alias = str(group.alias)
-                    self.map_id_to_name[group_id] = group_name
-                    self.map_name_to_id[group_name] = group_id
-                    self.map_alias_to_id[group_alias] = group_id
+            groups = self.client.group.get_all()
+            if not groups:
+                return
+            self.map_id_to_name = {}
+            self.map_name_to_id = {}
+            self.map_alias_to_id = {}
+            for group in groups:
+                group_id = str(group.id)
+                group_name = str(group.name)
+                group_alias = str(group.alias)
+                self.map_id_to_name[group_id] = group_name
+                self.map_name_to_id[group_name] = group_id
+                self.map_alias_to_id[group_alias] = group_id
 
     def _get_id_for_name(self, name: str) -> Optional[str]:
         """

pyatlan/client/atlan.py

@@ -44,6 +44,7 @@
 from pyatlan.cache.atlan_tag_cache import AtlanTagCache
 from pyatlan.cache.custom_metadata_cache import CustomMetadataCache
 from pyatlan.cache.enum_cache import EnumCache
+from pyatlan.cache.group_cache import GroupCache
 from pyatlan.client.admin import AdminClient
 from pyatlan.client.asset import A, AssetClient, IndexSearchResults, LineageListResults
 from pyatlan.client.audit import AuditClient
@@ -175,6 +176,7 @@ class AtlanClient(BaseSettings):
     _open_lineage_client: Optional[OpenLineageClient] = PrivateAttr(default=None)
     _atlan_tag_cache: Optional[AtlanTagCache] = PrivateAttr(default=None)
     _enum_cache: Optional[EnumCache] = PrivateAttr(default=None)
+    _group_cache: Optional[GroupCache] = PrivateAttr(default=None)
     _custom_metadata_cache: Optional[CustomMetadataCache] = PrivateAttr(default=None)
 
     class Config:
@@ -337,6 +339,12 @@ def enum_cache(self) -> EnumCache:
             self._enum_cache = EnumCache.get_cache(client=self)
         return self._enum_cache
 
+    @property
+    def group_cache(self) -> GroupCache:
+        if self._group_cache is None:
+            self._group_cache = GroupCache.get_cache(client=self)
+        return self._group_cache
+
     @property
     def custom_metadata_cache(self) -> CustomMetadataCache:
         if self._custom_metadata_cache is None: