[refactor] Rebased, refactored (dq_template_config_cache) and did cleaned up

Author: Aryamanz29

pyatlan/cache/aio/__init__.py

@@ -13,6 +13,7 @@
 from .atlan_tag_cache import AsyncAtlanTagCache
 from .connection_cache import AsyncConnectionCache
 from .custom_metadata_cache import AsyncCustomMetadataCache
+from .dq_template_config_cache import AsyncDQTemplateConfigCache
 from .enum_cache import AsyncEnumCache
 from .group_cache import AsyncGroupCache
 from .role_cache import AsyncRoleCache
@@ -23,6 +24,7 @@
     "AsyncAtlanTagCache",
     "AsyncConnectionCache",
     "AsyncCustomMetadataCache",
+    "AsyncDQTemplateConfigCache",
     "AsyncEnumCache",
     "AsyncGroupCache",
     "AsyncRoleCache",

pyatlan/cache/aio/dq_template_config_cache.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 Atlan Pte. Ltd.
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING, Dict, Optional
+
+from pyatlan.cache.common.dq_template_config_cache import DQTemplateConfigCacheCommon
+
+if TYPE_CHECKING:
+    from pyatlan.client.aio.client import AsyncAtlanClient
+
+
+class AsyncDQTemplateConfigCache:
+    """
+    Lazily-loaded async cache for DQ rule template configurations to avoid multiple API calls.
+    """
+
+    def __init__(self, client: AsyncAtlanClient):
+        self.client: AsyncAtlanClient = client
+        self._cache: Dict[str, Dict] = {}
+        self._lock: asyncio.Lock = asyncio.Lock()
+        self._initialized: bool = False
+
+    async def refresh_cache(self) -> None:
+        """
+        Refreshes the cache of DQ template configurations by requesting the full set from Atlan.
+        """
+        await self._refresh_cache()
+
+    async def get_template_config(self, rule_type: str) -> Optional[Dict]:
+        """
+        Get template configuration for a specific rule type.
+
+        :param rule_type: The display name of the rule type
+        :returns: Template configuration dict or None if not found
+        """
+        if not self._initialized:
+            await self._refresh_cache()
+
+        return self._cache.get(rule_type)
+
+    async def _refresh_cache(self) -> None:
+        """Refresh the cache by fetching all template configurations."""
+        async with self._lock:
+            if self._initialized:
+                return
+
+            try:
+                search_request = DQTemplateConfigCacheCommon.prepare_search_request()
+                request = search_request.to_request()
+                results = await self.client.asset.search(request)
+
+                success, error = DQTemplateConfigCacheCommon.process_search_results(
+                    results, self._cache
+                )
+
+                if success:
+                    self._initialized = True
+                else:
+                    # If cache refresh fails, mark as initialized to prevent infinite retries
+                    self._initialized = True
+                    if error:
+                        raise error
+            except Exception:
+                # If cache refresh fails, mark as initialized to prevent infinite retries
+                self._initialized = True
+                raise

pyatlan/cache/common/__init__.py

@@ -15,6 +15,7 @@
 # Cache shared logic classes
 from .atlan_tag_cache import AtlanTagCacheCommon
 from .custom_metadata_cache import CustomMetadataCacheCommon
+from .dq_template_config_cache import DQTemplateConfigCacheCommon
 from .enum_cache import EnumCacheCommon
 from .group_cache import GroupCacheCommon
 from .role_cache import RoleCacheCommon
@@ -24,6 +25,7 @@
     # Cache shared logic classes
     "AtlanTagCacheCommon",
     "CustomMetadataCacheCommon",
+    "DQTemplateConfigCacheCommon",
     "EnumCacheCommon",
     "GroupCacheCommon",
     "RoleCacheCommon",

pyatlan/cache/common/dq_template_config_cache.py

