Merge pull request #643 from atlanhq/APP-6766

APP-6766: Added default `typeName=Referenceable` for must clauses

Author: Aryamanz29

pyatlan/client/asset.py

@@ -172,30 +172,39 @@ def _get_bulk_search_log_message(self, bulk):
     @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'.
+        Ensures that at least one 'typeName' filter is present in both 'must' and 'filter' clauses.
+        If missing in either, appends a default filter for 'Referenceable' to that clause.
         """
         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
-        )
+        query = criteria.dsl.query
+        default_filter = Term.with_super_type_names(Referenceable.__name__)
+        type_field = Referenceable.TYPE_NAME.keyword_field_name
 
-        if not has_type_filter:
-            criteria.dsl.query.filter.append(
-                Term.with_super_type_names(Referenceable.__name__)
+        def needs_type_filter(clause: Optional[List]) -> bool:
+            return not any(
+                isinstance(f, (Term, Terms)) and f.field == type_field
+                for f in clause or []
             )
 
+        # Update 'filter' clause if needed
+        if needs_type_filter(query.filter):
+            if query.filter is None:
+                query.filter = []
+            query.filter.append(default_filter)
+
+        # Update 'must' clause if needed
+        if needs_type_filter(query.must):
+            if query.must is None:
+                query.must = []
+            query.must.append(default_filter)
+
     # 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

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, TermAttributes
+from pyatlan.model.search import DSL, Bool, IndexSearchRequest, Term, TermAttributes
 from pyatlan.model.search_log import SearchLogRequest
 from pyatlan.model.typedef import EnumDef
 from pyatlan.model.user import AtlanUser, UserRequest
@@ -1606,7 +1606,7 @@ def test_index_search_with_no_aggregation_results(
     mock_api_caller.reset_mock()
 
 
-def test_type_name_in_asset_search(mock_api_caller):
+def test_type_name_in_asset_search_bool_filter(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)
@@ -1656,6 +1656,58 @@ def test_type_name_in_asset_search(mock_api_caller):
     assert has_type_filter is False
 
 
+def test_type_name_in_asset_search_bool_must(mock_api_caller):
+    # When the type name is not present in the request
+    query = Bool(must=[Term.with_state("ACTIVE")])
+    request = IndexSearchRequest(dsl=DSL(query=query))
+
+    client = AssetClient(mock_api_caller)
+    client._ensure_type_filter_present(request)
+
+    assert request.dsl.query and request.dsl.query.must
+    assert isinstance(request.dsl.query.must, list)
+
+    has_type_filter = any(
+        isinstance(f, Term) and f.field == TermAttributes.SUPER_TYPE_NAMES.value
+        for f in request.dsl.query.must
+    )
+    assert has_type_filter is True
+
+    # When the type name is present in the request (no need to add super type filter)
+    query = Bool(must=[Term.with_state("ACTIVE"), Term.with_type_name("AtlasGlossary")])
+    request = IndexSearchRequest(dsl=DSL(query=query))
+    client._ensure_type_filter_present(request)
+
+    assert request.dsl.query and request.dsl.query.must
+    assert isinstance(request.dsl.query.must, list)
+
+    has_type_filter = any(
+        isinstance(f, Term) and f.field == TermAttributes.SUPER_TYPE_NAMES.value
+        for f in request.dsl.query.must
+    )
+    assert has_type_filter is False
+
+    # When multiple type name(s) is present in the request (no need to add super type filter)
+    query = Bool(
+        must=[
+            Term.with_state("ACTIVE"),
+            Term.with_type_name("AtlasGlossary"),
+            Term.with_type_name("AtlasGlossaryTerm"),
+        ]
+    )
+    request = IndexSearchRequest(dsl=DSL(query=query))
+    client._ensure_type_filter_present(request)
+
+    assert request.dsl.query and request.dsl.query.must
+    assert isinstance(request.dsl.query.must, list)
+
+    has_type_filter = any(
+        isinstance(f, Term) and f.field == TermAttributes.SUPER_TYPE_NAMES.value
+        for f in request.dsl.query.must
+    )
+    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]