-
Notifications
You must be signed in to change notification settings - Fork 121
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[Internal] Update Jobs GetRun API to support paginated responses for jobs and ForEach tasks with >100 runs #725
base: main
Are you sure you want to change the base?
Changes from all commits
4ff5852
895e5d0
4e7ed82
0abf0de
fab4fab
0465c5f
12bd298
ab675c5
518ebf0
a0c8c2b
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,51 @@ | ||
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 |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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") | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Will users ever specify an initial token? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. At this point users cannot get initial page token from any SDK function. So this parameter is redundant. The function that is overloaded is generated from openApi spec and it also comes with the page_token param. So I assumed that it's ok to repeat the same function contract in the mixin as well.
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Discussed offline. Let's leave the parameter there so the mixin matches the underlying API. We'll clean up token parameters separately and can remove it here when we do that. |
||
|
||
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") | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. ditto about the initial token? |
||
|
||
assert run.as_dict() == { | ||
"tasks": [{ | ||
'run_id': 1337 | ||
}], | ||
"iterations": [{ | ||
'run_id': 0 | ||
}, { | ||
'run_id': 1 | ||
}, { | ||
'run_id': 2 | ||
}, { | ||
'run_id': 3 | ||
}, { | ||
'run_id': 4 | ||
}], | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Should we have a doccomment here explaining that this makes multiple requests to return all tasks/iterations when the underlying has more than 100?