@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 Atlan Pte. Ltd.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Dict, Optional, Tuple
+
+from pyatlan.model.assets import Asset
+from pyatlan.model.fluent_search import FluentSearch
+
+if TYPE_CHECKING:
+    pass
+
+
+class DQTemplateConfigCacheCommon:
+    """
+    Common logic for DQ rule template configuration cache operations.
+    Provides shared functionality between sync and async implementations.
+    """
+
+    @classmethod
+    def prepare_search_request(cls) -> FluentSearch:
+        """
+        Prepare the search request for fetching DQ rule template configurations.
+
+        :returns: FluentSearch configured for DQ rule templates
+        """
+        try:
+            from pyatlan.model.assets.core.alpha__d_q_rule_template import (
+                alpha_DQRuleTemplate,
+            )
+
+            return (
+                FluentSearch()
+                .where(Asset.TYPE_NAME.eq(alpha_DQRuleTemplate.__name__))
+                .include_on_results(alpha_DQRuleTemplate.NAME)
+                .include_on_results(alpha_DQRuleTemplate.QUALIFIED_NAME)
+                .include_on_results(alpha_DQRuleTemplate.DISPLAY_NAME)
+                .include_on_results(
+                    alpha_DQRuleTemplate.ALPHADQ_RULE_TEMPLATE_DIMENSION
+                )
+                .include_on_results(alpha_DQRuleTemplate.ALPHADQ_RULE_TEMPLATE_CONFIG)
+            )
+        except ImportError:
+            # If the alpha_DQRuleTemplate is not available, return empty search
+            return FluentSearch()
+
+    @classmethod
+    def process_search_results(
+        cls, search_results, cache: Dict[str, Dict]
+    ) -> Tuple[bool, Optional[Exception]]:
+        """
+        Process search results and populate the cache.
+
+        :param search_results: Iterator of search results
+        :param cache: Cache dictionary to populate
+        :returns: Tuple of (success, exception if any)
+        """
+        try:
+            for result in search_results:
+                template_config = {
+                    "name": result.name,
+                    "qualified_name": result.qualified_name,
+                    "display_name": result.display_name,
+                    "dimension": result.alpha_dq_rule_template_dimension,  # type: ignore
+                    "config": result.alpha_dq_rule_template_config,  # type: ignore
+                }
+                cache[result.display_name] = template_config  # type: ignore
+            return True, None
+        except Exception as e:
+            return False, e

pyatlan/cache/dq_template_config_cache.py

@@ -5,8 +5,7 @@
 import threading
 from typing import TYPE_CHECKING, Dict, Optional
 
-from pyatlan.model.assets import Asset
-from pyatlan.model.fluent_search import FluentSearch
+from pyatlan.cache.common.dq_template_config_cache import DQTemplateConfigCacheCommon
 
 if TYPE_CHECKING:
     from pyatlan.client.atlan import AtlanClient
@@ -23,6 +22,12 @@ def __init__(self, client: AtlanClient):
         self._lock: threading.Lock = threading.Lock()
         self._initialized: bool = False
 
+    def refresh_cache(self) -> None:
+        """
+        Refreshes the cache of DQ template configurations by requesting the full set from Atlan.
+        """
+        self._refresh_cache()
+
     def get_template_config(self, rule_type: str) -> Optional[Dict]:
         """
         Get template configuration for a specific rule type.
@@ -42,36 +47,21 @@ def _refresh_cache(self) -> None:
                 return
 
             try:
-                from pyatlan.model.assets.core.alpha__d_q_rule_template import (
-                    alpha_DQRuleTemplate,
-                )
-
-                request = (
-                    FluentSearch()
-                    .where(Asset.TYPE_NAME.eq(alpha_DQRuleTemplate.__name__))
-                    .include_on_results(alpha_DQRuleTemplate.NAME)
-                    .include_on_results(alpha_DQRuleTemplate.QUALIFIED_NAME)
-                    .include_on_results(alpha_DQRuleTemplate.DISPLAY_NAME)
-                    .include_on_results(
-                        alpha_DQRuleTemplate.ALPHADQ_RULE_TEMPLATE_DIMENSION
-                    )
-                    .include_on_results(
-                        alpha_DQRuleTemplate.ALPHADQ_RULE_TEMPLATE_CONFIG
-                    )
-                ).to_request()
-
+                search_request = DQTemplateConfigCacheCommon.prepare_search_request()
+                request = search_request.to_request()
                 results = self.client.asset.search(request)
-                for result in results:
-                    template_config = {
-                        "name": result.name,
-                        "qualified_name": result.qualified_name,
-                        "display_name": result.display_name,
-                        "dimension": result.alpha_dq_rule_template_dimension,  # type: ignore
-                        "config": result.alpha_dq_rule_template_config,  # type: ignore
-                    }
-                    self._cache[result.display_name] = template_config  # type: ignore
 
-                self._initialized = True
+                success, error = DQTemplateConfigCacheCommon.process_search_results(
+                    results, self._cache
+                )
+
+                if success:
+                    self._initialized = True
+                else:
+                    # If cache refresh fails, mark as initialized to prevent infinite retries
+                    self._initialized = True
+                    if error:
+                        raise error
             except Exception:
                 # If cache refresh fails, mark as initialized to prevent infinite retries
                 self._initialized = True

pyatlan/client/aio/asset.py

@@ -108,29 +108,21 @@ async def search(
         self,
         criteria: IndexSearchRequest,
         bulk=False,
-        concurrent_pages: bool = False,
-        prefetch_pages: int = 1,
     ) -> AsyncIndexSearchResults:
         """
         Async search that reuses shared business logic via Search.
 
         :param criteria: search criteria
         :param bulk: whether to use bulk search mode
