Merge pull request #637 from atlanhq/APP-6827

APP-6827 Add utility method to facilitate processing assets while iterating thru results of a fluent search.

Author: Aryamanz29

pyatlan/client/asset.py

@@ -9,11 +9,13 @@
 from enum import Enum
 from typing import (
     TYPE_CHECKING,
+    Callable,
     Dict,
     Generator,
     Iterable,
     List,
     Optional,
+    Protocol,
     Set,
     Type,
     TypeVar,
@@ -134,6 +136,11 @@
 LOGGER = logging.getLogger(__name__)
 
 
+class IndexSearchRequestProvider(Protocol):
+    def to_request(self) -> IndexSearchRequest:
+        pass
+
+
 class AssetClient:
     """
     This class can be used to retrieve information about assets. This class does not need to be instantiated
@@ -1861,6 +1868,45 @@ def get_hierarchy(
             )
         return CategoryHierarchy(top_level=top_categories, stub_dict=category_dict)
 
+    def process_assets(
+        self, search: IndexSearchRequestProvider, func: Callable[[Asset], None]
+    ) -> int:
+        """
+        Process assets matching a search query and apply a processing function to each unique asset.
+
+        This function iteratively searches for assets using the search provider and processes each
+        unique asset using the provided callable function. The uniqueness of assets is determined
+        based on their GUIDs. If new assets are found in later iterations that haven't been
+        processed yet, the process continues until no more new assets are available to process.
+
+        Arguments:
+            search: IndexSearchRequestProvider
+                The search provider that generates search queries and contains the criteria for
+                searching the assets such as a FluentSearch.
+            func: Callable[[Asset], None]
+                A callable function that receives each unique asset as its parameter and performs
+                the required operations on it.
+
+        Returns:
+            int: The total number of unique assets that have been processed.
+        """
+        guids_processed: set[str] = set()
+        has_assets_to_process: bool = True
+        iteration_count = 0
+        while has_assets_to_process:
+            iteration_count += 1
+            has_assets_to_process = False
+            response = self.search(search.to_request())
+            LOGGER.debug(
+                "Iteration %d found %d assets.", iteration_count, response.count
+            )
+            for asset in response:
+                if asset.guid not in guids_processed:
+                    guids_processed.add(asset.guid)
+                    has_assets_to_process = True
+                    func(asset)
+        return len(guids_processed)
+
 
 class SearchResults(ABC, Iterable):
     """

tests/integration/test_client.py

@@ -54,6 +54,7 @@
 MODULE_NAME = TestId.make_unique("Client")
 TEST_USER_DESCRIPTION = "Automated testing of the Python SDK. (USER)"
 TEST_SYSTEM_DESCRIPTION = "Automated testing of the Python SDK. (SYSTEM)"
+call_count = 0
 
 
 @pytest.fixture(scope="module")
@@ -1536,3 +1537,38 @@ def test_client_401_token_refresh(
     # Confirm the API key has been updated and results are returned
     assert client.api_key != expired_api_token
     assert results and results.count >= 1
+
+
+def test_process_assets_when_no_assets_found(client: AtlanClient):
+    def should_never_be_called(_: Asset):
+        pytest.fail("Should not be called")
+
+    search = (
+        FluentSearch()
+        .where(Term.with_state("ACTIVE"))
+        .where(Asset.NAME.startswith("zXZ"))
+    )
+
+    processed_count = client.asset.process_assets(
+        search=search, func=should_never_be_called
+    )
+    assert processed_count == 0
+
+
+def test_process_assets_when_assets_found(client: AtlanClient):
+    def doit(asset: Asset):
+        global call_count
+        call_count += 1
+
+    search = (
+        FluentSearch()
+        .where(Term.with_state("ACTIVE"))
+        .where(Asset.TYPE_NAME.eq("Table"))
+        .where(Asset.NAME.startswith("B"))
+    )
+    expected_count = client.asset.search(search.to_request()).count
+
+    processed_count = client.asset.process_assets(search=search, func=doit)
+
+    assert call_count == expected_count
+    assert processed_count == expected_count