Skip to content

Commit 75bcfd3

Browse files
authored
[Internal] GetRun logic paginates more arrays (#867)
## What changes are proposed in this pull request? The existing code only paginates tasks and iterations. With this PR we update the logic to also paginate job_clusters, job_parameters and repair_history. This changes are needed for [Jobs API 2.2](https://docs.databricks.com/api/workspace/jobs/getrun) compatibility. ## How is this tested? I enabled API 2.2 calls by modifying URL string `/api/2.2/jobs/runs/get` in databricks/sdk/service/jobs.py. Then I ran unit test from tests/test_jobs_mixin.py
1 parent 41f5f4b commit 75bcfd3

File tree

2 files changed

+46
-22
lines changed

2 files changed

+46
-22
lines changed

databricks/sdk/mixins/jobs.py

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,18 +11,20 @@ def get_run(self,
1111
include_history: Optional[bool] = None,
1212
include_resolved_values: Optional[bool] = None,
1313
page_token: Optional[str] = None) -> jobs.Run:
14-
"""
15-
This method fetches the details of a run identified by `run_id`. If the run has multiple pages of tasks or iterations,
16-
it will paginate through all pages and aggregate the results.
14+
"""Get a single job run.
15+
16+
Retrieve the metadata of a run. If a run has multiple pages of tasks, it will paginate through all pages of tasks, iterations, job_clusters, job_parameters, and repair history.
17+
1718
:param run_id: int
1819
The canonical identifier of the run for which to retrieve the metadata. This field is required.
1920
:param include_history: bool (optional)
2021
Whether to include the repair history in the response.
2122
:param include_resolved_values: bool (optional)
2223
Whether to include resolved parameter values in the response.
2324
:param page_token: str (optional)
24-
To list the next page or the previous page of job tasks, set this field to the value of the
25-
`next_page_token` or `prev_page_token` returned in the GetJob response.
25+
To list the next page of job tasks, set this field to the value of the `next_page_token` returned in
26+
the GetJob response.
27+
2628
:returns: :class:`Run`
2729
"""
2830
run = super().get_run(run_id,
@@ -34,6 +36,7 @@ def get_run(self,
3436
# 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.
3537
is_paginating_iterations = run.iterations is not None and len(run.iterations) > 0
3638

39+
# runs/get response includes next_page_token as long as there are more pages to fetch.
3740
while run.next_page_token is not None:
3841
next_run = super().get_run(run_id,
3942
include_history=include_history,
@@ -43,7 +46,10 @@ def get_run(self,
4346
run.iterations.extend(next_run.iterations)
4447
else:
4548
run.tasks.extend(next_run.tasks)
49+
# Each new page of runs/get response includes the next page of the job_clusters, job_parameters, and repair history.
50+
run.job_clusters.extend(next_run.job_clusters)
51+
run.job_parameters.extend(next_run.job_parameters)
52+
run.repair_history.extend(next_run.repair_history)
4653
run.next_page_token = next_run.next_page_token
4754

48-
run.prev_page_token = None
4955
return run

tests/test_jobs_mixin.py

Lines changed: 34 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -22,25 +22,43 @@ def test_get_run_with_no_pagination(config, requests_mock):
2222

2323

2424
def test_get_run_pagination_with_tasks(config, requests_mock):
25+
from databricks.sdk.service import compute, jobs
26+
cluster_spec = compute.ClusterSpec(spark_version="11.3.x-scala2.12",
27+
custom_tags={"ResourceClass": "SingleNode"},
28+
num_workers=0,
29+
node_type_id="Standard_DS3_v2",
30+
)
31+
cluster1 = jobs.JobCluster(job_cluster_key="cluster1", new_cluster=cluster_spec)
32+
cluster2 = jobs.JobCluster(job_cluster_key="cluster2", new_cluster=cluster_spec)
33+
cluster3 = jobs.JobCluster(job_cluster_key="cluster3", new_cluster=cluster_spec)
34+
cluster4 = jobs.JobCluster(job_cluster_key="cluster4", new_cluster=cluster_spec)
2535
run1 = {
2636
"tasks": [{
2737
"run_id": 0
2838
}, {
2939
"run_id": 1
3040
}],
41+
"job_clusters": [cluster1.as_dict(), cluster2.as_dict(), ],
42+
"job_parameters": [{
43+
"name": "param1",
44+
"value": "value1"
45+
}],
3146
"next_page_token": "tokenToSecondPage",
32-
"prev_page_token": "tokenToPreviousPage"
3347
}
3448
run2 = {
3549
"tasks": [{
3650
"run_id": 2
3751
}, {
3852
"run_id": 3
3953
}],
54+
"job_clusters": [cluster3.as_dict(), cluster4.as_dict(), ],
55+
"job_parameters": [{
56+
"name": "param2",
57+
"value": "value2"
58+
}],
4059
"next_page_token": "tokenToThirdPage",
41-
"prev_page_token": "initialToken"
4260
}
43-
run3 = {"tasks": [{"run_id": 4}], "next_page_token": None, "prev_page_token": "tokenToSecondPage"}
61+
run3 = {"tasks": [{"run_id": 4}]}
4462
requests_mock.get(make_path_pattern(1337, "initialToken"), text=json.dumps(run1))
4563
requests_mock.get(make_path_pattern(1337, "tokenToSecondPage"), text=json.dumps(run2))
4664
requests_mock.get(make_path_pattern(1337, "tokenToThirdPage"), text=json.dumps(run3))
@@ -60,6 +78,17 @@ def test_get_run_pagination_with_tasks(config, requests_mock):
6078
}, {
6179
'run_id': 4
6280
}],
81+
"job_clusters": [cluster1.as_dict(),
82+
cluster2.as_dict(),
83+
cluster3.as_dict(),
84+
cluster4.as_dict()],
85+
"job_parameters": [{
86+
"name": "param1",
87+
"value": "value1"
88+
}, {
89+
"name": "param2",
90+
"value": "value2"
91+
}],
6392
}
6493

6594

@@ -74,7 +103,6 @@ def test_get_run_pagination_with_iterations(config, requests_mock):
74103
"run_id": 1
75104
}],
76105
"next_page_token": "tokenToSecondPage",
77-
"prev_page_token": "tokenToPreviousPage"
78106
}
79107
run2 = {
80108
"tasks": [{
@@ -86,18 +114,8 @@ def test_get_run_pagination_with_iterations(config, requests_mock):
86114
"run_id": 3
87115
}],
88116
"next_page_token": "tokenToThirdPage",
89-
"prev_page_token": "initialToken"
90-
}
91-
run3 = {
92-
"tasks": [{
93-
"run_id": 1337
94-
}],
95-
"iterations": [{
96-
"run_id": 4
97-
}],
98-
"next_page_token": None,
99-
"prev_page_token": "tokenToSecondPage"
100117
}
118+
run3 = {"tasks": [{"run_id": 1337}], "iterations": [{"run_id": 4}], }
101119
requests_mock.get(make_path_pattern(1337, "initialToken"), text=json.dumps(run1))
102120
requests_mock.get(make_path_pattern(1337, "tokenToSecondPage"), text=json.dumps(run2))
103121
requests_mock.get(make_path_pattern(1337, "tokenToThirdPage"), text=json.dumps(run3))
@@ -120,4 +138,4 @@ def test_get_run_pagination_with_iterations(config, requests_mock):
120138
}, {
121139
'run_id': 4
122140
}],
123-
}
141+
}

0 commit comments

Comments
 (0)