FT-853: Implement query rewrite for Searchlog responses

pyatlan/client/search_log.py

@@ -1,10 +1,12 @@
+import logging
 from typing import List, Union
 
 from pydantic.v1 import ValidationError, parse_obj_as, validate_arguments
 
 from pyatlan.client.common import ApiCaller
 from pyatlan.client.constants import SEARCH_LOG
 from pyatlan.errors import ErrorCode
+from pyatlan.model.search import SortItem
 from pyatlan.model.search_log import (
     AssetViews,
     SearchLogEntry,
@@ -16,6 +18,7 @@
 
 UNIQUE_USERS = "uniqueUsers"
 UNIQUE_ASSETS = "uniqueAssets"
+LOGGER = logging.getLogger(__name__)
 
 
 class SearchLogClient:
@@ -67,27 +70,69 @@ def _call_search_api(self, criteria: SearchLogRequest) -> dict:
         :param criteria: An instance of SearchLogRequest detailing the search query, parameters, etc.
         :return: A dictionary representing the raw JSON response from the search API.
         """
-        return self._client._call_api(
-            SEARCH_LOG,
-            request_obj=criteria,
+        return self._client._call_api(SEARCH_LOG, request_obj=criteria)
+
+    @staticmethod
+    def _prepare_sorts_for_sl_bulk_search(
+        sorts: List[SortItem],
+    ) -> List[SortItem]:
+        """
+        Ensures that sorting by creation timestamp is prioritized for search log bulk searches.
+
+        :param sorts: list of existing sorting options.
+        :returns: a modified list of sorting options with creation timestamp as the top priority.
+        """
+        if not SearchLogResults.presorted_by_timestamp(sorts):
+            return SearchLogResults.sort_by_timestamp_first(sorts)
+        return sorts
+
+    def _get_bulk_search_log_message(self, bulk):
+        return (
+            (
+                "Search log bulk search option is enabled. "
+                if bulk
+                else "Result size (%s) exceeds threshold (%s). "
+            )
+            + "Ignoring requests for offset-based paging and using timestamp-based paging instead."
         )
 
     @validate_arguments
     def search(
-        self, criteria: SearchLogRequest
+        self, criteria: SearchLogRequest, bulk=False
     ) -> Union[SearchLogViewResults, SearchLogResults]:
         """
-        Search for assets using the provided criteria.
+        Search for search logs using the provided criteria.
+        `Note:` if the number of results exceeds the predefined threshold
+        (10,000 search logs) this will be automatically converted into an search log `bulk` search.
 
         :param criteria: detailing the search query, parameters, and so on to run
-        :returns: the results of the search
+        :param bulk: whether to run the search to retrieve search logs that match the supplied criteria,
+        for large numbers of results (> `10,000`), defaults to `False`. Note: this will reorder the results
+        (based on creation timestamp) in order to iterate through a large number (more than `10,000`) results.
+        :raises InvalidRequestError:
+
+            - if search log bulk search is enabled (`bulk=True`) and any
+              user-specified sorting options are found in the search request.
+            - if search log bulk search is disabled (`bulk=False`) and the number of results
+              exceeds the predefined threshold (i.e: `10,000` assets)
+              and any user-specified sorting options are found in the search request.
+
         :raises AtlanError: on any API communication issue
+        :returns: the results of the search
         """
+        if bulk:
+            if criteria.dsl.sort and len(criteria.dsl.sort) > 2:
+                raise ErrorCode.UNABLE_TO_RUN_SEARCH_LOG_BULK_WITH_SORTS.exception_with_parameters()
+            criteria.dsl.sort = self._prepare_sorts_for_sl_bulk_search(
+                criteria.dsl.sort
+            )
+            LOGGER.debug(self._get_bulk_search_log_message(bulk))
         user_views = []
         asset_views = []
         log_entries = []
         raw_json = self._call_search_api(criteria)
         count = raw_json.get("approximateCount", 0)
+
         if "aggregations" in raw_json and UNIQUE_USERS in raw_json.get(
             "aggregations", {}
         ):
@@ -140,6 +185,21 @@ def search(
                 raise ErrorCode.JSON_ERROR.exception_with_parameters(
                     raw_json, 200, str(err)
                 ) from err
+        if (
+            count > SearchLogResults._MASS_EXTRACT_THRESHOLD
+            and not SearchLogResults.presorted_by_timestamp(criteria.dsl.sort)
+        ):
+            if criteria.dsl.sort and len(criteria.dsl.sort) > 2:
+                raise ErrorCode.UNABLE_TO_RUN_SEARCH_LOG_BULK_WITH_SORTS.exception_with_parameters()
+            criteria.dsl.sort = self._prepare_sorts_for_sl_bulk_search(
+                criteria.dsl.sort
+            )
+            LOGGER.debug(
+                self._get_bulk_search_log_message(bulk),
+                count,
+                SearchLogResults._MASS_EXTRACT_THRESHOLD,
+            )
+            return self.search(criteria)
         return SearchLogResults(
             client=self._client,
             criteria=criteria,
@@ -148,4 +208,5 @@ def search(
             count=count,
             log_entries=log_entries,
             aggregations={},
+            bulk=bulk,
         )

pyatlan/errors.py

@@ -579,6 +579,14 @@ class ErrorCode(Enum):
         + "audit search request when performing a bulk search.",
         InvalidRequestError,
     )
+    UNABLE_TO_RUN_SEARCH_LOG_BULK_WITH_SORTS = (
+        400,
+        "ATLAN-PYTHON-400-067",
+        "Unable to execute search log bulk search with user-defined sorting options.",
+        "Please ensure that no sorting options are included in your "
+        + "search log search request when performing a bulk search.",
+        InvalidRequestError,
+    )
     MISSING_NAME = (
         400,
         "ATLAN-PYTHON-400-067",