Merge pull request #633 from atlanhq/APP-6766

APP-6766: Prevent pagination loop break by filtering invalid assets

Author: Aryamanz29

pyatlan/client/asset.py

@@ -99,6 +99,7 @@
     Range,
     SortItem,
     Term,
+    Terms,
     with_active_category,
     with_active_glossary,
     with_active_term,
@@ -171,6 +172,33 @@ def _get_bulk_search_log_message(self, bulk):
             + "Ignoring requests for offset-based paging and using timestamp-based paging instead."
         )
 
+    @staticmethod
+    def _ensure_type_filter_present(criteria: IndexSearchRequest) -> None:
+        """
+        Ensures that at least one 'typeName' filter is present in the given search criteria.
+        If no such filter exists, appends a default filter for 'Referenceable'.
+        """
+        if not (
+            criteria
+            and criteria.dsl
+            and criteria.dsl.query
+            and isinstance(criteria.dsl.query, Bool)
+            and criteria.dsl.query.filter
+            and isinstance(criteria.dsl.query.filter, list)
+        ):
+            return
+
+        has_type_filter = any(
+            isinstance(f, (Term, Terms))
+            and f.field == Referenceable.TYPE_NAME.keyword_field_name
+            for f in criteria.dsl.query.filter
+        )
+
+        if not has_type_filter:
+            criteria.dsl.query.filter.append(
+                Term.with_super_type_names(Referenceable.__name__)
+            )
+
     # TODO: Try adding @validate_arguments to this method once
     # the issue below is fixed or when we switch to pydantic v2
     # https://github.com/atlanhq/atlan-python/pull/88#discussion_r1260892704
@@ -201,6 +229,7 @@ def search(self, criteria: IndexSearchRequest, bulk=False) -> IndexSearchResults
                 raise ErrorCode.UNABLE_TO_RUN_BULK_WITH_SORTS.exception_with_parameters()
             criteria.dsl.sort = self._prepare_sorts_for_bulk_search(criteria.dsl.sort)
             LOGGER.debug(self._get_bulk_search_log_message(bulk))
+        self._ensure_type_filter_present(criteria)
         raw_json = self._client._call_api(
             INDEX_SEARCH,
             request_obj=criteria,
@@ -1975,8 +2004,8 @@ def _get_next_page_json(self, is_bulk_search: bool = False):
         try:
             self._process_entities(raw_json["entities"])
             if is_bulk_search:
-                self._update_first_last_record_creation_times()
                 self._filter_processed_assets()
+                self._update_first_last_record_creation_times()
             return raw_json
 
         except ValidationError as err:

pyatlan/model/audit.py

@@ -341,8 +341,8 @@ def _get_next_page_json(self, is_bulk_search: bool = False):
                 List[EntityAudit], raw_json[ENTITY_AUDITS]
             )
             if is_bulk_search:
-                self._update_first_last_record_creation_times()
                 self._filter_processed_entities()
+                self._update_first_last_record_creation_times()
             return raw_json
         except ValidationError as err:
             raise ErrorCode.JSON_ERROR.exception_with_parameters(

pyatlan/model/search_log.py

@@ -502,8 +502,8 @@ def _get_next_page_json(self, is_bulk_search: bool = False):
         try:
             self._log_entries = parse_obj_as(List[SearchLogEntry], raw_json["logs"])
             if is_bulk_search:
-                self._update_first_last_record_creation_times()
                 self._filter_processed_entities()
+                self._update_first_last_record_creation_times()
             return raw_json
         except ValidationError as err:
             raise ErrorCode.JSON_ERROR.exception_with_parameters(

tests/unit/test_client.py

@@ -55,7 +55,7 @@
 from pyatlan.model.group import GroupRequest
 from pyatlan.model.lineage import LineageListRequest
 from pyatlan.model.response import AssetMutationResponse
-from pyatlan.model.search import Bool, Term
+from pyatlan.model.search import Bool, Term, TermAttributes
 from pyatlan.model.search_log import SearchLogRequest
 from pyatlan.model.typedef import EnumDef
 from pyatlan.model.user import AtlanUser, UserRequest
@@ -1606,6 +1606,56 @@ def test_index_search_with_no_aggregation_results(
     mock_api_caller.reset_mock()
 
 
+def test_type_name_in_asset_search(mock_api_caller):
+    # When the type name is not present in the request
+    request = (FluentSearch().where(CompoundQuery.active_assets())).to_request()
+    client = AssetClient(mock_api_caller)
+    client._ensure_type_filter_present(request)
+
+    assert request.dsl.query and request.dsl.query.filter
+    assert isinstance(request.dsl.query.filter, list)
+
+    has_type_filter = any(
+        isinstance(f, Term) and f.field == TermAttributes.SUPER_TYPE_NAMES.value
+        for f in request.dsl.query.filter
+    )
+    assert has_type_filter is True
+
+    # When the type name is present in the request (no need to add super type filter)
+    request = (
+        FluentSearch()
+        .where(CompoundQuery.active_assets())
+        .where(CompoundQuery.asset_type(AtlasGlossary))
+    ).to_request()
+    client._ensure_type_filter_present(request)
+
+    assert request.dsl.query and request.dsl.query.filter
+    assert isinstance(request.dsl.query.filter, list)
+
+    has_type_filter = any(
+        isinstance(f, Term) and f.field == TermAttributes.SUPER_TYPE_NAMES.value
+        for f in request.dsl.query.filter
+    )
+    assert has_type_filter is False
+
+    # When multiple type name(s) is present in the request (no need to add super type filter)
+    request = (
+        FluentSearch()
+        .where(CompoundQuery.active_assets())
+        .where(CompoundQuery.asset_types([AtlasGlossary, AtlasGlossaryTerm]))
+    ).to_request()
+    client._ensure_type_filter_present(request)
+
+    assert request.dsl.query and request.dsl.query.filter
+    assert isinstance(request.dsl.query.filter, list)
+
+    has_type_filter = any(
+        isinstance(f, Term) and f.field == TermAttributes.SUPER_TYPE_NAMES.value
+        for f in request.dsl.query.filter
+    )
+    assert has_type_filter is False
+
+
 def _assert_search_results(results, response_json, sorts, bulk=False):
     for i, result in enumerate(results):
         assert result and response_json["entities"][i]
@@ -2553,7 +2603,7 @@ def test_atlan_client_headers(client: AtlanClient):
     assert expected == client._session.headers
 
 
-def test_get_all_pagation(group_client, mock_api_caller):
+def test_get_all_pagination(group_client, mock_api_caller):
     mock_page_1 = [
         {"id": "1", "alias": "Group3"},
         {"id": "2", "alias": "Group4"},