Made major changes to workflow iterative pagination

Author: vaibhavatlan

pyatlan/client/workflow.py

@@ -167,7 +167,7 @@ def find_runs_by_status_and_time_range(
         finished_at: Optional[str] = None,
         from_: int = 0,
         size: int = 100,
-    ) -> List[WorkflowSearchResult]:
+    ) -> WorkflowSearchResponse:
         """
         Find workflow runs based on their status and time range.
 
@@ -205,7 +205,7 @@ def find_runs_by_status_and_time_range(
         run_lookup_results = self._find_runs(
             query=run_lookup_query, from_=from_, size=size
         )
-        return run_lookup_results.hits and run_lookup_results.hits.hits or []
+        return run_lookup_results
 
     @validate_arguments
     def _find_latest_run(self, workflow_name: str) -> Optional[WorkflowSearchResult]:
@@ -280,7 +280,16 @@ def _find_runs(
             WORKFLOW_INDEX_RUN_SEARCH,
             request_obj=request,
         )
-        return WorkflowSearchResponse(**raw_json)
+        return WorkflowSearchResponse(
+            client=self._client,
+            endpoint=WORKFLOW_INDEX_RUN_SEARCH,
+            criteria=query,
+            start=request.from_,
+            size=request.size,
+            took=raw_json.get("took"),
+            hits=raw_json.get("hits"),
+            shards=raw_json.get("_shards"),
+        )
 
     def _add_schedule(
         self,
@@ -517,7 +526,7 @@ def get_runs(
         workflow_phase: AtlanWorkflowPhase,
         from_: int = 0,
         size: int = 100,
-    ) -> Optional[List[WorkflowSearchResult]]:
+    ) -> Optional[WorkflowSearchResponse]:
         """
         Retrieves all workflow runs.
 
@@ -542,7 +551,7 @@ def get_runs(
             filter=[Term(field="status.phase.keyword", value=workflow_phase.value)],
         )
         response = self._find_runs(query, from_=from_, size=size)
-        return results if (results := response.hits and response.hits.hits) else None
+        return response
 
     @validate_arguments
     def stop(

pyatlan/model/workflow.py

@@ -215,17 +215,6 @@ def __init__(__pydantic_self__, **data: Any) -> None:
             ["from_", "size", "track_total_hits", "sort"]
         )
 
-    @property
-    def query_params(self) -> dict:
-        qp: Dict[str, object] = {}
-        if self.post_filter:
-            qp["filter"] = self.post_filter
-        if self.sort:
-            qp["sort"] = self.sort
-        qp["from_"] = self.from_
-        qp["size"] = self.size
-        return qp
-
 
 class WorkflowSearchResponse(AtlanObject):
     _size: int = PrivateAttr()
@@ -243,30 +232,31 @@ def __init__(self, **data: Any):
         self._client = data.get("client")  # type: ignore[assignment]
         self._criteria = data.get("criteria")  # type: ignore[assignment]
         self._size = data.get("size")  # type: ignore[assignment]
-        self._start = data.get("from")  # type: ignore[assignment]
+        self._start = data.get("start")  # type: ignore[assignment]
 
     def current_page(self) -> Optional[List[WorkflowSearchResult]]:
-        return self.records
+        return self.hits.hits  # type: ignore
 
     def next_page(self, start=None, size=None) -> bool:
         self._start = start or self._start + self._size
         if size:
             self._size = size
-        return self._get_next_page() if self.records else False
+        return self._get_next_page() if self.hits.hits else False  # type: ignore
 
     def _get_next_page(self):
-        self._criteria.from_ = self._start
-        self._criteria.size = self._size
+        request = WorkflowSearchRequest(
+            query=self._criteria, from_=self._start, size=self._size
+        )
         raw_json = self._client._call_api(
-            api=self._endpoint.format_path_with_params(),
-            query_params=self._criteria.query_params,
+            api=self._endpoint,
+            request_obj=request,
         )
-        if not raw_json.get("records"):
-            self.records = []
+        if not raw_json.get("hits", {}).get("hits"):
+            self.hits.hits = []
             return False
         try:
-            self.records = parse_obj_as(
-                List[WorkflowSearchResult], raw_json.get("records")
+            self.hits.hits = parse_obj_as(
+                List[WorkflowSearchResult], raw_json["hits"]["hits"]
             )
         except ValidationError as err:
             raise ErrorCode.JSON_ERROR.exception_with_parameters(

tests/integration/test_workflow_client.py

@@ -135,8 +135,9 @@ def test_workflow_get_runs_and_stop(client: AtlanClient, workflow: WorkflowRespo
         workflow_name=workflow.metadata.name, workflow_phase=AtlanWorkflowPhase.RUNNING
     )
     assert runs
-    assert len(runs) == 1
-    run = runs[0]
+    total_runs = runs.hits.hits  # type: ignore
+    assert len(total_runs) == 1  # type: ignore
+    run = total_runs[0]  # type: ignore
     assert run and run.id
     assert workflow.metadata.name and (workflow.metadata.name in run.id)
 
@@ -177,7 +178,7 @@ def test_workflow_get_runs_and_stop(client: AtlanClient, workflow: WorkflowRespo
         [AtlanWorkflowPhase.FAILED], started_at="now-1h"
     )
     assert runs_status
-    workflow_run_status = runs_status[0]
+    workflow_run_status = runs_status.hits.hits[0]  # type: ignore
     start_time = workflow_run_status.source.status.startedAt  # type: ignore
     start_datetime = datetime.strptime(start_time, "%Y-%m-%dT%H:%M:%SZ")  # type: ignore
     start_datetime = start_datetime.replace(tzinfo=timezone.utc)