[refactor] Migrated caches to use common operations

- Fixed local imports due to recent refactoring.

Author: Aryamanz29

pyatlan/cache/aio/__init__.py

@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 Atlan Pte. Ltd.
+"""
+Async cache modules for Atlan.
+
+This module provides async versions of all cache functionality
+with the same API as the sync versions, just requiring await.
+
+Pattern: All async cache methods reuse shared business logic from pyatlan.cache.common
+to ensure identical behavior with sync cache implementations.
+"""
+
+from .atlan_tag_cache import AsyncAtlanTagCache
+from .connection_cache import AsyncConnectionCache
+from .custom_metadata_cache import AsyncCustomMetadataCache
+from .enum_cache import AsyncEnumCache
+from .group_cache import AsyncGroupCache
+from .role_cache import AsyncRoleCache
+from .source_tag_cache import AsyncSourceTagCache
+from .user_cache import AsyncUserCache
+
+__all__ = [
+    "AsyncAtlanTagCache",
+    "AsyncConnectionCache",
+    "AsyncCustomMetadataCache",
+    "AsyncEnumCache",
+    "AsyncGroupCache",
+    "AsyncRoleCache",
+    "AsyncSourceTagCache",
+    "AsyncUserCache",
+]

pyatlan/cache/aio/abstract_asset_cache.py

@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 Atlan Pte. Ltd.
+from __future__ import annotations
+
+import asyncio
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Dict
+
+from pyatlan.errors import ErrorCode
+from pyatlan.model.assets import Asset
+
+if TYPE_CHECKING:
+    from pyatlan.client.aio import AsyncAtlanClient
+
+
+class AsyncAbstractAssetCache(ABC):
+    """
+    Async base class for reusable components that are common
+    to all caches, where a cache is populated entry-by-entry.
+    """
+
+    def __init__(self, client: AsyncAtlanClient):
+        self.client = client
+        self.lock = asyncio.Lock()
+        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
+    async def lookup_by_guid(self, guid: str):
+        """Abstract method to lookup asset by guid."""
+
+    @abstractmethod
+    async def lookup_by_qualified_name(self, qualified_name: str):
+        """Abstract method to lookup asset by qualified name."""
+
+    @abstractmethod
+    async def lookup_by_name(self, name: Any):
+        """Abstract method to lookup asset by name."""
+
+    @abstractmethod
+    def get_name(self, asset: Asset):
+        """Abstract method to get name from asset."""
+
+    def is_guid_known(self, guid: str) -> bool:
+        """
+        Checks whether the provided Atlan-internal UUID is known.
+        NOTE: will not refresh the cache itself to determine this.
+
+        :param guid: Atlan-internal UUID of the object
+        :returns: `True` if the object is known, `False` otherwise
+        """
+        return guid in self.guid_to_asset
+
+    def is_qualified_name_known(self, qualified_name: str) -> bool:
+        """
+        Checks whether the provided Atlan-internal ID string is known.
+        NOTE: will not refresh the cache itself to determine this.
+
+        :param qualified_name: Atlan-internal ID string of the object
+        :returns: `True` if the object is known, `False` otherwise
+        """
+        return qualified_name in self.qualified_name_to_guid
+
+    def is_name_known(self, name: str) -> bool:
+        """
+        Checks whether the provided Atlan-internal ID string is known.
+        NOTE: will not refresh the cache itself to determine this.
+
+        :param name: human-constructable name of the object
+        :returns: `True` if the object is known, `False` otherwise
+        """
+        return name in self.name_to_guid
+
+    def cache(self, asset: Asset):
+        """
+        Add an entry to the cache.
+
+        :param asset: to be cached
+        """
+        name = asset and self.get_name(asset)
+        if not all([name, asset.guid, asset.qualified_name]):
+            return
+        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]
+
+    async def _get_by_guid(self, guid: str, allow_refresh: bool = True):
+        """
+        Retrieve an asset from the cache by its UUID.
+
+        :param guid: UUID of the asset in Atlan
+        :param allow_refresh: whether to allow a refresh of the cache (`True`) or not (`False`)
+        :returns: the asset (if found)
+        :raises AtlanError: on any API communication problem if the cache needs to be refreshed
+        :raises NotFoundError: if the object cannot be found (does not exist) in Atlan
+        :raises InvalidRequestError: if no UUID was provided for the object to retrieve
+        """
+        if not guid:
+            raise ErrorCode.MISSING_GUID.exception_with_parameters()
+        asset = self.guid_to_asset.get(guid)
+        if not asset and allow_refresh:
+            await self.lookup_by_guid(guid)
+            asset = self.guid_to_asset.get(guid)
+        if not asset:
+            raise ErrorCode.ASSET_NOT_FOUND_BY_GUID.exception_with_parameters(guid)
+        return asset
+
+    async def _get_by_qualified_name(
+        self, qualified_name: str, allow_refresh: bool = True
+    ):
+        """
+        Retrieve an asset from the cache by its qualifiedName.
+
+        :param qualified_name: qualifiedName of the asset in Atlan
+        :param allow_refresh: whether to allow a refresh of the cache (`True`) or not (`False`)
+        :returns: the asset (if found)
+        :raises AtlanError: on any API communication problem if the cache needs to be refreshed
+        :raises NotFoundError: if the object cannot be found (does not exist) in Atlan
+        :raises InvalidRequestError: if no qualified name was provided for the object to retrieve
+        """
+        if not qualified_name:
+            raise ErrorCode.MISSING_QUALIFIED_NAME.exception_with_parameters()
+        guid = self.qualified_name_to_guid.get(qualified_name)
+        if not guid and allow_refresh:
+            await self.lookup_by_qualified_name(qualified_name)
+            guid = self.qualified_name_to_guid.get(qualified_name)
+        if not guid:
+            raise ErrorCode.ASSET_NOT_FOUND_BY_QN.exception_with_parameters(
+                qualified_name
+            )
+
+        return await self._get_by_guid(guid=guid, allow_refresh=False)
+
+    async def _get_by_name(self, name, allow_refresh: bool = True):
+        """
+        Retrieve an asset from the cache by its uniquely identifiable name.
+
+        :param name: uniquely identifiable name of the asset in Atlan
+        :param allow_refresh: whether to allow a refresh of the cache (`True`) or not (`False`)
+        :returns: the asset (if found)
+        :raises AtlanError: on any API communication problem if the cache needs to be refreshed
+        :raises NotFoundError: if the object cannot be found (does not exist) in Atlan
+        :raises InvalidRequestError: if no name was provided for the object to retrieve
+        """
+        from pyatlan.cache.abstract_asset_cache import AbstractAssetName
+
+        if not isinstance(name, AbstractAssetName):
+            raise ErrorCode.MISSING_NAME.exception_with_parameters()
+        guid = self.name_to_guid.get(str(name))
+        if not guid and allow_refresh:
+            await self.lookup_by_name(name)
+            guid = self.name_to_guid.get(str(name))
+        if not guid:
+            raise ErrorCode.ASSET_NOT_FOUND_BY_NAME.exception_with_parameters(
+                name._TYPE_NAME, name
+            )
+
+        return await self._get_by_guid(guid=guid, allow_refresh=False)

