Skip to content

Commit 2173caa

Browse files
committed
[refactor] Append Referenceable type filter if no typeName is found in search criteria
1 parent c465579 commit 2173caa

File tree

1 file changed

+28
-0
lines changed

1 file changed

+28
-0
lines changed

pyatlan/client/asset.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,33 @@ def _get_bulk_search_log_message(self, bulk):
164164
+ "Ignoring requests for offset-based paging and using timestamp-based paging instead."
165165
)
166166

167+
@staticmethod
168+
def _ensure_type_filter_present(criteria: IndexSearchRequest) -> None:
169+
"""
170+
Ensures that at least one 'typeName' filter is present in the given search criteria.
171+
If no such filter exists, appends a default filter for 'Referenceable'.
172+
"""
173+
if not (
174+
criteria
175+
and criteria.dsl
176+
and criteria.dsl.query
177+
and isinstance(criteria.dsl.query, Bool)
178+
and criteria.dsl.query.filter
179+
and isinstance(criteria.dsl.query.filter, list)
180+
):
181+
return
182+
183+
has_type_filter = any(
184+
isinstance(f, Term)
185+
and f.field == Referenceable.TYPE_NAME.keyword_field_name
186+
for f in criteria.dsl.query.filter
187+
)
188+
189+
if not has_type_filter:
190+
criteria.dsl.query.filter.append(
191+
Term.with_super_type_names(Referenceable.__name__)
192+
)
193+
167194
# TODO: Try adding @validate_arguments to this method once
168195
# the issue below is fixed or when we switch to pydantic v2
169196
# https://github.com/atlanhq/atlan-python/pull/88#discussion_r1260892704
@@ -194,6 +221,7 @@ def search(self, criteria: IndexSearchRequest, bulk=False) -> IndexSearchResults
194221
raise ErrorCode.UNABLE_TO_RUN_BULK_WITH_SORTS.exception_with_parameters()
195222
criteria.dsl.sort = self._prepare_sorts_for_bulk_search(criteria.dsl.sort)
196223
LOGGER.debug(self._get_bulk_search_log_message(bulk))
224+
self._ensure_type_filter_present(criteria)
197225
raw_json = self._client._call_api(
198226
INDEX_SEARCH,
199227
request_obj=criteria,

0 commit comments

Comments
 (0)