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

Author: Aryamanz29

pyatlan/cache/enum_cache.py

@@ -1,33 +1,48 @@
 # SPDX-License-Identifier: Apache-2.0
-# Copyright 2023 Atlan Pte. Ltd.
-from threading import Lock
-from typing import Dict, Optional
+# Copyright 2025 Atlan Pte. Ltd.
+from __future__ import annotations
+
+from threading import Lock, local
+from typing import TYPE_CHECKING, Dict, Optional
 
-from pyatlan.client.typedef import TypeDefClient
 from pyatlan.errors import ErrorCode
 from pyatlan.model.enums import AtlanTypeCategory
 from pyatlan.model.typedef import EnumDef
 
+if TYPE_CHECKING:
+    from pyatlan.client.atlan import AtlanClient
+
 lock: Lock = Lock()
+enum_cache_tls = local()  # Thread-local storage (TLS)
 
 
 class EnumCache:
     """
     Lazily-loaded cache for accessing details of an enumeration.
     """
 
-    caches: Dict[int, "EnumCache"] = {}
+    def __init__(self, client: AtlanClient):
+        self.client: AtlanClient = client
+        self.cache_by_name: Dict[str, EnumDef] = {}
+        self.lock: Lock = Lock()
 
     @classmethod
-    def get_cache(cls) -> "EnumCache":
+    def get_cache(cls, client: Optional[AtlanClient] = None) -> EnumCache:
         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] = EnumCache(typedef_client=client.typedef)
-            return cls.caches[cache_key]
+
+            if not hasattr(enum_cache_tls, "caches"):
+                enum_cache_tls.caches = {}
+
+            if cache_key not in enum_cache_tls.caches:
+                cache_instance = EnumCache(client=client)
+                cache_instance._refresh_cache()  # Refresh on new cache instance
+                enum_cache_tls.caches[cache_key] = cache_instance
+
+            return enum_cache_tls.caches[cache_key]
 
     @classmethod
     def refresh_cache(cls) -> None:
@@ -49,17 +64,12 @@ def get_by_name(cls, name: str) -> EnumDef:
             raise ErrorCode.ENUM_NOT_FOUND.exception_with_parameters(name)
         return enum
 
-    def __init__(self, typedef_client: TypeDefClient):
-        self.typedef_client: TypeDefClient = typedef_client
-        self.cache_by_name: Dict[str, EnumDef] = {}
-        self.lock: Lock = Lock()
-
     def _refresh_cache(self) -> None:
         """
         Refreshes the cache of enumerations by requesting the full set of enumerations from Atlan.
         """
         with self.lock:
-            response = self.typedef_client.get(type_category=AtlanTypeCategory.ENUM)
+            response = self.client.typedef.get(type_category=AtlanTypeCategory.ENUM)
             if not response or not response.enum_defs:
                 raise ErrorCode.EXPIRED_API_TOKEN.exception_with_parameters()
             self.cache_by_name = {}

pyatlan/client/atlan.py

@@ -43,6 +43,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.client.admin import AdminClient
 from pyatlan.client.asset import A, AssetClient, IndexSearchResults, LineageListResults
 from pyatlan.client.audit import AuditClient
@@ -144,7 +145,7 @@ def get_session():
 
 
 class AtlanClient(BaseSettings):
-    _default_client_storage: ClassVar[local] = local()  # Thread-local storage
+    _default_client_tls: ClassVar[local] = local()  # Thread-local storage (TLS)
     base_url: Union[Literal["INTERNAL"], HttpUrl]
     api_key: str
     connect_timeout: float = 30.0  # 30 secs
@@ -173,6 +174,7 @@ class AtlanClient(BaseSettings):
     _contract_client: Optional[ContractClient] = PrivateAttr(default=None)
     _open_lineage_client: Optional[OpenLineageClient] = PrivateAttr(default=None)
     _atlan_tag_cache: Optional[AtlanTagCache] = PrivateAttr(default=None)
+    _enum_cache: Optional[EnumCache] = PrivateAttr(default=None)
     _custom_metadata_cache: Optional[CustomMetadataCache] = PrivateAttr(default=None)
 
     class Config:
@@ -185,7 +187,7 @@ def set_default_client(cls, client: "AtlanClient"):
         """
         if not isinstance(client, AtlanClient):
             raise ErrorCode.MISSING_ATLAN_CLIENT.exception_with_parameters()
-        cls._default_client_storage.client = client
+        cls._default_client_tls.client = client
 
     @classmethod
     def get_default_client(cls) -> AtlanClient:
@@ -194,9 +196,9 @@ def get_default_client(cls) -> AtlanClient:
 
         :returns: the default client
         """
-        if not hasattr(cls._default_client_storage, "client"):
+        if not hasattr(cls._default_client_tls, "client"):
             raise ErrorCode.NO_ATLAN_CLIENT_AVAILABLE.exception_with_parameters()
-        return cls._default_client_storage.client
+        return cls._default_client_tls.client
 
     def __init__(self, **data):
         super().__init__(**data)
@@ -329,6 +331,12 @@ def atlan_tag_cache(self) -> AtlanTagCache:
             self._atlan_tag_cache = AtlanTagCache.get_cache(client=self)
         return self._atlan_tag_cache
 
+    @property
+    def enum_cache(self) -> EnumCache:
+        if self._enum_cache is None:
+            self._enum_cache = EnumCache.get_cache(client=self)
+        return self._enum_cache
+
     @property
     def custom_metadata_cache(self) -> CustomMetadataCache:
         if self._custom_metadata_cache is None: