Merge pull request #573 from atlanhq/APP-5519

APP-5519: Migrated `AtlanClient._default_client` to thread-local storage + caches now bound to optional client

Author: Aryamanz29

pyatlan/cache/abstract_asset_cache.py

@@ -4,30 +4,28 @@
 
 import threading
 from abc import ABC, abstractmethod
-from typing import Any
+from typing import TYPE_CHECKING, Any, Dict
 
 from pyatlan.errors import ErrorCode
 from pyatlan.model.assets import Asset
 from pyatlan.model.enums import AtlanConnectorType
 
+if TYPE_CHECKING:
+    from pyatlan.client.atlan import AtlanClient
+
 
 class AbstractAssetCache(ABC):
     """
     Base class for reusable components that are common
     to all caches, where a cache is populated entry-by-entry.
     """
 
-    def __init__(self, client):
+    def __init__(self, client: AtlanClient):
         self.client = client
         self.lock = threading.Lock()
-        self.name_to_guid = dict()
-        self.guid_to_asset = dict()
-        self.qualified_name_to_guid = dict()
-
-    @classmethod
-    @abstractmethod
-    def get_cache(cls):
-        """Abstract method to retreive cache."""
+        self.name_to_guid: Dict[str, str] = dict()
+        self.guid_to_asset: Dict[str, Asset] = dict()
+        self.qualified_name_to_guid: Dict[str, str] = dict()
 
     @abstractmethod
     def lookup_by_guid(self, guid: str):
@@ -84,9 +82,9 @@ def cache(self, asset: Asset):
         name = asset and self.get_name(asset)
         if not all([name, asset.guid, asset.qualified_name]):
             return
-        self.name_to_guid[name] = asset.guid
-        self.guid_to_asset[asset.guid] = asset
-        self.qualified_name_to_guid[asset.qualified_name] = asset.guid
+        self.name_to_guid[name] = asset.guid  # type: ignore[index]
+        self.guid_to_asset[asset.guid] = asset  # type: ignore[index]
+        self.qualified_name_to_guid[asset.qualified_name] = asset.guid  # type: ignore[index]
 
     def _get_by_guid(self, guid: str, allow_refresh: bool = True):
         """

pyatlan/cache/atlan_tag_cache.py

@@ -1,13 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
-# Copyright 2022 Atlan Pte. Ltd.
+# Copyright 2025 Atlan Pte. Ltd.
+from __future__ import annotations
+
 from threading import Lock
-from typing import Dict, Optional, Set
+from typing import TYPE_CHECKING, Dict, Optional, Set
 
-from pyatlan.client.typedef import TypeDefClient
 from pyatlan.errors import ErrorCode
 from pyatlan.model.enums import AtlanTypeCategory
 from pyatlan.model.typedef import AtlanTagDef
 
+if TYPE_CHECKING:
+    from pyatlan.client.atlan import AtlanClient
+
 lock: Lock = Lock()
 
 
@@ -17,73 +21,56 @@ class AtlanTagCache:
     for Atlan tags.
     """
 
-    caches: Dict[int, "AtlanTagCache"] = {}
-
-    @classmethod
-    def get_cache(cls) -> "AtlanTagCache":
-        from pyatlan.client.atlan import AtlanClient
-
-        with lock:
-            client = AtlanClient.get_default_client()
-            cache_key = client.cache_key
-            if cache_key not in cls.caches:
-                cls.caches[cache_key] = AtlanTagCache(typedef_client=client.typedef)
-            return cls.caches[cache_key]
+    def __init__(self, client: AtlanClient):
+        self.client: AtlanClient = client
+        self.cache_by_id: Dict[str, AtlanTagDef] = {}
+        self.map_id_to_name: Dict[str, str] = {}
+        self.map_name_to_id: Dict[str, str] = {}
+        self.deleted_ids: Set[str] = set()
+        self.deleted_names: Set[str] = set()
+        self.map_id_to_source_tags_attr_id: Dict[str, str] = {}
+        self.lock: Lock = Lock()
 
-    @classmethod
-    def refresh_cache(cls) -> None:
+    def refresh_cache(self) -> None:
         """
         Refreshes the cache of Atlan tags by requesting the full set of Atlan tags from Atlan.
         """