pyatlan/cache/aio/atlan_tag_cache.py

@@ -0,0 +1,147 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 Atlan Pte. Ltd.
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING, Dict, Optional, Set
+
+from pyatlan.cache.common import AtlanTagCacheCommon
+from pyatlan.errors import ErrorCode
+from pyatlan.model.enums import AtlanTypeCategory
+from pyatlan.model.typedef import AtlanTagDef
+
+if TYPE_CHECKING:
+    from pyatlan.client.aio import AsyncAtlanClient
+
+
+class AsyncAtlanTagCache:
+    """
+    Async lazily-loaded cache for translating between Atlan-internal ID strings and human-readable names
+    for Atlan tags.
+    """
+
+    def __init__(self, client: AsyncAtlanClient):
+        self.client: AsyncAtlanClient = 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: asyncio.Lock = asyncio.Lock()
+
+    async def refresh_cache(self) -> None:
+        """
+        Refreshes the cache of Atlan tags by requesting the full set of Atlan tags from Atlan.
+        """
+        await self._refresh_cache()
+
+    async 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 await self._get_id_for_name(name=name)
+
+    async 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 await self._get_name_for_id(idstr=idstr)
+
+    async 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 await self._get_source_tags_attr_id(id)
+
+    async def _refresh_cache(self) -> None:
+        """
+        Refreshes the cache of Atlan tags by requesting the full set of Atlan tags from Atlan.
+        """
+        async with self.lock:
+            # Make async API call directly
+            response = await self.client.typedef.get(
+                type_category=[
+                    AtlanTypeCategory.CLASSIFICATION,
+                    AtlanTypeCategory.STRUCT,
+                ]
+            )
+
+            if not response or not response.struct_defs:
+                raise ErrorCode.EXPIRED_API_TOKEN.exception_with_parameters()
+
+            # Process response using shared logic
+            (
+                self.cache_by_id,
+                self.map_id_to_name,
+                self.map_name_to_id,
+                self.map_id_to_source_tags_attr_id,
+            ) = AtlanTagCacheCommon.refresh_cache_data(response)
+
+    async 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
+        """
+        if not self.cache_by_id:
+            await self._refresh_cache()
+        result, should_refresh = AtlanTagCacheCommon.get_id_for_name(
+            name, self.map_name_to_id, self.deleted_names
+        )
+        if should_refresh:
+            await self._refresh_cache()
+            return AtlanTagCacheCommon.get_id_for_name_after_refresh(
+                name, self.map_name_to_id, self.deleted_names
+            )
+        return result
+
+    async 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
+        """
+        if not self.cache_by_id:
+            await self._refresh_cache()
+        result, should_refresh = AtlanTagCacheCommon.get_name_for_id(
+            idstr, self.map_id_to_name, self.deleted_ids
+        )
+        if should_refresh:
+            await self._refresh_cache()
+            return AtlanTagCacheCommon.get_name_for_id_after_refresh(
+                idstr, self.map_id_to_name, self.deleted_ids
+            )
+        return result
+
+    async 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
+        """
+        if not self.cache_by_id:
+            await self._refresh_cache()
+        result, should_refresh = AtlanTagCacheCommon.get_source_tags_attr_id(
+            id, self.map_id_to_source_tags_attr_id, self.deleted_ids
+        )
+        if should_refresh:
+            await self._refresh_cache()
+            return AtlanTagCacheCommon.get_source_tags_attr_id_after_refresh(
+                id, self.map_id_to_source_tags_attr_id, self.deleted_ids
+            )
+        return result