databricks-sdk 0.44.1__py3-none-any.whl → 0.46.0__py3-none-any.whl

databricks/sdk/mixins/jobs.py

@@ -1,16 +1,159 @@
-from typing import Optional
+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 get_run(self,
-                run_id: int,
-                *,
-                include_history: Optional[bool] = None,
-                include_resolved_values: Optional[bool] = None,
-                page_token: Optional[str] = None) -> jobs.Run:
+    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.
@@ -27,10 +170,12 @@ class JobsExt(jobs.JobsAPI):
 
         :returns: :class:`Run`
         """
-        run = super().get_run(run_id,
-                              include_history=include_history,
-                              include_resolved_values=include_resolved_values,
-                              page_token=page_token)
+        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.
@@ -38,10 +183,12 @@ class JobsExt(jobs.JobsAPI):
 
         # 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)
+            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:
@@ -52,4 +199,32 @@ class JobsExt(jobs.JobsAPI):
             run.repair_history.extend(next_run.repair_history)
             run.next_page_token = next_run.next_page_token
 
-        return run
+        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

databricks/sdk/mixins/open_ai_client.py

@@ -40,8 +40,9 @@ class ServingEndpointsExt(ServingEndpointsAPI):
 
         return OpenAI(
             base_url=self._api._cfg.host + "/serving-endpoints",
-            api_key="no-token", # Passing in a placeholder to pass validations, this will not be used
-            http_client=self._get_authorized_http_client())
+            api_key="no-token",  # Passing in a placeholder to pass validations, this will not be used
+            http_client=self._get_authorized_http_client(),
+        )
 
     def get_langchain_chat_open_ai_client(self, model):
         try:
@@ -54,17 +55,20 @@ class ServingEndpointsExt(ServingEndpointsAPI):
         return ChatOpenAI(
             model=model,
             openai_api_base=self._api._cfg.host + "/serving-endpoints",
-            api_key="no-token", # Passing in a placeholder to pass validations, this will not be used
-            http_client=self._get_authorized_http_client())
-
-    def http_request(self,
-                     conn: str,
-                     method: ExternalFunctionRequestHttpMethod,
-                     path: str,
-                     *,
-                     headers: Optional[Dict[str, str]] = None,
-                     json: Optional[Dict[str, str]] = None,
-                     params: Optional[Dict[str, str]] = None) -> Response:
+            api_key="no-token",  # Passing in a placeholder to pass validations, this will not be used
+            http_client=self._get_authorized_http_client(),
+        )
+
+    def http_request(
+        self,
+        conn: str,
+        method: ExternalFunctionRequestHttpMethod,
+        path: str,
+        *,
+        headers: Optional[Dict[str, str]] = None,
+        json: Optional[Dict[str, str]] = None,
+        params: Optional[Dict[str, str]] = None,
+    ) -> Response:
         """Make external services call using the credentials stored in UC Connection.
         **NOTE:** Experimental: This API may change or be removed in a future release without warning.
         :param conn: str
@@ -84,16 +88,18 @@ class ServingEndpointsExt(ServingEndpointsAPI):
         """
         response = Response()
         response.status_code = 200
-        server_response = super().http_request(connection_name=conn,
-                                               method=method,
-                                               path=path,
-                                               headers=js.dumps(headers) if headers is not None else None,
-                                               json=js.dumps(json) if json is not None else None,
-                                               params=js.dumps(params) if params is not None else None)
+        server_response = super().http_request(
+            connection_name=conn,
+            method=method,
+            path=path,
+            headers=js.dumps(headers) if headers is not None else None,
+            json=js.dumps(json) if json is not None else None,
+            params=js.dumps(params) if params is not None else None,
+        )
 
         # Read the content from the HttpRequestResponse object
         if hasattr(server_response, "contents") and hasattr(server_response.contents, "read"):
-            raw_content = server_response.contents.read() # Read the bytes
+            raw_content = server_response.contents.read()  # Read the bytes
         else:
             raise ValueError("Invalid response from the server.")
 

databricks/sdk/mixins/workspace.py

@@ -1,23 +1,25 @@
-from typing import BinaryIO, Iterator, Optional, Union
+from typing import Any, BinaryIO, Iterator, Optional, Union
 
 from ..core import DatabricksError
 from ..service.workspace import (ExportFormat, ImportFormat, Language,
                                  ObjectInfo, ObjectType, WorkspaceAPI)
 
 
-def _fqcn(x: any) -> str:
-    return f'{x.__module__}.{x.__name__}'
+def _fqcn(x: Any) -> str:
+    return f"{x.__module__}.{x.__name__}"
 
 
 class WorkspaceExt(WorkspaceAPI):
     __doc__ = WorkspaceAPI.__doc__
 
-    def list(self,
-             path: str,
-             *,
-             notebooks_modified_after: Optional[int] = None,
-             recursive: Optional[bool] = False,
-             **kwargs) -> Iterator[ObjectInfo]:
+    def list(
+        self,
+        path: str,
+        *,
+        notebooks_modified_after: Optional[int] = None,
+        recursive: Optional[bool] = False,
+        **kwargs,
+    ) -> Iterator[ObjectInfo]:
         """List workspace objects
 
         :param recursive: bool