-        cls.get_cache()._refresh_cache()
+        self._refresh_cache()
 
-    @classmethod
-    def get_id_for_name(cls, name: str) -> Optional[str]:
+    def get_id_for_name(self, name: str) -> Optional[str]:
         """
         Translate the provided human-readable Atlan tag name to its Atlan-internal ID string.
 
         :param name: human-readable name of the Atlan tag
         :returns: Atlan-internal ID string of the Atlan tag
         """
-        return cls.get_cache()._get_id_for_name(name=name)
+        return self._get_id_for_name(name=name)
 
-    @classmethod
-    def get_name_for_id(cls, idstr: str) -> Optional[str]:
+    def get_name_for_id(self, idstr: str) -> Optional[str]:
         """
         Translate the provided Atlan-internal classification ID string to the human-readable Atlan tag name.
 
         :param idstr: Atlan-internal ID string of the Atlan tag
         :returns: human-readable name of the Atlan tag
         """
-        return cls.get_cache()._get_name_for_id(idstr=idstr)
+        return self._get_name_for_id(idstr=idstr)
 
-    @classmethod
-    def get_source_tags_attr_id(cls, id: str) -> Optional[str]:
+    def get_source_tags_attr_id(self, id: str) -> Optional[str]:
         """
         Translate the provided Atlan-internal Atlan tag ID string to the Atlan-internal name of the attribute that
         captures tag attachment details (for source-synced tags).
 
         :param id: Atlan-internal ID string of the Atlan tag
         :returns: Atlan-internal ID string of the attribute containing source-synced tag attachment details
         """
-        return cls.get_cache()._get_source_tags_attr_id(id)
-
-    def __init__(self, typedef_client: TypeDefClient):
-        self.typdef_client: TypeDefClient = typedef_client
-        self.cache_by_id: Dict[str, AtlanTagDef] = {}
-        self.map_id_to_name: Dict[str, str] = {}
-        self.map_name_to_id: Dict[str, str] = {}
-        self.deleted_ids: Set[str] = set()
-        self.deleted_names: Set[str] = set()
-        self.map_id_to_source_tags_attr_id: Dict[str, str] = {}
-        self.lock: Lock = Lock()
+        return self._get_source_tags_attr_id(id)
 
     def _refresh_cache(self) -> None:
         """
         Refreshes the cache of Atlan tags by requesting the full set of Atlan tags from Atlan.
         """
         with self.lock:
-            response = self.typdef_client.get(
+            response = self.client.typedef.get(
                 type_category=[
                     AtlanTypeCategory.CLASSIFICATION,
                     AtlanTypeCategory.STRUCT,

pyatlan/cache/connection_cache.py

@@ -1,21 +1,22 @@
 # SPDX-License-Identifier: Apache-2.0
-# Copyright 2024 Atlan Pte. Ltd.
+# Copyright 2025 Atlan Pte. Ltd.
 from __future__ import annotations
 
 import logging
 import threading
-from typing import Dict, Optional, Union
+from typing import TYPE_CHECKING, Optional, Union
 
 from pyatlan.cache.abstract_asset_cache import AbstractAssetCache, AbstractAssetName
-from pyatlan.client.atlan import AtlanClient
 from pyatlan.model.assets import Asset, Connection
 from pyatlan.model.enums import AtlanConnectorType
 from pyatlan.model.fluent_search import FluentSearch
 from pyatlan.model.search import Term
 
-LOGGER = logging.getLogger(__name__)
+if TYPE_CHECKING:
+    from pyatlan.client.atlan import AtlanClient
 
 lock = threading.Lock()
+LOGGER = logging.getLogger(__name__)
 
 
 class ConnectionCache(AbstractAssetCache):
@@ -37,24 +38,11 @@ class ConnectionCache(AbstractAssetCache):
         Connection.CONNECTOR_NAME,
     ]
     SEARCH_ATTRIBUTES = [field.atlan_field_name for field in _SEARCH_FIELDS]
-    caches: Dict[int, ConnectionCache] = dict()
 
     def __init__(self, client: AtlanClient):
         super().__init__(client)
 
-    @classmethod
-    def get_cache(cls) -> ConnectionCache:
-        from pyatlan.client.atlan import AtlanClient
-
-        with lock:
-            default_client = AtlanClient.get_default_client()
-            cache_key = default_client.cache_key
-            if cache_key not in cls.caches:
-                cls.caches[cache_key] = ConnectionCache(client=default_client)
-            return cls.caches[cache_key]
-
-    @classmethod
-    def get_by_guid(cls, guid: str, allow_refresh: bool = True) -> Connection:
+    def get_by_guid(self, guid: str, allow_refresh: bool = True) -> Connection:
         """
         Retrieve a connection from the cache by its UUID.
         If the asset is not found, it will be looked up and added to the cache.
@@ -66,11 +54,10 @@ def get_by_guid(cls, guid: str, allow_refresh: bool = True) -> Connection:
         :raises NotFoundError: if the connection cannot be found (does not exist) in Atlan
         :raises InvalidRequestError: if no UUID was provided for the connection to retrieve
         """
-        return cls.get_cache()._get_by_guid(guid=guid, allow_refresh=allow_refresh)
+        return self._get_by_guid(guid=guid, allow_refresh=allow_refresh)
 
-    @classmethod
     def get_by_qualified_name(
-        cls, qualified_name: str, allow_refresh: bool = True
+        self, qualified_name: str, allow_refresh: bool = True
     ) -> Connection:
         """
         Retrieve a connection from the cache by its unique Atlan-internal name.
@@ -84,13 +71,12 @@ def get_by_qualified_name(
         :raises NotFoundError: if the connection cannot be found (does not exist) in Atlan
         :raises InvalidRequestError: if no qualified_name was provided for the connection to retrieve
         """
-        return cls.get_cache()._get_by_qualified_name(
+        return self._get_by_qualified_name(
             qualified_name=qualified_name, allow_refresh=allow_refresh
         )
 
-    @classmethod
     def get_by_name(
-        cls, name: ConnectionName, allow_refresh: bool = True
+        self, name: ConnectionName, allow_refresh: bool = True
     ) -> Connection:
         """
         Retrieve an connection from the cache by its uniquely identifiable name.
@@ -104,7 +90,7 @@ def get_by_name(
         :raises NotFoundError: if the connection cannot be found (does not exist) in Atlan
         :raises InvalidRequestError: if no name was provided for the connection to retrieve
         """
-        return cls.get_cache()._get_by_name(name=name, allow_refresh=allow_refresh)
+        return self._get_by_name(name=name, allow_refresh=allow_refresh)
 
     def lookup_by_guid(self, guid: str) -> None:
         if not guid:
@@ -139,21 +125,22 @@ def lookup_by_qualified_name(self, connection_qn: str) -> None:
     def lookup_by_name(self, name: ConnectionName) -> None:
         if not isinstance(name, ConnectionName):
             return
-        results = self.client.asset.find_connections_by_name(
-            name=name.name,
-            connector_type=name.type,
-            attributes=self.SEARCH_ATTRIBUTES,
-        )
-        if not results:
-            return
-        if len(results) > 1:
-            LOGGER.warning(
-                (
-                    "Found multiple connections of the same type with the same name, caching only the first: %s"
-                ),
-                name,
+        with self.lock:
+            results = self.client.asset.find_connections_by_name(
+                name=name.name,  # type: ignore[arg-type]
+                connector_type=name.type,  # type: ignore[arg-type]
+                attributes=self.SEARCH_ATTRIBUTES,
             )
-        self.cache(results[0])
+            if not results:
+                return
+            if len(results) > 1:
+                LOGGER.warning(
+                    (
+                        "Found multiple connections of the same type with the same name, caching only the first: %s"
+                    ),
+                    name,
+                )
+            self.cache(results[0])
 
     def get_name(self, asset: Asset):
         if not isinstance(asset, Connection):
@@ -188,7 +175,7 @@ def __init__(
         elif isinstance(connection, str):
             tokens = connection.split("/")
             if len(tokens) > 1:
-                self.type = AtlanConnectorType(tokens[0])  # type: ignore[call-arg]
+                self.type = AtlanConnectorType(tokens[0]).value  # type: ignore[call-arg]
                 self.name = connection[len(tokens[0]) + 1 :]  # noqa
 
     def __hash__(self):