[Internal] Update Jobs GetRun API to support paginated responses for …

…jobs and ForEach tasks (#819)

## What changes are proposed in this pull request?

Introduces extension for jobs get_run call that paginates tasks and
iterations arrays in the response and returns aggregated response to the
caller. This change is necessary to prepare for jobs API 2.2 release
that serves paginated response. Pagination is over once the
next_page_token is absent from the response. The pagination logic is not
exposed to the customer.

## How is this tested?

Unit tests and manual test

databricks/sdk/mixins/jobs.py

@@ -0,0 +1,49 @@
+from typing import Optional
+
+from databricks.sdk.service import jobs
+
+
+class JobsExt(jobs.JobsAPI):
+
+    def get_run(self,
+                run_id: int,
+                *,
+                include_history: Optional[bool] = None,
+                include_resolved_values: Optional[bool] = None,
+                page_token: Optional[str] = None) -> jobs.Run:
+        """
+        This method fetches the details of a run identified by `run_id`. If the run has multiple pages of tasks or iterations,
+        it will paginate through all pages and aggregate the results.
+        :param run_id: int
+          The canonical identifier of the run for which to retrieve the metadata. This field is required.
+        :param include_history: bool (optional)
+          Whether to include the repair history in the response.
+        :param include_resolved_values: bool (optional)
+          Whether to include resolved parameter values in the response.
+        :param page_token: str (optional)
+          To list the next page or the previous page of job tasks, set this field to the value of the
+          `next_page_token` or `prev_page_token` returned in the GetJob response.
+        :returns: :class:`Run`
+        """
+        run = super().get_run(run_id,
+                              include_history=include_history,
+                              include_resolved_values=include_resolved_values,
+                              page_token=page_token)
+
+        # When querying a Job run, a page token is returned when there are more than 100 tasks. No iterations are defined for a Job run. Therefore, the next page in the response only includes the next page of tasks.
+        # When querying a ForEach task run, a page token is returned when there are more than 100 iterations. Only a single task is returned, corresponding to the ForEach task itself. Therefore, the client only reads the iterations from the next page and not the tasks.
+        is_paginating_iterations = run.iterations is not None and len(run.iterations) > 0
+
+        while run.next_page_token is not None:
+            next_run = super().get_run(run_id,
+                                       include_history=include_history,
+                                       include_resolved_values=include_resolved_values,
+                                       page_token=run.next_page_token)
+            if is_paginating_iterations:
+                run.iterations.extend(next_run.iterations)
+            else:
+                run.tasks.extend(next_run.tasks)
+            run.next_page_token = next_run.next_page_token
+
+        run.prev_page_token = None
+        return run

tests/test_jobs_mixin.py

@@ -0,0 +1,123 @@
+import json
+import re
+from typing import Pattern
+
+from databricks.sdk import WorkspaceClient
+
+
+def make_path_pattern(run_id: int, page_token: str) -> Pattern[str]:
+    return re.compile(
+        rf'{re.escape("http://localhost/api/")}2.\d{re.escape(f"/jobs/runs/get?page_token={page_token}&run_id={run_id}")}'
+    )
+
+
+def test_get_run_with_no_pagination(config, requests_mock):
+    run1 = {"tasks": [{"run_id": 0}, {"run_id": 1}], }
+    requests_mock.get(make_path_pattern(1337, "initialToken"), text=json.dumps(run1))
+    w = WorkspaceClient(config=config)
+
+    run = w.jobs.get_run(1337, page_token="initialToken")
+
+    assert run.as_dict() == {"tasks": [{'run_id': 0}, {'run_id': 1}], }
+
+
+def test_get_run_pagination_with_tasks(config, requests_mock):
+    run1 = {
+        "tasks": [{
+            "run_id": 0
+        }, {
+            "run_id": 1
+        }],
+        "next_page_token": "tokenToSecondPage",
+        "prev_page_token": "tokenToPreviousPage"
+    }
+    run2 = {
+        "tasks": [{
+            "run_id": 2
+        }, {
+            "run_id": 3
+        }],
+        "next_page_token": "tokenToThirdPage",
+        "prev_page_token": "initialToken"
+    }
+    run3 = {"tasks": [{"run_id": 4}], "next_page_token": None, "prev_page_token": "tokenToSecondPage"}
+    requests_mock.get(make_path_pattern(1337, "initialToken"), text=json.dumps(run1))
+    requests_mock.get(make_path_pattern(1337, "tokenToSecondPage"), text=json.dumps(run2))
+    requests_mock.get(make_path_pattern(1337, "tokenToThirdPage"), text=json.dumps(run3))
+    w = WorkspaceClient(config=config)
+
+    run = w.jobs.get_run(1337, page_token="initialToken")
+
+    assert run.as_dict() == {
+        "tasks": [{
+            'run_id': 0
+        }, {
+            'run_id': 1
+        }, {
+            'run_id': 2
+        }, {
+            'run_id': 3
+        }, {
+            'run_id': 4
+        }],
+    }
+
+
+def test_get_run_pagination_with_iterations(config, requests_mock):
+    run1 = {
+        "tasks": [{
+            "run_id": 1337
+        }],
+        "iterations": [{
+            "run_id": 0
+        }, {
+            "run_id": 1
+        }],
+        "next_page_token": "tokenToSecondPage",
+        "prev_page_token": "tokenToPreviousPage"
+    }
+    run2 = {
+        "tasks": [{
+            "run_id": 1337
+        }],
+        "iterations": [{
+            "run_id": 2
+        }, {
+            "run_id": 3
+        }],
+        "next_page_token": "tokenToThirdPage",
+        "prev_page_token": "initialToken"
+    }
+    run3 = {
+        "tasks": [{
+            "run_id": 1337
+        }],
+        "iterations": [{
+            "run_id": 4
+        }],
+        "next_page_token": None,
+        "prev_page_token": "tokenToSecondPage"
+    }
+    requests_mock.get(make_path_pattern(1337, "initialToken"), text=json.dumps(run1))
+    requests_mock.get(make_path_pattern(1337, "tokenToSecondPage"), text=json.dumps(run2))
+    requests_mock.get(make_path_pattern(1337, "tokenToThirdPage"), text=json.dumps(run3))
+    w = WorkspaceClient(config=config)
+
+    run = w.jobs.get_run(1337, page_token="initialToken")
+
+    assert run.as_dict() == {
+        "tasks": [{
+            'run_id': 1337
+        }],
+        "iterations": [{
+            'run_id': 0
+        }, {
+            'run_id': 1
+        }, {
+            'run_id': 2
+        }, {
+            'run_id': 3
+        }, {
+            'run_id': 4
+        }],
+    }