-        :param concurrent_pages: whether to enable concurrent page fetching (experimental)
-        :param prefetch_pages: number of pages to prefetch in background (only used with concurrent_pages=True)
-        :returns: AsyncIndexSearchResults or ConcurrentAsyncIndexSearchResults
+        :returns: AsyncIndexSearchResults
         """
         INDEX_SEARCH, request_obj = Search.prepare_request(criteria, bulk)
         raw_json = await self._async_client._call_api(
             INDEX_SEARCH, request_obj=request_obj
         )
-        response = Search.process_response(raw_json, criteria, bulk, self._async_client)
+        response = Search.process_response(raw_json, criteria)
         if Search._check_for_bulk_search(criteria, response["count"], bulk):
-            return await self.search(
-                criteria,
-                concurrent_pages=concurrent_pages,
-                prefetch_pages=prefetch_pages,
-            )
+            return await self.search(criteria, bulk)
 
         return AsyncIndexSearchResults(
             self._async_client,

pyatlan/client/aio/atlan.py

@@ -1,15 +1,18 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright 2025 Atlan Pte. Ltd.
+from __future__ import annotations
 
 import contextlib
-from typing import AsyncGenerator, Optional
+from typing import TYPE_CHECKING, AsyncGenerator, Optional
 
 from httpx_retries import Retry
 from pydantic.v1 import HttpUrl
 
-from pyatlan.client.aio.client import AsyncAtlanClient
 from pyatlan.client.atlan import DEFAULT_RETRY
 
+if TYPE_CHECKING:
+    from pyatlan.client.aio.client import AsyncAtlanClient
+
 
 @contextlib.asynccontextmanager
 async def client_connection(

pyatlan/client/aio/audit.py

@@ -57,7 +57,7 @@ async def search(
         raw_json = await self._client._call_api(endpoint, request_obj=request_obj)
 
         # Process response using shared logic
-        response = AuditSearch.process_response(raw_json, criteria, bulk, self._client)
+        response = AuditSearch.process_response(raw_json)
 
         # Check if we need to convert to bulk search using shared logic
         if AuditSearch.check_for_bulk_search(response["count"], criteria, bulk):

pyatlan/client/aio/client.py

@@ -17,6 +17,7 @@
     AsyncAtlanTagCache,
     AsyncConnectionCache,
     AsyncCustomMetadataCache,
+    AsyncDQTemplateConfigCache,
     AsyncEnumCache,
     AsyncGroupCache,
     AsyncRoleCache,
@@ -95,6 +96,9 @@ class AsyncAtlanClient(AtlanClient):
     _async_custom_metadata_cache: Optional[AsyncCustomMetadataCache] = PrivateAttr(
         default=None
     )
+    _async_dq_template_config_cache: Optional[AsyncDQTemplateConfigCache] = PrivateAttr(
+        default=None
+    )
     _async_enum_cache: Optional[AsyncEnumCache] = PrivateAttr(default=None)
     _async_group_cache: Optional[AsyncGroupCache] = PrivateAttr(default=None)
     _async_role_cache: Optional[AsyncRoleCache] = PrivateAttr(default=None)
@@ -252,6 +256,15 @@ def custom_metadata_cache(self) -> AsyncCustomMetadataCache:
             self._async_custom_metadata_cache = AsyncCustomMetadataCache(client=self)
         return self._async_custom_metadata_cache
 
+    @property
+    def dq_template_config_cache(self) -> AsyncDQTemplateConfigCache:
+        """Get async DQ template config cache with same API as sync"""
+        if self._async_dq_template_config_cache is None:
+            self._async_dq_template_config_cache = AsyncDQTemplateConfigCache(
+                client=self
+            )
+        return self._async_dq_template_config_cache
+
     @property
     def enum_cache(self) -> AsyncEnumCache:
         """Get async enum cache with same API as sync"""

pyatlan/client/aio/user.py

@@ -33,7 +33,7 @@
 from pyatlan.model.user import AtlanUser, UserMinimalResponse, UserRequest
 
 if TYPE_CHECKING:
-    from .client import AsyncAtlanClient
+    from pyatlan.client.aio.client import AsyncAtlanClient
 
 
 class AsyncUserClient: