star

venv/lib/python3.9/site-packages/databricks/sdk/mixins/jobs.py

@@ -0,0 +1,230 @@
+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