From 3f391459fb572c17db080769702c26c7271047a5 Mon Sep 17 00:00:00 2001 From: Parth Bansal Date: Wed, 8 Jan 2025 17:47:53 +0100 Subject: [PATCH 1/6] update --- .codegen/_openapi_sha | 2 +- .vscode/settings.json | 3 +- databricks/sdk/__init__.py | 11 +-- databricks/sdk/data_plane.py | 21 +++++- databricks/sdk/service/apps.py | 14 ++-- databricks/sdk/service/catalog.py | 1 + databricks/sdk/service/jobs.py | 99 ++++++++++++++++++++++----- databricks/sdk/service/oauth2.py | 78 +++++++++------------ databricks/sdk/service/pipelines.py | 101 ++++++++++++++++++++++++---- databricks/sdk/service/serving.py | 35 ++++++++-- tests/test_data_plane.py | 2 +- 11 files changed, 273 insertions(+), 94 deletions(-) diff --git a/.codegen/_openapi_sha b/.codegen/_openapi_sha index 8622b29ca..dfe78790a 100644 --- a/.codegen/_openapi_sha +++ b/.codegen/_openapi_sha @@ -1 +1 @@ -a6a317df8327c9b1e5cb59a03a42ffa2aabeef6d \ No newline at end of file +779817ed8d63031f5ea761fbd25ee84f38feec0d \ No newline at end of file diff --git a/.vscode/settings.json b/.vscode/settings.json index c36b4db6c..ad2eb823b 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -7,5 +7,6 @@ "python.envFile": "${workspaceFolder}/.databricks/.databricks.env", "databricks.python.envFile": "${workspaceFolder}/.env", "jupyter.interactiveWindow.cellMarker.codeRegex": "^# COMMAND ----------|^# Databricks notebook source|^(#\\s*%%|#\\s*\\|#\\s*In\\[\\d*?\\]|#\\s*In\\[ \\])", - "jupyter.interactiveWindow.cellMarker.default": "# COMMAND ----------" + "jupyter.interactiveWindow.cellMarker.default": "# COMMAND ----------", + "python.analysis.autoImportCompletions": true } diff --git a/databricks/sdk/__init__.py b/databricks/sdk/__init__.py index d27110b86..068069f04 100755 --- a/databricks/sdk/__init__.py +++ b/databricks/sdk/__init__.py @@ -5,7 +5,7 @@ from databricks.sdk import azure from databricks.sdk.credentials_provider import CredentialsStrategy from databricks.sdk.mixins.compute import ClustersExt -from databricks.sdk.mixins.files import DbfsExt, FilesExt +from databricks.sdk.mixins.files import DbfsExt from databricks.sdk.mixins.jobs import JobsExt from databricks.sdk.mixins.open_ai_client import ServingEndpointsExt from databricks.sdk.mixins.workspace import WorkspaceExt @@ -114,13 +114,6 @@ def _make_dbutils(config: client.Config): return runtime_dbutils -def _make_files_client(apiClient: client.ApiClient, config: client.Config): - if config.enable_experimental_files_api_client: - return FilesExt(apiClient, config) - else: - return FilesAPI(apiClient) - - class WorkspaceClient: """ The WorkspaceClient is a client for the workspace-level Databricks REST API. @@ -210,7 +203,7 @@ def __init__(self, self._dbsql_permissions = DbsqlPermissionsAPI(self._api_client) self._experiments = ExperimentsAPI(self._api_client) self._external_locations = ExternalLocationsAPI(self._api_client) - self._files = _make_files_client(self._api_client, self._config) + self._files = FilesAPI(self._api_client) self._functions = FunctionsAPI(self._api_client) self._genie = GenieAPI(self._api_client) self._git_credentials = GitCredentialsAPI(self._api_client) diff --git a/databricks/sdk/data_plane.py b/databricks/sdk/data_plane.py index 6f6ddf80c..2a34a9f43 100644 --- a/databricks/sdk/data_plane.py +++ b/databricks/sdk/data_plane.py @@ -1,9 +1,26 @@ import threading +from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import Callable, List +from typing import Callable, List, Optional from databricks.sdk.oauth import Token -from databricks.sdk.service.oauth2 import DataPlaneInfo + + +@dataclass +class DataPlaneInfo(ABC): + """ + Abstract base class for DataPlane information. + """ + + @property + @abstractmethod + def url(self) -> str: + """Gets the URL for the DataPlane endpoint.""" + + @property + @abstractmethod + def token(self) -> Optional[str]: + """Gets the token for the DataPlane endpoint.""" @dataclass diff --git a/databricks/sdk/service/apps.py b/databricks/sdk/service/apps.py index eee49a212..37af1011d 100755 --- a/databricks/sdk/service/apps.py +++ b/databricks/sdk/service/apps.py @@ -967,12 +967,14 @@ def wait_get_app_stopped(self, attempt += 1 raise TimeoutError(f'timed out after {timeout}: {status_message}') - def create(self, *, app: Optional[App] = None) -> Wait[App]: + def create(self, *, app: Optional[App] = None, no_compute: Optional[bool] = None) -> Wait[App]: """Create an app. Creates a new app. :param app: :class:`App` (optional) + :param no_compute: bool (optional) + If true, the app will not be started after creation. :returns: Long-running operation waiter for :class:`App`. @@ -981,11 +983,15 @@ def create(self, *, app: Optional[App] = None) -> Wait[App]: body = app.as_dict() headers = {'Accept': 'application/json', 'Content-Type': 'application/json', } - op_response = self._api.do('POST', '/api/2.0/apps', body=body, headers=headers) + op_response = self._api.do('POST', '/api/2.0/apps', query=query, body=body, headers=headers) return Wait(self.wait_get_app_active, response=App.from_dict(op_response), name=op_response['name']) - def create_and_wait(self, *, app: Optional[App] = None, timeout=timedelta(minutes=20)) -> App: - return self.create(app=app).result(timeout=timeout) + def create_and_wait(self, + *, + app: Optional[App] = None, + no_compute: Optional[bool] = None, + timeout=timedelta(minutes=20)) -> App: + return self.create(app=app, no_compute=no_compute).result(timeout=timeout) def delete(self, name: str) -> App: """Delete an app. diff --git a/databricks/sdk/service/catalog.py b/databricks/sdk/service/catalog.py index f1b549339..c56acce32 100755 --- a/databricks/sdk/service/catalog.py +++ b/databricks/sdk/service/catalog.py @@ -5810,6 +5810,7 @@ def from_dict(cls, d: Dict[str, any]) -> ProvisioningInfo: class ProvisioningInfoState(Enum): ACTIVE = 'ACTIVE' + DEGRADED = 'DEGRADED' DELETING = 'DELETING' FAILED = 'FAILED' PROVISIONING = 'PROVISIONING' diff --git a/databricks/sdk/service/jobs.py b/databricks/sdk/service/jobs.py index 105c7cd22..c5fdb8393 100755 --- a/databricks/sdk/service/jobs.py +++ b/databricks/sdk/service/jobs.py @@ -35,6 +35,11 @@ class BaseJob: Jobs UI in the job details page and Jobs API using `budget_policy_id` 3. Inferred default based on accessible budget policies of the run_as identity on job creation or modification.""" + has_more: Optional[bool] = None + """Indicates if the job has more sub-resources (`tasks`, `job_clusters`) that are not shown. They + can be accessed via :method:jobs/get endpoint. It is only relevant for API 2.2 :method:jobs/list + requests with `expand_tasks=true`.""" + job_id: Optional[int] = None """The canonical identifier for this job.""" @@ -49,6 +54,7 @@ def as_dict(self) -> dict: if self.creator_user_name is not None: body['creator_user_name'] = self.creator_user_name if self.effective_budget_policy_id is not None: body['effective_budget_policy_id'] = self.effective_budget_policy_id + if self.has_more is not None: body['has_more'] = self.has_more if self.job_id is not None: body['job_id'] = self.job_id if self.settings: body['settings'] = self.settings.as_dict() return body @@ -60,6 +66,7 @@ def as_shallow_dict(self) -> dict: if self.creator_user_name is not None: body['creator_user_name'] = self.creator_user_name if self.effective_budget_policy_id is not None: body['effective_budget_policy_id'] = self.effective_budget_policy_id + if self.has_more is not None: body['has_more'] = self.has_more if self.job_id is not None: body['job_id'] = self.job_id if self.settings: body['settings'] = self.settings return body @@ -70,6 +77,7 @@ def from_dict(cls, d: Dict[str, any]) -> BaseJob: return cls(created_time=d.get('created_time', None), creator_user_name=d.get('creator_user_name', None), effective_budget_policy_id=d.get('effective_budget_policy_id', None), + has_more=d.get('has_more', None), job_id=d.get('job_id', None), settings=_from_dict(d, 'settings', JobSettings)) @@ -124,10 +132,16 @@ class BaseRun: Note: dbt and SQL File tasks support only version-controlled sources. If dbt or SQL File tasks are used, `git_source` must be defined on the job.""" + has_more: Optional[bool] = None + """Indicates if the run has more sub-resources (`tasks`, `job_clusters`) that are not shown. They + can be accessed via :method:jobs/getrun endpoint. It is only relevant for API 2.2 + :method:jobs/listruns requests with `expand_tasks=true`.""" + job_clusters: Optional[List[JobCluster]] = None """A list of job cluster specifications that can be shared and reused by tasks of this job. Libraries cannot be declared in a shared job cluster. You must declare dependent libraries in - task settings.""" + task settings. If more than 100 job clusters are available, you can paginate through them using + :method:jobs/getrun.""" job_id: Optional[int] = None """The canonical identifier of the job that contains this run.""" @@ -198,7 +212,9 @@ class BaseRun: tasks: Optional[List[RunTask]] = None """The list of tasks performed by the run. Each task has its own `run_id` which you can use to call - `JobsGetOutput` to retrieve the run resutls.""" + `JobsGetOutput` to retrieve the run resutls. If more than 100 tasks are available, you can + paginate through them using :method:jobs/getrun. Use the `next_page_token` field at the object + root to determine if more results are available.""" trigger: Optional[TriggerType] = None """The type of trigger that fired this run. @@ -227,6 +243,7 @@ def as_dict(self) -> dict: if self.end_time is not None: body['end_time'] = self.end_time if self.execution_duration is not None: body['execution_duration'] = self.execution_duration if self.git_source: body['git_source'] = self.git_source.as_dict() + if self.has_more is not None: body['has_more'] = self.has_more if self.job_clusters: body['job_clusters'] = [v.as_dict() for v in self.job_clusters] if self.job_id is not None: body['job_id'] = self.job_id if self.job_parameters: body['job_parameters'] = [v.as_dict() for v in self.job_parameters] @@ -264,6 +281,7 @@ def as_shallow_dict(self) -> dict: if self.end_time is not None: body['end_time'] = self.end_time if self.execution_duration is not None: body['execution_duration'] = self.execution_duration if self.git_source: body['git_source'] = self.git_source + if self.has_more is not None: body['has_more'] = self.has_more if self.job_clusters: body['job_clusters'] = self.job_clusters if self.job_id is not None: body['job_id'] = self.job_id if self.job_parameters: body['job_parameters'] = self.job_parameters @@ -301,6 +319,7 @@ def from_dict(cls, d: Dict[str, any]) -> BaseRun: end_time=d.get('end_time', None), execution_duration=d.get('execution_duration', None), git_source=_from_dict(d, 'git_source', GitSource), + has_more=d.get('has_more', None), job_clusters=_repeated_dict(d, 'job_clusters', JobCluster), job_id=d.get('job_id', None), job_parameters=_repeated_dict(d, 'job_parameters', JobParameter), @@ -754,7 +773,8 @@ class CreateJob: job_clusters: Optional[List[JobCluster]] = None """A list of job cluster specifications that can be shared and reused by tasks of this job. Libraries cannot be declared in a shared job cluster. You must declare dependent libraries in - task settings.""" + task settings. If more than 100 job clusters are available, you can paginate through them using + :method:jobs/get.""" max_concurrent_runs: Optional[int] = None """An optional maximum allowed number of concurrent runs of the job. Set this value if you want to @@ -795,7 +815,9 @@ class CreateJob: be added to the job.""" tasks: Optional[List[Task]] = None - """A list of task specifications to be executed by this job.""" + """A list of task specifications to be executed by this job. If more than 100 tasks are available, + you can paginate through them using :method:jobs/get. Use the `next_page_token` field at the + object root to determine if more results are available.""" timeout_seconds: Optional[int] = None """An optional timeout applied to each run of this job. A value of `0` means no timeout.""" @@ -1680,9 +1702,17 @@ class Job: Jobs UI in the job details page and Jobs API using `budget_policy_id` 3. Inferred default based on accessible budget policies of the run_as identity on job creation or modification.""" + has_more: Optional[bool] = None + """Indicates if the job has more sub-resources (`tasks`, `job_clusters`) that are not shown. They + can be accessed via :method:jobs/get endpoint. It is only relevant for API 2.2 :method:jobs/list + requests with `expand_tasks=true`.""" + job_id: Optional[int] = None """The canonical identifier for this job.""" + next_page_token: Optional[str] = None + """A token that can be used to list the next page of sub-resources.""" + run_as_user_name: Optional[str] = None """The email of an active workspace user or the application ID of a service principal that the job runs as. This value can be changed by setting the `run_as` field when creating or updating a @@ -1703,7 +1733,9 @@ def as_dict(self) -> dict: if self.creator_user_name is not None: body['creator_user_name'] = self.creator_user_name if self.effective_budget_policy_id is not None: body['effective_budget_policy_id'] = self.effective_budget_policy_id + if self.has_more is not None: body['has_more'] = self.has_more if self.job_id is not None: body['job_id'] = self.job_id + if self.next_page_token is not None: body['next_page_token'] = self.next_page_token if self.run_as_user_name is not None: body['run_as_user_name'] = self.run_as_user_name if self.settings: body['settings'] = self.settings.as_dict() return body @@ -1715,7 +1747,9 @@ def as_shallow_dict(self) -> dict: if self.creator_user_name is not None: body['creator_user_name'] = self.creator_user_name if self.effective_budget_policy_id is not None: body['effective_budget_policy_id'] = self.effective_budget_policy_id + if self.has_more is not None: body['has_more'] = self.has_more if self.job_id is not None: body['job_id'] = self.job_id + if self.next_page_token is not None: body['next_page_token'] = self.next_page_token if self.run_as_user_name is not None: body['run_as_user_name'] = self.run_as_user_name if self.settings: body['settings'] = self.settings return body @@ -1726,7 +1760,9 @@ def from_dict(cls, d: Dict[str, any]) -> Job: return cls(created_time=d.get('created_time', None), creator_user_name=d.get('creator_user_name', None), effective_budget_policy_id=d.get('effective_budget_policy_id', None), + has_more=d.get('has_more', None), job_id=d.get('job_id', None), + next_page_token=d.get('next_page_token', None), run_as_user_name=d.get('run_as_user_name', None), settings=_from_dict(d, 'settings', JobSettings)) @@ -2366,7 +2402,8 @@ class JobSettings: job_clusters: Optional[List[JobCluster]] = None """A list of job cluster specifications that can be shared and reused by tasks of this job. Libraries cannot be declared in a shared job cluster. You must declare dependent libraries in - task settings.""" + task settings. If more than 100 job clusters are available, you can paginate through them using + :method:jobs/get.""" max_concurrent_runs: Optional[int] = None """An optional maximum allowed number of concurrent runs of the job. Set this value if you want to @@ -2407,7 +2444,9 @@ class JobSettings: be added to the job.""" tasks: Optional[List[Task]] = None - """A list of task specifications to be executed by this job.""" + """A list of task specifications to be executed by this job. If more than 100 tasks are available, + you can paginate through them using :method:jobs/get. Use the `next_page_token` field at the + object root to determine if more results are available.""" timeout_seconds: Optional[int] = None """An optional timeout applied to each run of this job. A value of `0` means no timeout.""" @@ -3663,13 +3702,19 @@ class Run: Note: dbt and SQL File tasks support only version-controlled sources. If dbt or SQL File tasks are used, `git_source` must be defined on the job.""" + has_more: Optional[bool] = None + """Indicates if the run has more sub-resources (`tasks`, `job_clusters`) that are not shown. They + can be accessed via :method:jobs/getrun endpoint. It is only relevant for API 2.2 + :method:jobs/listruns requests with `expand_tasks=true`.""" + iterations: Optional[List[RunTask]] = None """Only populated by for-each iterations. The parent for-each task is located in tasks array.""" job_clusters: Optional[List[JobCluster]] = None """A list of job cluster specifications that can be shared and reused by tasks of this job. Libraries cannot be declared in a shared job cluster. You must declare dependent libraries in - task settings.""" + task settings. If more than 100 job clusters are available, you can paginate through them using + :method:jobs/getrun.""" job_id: Optional[int] = None """The canonical identifier of the job that contains this run.""" @@ -3743,7 +3788,9 @@ class Run: tasks: Optional[List[RunTask]] = None """The list of tasks performed by the run. Each task has its own `run_id` which you can use to call - `JobsGetOutput` to retrieve the run resutls.""" + `JobsGetOutput` to retrieve the run resutls. If more than 100 tasks are available, you can + paginate through them using :method:jobs/getrun. Use the `next_page_token` field at the object + root to determine if more results are available.""" trigger: Optional[TriggerType] = None """The type of trigger that fired this run. @@ -3772,6 +3819,7 @@ def as_dict(self) -> dict: if self.end_time is not None: body['end_time'] = self.end_time if self.execution_duration is not None: body['execution_duration'] = self.execution_duration if self.git_source: body['git_source'] = self.git_source.as_dict() + if self.has_more is not None: body['has_more'] = self.has_more if self.iterations: body['iterations'] = [v.as_dict() for v in self.iterations] if self.job_clusters: body['job_clusters'] = [v.as_dict() for v in self.job_clusters] if self.job_id is not None: body['job_id'] = self.job_id @@ -3811,6 +3859,7 @@ def as_shallow_dict(self) -> dict: if self.end_time is not None: body['end_time'] = self.end_time if self.execution_duration is not None: body['execution_duration'] = self.execution_duration if self.git_source: body['git_source'] = self.git_source + if self.has_more is not None: body['has_more'] = self.has_more if self.iterations: body['iterations'] = self.iterations if self.job_clusters: body['job_clusters'] = self.job_clusters if self.job_id is not None: body['job_id'] = self.job_id @@ -3850,6 +3899,7 @@ def from_dict(cls, d: Dict[str, any]) -> Run: end_time=d.get('end_time', None), execution_duration=d.get('execution_duration', None), git_source=_from_dict(d, 'git_source', GitSource), + has_more=d.get('has_more', None), iterations=_repeated_dict(d, 'iterations', RunTask), job_clusters=_repeated_dict(d, 'job_clusters', JobCluster), job_id=d.get('job_id', None), @@ -7066,6 +7116,7 @@ def create(self, :param job_clusters: List[:class:`JobCluster`] (optional) A list of job cluster specifications that can be shared and reused by tasks of this job. Libraries cannot be declared in a shared job cluster. You must declare dependent libraries in task settings. + If more than 100 job clusters are available, you can paginate through them using :method:jobs/get. :param max_concurrent_runs: int (optional) An optional maximum allowed number of concurrent runs of the job. Set this value if you want to be able to execute multiple runs of the same job concurrently. This is useful for example if you @@ -7097,7 +7148,9 @@ def create(self, clusters, and are subject to the same limitations as cluster tags. A maximum of 25 tags can be added to the job. :param tasks: List[:class:`Task`] (optional) - A list of task specifications to be executed by this job. + A list of task specifications to be executed by this job. If more than 100 tasks are available, you + can paginate through them using :method:jobs/get. Use the `next_page_token` field at the object root + to determine if more results are available. :param timeout_seconds: int (optional) An optional timeout applied to each run of this job. A value of `0` means no timeout. :param trigger: :class:`TriggerSettings` (optional) @@ -7193,19 +7246,28 @@ def export_run(self, run_id: int, *, views_to_export: Optional[ViewsToExport] = res = self._api.do('GET', '/api/2.1/jobs/runs/export', query=query, headers=headers) return ExportRunOutput.from_dict(res) - def get(self, job_id: int) -> Job: + def get(self, job_id: int, *, page_token: Optional[str] = None) -> Job: """Get a single job. Retrieves the details for a single job. + In Jobs API 2.2, requests for a single job support pagination of `tasks` and `job_clusters` when + either exceeds 100 elements. Use the `next_page_token` field to check for more results and pass its + value as the `page_token` in subsequent requests. Arrays with fewer than 100 elements in a page will + be empty on later pages. + :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` """ query = {} if job_id is not None: query['job_id'] = job_id + if page_token is not None: query['page_token'] = page_token headers = {'Accept': 'application/json', } res = self._api.do('GET', '/api/2.1/jobs/get', query=query, headers=headers) @@ -7251,7 +7313,12 @@ def get_run(self, page_token: Optional[str] = None) -> Run: """Get a single job run. - Retrieve the metadata of a run. + Retrieves the metadata of a run. + + In Jobs API 2.2, requests for a single job run support pagination of `tasks` and `job_clusters` when + either exceeds 100 elements. Use the `next_page_token` field to check for more results and pass its + value as the `page_token` in subsequent requests. Arrays with fewer than 100 elements in a page will + be empty on later pages. :param run_id: int The canonical identifier of the run for which to retrieve the metadata. This field is required. @@ -7260,8 +7327,8 @@ def get_run(self, :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. + Use `next_page_token` returned from the previous GetRun to request the next page of the run's + sub-resources. :returns: :class:`Run` """ @@ -7313,7 +7380,8 @@ def list(self, Retrieves a list of jobs. :param expand_tasks: bool (optional) - Whether to include task and cluster details in the response. + 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. @@ -7370,7 +7438,8 @@ def list_runs(self, 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. + 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) diff --git a/databricks/sdk/service/oauth2.py b/databricks/sdk/service/oauth2.py index f7df5a25e..1aac8bc1c 100755 --- a/databricks/sdk/service/oauth2.py +++ b/databricks/sdk/service/oauth2.py @@ -202,35 +202,6 @@ def from_dict(cls, d: Dict[str, any]) -> CreateServicePrincipalSecretResponse: update_time=d.get('update_time', None)) -@dataclass -class DataPlaneInfo: - authorization_details: Optional[str] = None - """Authorization details as a string.""" - - endpoint_url: Optional[str] = None - """The URL of the endpoint for this operation in the dataplane.""" - - def as_dict(self) -> dict: - """Serializes the DataPlaneInfo into a dictionary suitable for use as a JSON request body.""" - body = {} - if self.authorization_details is not None: body['authorization_details'] = self.authorization_details - if self.endpoint_url is not None: body['endpoint_url'] = self.endpoint_url - return body - - def as_shallow_dict(self) -> dict: - """Serializes the DataPlaneInfo into a shallow dictionary of its immediate attributes.""" - body = {} - if self.authorization_details is not None: body['authorization_details'] = self.authorization_details - if self.endpoint_url is not None: body['endpoint_url'] = self.endpoint_url - return body - - @classmethod - def from_dict(cls, d: Dict[str, any]) -> DataPlaneInfo: - """Deserializes the DataPlaneInfo from a dictionary.""" - return cls(authorization_details=d.get('authorization_details', None), - endpoint_url=d.get('endpoint_url', None)) - - @dataclass class DeleteCustomAppIntegrationOutput: @@ -297,8 +268,13 @@ class FederationPolicy: """Description of the federation policy.""" name: Optional[str] = None - """Name of the federation policy. The name must contain only lowercase alphanumeric characters, - numbers, and hyphens. It must be unique within the account.""" + """Resource name for the federation policy. Example values include + `accounts//federationPolicies/my-federation-policy` for Account Federation Policies, + and + `accounts//servicePrincipals//federationPolicies/my-federation-policy` + for Service Principal Federation Policies. Typically an output parameter, which does not need to + be specified in create or update requests. If specified in a request, must match the value in + the request URL.""" oidc_policy: Optional[OidcFederationPolicy] = None """Specifies the policy to use for validating OIDC claims in your federated tokens.""" @@ -961,7 +937,8 @@ def create(self, :param policy: :class:`FederationPolicy` (optional) :param policy_id: str (optional) - The identifier for the federation policy. If unspecified, the id will be assigned by Databricks. + The identifier for the federation policy. The identifier must contain only lowercase alphanumeric + characters, numbers, hyphens, and slashes. If unspecified, the id will be assigned by Databricks. :returns: :class:`FederationPolicy` """ @@ -979,6 +956,7 @@ def delete(self, policy_id: str): """Delete account federation policy. :param policy_id: str + The identifier for the federation policy. """ @@ -993,6 +971,7 @@ def get(self, policy_id: str) -> FederationPolicy: """Get account federation policy. :param policy_id: str + The identifier for the federation policy. :returns: :class:`FederationPolicy` """ @@ -1035,17 +1014,20 @@ def list(self, def update(self, policy_id: str, - update_mask: str, *, - policy: Optional[FederationPolicy] = None) -> FederationPolicy: + policy: Optional[FederationPolicy] = None, + update_mask: Optional[str] = None) -> FederationPolicy: """Update account federation policy. :param policy_id: str - :param update_mask: str - Field mask is required to be passed into the PATCH request. Field mask specifies which fields of the - setting payload will be updated. The field mask needs to be supplied as single string. To specify - multiple fields in the field mask, use comma as the separator (no space). + The identifier for the federation policy. :param policy: :class:`FederationPolicy` (optional) + :param update_mask: str (optional) + The field mask specifies which fields of the policy to update. To specify multiple fields in the + field mask, use comma as the separator (no space). The special value '*' indicates that all fields + should be updated (full replacement). If unspecified, all fields that are set in the policy provided + in the update request will overwrite the corresponding fields in the existing policy. Example value: + 'description,oidc_policy.audiences'. :returns: :class:`FederationPolicy` """ @@ -1433,7 +1415,8 @@ def create(self, The service principal id for the federation policy. :param policy: :class:`FederationPolicy` (optional) :param policy_id: str (optional) - The identifier for the federation policy. If unspecified, the id will be assigned by Databricks. + The identifier for the federation policy. The identifier must contain only lowercase alphanumeric + characters, numbers, hyphens, and slashes. If unspecified, the id will be assigned by Databricks. :returns: :class:`FederationPolicy` """ @@ -1454,6 +1437,7 @@ def delete(self, service_principal_id: int, policy_id: str): :param service_principal_id: int The service principal id for the federation policy. :param policy_id: str + The identifier for the federation policy. """ @@ -1471,6 +1455,7 @@ def get(self, service_principal_id: int, policy_id: str) -> FederationPolicy: :param service_principal_id: int The service principal id for the federation policy. :param policy_id: str + The identifier for the federation policy. :returns: :class:`FederationPolicy` """ @@ -1519,19 +1504,22 @@ def list(self, def update(self, service_principal_id: int, policy_id: str, - update_mask: str, *, - policy: Optional[FederationPolicy] = None) -> FederationPolicy: + policy: Optional[FederationPolicy] = None, + update_mask: Optional[str] = None) -> FederationPolicy: """Update service principal federation policy. :param service_principal_id: int The service principal id for the federation policy. :param policy_id: str - :param update_mask: str - Field mask is required to be passed into the PATCH request. Field mask specifies which fields of the - setting payload will be updated. The field mask needs to be supplied as single string. To specify - multiple fields in the field mask, use comma as the separator (no space). + The identifier for the federation policy. :param policy: :class:`FederationPolicy` (optional) + :param update_mask: str (optional) + The field mask specifies which fields of the policy to update. To specify multiple fields in the + field mask, use comma as the separator (no space). The special value '*' indicates that all fields + should be updated (full replacement). If unspecified, all fields that are set in the policy provided + in the update request will overwrite the corresponding fields in the existing policy. Example value: + 'description,oidc_policy.audiences'. :returns: :class:`FederationPolicy` """ diff --git a/databricks/sdk/service/pipelines.py b/databricks/sdk/service/pipelines.py index 8f8b015c5..db5d698d6 100755 --- a/databricks/sdk/service/pipelines.py +++ b/databricks/sdk/service/pipelines.py @@ -85,6 +85,14 @@ class CreatePipeline: restart_window: Optional[RestartWindow] = None """Restart window of this pipeline.""" + run_as: Optional[RunAs] = None + """Write-only setting, available only in Create/Update calls. Specifies the user or service + principal that the pipeline runs as. If not specified, the pipeline runs as the user who created + the pipeline. + + Only `user_name` or `service_principal_name` can be specified. If both are specified, an error + is thrown.""" + schema: Optional[str] = None """The default schema (database) where tables are read from or published to. The presence of this field implies that the pipeline is in direct publishing mode.""" @@ -126,6 +134,7 @@ def as_dict(self) -> dict: if self.notifications: body['notifications'] = [v.as_dict() for v in self.notifications] if self.photon is not None: body['photon'] = self.photon if self.restart_window: body['restart_window'] = self.restart_window.as_dict() + if self.run_as: body['run_as'] = self.run_as.as_dict() if self.schema is not None: body['schema'] = self.schema if self.serverless is not None: body['serverless'] = self.serverless if self.storage is not None: body['storage'] = self.storage @@ -156,6 +165,7 @@ def as_shallow_dict(self) -> dict: if self.notifications: body['notifications'] = self.notifications if self.photon is not None: body['photon'] = self.photon if self.restart_window: body['restart_window'] = self.restart_window + if self.run_as: body['run_as'] = self.run_as if self.schema is not None: body['schema'] = self.schema if self.serverless is not None: body['serverless'] = self.serverless if self.storage is not None: body['storage'] = self.storage @@ -186,6 +196,7 @@ def from_dict(cls, d: Dict[str, any]) -> CreatePipeline: notifications=_repeated_dict(d, 'notifications', Notifications), photon=d.get('photon', None), restart_window=_from_dict(d, 'restart_window', RestartWindow), + run_as=_from_dict(d, 'run_as', RunAs), schema=d.get('schema', None), serverless=d.get('serverless', None), storage=d.get('storage', None), @@ -277,6 +288,19 @@ def from_dict(cls, d: Dict[str, any]) -> DataPlaneId: return cls(instance=d.get('instance', None), seq_no=d.get('seq_no', None)) +class DayOfWeek(Enum): + """Days of week in which the restart is allowed to happen (within a five-hour window starting at + start_hour). If not specified all days of the week will be used.""" + + FRIDAY = 'FRIDAY' + MONDAY = 'MONDAY' + SATURDAY = 'SATURDAY' + SUNDAY = 'SUNDAY' + THURSDAY = 'THURSDAY' + TUESDAY = 'TUESDAY' + WEDNESDAY = 'WEDNESDAY' + + @dataclass class DeletePipelineResponse: @@ -373,6 +397,14 @@ class EditPipeline: restart_window: Optional[RestartWindow] = None """Restart window of this pipeline.""" + run_as: Optional[RunAs] = None + """Write-only setting, available only in Create/Update calls. Specifies the user or service + principal that the pipeline runs as. If not specified, the pipeline runs as the user who created + the pipeline. + + Only `user_name` or `service_principal_name` can be specified. If both are specified, an error + is thrown.""" + schema: Optional[str] = None """The default schema (database) where tables are read from or published to. The presence of this field implies that the pipeline is in direct publishing mode.""" @@ -416,6 +448,7 @@ def as_dict(self) -> dict: if self.photon is not None: body['photon'] = self.photon if self.pipeline_id is not None: body['pipeline_id'] = self.pipeline_id if self.restart_window: body['restart_window'] = self.restart_window.as_dict() + if self.run_as: body['run_as'] = self.run_as.as_dict() if self.schema is not None: body['schema'] = self.schema if self.serverless is not None: body['serverless'] = self.serverless if self.storage is not None: body['storage'] = self.storage @@ -448,6 +481,7 @@ def as_shallow_dict(self) -> dict: if self.photon is not None: body['photon'] = self.photon if self.pipeline_id is not None: body['pipeline_id'] = self.pipeline_id if self.restart_window: body['restart_window'] = self.restart_window + if self.run_as: body['run_as'] = self.run_as if self.schema is not None: body['schema'] = self.schema if self.serverless is not None: body['serverless'] = self.serverless if self.storage is not None: body['storage'] = self.storage @@ -479,6 +513,7 @@ def from_dict(cls, d: Dict[str, any]) -> EditPipeline: photon=d.get('photon', None), pipeline_id=d.get('pipeline_id', None), restart_window=_from_dict(d, 'restart_window', RestartWindow), + run_as=_from_dict(d, 'run_as', RunAs), schema=d.get('schema', None), serverless=d.get('serverless', None), storage=d.get('storage', None), @@ -2105,7 +2140,7 @@ class RestartWindow: """An integer between 0 and 23 denoting the start hour for the restart window in the 24-hour day. Continuous pipeline restart is triggered only within a five-hour window starting at this hour.""" - days_of_week: Optional[List[RestartWindowDaysOfWeek]] = None + days_of_week: Optional[List[DayOfWeek]] = None """Days of week in which the restart is allowed to happen (within a five-hour window starting at start_hour). If not specified all days of the week will be used.""" @@ -2133,22 +2168,48 @@ def as_shallow_dict(self) -> dict: @classmethod def from_dict(cls, d: Dict[str, any]) -> RestartWindow: """Deserializes the RestartWindow from a dictionary.""" - return cls(days_of_week=_repeated_enum(d, 'days_of_week', RestartWindowDaysOfWeek), + return cls(days_of_week=_repeated_enum(d, 'days_of_week', DayOfWeek), start_hour=d.get('start_hour', None), time_zone_id=d.get('time_zone_id', None)) -class RestartWindowDaysOfWeek(Enum): - """Days of week in which the restart is allowed to happen (within a five-hour window starting at - start_hour). If not specified all days of the week will be used.""" +@dataclass +class RunAs: + """Write-only setting, available only in Create/Update calls. Specifies the user or service + principal that the pipeline runs as. If not specified, the pipeline runs as the user who created + the pipeline. + + Only `user_name` or `service_principal_name` can be specified. If both are specified, an error + is thrown.""" - FRIDAY = 'FRIDAY' - MONDAY = 'MONDAY' - SATURDAY = 'SATURDAY' - SUNDAY = 'SUNDAY' - THURSDAY = 'THURSDAY' - TUESDAY = 'TUESDAY' - WEDNESDAY = 'WEDNESDAY' + service_principal_name: Optional[str] = None + """Application ID of an active service principal. Setting this field requires the + `servicePrincipal/user` role.""" + + user_name: Optional[str] = None + """The email of an active workspace user. Users can only set this field to their own email.""" + + def as_dict(self) -> dict: + """Serializes the RunAs into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.service_principal_name is not None: + body['service_principal_name'] = self.service_principal_name + if self.user_name is not None: body['user_name'] = self.user_name + return body + + def as_shallow_dict(self) -> dict: + """Serializes the RunAs into a shallow dictionary of its immediate attributes.""" + body = {} + if self.service_principal_name is not None: + body['service_principal_name'] = self.service_principal_name + if self.user_name is not None: body['user_name'] = self.user_name + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> RunAs: + """Deserializes the RunAs from a dictionary.""" + return cls(service_principal_name=d.get('service_principal_name', None), + user_name=d.get('user_name', None)) @dataclass @@ -2791,6 +2852,7 @@ def create(self, notifications: Optional[List[Notifications]] = None, photon: Optional[bool] = None, restart_window: Optional[RestartWindow] = None, + run_as: Optional[RunAs] = None, schema: Optional[str] = None, serverless: Optional[bool] = None, storage: Optional[str] = None, @@ -2843,6 +2905,12 @@ def create(self, Whether Photon is enabled for this pipeline. :param restart_window: :class:`RestartWindow` (optional) Restart window of this pipeline. + :param run_as: :class:`RunAs` (optional) + Write-only setting, available only in Create/Update calls. Specifies the user or service principal + that the pipeline runs as. If not specified, the pipeline runs as the user who created the pipeline. + + Only `user_name` or `service_principal_name` can be specified. If both are specified, an error is + thrown. :param schema: str (optional) The default schema (database) where tables are read from or published to. The presence of this field implies that the pipeline is in direct publishing mode. @@ -2879,6 +2947,7 @@ def create(self, if notifications is not None: body['notifications'] = [v.as_dict() for v in notifications] if photon is not None: body['photon'] = photon if restart_window is not None: body['restart_window'] = restart_window.as_dict() + if run_as is not None: body['run_as'] = run_as.as_dict() if schema is not None: body['schema'] = schema if serverless is not None: body['serverless'] = serverless if storage is not None: body['storage'] = storage @@ -3213,6 +3282,7 @@ def update(self, notifications: Optional[List[Notifications]] = None, photon: Optional[bool] = None, restart_window: Optional[RestartWindow] = None, + run_as: Optional[RunAs] = None, schema: Optional[str] = None, serverless: Optional[bool] = None, storage: Optional[str] = None, @@ -3268,6 +3338,12 @@ def update(self, Whether Photon is enabled for this pipeline. :param restart_window: :class:`RestartWindow` (optional) Restart window of this pipeline. + :param run_as: :class:`RunAs` (optional) + Write-only setting, available only in Create/Update calls. Specifies the user or service principal + that the pipeline runs as. If not specified, the pipeline runs as the user who created the pipeline. + + Only `user_name` or `service_principal_name` can be specified. If both are specified, an error is + thrown. :param schema: str (optional) The default schema (database) where tables are read from or published to. The presence of this field implies that the pipeline is in direct publishing mode. @@ -3304,6 +3380,7 @@ def update(self, if notifications is not None: body['notifications'] = [v.as_dict() for v in notifications] if photon is not None: body['photon'] = photon if restart_window is not None: body['restart_window'] = restart_window.as_dict() + if run_as is not None: body['run_as'] = run_as.as_dict() if schema is not None: body['schema'] = schema if serverless is not None: body['serverless'] = serverless if storage is not None: body['storage'] = storage diff --git a/databricks/sdk/service/serving.py b/databricks/sdk/service/serving.py index cb7861e88..39ff002fb 100755 --- a/databricks/sdk/service/serving.py +++ b/databricks/sdk/service/serving.py @@ -18,8 +18,6 @@ _LOG = logging.getLogger('databricks.sdk') -from databricks.sdk.service import oauth2 - # all definitions in this file are in alphabetical order @@ -712,6 +710,35 @@ def from_dict(cls, d: Dict[str, any]) -> CreateServingEndpoint: tags=_repeated_dict(d, 'tags', EndpointTag)) +@dataclass +class DataPlaneInfo: + authorization_details: Optional[str] = None + """Authorization details as a string.""" + + endpoint_url: Optional[str] = None + """The URL of the endpoint for this operation in the dataplane.""" + + def as_dict(self) -> dict: + """Serializes the DataPlaneInfo into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.authorization_details is not None: body['authorization_details'] = self.authorization_details + if self.endpoint_url is not None: body['endpoint_url'] = self.endpoint_url + return body + + def as_shallow_dict(self) -> dict: + """Serializes the DataPlaneInfo into a shallow dictionary of its immediate attributes.""" + body = {} + if self.authorization_details is not None: body['authorization_details'] = self.authorization_details + if self.endpoint_url is not None: body['endpoint_url'] = self.endpoint_url + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> DataPlaneInfo: + """Deserializes the DataPlaneInfo from a dictionary.""" + return cls(authorization_details=d.get('authorization_details', None), + endpoint_url=d.get('endpoint_url', None)) + + @dataclass class DatabricksModelServingConfig: databricks_workspace_url: str @@ -1444,7 +1471,7 @@ def from_dict(cls, d: Dict[str, any]) -> ListEndpointsResponse: @dataclass class ModelDataPlaneInfo: - query_info: Optional[oauth2.DataPlaneInfo] = None + query_info: Optional[DataPlaneInfo] = None """Information required to query DataPlane API 'query' endpoint.""" def as_dict(self) -> dict: @@ -1462,7 +1489,7 @@ def as_shallow_dict(self) -> dict: @classmethod def from_dict(cls, d: Dict[str, any]) -> ModelDataPlaneInfo: """Deserializes the ModelDataPlaneInfo from a dictionary.""" - return cls(query_info=_from_dict(d, 'query_info', oauth2.DataPlaneInfo)) + return cls(query_info=_from_dict(d, 'query_info', DataPlaneInfo)) @dataclass diff --git a/tests/test_data_plane.py b/tests/test_data_plane.py index a74658964..1eac92382 100644 --- a/tests/test_data_plane.py +++ b/tests/test_data_plane.py @@ -2,7 +2,7 @@ from databricks.sdk.data_plane import DataPlaneService from databricks.sdk.oauth import Token -from databricks.sdk.service.oauth2 import DataPlaneInfo +from databricks.sdk.service.serving import DataPlaneInfo info = DataPlaneInfo(authorization_details="authDetails", endpoint_url="url") From 496dadaaf67e28622034cdf1904aeecf33478978 Mon Sep 17 00:00:00 2001 From: Parth Bansal Date: Wed, 8 Jan 2025 18:22:14 +0100 Subject: [PATCH 2/6] update --- databricks/sdk/__init__.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/databricks/sdk/__init__.py b/databricks/sdk/__init__.py index 068069f04..d27110b86 100755 --- a/databricks/sdk/__init__.py +++ b/databricks/sdk/__init__.py @@ -5,7 +5,7 @@ from databricks.sdk import azure from databricks.sdk.credentials_provider import CredentialsStrategy from databricks.sdk.mixins.compute import ClustersExt -from databricks.sdk.mixins.files import DbfsExt +from databricks.sdk.mixins.files import DbfsExt, FilesExt from databricks.sdk.mixins.jobs import JobsExt from databricks.sdk.mixins.open_ai_client import ServingEndpointsExt from databricks.sdk.mixins.workspace import WorkspaceExt @@ -114,6 +114,13 @@ def _make_dbutils(config: client.Config): return runtime_dbutils +def _make_files_client(apiClient: client.ApiClient, config: client.Config): + if config.enable_experimental_files_api_client: + return FilesExt(apiClient, config) + else: + return FilesAPI(apiClient) + + class WorkspaceClient: """ The WorkspaceClient is a client for the workspace-level Databricks REST API. @@ -203,7 +210,7 @@ def __init__(self, self._dbsql_permissions = DbsqlPermissionsAPI(self._api_client) self._experiments = ExperimentsAPI(self._api_client) self._external_locations = ExternalLocationsAPI(self._api_client) - self._files = FilesAPI(self._api_client) + self._files = _make_files_client(self._api_client, self._config) self._functions = FunctionsAPI(self._api_client) self._genie = GenieAPI(self._api_client) self._git_credentials = GitCredentialsAPI(self._api_client) From d056712f4dd8bee9e8434e0d26ca7150b73f1c3d Mon Sep 17 00:00:00 2001 From: Parth Bansal Date: Thu, 9 Jan 2025 11:12:29 +0100 Subject: [PATCH 3/6] update --- databricks/sdk/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/databricks/sdk/__init__.py b/databricks/sdk/__init__.py index d27110b86..80fe188b8 100755 --- a/databricks/sdk/__init__.py +++ b/databricks/sdk/__init__.py @@ -1,3 +1,5 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. + from typing import Optional import databricks.sdk.core as client From a19d7785462bed2e4917ed4dee59c379e2b2bb0e Mon Sep 17 00:00:00 2001 From: Parth Bansal Date: Thu, 9 Jan 2025 14:51:47 +0100 Subject: [PATCH 4/6] update --- databricks/sdk/data_plane.py | 21 ++------------------- databricks/sdk/service/serving.py | 3 ++- 2 files changed, 4 insertions(+), 20 deletions(-) diff --git a/databricks/sdk/data_plane.py b/databricks/sdk/data_plane.py index 2a34a9f43..5ad9b79ad 100644 --- a/databricks/sdk/data_plane.py +++ b/databricks/sdk/data_plane.py @@ -1,28 +1,10 @@ import threading -from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import Callable, List, Optional +from typing import Callable, List from databricks.sdk.oauth import Token -@dataclass -class DataPlaneInfo(ABC): - """ - Abstract base class for DataPlane information. - """ - - @property - @abstractmethod - def url(self) -> str: - """Gets the URL for the DataPlane endpoint.""" - - @property - @abstractmethod - def token(self) -> Optional[str]: - """Gets the token for the DataPlane endpoint.""" - - @dataclass class DataPlaneDetails: """ @@ -36,6 +18,7 @@ class DataPlaneDetails: class DataPlaneService: """Helper class to fetch and manage DataPlane details.""" + from .service.serving import DataPlaneInfo def __init__(self): self._data_plane_info = {} diff --git a/databricks/sdk/service/serving.py b/databricks/sdk/service/serving.py index 39ff002fb..6541f700e 100755 --- a/databricks/sdk/service/serving.py +++ b/databricks/sdk/service/serving.py @@ -12,7 +12,6 @@ import requests -from ..data_plane import DataPlaneService from ..errors import OperationFailed from ._internal import Wait, _enum, _from_dict, _repeated_dict @@ -3750,6 +3749,8 @@ class ServingEndpointsDataPlaneAPI: endpoints service.""" def __init__(self, api_client, control_plane): + from ..data_plane import DataPlaneService + self._api = api_client self._control_plane = control_plane self._data_plane_service = DataPlaneService() From 9ca6a395ce97267320bec59d0d38166f28b36425 Mon Sep 17 00:00:00 2001 From: Parth Bansal Date: Thu, 9 Jan 2025 15:05:11 +0100 Subject: [PATCH 5/6] update --- databricks/sdk/service/serving.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/databricks/sdk/service/serving.py b/databricks/sdk/service/serving.py index 6541f700e..1ada305cd 100755 --- a/databricks/sdk/service/serving.py +++ b/databricks/sdk/service/serving.py @@ -3749,10 +3749,9 @@ class ServingEndpointsDataPlaneAPI: endpoints service.""" def __init__(self, api_client, control_plane): - from ..data_plane import DataPlaneService - self._api = api_client self._control_plane = control_plane + from ..data_plane import DataPlaneService self._data_plane_service = DataPlaneService() def query(self, From 1b70cba6e903bb8d7568daa5ce84c10bb41f22eb Mon Sep 17 00:00:00 2001 From: Parth Bansal Date: Thu, 9 Jan 2025 15:11:10 +0100 Subject: [PATCH 6/6] update --- .vscode/settings.json | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/.vscode/settings.json b/.vscode/settings.json index ad2eb823b..c36b4db6c 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -7,6 +7,5 @@ "python.envFile": "${workspaceFolder}/.databricks/.databricks.env", "databricks.python.envFile": "${workspaceFolder}/.env", "jupyter.interactiveWindow.cellMarker.codeRegex": "^# COMMAND ----------|^# Databricks notebook source|^(#\\s*%%|#\\s*\\|#\\s*In\\[\\d*?\\]|#\\s*In\\[ \\])", - "jupyter.interactiveWindow.cellMarker.default": "# COMMAND ----------", - "python.analysis.autoImportCompletions": true + "jupyter.interactiveWindow.cellMarker.default": "# COMMAND ----------" }