@@ -35,13 +37,15 @@ class WorkspaceExt(WorkspaceAPI):
                     continue
                 yield object_info
 
-    def upload(self,
-               path: str,
-               content: Union[bytes, BinaryIO],
-               *,
-               format: Optional[ImportFormat] = None,
-               language: Optional[Language] = None,
-               overwrite: Optional[bool] = False) -> None:
+    def upload(
+        self,
+        path: str,
+        content: Union[bytes, BinaryIO],
+        *,
+        format: Optional[ImportFormat] = None,
+        language: Optional[Language] = None,
+        overwrite: Optional[bool] = False,
+    ) -> None:
         """
         Uploads a workspace object (for example, a notebook or file) or the contents of an entire
         directory (`DBC` format).
@@ -60,31 +64,37 @@ class WorkspaceExt(WorkspaceAPI):
         :param language: Only required if using `ExportFormat.SOURCE`.
         """
         if format is not None and not isinstance(format, ImportFormat):
-            raise ValueError(
-                f'format is expected to be {_fqcn(ImportFormat)}, but got {_fqcn(format.__class__)}')
+            raise ValueError(f"format is expected to be {_fqcn(ImportFormat)}, but got {_fqcn(format.__class__)}")
         if (not format or format == ImportFormat.SOURCE) and not language:
             suffixes = {
-                '.py': Language.PYTHON,
-                '.sql': Language.SQL,
-                '.scala': Language.SCALA,
-                '.R': Language.R
+                ".py": Language.PYTHON,
+                ".sql": Language.SQL,
+                ".scala": Language.SCALA,
+                ".R": Language.R,
             }
             for sfx, lang in suffixes.items():
                 if path.endswith(sfx):
                     language = lang
                     break
         if language is not None and not isinstance(language, Language):
-            raise ValueError(
-                f'language is expected to be {_fqcn(Language)}, but got {_fqcn(language.__class__)}')
-        data = {'path': path}
-        if format: data['format'] = format.value
-        if language: data['language'] = language.value
-        if overwrite: data['overwrite'] = 'true'
+            raise ValueError(f"language is expected to be {_fqcn(Language)}, but got {_fqcn(language.__class__)}")
+        data = {"path": path}
+        if format:
+            data["format"] = format.value
+        if language:
+            data["language"] = language.value
+        if overwrite:
+            data["overwrite"] = "true"
         try:
-            return self._api.do('POST', '/api/2.0/workspace/import', files={'content': content}, data=data)
+            return self._api.do(
+                "POST",
+                "/api/2.0/workspace/import",
+                files={"content": content},
+                data=data,
+            )
         except DatabricksError as e:
-            if e.error_code == 'INVALID_PARAMETER_VALUE':
-                msg = f'Perhaps you forgot to specify the `format=ImportFormat.AUTO`. {e}'
+            if e.error_code == "INVALID_PARAMETER_VALUE":
+                msg = f"Perhaps you forgot to specify the `format=ImportFormat.AUTO`. {e}"
                 raise DatabricksError(message=msg, error_code=e.error_code)
             else:
                 raise e
@@ -100,7 +110,8 @@ class WorkspaceExt(WorkspaceAPI):
                          the request.
         :return:         file-like `io.BinaryIO` of the `path` contents.
         """
-        query = {'path': path, 'direct_download': 'true'}
-        if format: query['format'] = format.value
-        response = self._api.do('GET', '/api/2.0/workspace/export', query=query, raw=True)
+        query = {"path": path, "direct_download": "true"}
+        if format:
+            query["format"] = format.value
+        response = self._api.do("GET", "/api/2.0/workspace/export", query=query, raw=True)
         return response["contents"]