231 lines
11 KiB
Python
231 lines
11 KiB
Python
from typing import Iterator, Optional
|
|
|
|
from databricks.sdk.service import jobs
|
|
from databricks.sdk.service.jobs import BaseJob, BaseRun, Job, RunType
|
|
|
|
|
|
class JobsExt(jobs.JobsAPI):
|
|
|
|
def list(
|
|
self,
|
|
*,
|
|
expand_tasks: Optional[bool] = None,
|
|
limit: Optional[int] = None,
|
|
name: Optional[str] = None,
|
|
offset: Optional[int] = None,
|
|
page_token: Optional[str] = None,
|
|
) -> Iterator[BaseJob]:
|
|
"""List jobs.
|
|
|
|
Retrieves a list of jobs. If the job has multiple pages of tasks, job_clusters, parameters or environments,
|
|
it will paginate through all pages and aggregate the results.
|
|
|
|
:param expand_tasks: bool (optional)
|
|
Whether to include task and cluster details in the response. Note that in API 2.2, only the first
|
|
100 elements will be shown. Use :method:jobs/get to paginate through all tasks and clusters.
|
|
:param limit: int (optional)
|
|
The number of jobs to return. This value must be greater than 0 and less or equal to 100. The
|
|
default value is 20.
|
|
:param name: str (optional)
|
|
A filter on the list based on the exact (case insensitive) job name.
|
|
:param offset: int (optional)
|
|
The offset of the first job to return, relative to the most recently created job. Deprecated since
|
|
June 2023. Use `page_token` to iterate through the pages instead.
|
|
:param page_token: str (optional)
|
|
Use `next_page_token` or `prev_page_token` returned from the previous request to list the next or
|
|
previous page of jobs respectively.
|
|
|
|
:returns: Iterator over :class:`BaseJob`
|
|
"""
|
|
# fetch jobs with limited elements in top level arrays
|
|
jobs_list = super().list(
|
|
expand_tasks=expand_tasks,
|
|
limit=limit,
|
|
name=name,
|
|
offset=offset,
|
|
page_token=page_token,
|
|
)
|
|
if not expand_tasks:
|
|
yield from jobs_list
|
|
|
|
# fully fetch all top level arrays for each job in the list
|
|
for job in jobs_list:
|
|
if job.has_more:
|
|
job_from_get_call = self.get(job.job_id)
|
|
job.settings.tasks = job_from_get_call.settings.tasks
|
|
job.settings.job_clusters = job_from_get_call.settings.job_clusters
|
|
job.settings.parameters = job_from_get_call.settings.parameters
|
|
job.settings.environments = job_from_get_call.settings.environments
|
|
# Remove has_more fields for each job in the list.
|
|
# This field in Jobs API 2.2 is useful for pagination. It indicates if there are more than 100 tasks or job_clusters in the job.
|
|
# This function hides pagination details from the user. So the field does not play useful role here.
|
|
if hasattr(job, "has_more"):
|
|
delattr(job, "has_more")
|
|
yield job
|
|
|
|
def list_runs(
|
|
self,
|
|
*,
|
|
active_only: Optional[bool] = None,
|
|
completed_only: Optional[bool] = None,
|
|
expand_tasks: Optional[bool] = None,
|
|
job_id: Optional[int] = None,
|
|
limit: Optional[int] = None,
|
|
offset: Optional[int] = None,
|
|
page_token: Optional[str] = None,
|
|
run_type: Optional[RunType] = None,
|
|
start_time_from: Optional[int] = None,
|
|
start_time_to: Optional[int] = None,
|
|
) -> Iterator[BaseRun]:
|
|
"""List job runs.
|
|
|
|
List runs in descending order by start time. If the job has multiple pages of tasks, job_clusters, parameters or repair history,
|
|
it will paginate through all pages and aggregate the results.
|
|
|
|
:param active_only: bool (optional)
|
|
If active_only is `true`, only active runs are included in the results; otherwise, lists both active
|
|
and completed runs. An active run is a run in the `QUEUED`, `PENDING`, `RUNNING`, or `TERMINATING`.
|
|
This field cannot be `true` when completed_only is `true`.
|
|
:param completed_only: bool (optional)
|
|
If completed_only is `true`, only completed runs are included in the results; otherwise, lists both
|
|
active and completed runs. This field cannot be `true` when active_only is `true`.
|
|
:param expand_tasks: bool (optional)
|
|
Whether to include task and cluster details in the response. Note that in API 2.2, only the first
|
|
100 elements will be shown. Use :method:jobs/getrun to paginate through all tasks and clusters.
|
|
:param job_id: int (optional)
|
|
The job for which to list runs. If omitted, the Jobs service lists runs from all jobs.
|
|
:param limit: int (optional)
|
|
The number of runs to return. This value must be greater than 0 and less than 25. The default value
|
|
is 20. If a request specifies a limit of 0, the service instead uses the maximum limit.
|
|
:param offset: int (optional)
|
|
The offset of the first run to return, relative to the most recent run. Deprecated since June 2023.
|
|
Use `page_token` to iterate through the pages instead.
|
|
:param page_token: str (optional)
|
|
Use `next_page_token` or `prev_page_token` returned from the previous request to list the next or
|
|
previous page of runs respectively.
|
|
:param run_type: :class:`RunType` (optional)
|
|
The type of runs to return. For a description of run types, see :method:jobs/getRun.
|
|
:param start_time_from: int (optional)
|
|
Show runs that started _at or after_ this value. The value must be a UTC timestamp in milliseconds.
|
|
Can be combined with _start_time_to_ to filter by a time range.
|
|
:param start_time_to: int (optional)
|
|
Show runs that started _at or before_ this value. The value must be a UTC timestamp in milliseconds.
|
|
Can be combined with _start_time_from_ to filter by a time range.
|
|
|
|
:returns: Iterator over :class:`BaseRun`
|
|
"""
|
|
# fetch runs with limited elements in top level arrays
|
|
runs_list = super().list_runs(
|
|
active_only=active_only,
|
|
completed_only=completed_only,
|
|
expand_tasks=expand_tasks,
|
|
job_id=job_id,
|
|
limit=limit,
|
|
offset=offset,
|
|
page_token=page_token,
|
|
run_type=run_type,
|
|
start_time_from=start_time_from,
|
|
start_time_to=start_time_to,
|
|
)
|
|
|
|
if not expand_tasks:
|
|
yield from runs_list
|
|
|
|
# fully fetch all top level arrays for each run in the list
|
|
for run in runs_list:
|
|
if run.has_more:
|
|
run_from_get_call = self.get_run(run.run_id)
|
|
run.tasks = run_from_get_call.tasks
|
|
run.job_clusters = run_from_get_call.job_clusters
|
|
run.job_parameters = run_from_get_call.job_parameters
|
|
run.repair_history = run_from_get_call.repair_history
|
|
# Remove has_more fields for each run in the list.
|
|
# This field in Jobs API 2.2 is useful for pagination. It indicates if there are more than 100 tasks or job_clusters in the run.
|
|
# This function hides pagination details from the user. So the field does not play useful role here.
|
|
if hasattr(run, "has_more"):
|
|
delattr(run, "has_more")
|
|
yield run
|
|
|
|
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:
|
|
"""Get a single job run.
|
|
|
|
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.
|
|
|
|
: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 of job tasks, set this field to the value of the `next_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
|
|
|
|
# runs/get response includes next_page_token as long as there are more pages to fetch.
|
|
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)
|
|
# Each new page of runs/get response includes the next page of the job_clusters, job_parameters, and repair history.
|
|
run.job_clusters.extend(next_run.job_clusters)
|
|
run.job_parameters.extend(next_run.job_parameters)
|
|
run.repair_history.extend(next_run.repair_history)
|
|
run.next_page_token = next_run.next_page_token
|
|
|
|
return run
|
|
|
|
def get(self, job_id: int, *, page_token: Optional[str] = None) -> Job:
|
|
"""Get a single job.
|
|
|
|
Retrieves the details for a single job. If the job has multiple pages of tasks, job_clusters, parameters or environments,
|
|
it will paginate through all pages and aggregate the results.
|
|
|
|
:param job_id: int
|
|
The canonical identifier of the job to retrieve information about. This field is required.
|
|
:param page_token: str (optional)
|
|
Use `next_page_token` returned from the previous GetJob to request the next page of the job's
|
|
sub-resources.
|
|
|
|
:returns: :class:`Job`
|
|
"""
|
|
job = super().get(job_id, page_token=page_token)
|
|
|
|
# jobs/get response includes next_page_token as long as there are more pages to fetch.
|
|
while job.next_page_token is not None:
|
|
next_job = super().get(job_id, page_token=job.next_page_token)
|
|
# Each new page of jobs/get response includes the next page of the tasks, job_clusters, job_parameters, and environments.
|
|
job.settings.tasks.extend(next_job.settings.tasks)
|
|
job.settings.job_clusters.extend(next_job.settings.job_clusters)
|
|
job.settings.parameters.extend(next_job.settings.parameters)
|
|
job.settings.environments.extend(next_job.settings.environments)
|
|
job.next_page_token = next_job.next_page_token
|
|
|
|
return job
|