diff --git a/.codegen/_last_sha b/.codegen/_last_sha index 862f7f3ac..249e78289 100644 --- a/.codegen/_last_sha +++ b/.codegen/_last_sha @@ -1 +1 @@ -c3b5743c6a0e74fa11ab0079f83fbdbc349a8861 \ No newline at end of file +07bd11eefe3ed30eebcfa846546af8b56b655973 \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md index 6ecadb34d..c9b52a592 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,141 @@ # Version changelog +## Release v0.159.0 (2026-11-02) + +### API Changes +* Add `databricks.sdk.service.computemigrations` and `databricks.sdk.service.genieagents` packages. +* Add [w.compute_migrations](https://databricks-sdk-py.readthedocs.io/en/latest/workspace/computemigrations/compute_migrations.html) workspace-level service. +* Add [w.genie_agents](https://databricks-sdk-py.readthedocs.io/en/latest/workspace/genieagents/genie_agents.html) workspace-level service. +* Add `download_file()` and `upload_file()` methods for [w.sandbox](https://databricks-sdk-py.readthedocs.io/en/latest/workspace/sandbox/sandbox.html) workspace-level service. +* Add `update_network_connectivity_config()` method for [a.network_connectivity](https://databricks-sdk-py.readthedocs.io/en/latest/account/settings/network_connectivity.html) account-level service. +* [Breaking] Add `parent` field for `databricks.sdk.service.iamv2.CreateDirectGroupMemberProxyRequest`. +* [Breaking] Add `parent` field for `databricks.sdk.service.iamv2.CreateDirectGroupMemberRequest`. +* Add `parent` field for `databricks.sdk.service.iamv2.CreateGroupProxyRequest`. +* [Breaking] Add `parent` field for `databricks.sdk.service.iamv2.CreateGroupRequest`. +* Add `parent` field for `databricks.sdk.service.iamv2.CreateServicePrincipalProxyRequest`. +* [Breaking] Add `parent` field for `databricks.sdk.service.iamv2.CreateServicePrincipalRequest`. +* Add `parent` field for `databricks.sdk.service.iamv2.CreateUserProxyRequest`. +* [Breaking] Add `parent` field for `databricks.sdk.service.iamv2.CreateUserRequest`. +* Add `parent` field for `databricks.sdk.service.iamv2.CreateWorkspaceAssignmentProxyRequest`. +* [Breaking] Add `parent` field for `databricks.sdk.service.iamv2.CreateWorkspaceAssignmentRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.DeleteDirectGroupMemberProxyRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.DeleteDirectGroupMemberRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.DeleteGroupProxyRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.DeleteGroupRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.DeleteServicePrincipalProxyRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.DeleteServicePrincipalRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.DeleteUserProxyRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.DeleteUserRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.DeleteWorkspaceAssignmentProxyRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.DeleteWorkspaceAssignmentRequest`. +* Add `name` and `parent` fields for `databricks.sdk.service.iamv2.DirectGroupMember`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.GetDirectGroupMemberProxyRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.GetDirectGroupMemberRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.GetGroupProxyRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.GetGroupRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.GetServicePrincipalProxyRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.GetServicePrincipalRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.GetUserProxyRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.GetUserRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.GetWorkspaceAccessDetailLocalRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.GetWorkspaceAccessDetailRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.GetWorkspaceAssignmentProxyRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.GetWorkspaceAssignmentRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.GetWorkspaceIdentityDetailRequest`. +* Add `name` and `parent` fields for `databricks.sdk.service.iamv2.Group`. +* [Breaking] Add `parent` field for `databricks.sdk.service.iamv2.ListDirectGroupMembersProxyRequest`. +* [Breaking] Add `parent` field for `databricks.sdk.service.iamv2.ListDirectGroupMembersRequest`. +* Add `parent` field for `databricks.sdk.service.iamv2.ListGroupsProxyRequest`. +* [Breaking] Add `parent` field for `databricks.sdk.service.iamv2.ListGroupsRequest`. +* Add `parent` field for `databricks.sdk.service.iamv2.ListServicePrincipalsProxyRequest`. +* [Breaking] Add `parent` field for `databricks.sdk.service.iamv2.ListServicePrincipalsRequest`. +* [Breaking] Add `parent` field for `databricks.sdk.service.iamv2.ListTransitiveParentGroupsProxyRequest`. +* [Breaking] Add `parent` field for `databricks.sdk.service.iamv2.ListTransitiveParentGroupsRequest`. +* Add `parent` field for `databricks.sdk.service.iamv2.ListUsersProxyRequest`. +* [Breaking] Add `parent` field for `databricks.sdk.service.iamv2.ListUsersRequest`. +* Add `parent` field for `databricks.sdk.service.iamv2.ListWorkspaceAssignmentsProxyRequest`. +* [Breaking] Add `parent` field for `databricks.sdk.service.iamv2.ListWorkspaceAssignmentsRequest`. +* Add `name` and `parent` fields for `databricks.sdk.service.iamv2.ServicePrincipal`. +* Add `name` field for `databricks.sdk.service.iamv2.TransitiveParentGroup`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.UpdateGroupProxyRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.UpdateGroupRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.UpdateServicePrincipalProxyRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.UpdateServicePrincipalRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.UpdateUserProxyRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.UpdateUserRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.UpdateWorkspaceAssignmentProxyRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.UpdateWorkspaceAssignmentRequest`. +* [Breaking] Add `name` field for `databricks.sdk.service.iamv2.UpdateWorkspaceIdentityDetailRequest`. +* Add `name` and `parent` fields for `databricks.sdk.service.iamv2.User`. +* Add `name` field for `databricks.sdk.service.iamv2.WorkspaceAccessDetail`. +* Add `name` and `parent` fields for `databricks.sdk.service.iamv2.WorkspaceAssignment`. +* Add `name` field for `databricks.sdk.service.iamv2.WorkspaceIdentityDetail`. +* Add `serverless_compute` field for `databricks.sdk.service.jobs.JobCluster`. +* Add `job_completion` field for `databricks.sdk.service.jobs.PerTriggerState`. +* Add `ai_function_task` field for `databricks.sdk.service.jobs.ResolvedValues`. +* Add `ai_function_task_output` field for `databricks.sdk.service.jobs.RunOutput`. +* Add `ai_function_task` field for `databricks.sdk.service.jobs.RunTask`. +* Add `ai_function_task` field for `databricks.sdk.service.jobs.SubmitTask`. +* Add `ai_function_task` field for `databricks.sdk.service.jobs.Task`. +* Add `job_completion` field for `databricks.sdk.service.jobs.TriggerConfiguration`. +* Add `job_completion` field for `databricks.sdk.service.jobs.TriggerInfo`. +* Add `job_completion` field for `databricks.sdk.service.jobs.TriggerSettings`. +* Add `job_completion` field for `databricks.sdk.service.jobs.TriggerStateProto`. +* Add `billing_workspace_id` field for `databricks.sdk.service.settings.NetworkConnectivityConfiguration`. +* Add `create_feature` and `read_feature` enum values for `databricks.sdk.service.catalog.Privilege`. +* Add `feature` enum value for `databricks.sdk.service.catalog.SecurableType`. +* Add `job_completion` enum value for `databricks.sdk.service.jobs.TriggerType`. +* [Breaking] Change `create_direct_group_member()`, `create_group()`, `create_service_principal()`, `create_user()`, `create_workspace_assignment()`, `delete_direct_group_member()`, `delete_group()`, `delete_service_principal()`, `delete_user()`, `delete_workspace_assignment()`, `get_direct_group_member()`, `get_group()`, `get_service_principal()`, `get_user()`, `get_workspace_access_detail()`, `get_workspace_assignment()`, `list_direct_group_members()`, `list_groups()`, `list_service_principals()`, `list_transitive_parent_groups()`, `list_users()`, `list_workspace_assignments()`, `update_group()`, `update_service_principal()`, `update_user()` and `update_workspace_assignment()` methods for [a.account_iam_v2](https://databricks-sdk-py.readthedocs.io/en/latest/account/iamv2/account_iam_v2.html) account-level service with new required argument order. +* [Breaking] Change `create_direct_group_member()`, `create_group()`, `create_service_principal()`, `create_user()`, `create_workspace_assignment()`, `delete_direct_group_member()`, `delete_group()`, `delete_service_principal()`, `delete_user()`, `delete_workspace_assignment()`, `get_direct_group_member()`, `get_group()`, `get_service_principal()`, `get_user()`, `get_workspace_access_detail()`, `get_workspace_assignment()`, `list_direct_group_members()`, `list_groups()`, `list_service_principals()`, `list_transitive_parent_groups()`, `list_users()`, `list_workspace_assignments()`, `update_group()`, `update_service_principal()`, `update_user()` and `update_workspace_assignment()` methods for [a.account_iam_v2](https://databricks-sdk-py.readthedocs.io/en/latest/account/iamv2/account_iam_v2.html) account-level service. Method path has changed. +* [Breaking] Change `create_direct_group_member_proxy()`, `delete_direct_group_member_proxy()`, `delete_group_proxy()`, `delete_service_principal_proxy()`, `delete_user_proxy()`, `delete_workspace_assignment_proxy()`, `get_direct_group_member_proxy()`, `get_group_proxy()`, `get_service_principal_proxy()`, `get_user_proxy()`, `get_workspace_access_detail_local()`, `get_workspace_assignment_proxy()`, `get_workspace_identity_detail()`, `list_direct_group_members_proxy()`, `list_transitive_parent_groups_proxy()`, `update_group_proxy()`, `update_service_principal_proxy()`, `update_user_proxy()`, `update_workspace_assignment_proxy()` and `update_workspace_identity_detail()` methods for [w.workspace_iam_v2](https://databricks-sdk-py.readthedocs.io/en/latest/workspace/iamv2/workspace_iam_v2.html) workspace-level service. Method path has changed. +* [Breaking] Change `create_direct_group_member_proxy()`, `delete_direct_group_member_proxy()`, `delete_group_proxy()`, `delete_service_principal_proxy()`, `delete_user_proxy()`, `delete_workspace_assignment_proxy()`, `get_direct_group_member_proxy()`, `get_group_proxy()`, `get_service_principal_proxy()`, `get_user_proxy()`, `get_workspace_access_detail_local()`, `get_workspace_assignment_proxy()`, `get_workspace_identity_detail()`, `list_direct_group_members_proxy()`, `list_transitive_parent_groups_proxy()`, `update_group_proxy()`, `update_service_principal_proxy()`, `update_user_proxy()`, `update_workspace_assignment_proxy()` and `update_workspace_identity_detail()` methods for [w.workspace_iam_v2](https://databricks-sdk-py.readthedocs.io/en/latest/workspace/iamv2/workspace_iam_v2.html) workspace-level service with new required argument order. +* [Breaking] Remove `group_id` field for `databricks.sdk.service.iamv2.CreateDirectGroupMemberProxyRequest`. +* [Breaking] Remove `account_id` and `group_id` fields for `databricks.sdk.service.iamv2.CreateDirectGroupMemberRequest`. +* [Breaking] Remove `account_id` field for `databricks.sdk.service.iamv2.CreateGroupRequest`. +* [Breaking] Remove `account_id` field for `databricks.sdk.service.iamv2.CreateServicePrincipalRequest`. +* [Breaking] Remove `account_id` field for `databricks.sdk.service.iamv2.CreateUserRequest`. +* [Breaking] Remove `account_id` and `workspace_id` fields for `databricks.sdk.service.iamv2.CreateWorkspaceAssignmentRequest`. +* [Breaking] Remove `group_id` and `principal_id` fields for `databricks.sdk.service.iamv2.DeleteDirectGroupMemberProxyRequest`. +* [Breaking] Remove `account_id`, `group_id` and `principal_id` fields for `databricks.sdk.service.iamv2.DeleteDirectGroupMemberRequest`. +* [Breaking] Remove `group_id` field for `databricks.sdk.service.iamv2.DeleteGroupProxyRequest`. +* [Breaking] Remove `account_id` and `group_id` fields for `databricks.sdk.service.iamv2.DeleteGroupRequest`. +* [Breaking] Remove `service_principal_id` field for `databricks.sdk.service.iamv2.DeleteServicePrincipalProxyRequest`. +* [Breaking] Remove `account_id` and `service_principal_id` fields for `databricks.sdk.service.iamv2.DeleteServicePrincipalRequest`. +* [Breaking] Remove `user_id` field for `databricks.sdk.service.iamv2.DeleteUserProxyRequest`. +* [Breaking] Remove `account_id` and `user_id` fields for `databricks.sdk.service.iamv2.DeleteUserRequest`. +* [Breaking] Remove `principal_id` field for `databricks.sdk.service.iamv2.DeleteWorkspaceAssignmentProxyRequest`. +* [Breaking] Remove `account_id`, `workspace_id` and `principal_id` fields for `databricks.sdk.service.iamv2.DeleteWorkspaceAssignmentRequest`. +* [Breaking] Remove `group_id` and `principal_id` fields for `databricks.sdk.service.iamv2.GetDirectGroupMemberProxyRequest`. +* [Breaking] Remove `account_id`, `group_id` and `principal_id` fields for `databricks.sdk.service.iamv2.GetDirectGroupMemberRequest`. +* [Breaking] Remove `group_id` field for `databricks.sdk.service.iamv2.GetGroupProxyRequest`. +* [Breaking] Remove `account_id` and `group_id` fields for `databricks.sdk.service.iamv2.GetGroupRequest`. +* [Breaking] Remove `service_principal_id` field for `databricks.sdk.service.iamv2.GetServicePrincipalProxyRequest`. +* [Breaking] Remove `account_id` and `service_principal_id` fields for `databricks.sdk.service.iamv2.GetServicePrincipalRequest`. +* [Breaking] Remove `user_id` field for `databricks.sdk.service.iamv2.GetUserProxyRequest`. +* [Breaking] Remove `account_id` and `user_id` fields for `databricks.sdk.service.iamv2.GetUserRequest`. +* [Breaking] Remove `principal_id` field for `databricks.sdk.service.iamv2.GetWorkspaceAccessDetailLocalRequest`. +* [Breaking] Remove `account_id`, `workspace_id` and `principal_id` fields for `databricks.sdk.service.iamv2.GetWorkspaceAccessDetailRequest`. +* [Breaking] Remove `principal_id` field for `databricks.sdk.service.iamv2.GetWorkspaceAssignmentProxyRequest`. +* [Breaking] Remove `account_id`, `workspace_id` and `principal_id` fields for `databricks.sdk.service.iamv2.GetWorkspaceAssignmentRequest`. +* [Breaking] Remove `principal_id` field for `databricks.sdk.service.iamv2.GetWorkspaceIdentityDetailRequest`. +* [Breaking] Remove `group_id` field for `databricks.sdk.service.iamv2.ListDirectGroupMembersProxyRequest`. +* [Breaking] Remove `account_id` and `group_id` fields for `databricks.sdk.service.iamv2.ListDirectGroupMembersRequest`. +* [Breaking] Remove `account_id` field for `databricks.sdk.service.iamv2.ListGroupsRequest`. +* [Breaking] Remove `account_id` field for `databricks.sdk.service.iamv2.ListServicePrincipalsRequest`. +* [Breaking] Remove `principal_id` field for `databricks.sdk.service.iamv2.ListTransitiveParentGroupsProxyRequest`. +* [Breaking] Remove `account_id` and `principal_id` fields for `databricks.sdk.service.iamv2.ListTransitiveParentGroupsRequest`. +* [Breaking] Remove `account_id` field for `databricks.sdk.service.iamv2.ListUsersRequest`. +* [Breaking] Remove `account_id` and `workspace_id` fields for `databricks.sdk.service.iamv2.ListWorkspaceAssignmentsRequest`. +* [Breaking] Remove `group_id` field for `databricks.sdk.service.iamv2.UpdateGroupProxyRequest`. +* [Breaking] Remove `account_id` and `group_id` fields for `databricks.sdk.service.iamv2.UpdateGroupRequest`. +* [Breaking] Remove `service_principal_id` field for `databricks.sdk.service.iamv2.UpdateServicePrincipalProxyRequest`. +* [Breaking] Remove `account_id` and `service_principal_id` fields for `databricks.sdk.service.iamv2.UpdateServicePrincipalRequest`. +* [Breaking] Remove `user_id` field for `databricks.sdk.service.iamv2.UpdateUserProxyRequest`. +* [Breaking] Remove `account_id` and `user_id` fields for `databricks.sdk.service.iamv2.UpdateUserRequest`. +* [Breaking] Remove `principal_id` field for `databricks.sdk.service.iamv2.UpdateWorkspaceAssignmentProxyRequest`. +* [Breaking] Remove `account_id`, `workspace_id` and `principal_id` fields for `databricks.sdk.service.iamv2.UpdateWorkspaceAssignmentRequest`. +* [Breaking] Remove `principal_id` field for `databricks.sdk.service.iamv2.UpdateWorkspaceIdentityDetailRequest`. + ## Release v0.158.0 (2026-10-29) ### API Changes diff --git a/databricks/sdk/__init__.py b/databricks/sdk/__init__.py index 0500f1398..6bf05feb6 100644 --- a/databricks/sdk/__init__.py +++ b/databricks/sdk/__init__.py @@ -19,6 +19,7 @@ from databricks.sdk.service import catalog as pkg_catalog from databricks.sdk.service import cleanrooms as pkg_cleanrooms from databricks.sdk.service import compute as pkg_compute +from databricks.sdk.service import computemigrations as pkg_computemigrations from databricks.sdk.service import dashboards as pkg_dashboards from databricks.sdk.service import database as pkg_database from databricks.sdk.service import dataclassification as pkg_dataclassification @@ -27,6 +28,7 @@ from databricks.sdk.service import domains as pkg_domains from databricks.sdk.service import environments as pkg_environments from databricks.sdk.service import files as pkg_files +from databricks.sdk.service import genieagents as pkg_genieagents from databricks.sdk.service import iam as pkg_iam from databricks.sdk.service import iamv2 as pkg_iamv2 from databricks.sdk.service import jobs as pkg_jobs @@ -88,6 +90,7 @@ from databricks.sdk.service.compute import ClustersAPI from databricks.sdk.service.compute import CommandExecutionAPI from databricks.sdk.service.settings import ComplianceSecurityProfileAPI +from databricks.sdk.service.computemigrations import ComputeMigrationsAPI from databricks.sdk.service.catalog import ConnectionsAPI from databricks.sdk.service.marketplace import ConsumerFulfillmentsAPI from databricks.sdk.service.marketplace import ConsumerInstallationsAPI @@ -137,6 +140,7 @@ from databricks.sdk.service.ml import ForecastingAPI from databricks.sdk.service.catalog import FunctionsAPI from databricks.sdk.service.dashboards import GenieAPI +from databricks.sdk.service.genieagents import GenieAgentsAPI from databricks.sdk.service.workspace import GitCredentialsAPI from databricks.sdk.service.compute import GlobalInitScriptsAPI from databricks.sdk.service.catalog import GrantsAPI @@ -389,6 +393,7 @@ def __init__( self._cluster_policies = pkg_compute.ClusterPoliciesAPI(self._api_client) self._clusters = ClustersExt(self._api_client) self._command_execution = pkg_compute.CommandExecutionAPI(self._api_client) + self._compute_migrations = pkg_computemigrations.ComputeMigrationsAPI(self._api_client) self._connections = pkg_catalog.ConnectionsAPI(self._api_client) self._consumer_fulfillments = pkg_marketplace.ConsumerFulfillmentsAPI(self._api_client) self._consumer_installations = pkg_marketplace.ConsumerInstallationsAPI(self._api_client) @@ -419,6 +424,7 @@ def __init__( self._forecasting = pkg_ml.ForecastingAPI(self._api_client) self._functions = pkg_catalog.FunctionsAPI(self._api_client) self._genie = pkg_dashboards.GenieAPI(self._api_client) + self._genie_agents = pkg_genieagents.GenieAgentsAPI(self._api_client) self._git_credentials = pkg_workspace.GitCredentialsAPI(self._api_client) self._global_init_scripts = pkg_compute.GlobalInitScriptsAPI(self._api_client) self._grants = pkg_catalog.GrantsAPI(self._api_client) @@ -642,6 +648,11 @@ def command_execution(self) -> pkg_compute.CommandExecutionAPI: """This API allows execution of Python, Scala, SQL, or R commands on running Databricks Clusters.""" return self._command_execution + @property + def compute_migrations(self) -> pkg_computemigrations.ComputeMigrationsAPI: + """Manage migrations of Classic compute workloads to Serverless.""" + return self._compute_migrations + @property def connections(self) -> pkg_catalog.ConnectionsAPI: """A connection represents an external data source for use within Databricks.""" @@ -787,6 +798,11 @@ def genie(self) -> pkg_dashboards.GenieAPI: """Genie provides a no-code experience for business users, powered by AI/BI.""" return self._genie + @property + def genie_agents(self) -> pkg_genieagents.GenieAgentsAPI: + """Genie provides a no-code experience for business users, powered by AI/BI.""" + return self._genie_agents + @property def git_credentials(self) -> pkg_workspace.GitCredentialsAPI: """Registers personal access token for Databricks to do operations on behalf of the user.""" diff --git a/databricks/sdk/config.py b/databricks/sdk/config.py index 4a838000c..bf7f268bf 100644 --- a/databricks/sdk/config.py +++ b/databricks/sdk/config.py @@ -179,7 +179,7 @@ class Config: # File multipart upload/download parameters # ---------------------- - # Minimal input stream size (bytes) to use multipart / resumable uploads. + # Minimal input stream size (bytes) to use multipart uploads. # For small files it's more efficient to make one single-shot upload request. # When uploading a file, SDK will initially buffer this many bytes from input stream. # This parameter can be less or bigger than multipart_upload_chunk_size. @@ -247,12 +247,9 @@ class Config: # but a maximum time between consecutive data reception events (even 1 byte) from the server files_ext_network_transfer_inactivity_timeout_seconds: float = 60 - # Cap on the number of custom retries during incremental uploads: - # 1) multipart: upload part URL is expired, so new upload URLs must be requested to continue upload - # 2) resumable: chunk upload produced a retryable response (or exception), so upload status must be - # retrieved to continue the upload. - # In these two cases standard SDK retries (which are capped by the `retry_timeout_seconds` option) are not used. - # Note that retry counter is reset when upload is successfully resumed. + # Cap on the number of custom retries during multipart uploads, when an upload part URL is expired and new + # upload URLs must be requested to continue the upload. Standard SDK retries (capped by the + # `retry_timeout_seconds` option) are not used in this case. files_ext_multipart_upload_max_retries = 3 # Cap on the number of custom retries during parallel downloads. diff --git a/databricks/sdk/mixins/files.py b/databricks/sdk/mixins/files.py index 16173d1e4..b6f4b6afc 100644 --- a/databricks/sdk/mixins/files.py +++ b/databricks/sdk/mixins/files.py @@ -7,7 +7,6 @@ import os import pathlib import platform -import re import shutil import xml.etree.ElementTree as ET from abc import ABC, abstractmethod @@ -32,7 +31,7 @@ from .._base_client import _BaseClient, _RawResponse, _StreamingResponse from .._property import _cached_property from ..config import Config -from ..errors import AlreadyExists, InternalError, NotFound, PermissionDenied +from ..errors import InternalError, NotFound, PermissionDenied from ..errors.mapper import _error_mapper from ..retries import retried from databricks.sdk.service import files @@ -779,8 +778,8 @@ class _Endpoint: class _PresignedUrlRequestBuilder: """Builds presigned URL requests for file upload operations. - Encapsulates the coordination API calls (create-upload-part-urls, - create-resumable-upload-url) and response parsing into a single place. + Encapsulates the coordination API calls (create-upload-part-urls, create-download-url) and response parsing + into a single place. """ def __init__(self, api, hostname: str): @@ -817,27 +816,6 @@ def build_upload_part_urls( results.append(_PresignedUrl(url=url, headers=headers, part_number=part_number)) return results - def build_resumable_upload_url(self, path: str, session_token: str) -> _PresignedUrl: - """Fetches a presigned URL for resumable upload.""" - body = {"path": path, "session_token": session_token} - # The _api.do() method handles retries and will raise an exception in case of failure. - response = self._api.do( - "POST", - url=f"{self._hostname}/api/2.0/fs/create-resumable-upload-url", - headers={"Content-Type": "application/json"}, - body=body, - ) - url_node = response.get("resumable_upload_url") - if not url_node: - raise ValueError(f"Unexpected server response: {response}") - url = url_node.get("url") - if not url: - raise ValueError(f"Unexpected server response: {response}") - headers: dict[str, str] = {} - for h in url_node.get("headers", []): - headers[h["name"]] = h["value"] - return _PresignedUrl(url=url, headers=headers) - def build_download_url(self, path: str, expire_time: str) -> CreateDownloadUrlResponse: """Fetches a presigned URL for downloading a file.""" response = self._api.do( @@ -888,21 +866,6 @@ def build_upload_part_urls( ) return results - def build_resumable_upload_url(self, path: str, session_token: str) -> _PresignedUrl: - """Builds a URL for resumable upload directly to the storage proxy.""" - escaped = _escape_multi_segment_path_parameter(path) - base = f"{self._hostname}/api/2.0/fs/files{escaped}" - query = parse.urlencode( - { - "session_token": session_token, - "upload_type": "resumable", - } - ) - return _PresignedUrl( - url=f"{base}?{query}", - headers={"Content-Type": "application/octet-stream"}, - ) - def build_download_url(self, path: str, expire_time: str) -> CreateDownloadUrlResponse: """Builds a URL for downloading a file directly from the storage proxy.""" escaped = _escape_multi_segment_path_parameter(path) @@ -952,7 +915,6 @@ class _UploadContext: def __init__(self, api_client, config: Config): super().__init__(api_client) self._config = config.copy() - self._multipart_upload_read_ahead_bytes = 1 self._cached_cloud_provider_session: Optional[requests.Session] = None self._cloud_provider_session_lock = Lock() self._dp_hostname_available: Optional[bool] = None @@ -1713,16 +1675,6 @@ def _single_thread_multipart_upload(self, ctx: _UploadContext, contents: BinaryI # Rethrow the original exception raise e from None - elif initiate_upload_response.get("resumable_upload"): - session_token = initiate_upload_response["resumable_upload"]["session_token"] - - try: - self._perform_resumable_upload(ctx, contents, session_token, pre_read_buffer) - except FallbackToUploadUsingFilesApi as e: - _LOG.info(f"Falling back to single-shot upload with Files API: {e}") - # Concatenate the buffered part and the rest of the stream. - full_stream = _ConcatenatedInputStream(BytesIO(e.buffer), contents) - return self._single_thread_single_shot_upload(ctx, full_stream) else: raise ValueError(f"Unexpected server response: {initiate_upload_response}") @@ -1734,10 +1686,7 @@ def _parallel_upload_from_stream(self, ctx: _UploadContext, contents: BinaryIO) """ initiate_upload_response = self._initiate_multipart_upload(ctx) - if initiate_upload_response.get("resumable_upload"): - _LOG.warning("GCP does not support parallel resumable uploads, falling back to single-threaded upload") - return self._single_thread_multipart_upload(ctx, contents) - elif initiate_upload_response.get("multipart_upload"): + if initiate_upload_response.get("multipart_upload"): session_token = initiate_upload_response["multipart_upload"].get("session_token") if not session_token: raise ValueError(f"Unexpected server response: {initiate_upload_response}") @@ -1803,10 +1752,6 @@ def _parallel_upload_from_file( # Rethrow the original exception. raise e from None - elif initiate_upload_response.get("resumable_upload"): - _LOG.warning("GCP does not support parallel resumable uploads, falling back to single-threaded upload") - with open(ctx.source_file_path, "rb") as f: - return self._upload_single_thread_with_known_size(ctx, f) else: raise ValueError(f"Unexpected server response: {initiate_upload_response}") @@ -2339,212 +2284,6 @@ def _is_presigned_urls_network_zone_error(self, e: InternalError) -> bool: return True return False - def _perform_resumable_upload( - self, - ctx: _UploadContext, - input_stream: BinaryIO, - session_token: str, - pre_read_buffer: bytes, - ) -> None: - """ - Performs resumable upload on GCP: https://cloud.google.com/storage/docs/performing-resumable-uploads - """ - # Session URI we're using expires after a week - - # Why are we buffering the current chunk? - # When using resumable upload API we're uploading data in chunks. During chunk upload - # server responds with the "received offset" confirming how much data it stored so far, - # so we should continue uploading from that offset. (Note this is not a failure but an - # expected behaviour as per the docs.) But, input stream might be consumed beyond that - # offset, since server might have read more data than it confirmed received, or some data - # might have been pre-cached by e.g. OS or a proxy. So, to continue upload, we must rewind - # the input stream back to the byte next to "received offset". This is not possible - # for non-seekable input stream, so we must buffer the whole last chunk and seek inside - # the buffer. By always uploading from the buffer we fully support non-seekable streams. - - # Why are we doing read-ahead? - # It's not possible to upload an empty chunk as "Content-Range" header format does not - # support this. So if current chunk happens to finish exactly at the end of the stream, - # we need to know that and mark the chunk as last (by passing real file size in the - # "Content-Range" header) when uploading it. To detect if we're at the end of the stream - # we're reading "ahead" an extra bytes but not uploading them immediately. If - # nothing has been read ahead, it means we're at the end of the stream. - # On the contrary, in multipart upload we can decide to complete upload *after* - # last chunk has been sent. - - builder = self._create_request_builder() - - try: - presigned = builder.build_resumable_upload_url(ctx.target_path, session_token) - except Exception as e: - raise FallbackToUploadUsingFilesApi( - pre_read_buffer, f"Failed to obtain resumable upload URL: {e}, falling back to single shot upload" - ) from e - - resumable_upload_url = presigned.url - base_headers = presigned.headers - - try: - # We will buffer this many bytes: one chunk + read-ahead block. - # Note buffer may contain more data initially (from pre_read_buffer). - min_buffer_size = ctx.part_size + self._multipart_upload_read_ahead_bytes - - buffer = pre_read_buffer - - # How many bytes in the buffer were confirmed to be received by the server. - # All the remaining bytes in the buffer must be uploaded. - uploaded_bytes_count = 0 - - chunk_offset = 0 - - retry_count = 0 - while True: - # If needed, fill the buffer to contain at least min_buffer_size bytes - # (unless end of stream), discarding already uploaded bytes. - bytes_to_read = max(0, min_buffer_size - (len(buffer) - uploaded_bytes_count)) - next_buf = input_stream.read(bytes_to_read) - buffer = buffer[uploaded_bytes_count:] + next_buf - - if len(next_buf) < bytes_to_read: - # This is the last chunk in the stream. - # Let's upload all the remaining bytes in one go. - actual_chunk_length = len(buffer) - file_size = chunk_offset + actual_chunk_length - else: - # More chunks expected, let's upload current chunk (excluding read-ahead block). - actual_chunk_length = ctx.part_size - file_size = "*" - - headers: dict = {"Content-Type": "application/octet-stream", **base_headers} - - chunk_last_byte_offset = chunk_offset + actual_chunk_length - 1 - content_range_header = f"bytes {chunk_offset}-{chunk_last_byte_offset}/{file_size}" - _LOG.debug(f"Uploading chunk: {content_range_header}") - headers["Content-Range"] = content_range_header - - def retrieve_upload_status() -> Optional[requests.Response]: - def perform(): - return self._cloud_provider_session().request( - "PUT", - resumable_upload_url, - headers={"Content-Range": "bytes */*"}, - data=b"", - timeout=self._config.files_ext_network_transfer_inactivity_timeout_seconds, - ) - - try: - return self._retry_cloud_idempotent_operation(perform) - except RequestException: - _LOG.warning("Failed to retrieve upload status") - return None - - try: - upload_response = self._cloud_provider_session().request( - "PUT", - resumable_upload_url, - headers=headers, - data=BytesIO(buffer[:actual_chunk_length]), - timeout=self._config.files_ext_network_transfer_inactivity_timeout_seconds, - ) - - # https://cloud.google.com/storage/docs/performing-resumable-uploads#resume-upload - # If an upload request is terminated before receiving a response, or if you receive - # a 503 or 500 response, then you need to resume the interrupted upload from where it left off. - - # Let's follow that for all potentially retryable status codes. - if upload_response.status_code in self._RETRYABLE_STATUS_CODES: - if retry_count < self._config.files_ext_multipart_upload_max_retries: - retry_count += 1 - # let original upload_response be handled as an error - upload_response = retrieve_upload_status() or upload_response - else: - # we received non-retryable response, reset retry count - retry_count = 0 - - except RequestException as e: - # Let's do the same for retryable network errors. - if ( - _BaseClient._is_retryable(e) - and retry_count < self._config.files_ext_multipart_upload_max_retries - ): - retry_count += 1 - upload_response = retrieve_upload_status() - if not upload_response: - # rethrow original exception - raise e from None - else: - # rethrow original exception - raise e from None - - if upload_response.status_code in (200, 201): - if file_size == "*": - raise ValueError( - f"Received unexpected status {upload_response.status_code} before reaching end of stream" - ) - - # upload complete - break - - elif upload_response.status_code == 308: - # chunk accepted (or check-status succeeded), let's determine received offset to resume from there - range_string = upload_response.headers.get("Range") - confirmed_offset = self._extract_range_offset(range_string) - _LOG.debug(f"Received confirmed offset: {confirmed_offset}") - - if confirmed_offset: - if confirmed_offset < chunk_offset - 1 or confirmed_offset > chunk_last_byte_offset: - raise ValueError( - f"Unexpected received offset: {confirmed_offset} is outside of expected range, chunk offset: {chunk_offset}, chunk last byte offset: {chunk_last_byte_offset}" - ) - else: - if chunk_offset > 0: - raise ValueError( - f"Unexpected received offset: {confirmed_offset} is outside of expected range, chunk offset: {chunk_offset}, chunk last byte offset: {chunk_last_byte_offset}" - ) - - # We have just uploaded a part of chunk starting from offset "chunk_offset" and ending - # at offset "confirmed_offset" (inclusive), so the next chunk will start at - # offset "confirmed_offset + 1" - if confirmed_offset: - next_chunk_offset = confirmed_offset + 1 - else: - next_chunk_offset = chunk_offset - uploaded_bytes_count = next_chunk_offset - chunk_offset - chunk_offset = next_chunk_offset - - elif upload_response.status_code == 412 and not ctx.overwrite: - # Assuming this is only possible reason - # Full message in this case: "At least one of the pre-conditions you specified did not hold." - raise AlreadyExists("The file being created already exists.") - - else: - message = f"Unsuccessful chunk upload. Response status: {upload_response.status_code}, body: {upload_response.content}" - _LOG.warning(message) - mapped_error = _error_mapper(upload_response, {}) - raise mapped_error or ValueError(message) - - except Exception as e: - _LOG.info(f"Aborting resumable upload on error: {e}") - try: - self._abort_resumable_upload(resumable_upload_url, base_headers) - except BaseException as ex: - _LOG.warning(f"Failed to abort upload: {ex}") - # ignore, abort is a best-effort - finally: - # rethrow original exception - raise e from None - - @staticmethod - def _extract_range_offset(range_string: Optional[str]) -> Optional[int]: - """Parses the response range header to extract the last byte.""" - if not range_string: - return None # server did not yet confirm any bytes - - if match := re.match("bytes=0-(\\d+)", range_string): - return int(match.group(1)) - else: - raise ValueError(f"Cannot parse response header: Range: {range_string}") - def _get_rfc339_timestamp_with_future_offset(self, base_time: datetime.datetime, offset: timedelta) -> str: """Generates an offset timestamp in an RFC3339 format suitable for URL generation""" offset_timestamp = base_time + offset @@ -2584,19 +2323,6 @@ def _abort_multipart_upload(self, ctx: _UploadContext, session_token: str) -> No if not abort_response.ok: raise ValueError(abort_response) - def _abort_resumable_upload(self, resumable_upload_url: str, headers: dict[str, str]) -> None: - """Aborts ongoing resumable upload session to clean up incomplete file. Best-effort, so sent once.""" - abort_response = self._cloud_provider_session().request( - "DELETE", - resumable_upload_url, - headers=headers, - data=b"", - timeout=self._config.files_ext_network_transfer_inactivity_timeout_seconds, - ) - - if abort_response.status_code not in (200, 201): - raise ValueError(abort_response) - def _cloud_provider_session(self) -> requests.Session: """Returns a session for cloud storage operations. diff --git a/databricks/sdk/service/billing.py b/databricks/sdk/service/billing.py index 0dd77e03a..93de0c72e 100644 --- a/databricks/sdk/service/billing.py +++ b/databricks/sdk/service/billing.py @@ -195,8 +195,9 @@ class AlertSpendStatus: percent_consumed: Optional[str] = None """Percentage of the threshold consumed, as a decimal string rounded down to one decimal place (for example, "68.4" means 68.4%). May exceed 100 when spend is over the threshold. Returns "0.0" if - there's no spend, and is omitted if the threshold is $0. For shared thresholds, this is the only - value returned; the absolute ``spend`` and ``threshold`` dollar amounts are withheld.""" + there's no spend, except that a $0 threshold always returns "100.0". For shared thresholds, this + is the only value returned; the absolute ``spend`` and ``threshold`` dollar amounts are + withheld.""" scope_type: Optional[AlertConfigurationScopeType] = None """Whether the threshold is evaluated per individual user diff --git a/databricks/sdk/service/catalog.py b/databricks/sdk/service/catalog.py index 381aa1c6e..0a1ff1327 100644 --- a/databricks/sdk/service/catalog.py +++ b/databricks/sdk/service/catalog.py @@ -10857,6 +10857,7 @@ class Privilege(Enum): CREATE_EXTERNAL_LOCATION = "CREATE_EXTERNAL_LOCATION" CREATE_EXTERNAL_TABLE = "CREATE_EXTERNAL_TABLE" CREATE_EXTERNAL_VOLUME = "CREATE_EXTERNAL_VOLUME" + CREATE_FEATURE = "CREATE_FEATURE" CREATE_FOREIGN_CATALOG = "CREATE_FOREIGN_CATALOG" CREATE_FOREIGN_SECURABLE = "CREATE_FOREIGN_SECURABLE" CREATE_FUNCTION = "CREATE_FUNCTION" @@ -10880,6 +10881,7 @@ class Privilege(Enum): MANAGE_ALLOWLIST = "MANAGE_ALLOWLIST" MODIFY = "MODIFY" MODIFY_CLEAN_ROOM = "MODIFY_CLEAN_ROOM" + READ_FEATURE = "READ_FEATURE" READ_FILES = "READ_FILES" READ_METADATA = "READ_METADATA" READ_PRIVATE_FILES = "READ_PRIVATE_FILES" @@ -12010,6 +12012,7 @@ class SecurableType(Enum): CREDENTIAL = "CREDENTIAL" EXTERNAL_LOCATION = "EXTERNAL_LOCATION" EXTERNAL_METADATA = "EXTERNAL_METADATA" + FEATURE = "FEATURE" FUNCTION = "FUNCTION" MCP_SERVICE = "MCP_SERVICE" METASTORE = "METASTORE" @@ -16184,23 +16187,22 @@ def validate_credential( ) -> ValidateCredentialResponse: """Validates a credential. - For service credentials (purpose is **SERVICE**), either the **credential_name** or the cloud-specific - credential must be provided. - For storage credentials (purpose is **STORAGE**), at least one of **external_location_name** and **url** need to be provided. If only one of them is provided, it will be used for validation. And if both are provided, the **url** will be used for validation, and **external_location_name** will be - ignored when checking overlapping urls. Either the **credential_name** or the cloud-specific - credential must be provided. + ignored when checking overlapping urls. The caller must be a metastore admin or the credential owner or have the required permission on the metastore and the credential (e.g., **CREATE_EXTERNAL_LOCATION** when purpose is **STORAGE**). :param aws_iam_role: :class:`AwsIamRole` (optional) + Deprecated: create a credential and pass its credential_name instead. :param azure_managed_identity: :class:`AzureManagedIdentity` (optional) + Deprecated: create a credential and pass its credential_name instead. :param credential_name: str (optional) Required. The name of an existing credential or long-lived cloud credential to validate. :param databricks_gcp_service_account: :class:`DatabricksGcpServiceAccount` (optional) + Deprecated: create a credential and pass its credential_name instead. :param external_location_name: str (optional) The name of an existing external location to validate. Only applicable for storage credentials (purpose is **STORAGE**.) @@ -20507,21 +20509,24 @@ def validate( the **url** will be used for validation, and **external_location_name** will be ignored when checking overlapping urls. - Either the **storage_credential_name** or the cloud-specific credential must be provided. - The caller must be a metastore admin or the storage credential owner or have the **CREATE_EXTERNAL_LOCATION** privilege on the metastore and the storage credential. :param aws_iam_role: :class:`AwsIamRoleRequest` (optional) - The AWS IAM role configuration. + Deprecated: create a storage credential and pass its storage_credential_name instead. The AWS IAM + role configuration. :param azure_managed_identity: :class:`AzureManagedIdentityRequest` (optional) - The Azure managed identity configuration. + Deprecated: create a storage credential and pass its storage_credential_name instead. The Azure + managed identity configuration. :param azure_service_principal: :class:`AzureServicePrincipal` (optional) - The Azure service principal configuration. + Deprecated: create a storage credential and pass its storage_credential_name instead. The Azure + service principal configuration. :param cloudflare_api_token: :class:`CloudflareApiToken` (optional) - The Cloudflare API token configuration. + Deprecated: create a storage credential and pass its storage_credential_name instead. The Cloudflare + API token configuration. :param databricks_gcp_service_account: :class:`DatabricksGcpServiceAccountRequest` (optional) - The Databricks created GCP service account configuration. + Deprecated: create a storage credential and pass its storage_credential_name instead. The Databricks + created GCP service account configuration. :param external_location_name: str (optional) The name of an existing external location to validate. :param read_only: bool (optional) @@ -21534,7 +21539,7 @@ def read(self, name: str, *, include_browse: Optional[bool] = None) -> VolumeInf on the parent catalog and the **USE_SCHEMA** privilege on the parent schema. :param name: str - The three-level (fully qualified) name of the volume + The fully qualified name of the volume. :param include_browse: bool (optional) Whether to include volumes in the response for which the principal can only access selective metadata for diff --git a/databricks/sdk/service/compute.py b/databricks/sdk/service/compute.py index 579ad462c..117c06862 100644 --- a/databricks/sdk/service/compute.py +++ b/databricks/sdk/service/compute.py @@ -3658,38 +3658,41 @@ class Environment: base_environment: Optional[str] = None """The base environment this environment is built on top of. A base environment defines the environment version and a list of dependencies for serverless compute. The value can be a file - path to a custom ``env.yaml`` file (e.g., ``/Workspace/path/to/env.yaml``). Support for a - Databricks-provided base environment ID (e.g., ``workspace-base-environments/databricks_ai_v4``) - and workspace base environment ID (e.g., - ``workspace-base-environments/dbe_b849b66e-b31a-4cb5-b161-1f2b10877fb7``) is in Beta. Either - ``environment_version`` or ``base_environment`` can be provided. For more information about - Databricks-provided base environments, see the [list workspace base - environments](:method:Environments/ListWorkspaceBaseEnvironments) API. For more information, see""" + path to a custom ``env.yaml`` file (e.g., ``/Workspace/path/to/env.yaml``). + + Support for a Databricks-provided base environment ID (e.g., + ``workspace-base-environments/databricks_ai_v4``) and workspace base environment ID (e.g., + ``workspace-base-environments/dbe_b849b66e-b31a-4cb5-b161-1f2b10877fb7``) is in Beta. + + Either ``environment_version`` or ``base_environment`` can be provided. + + For more information about Databricks-provided base environments, see the [list workspace base + environments](:method:Environments/ListWorkspaceBaseEnvironments) API.""" client: Optional[str] = None """Use ``environment_version`` instead.""" dependencies: Optional[List[str]] = None """List of pip dependencies, as supported by the version of pip in this environment. Each - dependency is a valid pip requirements file line per - https://pip.pypa.io/en/stable/reference/requirements-file-format/. Allowed dependencies include - a requirement specifier, an archive URL, a local project path (such as WSFS or UC Volumes in - Databricks), or a VCS project URL.""" + dependency is a valid `pip requirements file line + `__. Allowed dependencies + include a requirement specifier, an archive URL, a local project path (such as WSFS or UC + Volumes in Databricks), or a VCS project URL.""" environment_version: Optional[str] = None - """Either ``environment_version`` or ``base_environment`` needs to be provided. Environment version - used by the environment. Each version comes with a specific Python version and a set of Python - packages. The version is a string, consisting of an integer.""" + """The environment version used by this environment. Either ``environment_version`` or + ``base_environment`` must be provided. Each version comes with a specific Python version and a + set of Python packages. The version is a string, consisting of an integer.""" java_dependencies: Optional[List[str]] = None - """List of java dependencies. Each dependency is a string representing a java library path. For + """List of Java dependencies. Each dependency is a string representing a Java library path. For example: ``/Volumes/path/to/test.jar``.""" project_environment: Optional[str] = None """Path to the ``pyproject.toml`` file that defines the project environment. This field cannot be combined with ``environment_version``, ``base_environment``, dependency fields, or any other environment-source fields. Set the environment version and dependencies in ``pyproject.toml`` - instead. For more information, see the serverless dependencies documentation.""" + instead.""" def as_dict(self) -> dict: """Serializes the Environment into a dictionary suitable for use as a JSON request body.""" diff --git a/databricks/sdk/service/computemigrations.py b/databricks/sdk/service/computemigrations.py new file mode 100644 index 000000000..4d0a2fe94 --- /dev/null +++ b/databricks/sdk/service/computemigrations.py @@ -0,0 +1,934 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. +# ruff: noqa: F811, F841 +# F401 is intentionally NOT covered: `make fmt` uses `ruff check --fix-only` +# to strip the fat-import header below; ignoring F401 would defeat that. + +from __future__ import annotations +from dataclasses import dataclass +from enum import Enum +from typing import Dict, List, Any, Iterator, Optional + +from google.protobuf.timestamp_pb2 import Timestamp + +import logging +import uuid + +from databricks.sdk.service._internal import ( + _enum, + _from_dict, + _int64, + _repeated_dict, + _timestamp, +) + + +_LOG = logging.getLogger("databricks.sdk") + + +# all definitions in this file are in alphabetical order + + +@dataclass +class BakingConfig: + """Configuration for the baking window (post-execution production monitoring).""" + + run_count: Optional[int] = None + """Number of successful production runs required before the migration reaches MIGRATION_COMPLETE. + Default applied by the server (configurable).""" + + time_window_hours: Optional[int] = None + """Maximum baking window duration in hours. If the required run_count has not been reached within + this window, the migration transitions to FAILED. Default applied by the server (configurable).""" + + def as_dict(self) -> dict: + """Serializes the BakingConfig into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.run_count is not None: + body["run_count"] = self.run_count + if self.time_window_hours is not None: + body["time_window_hours"] = self.time_window_hours + return body + + def as_shallow_dict(self) -> dict: + """Serializes the BakingConfig into a shallow dictionary of its immediate attributes.""" + body = {} + if self.run_count is not None: + body["run_count"] = self.run_count + if self.time_window_hours is not None: + body["time_window_hours"] = self.time_window_hours + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> BakingConfig: + """Deserializes the BakingConfig from a dictionary.""" + return cls(run_count=d.get("run_count", None), time_window_hours=d.get("time_window_hours", None)) + + +@dataclass +class BakingReport: + """Summary of baking metrics after the baking window completes.""" + + successful_runs: Optional[int] = None + """Number of successful runs.""" + + total_runs: Optional[int] = None + """Number of runs observed during the baking window.""" + + def as_dict(self) -> dict: + """Serializes the BakingReport into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.successful_runs is not None: + body["successful_runs"] = self.successful_runs + if self.total_runs is not None: + body["total_runs"] = self.total_runs + return body + + def as_shallow_dict(self) -> dict: + """Serializes the BakingReport into a shallow dictionary of its immediate attributes.""" + body = {} + if self.successful_runs is not None: + body["successful_runs"] = self.successful_runs + if self.total_runs is not None: + body["total_runs"] = self.total_runs + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> BakingReport: + """Deserializes the BakingReport from a dictionary.""" + return cls(successful_runs=d.get("successful_runs", None), total_runs=d.get("total_runs", None)) + + +@dataclass +class ListMigrationsResponse: + """Response containing a paginated list of migrations matching the request filters.""" + + migrations: Optional[List[Migration]] = None + """The migrations matching the request filters.""" + + next_page_token: Optional[str] = None + """Token for the next page. Omitted if there are no more results.""" + + def as_dict(self) -> dict: + """Serializes the ListMigrationsResponse into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.migrations: + body["migrations"] = [v.as_dict() for v in self.migrations] + if self.next_page_token is not None: + body["next_page_token"] = self.next_page_token + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ListMigrationsResponse into a shallow dictionary of its immediate attributes.""" + body = {} + if self.migrations: + body["migrations"] = self.migrations + if self.next_page_token is not None: + body["next_page_token"] = self.next_page_token + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ListMigrationsResponse: + """Deserializes the ListMigrationsResponse from a dictionary.""" + return cls( + migrations=_repeated_dict(d, "migrations", Migration), next_page_token=d.get("next_page_token", None) + ) + + +@dataclass +class Migration: + """A persistent object tracking the full lifecycle of migrating one Classic compute asset to + Serverless. Each Migration has a well-defined state machine and is the unit of orchestration, + approval, and rollback.""" + + asset_type: MigrationAssetType + """Type of Classic compute asset being migrated. Cannot be changed after creation.""" + + asset_uri: str + """URI of the source asset (e.g., "jobs/12345"). Cannot be changed after creation.""" + + auto_approve: Optional[bool] = None + """When true, the server skips the translation review (TRANSLATION_READY) instead of waiting for an + explicit ApproveMigrationTranslation. MIGRATION_REVIEW is never skipped. Unset/false means the + translation review requires explicit user action. Set at creation; cannot be changed afterward. + Does not suppress cancellation or terminal failures.""" + + baking_config: Optional[BakingConfig] = None + """Baking window configuration. Mutable via UpdateMigration before baking starts.""" + + baking_report: Optional[BakingReport] = None + """Summary of baking metrics after the baking window completes.""" + + create_time: Optional[Timestamp] = None + """Timestamp when the migration was created.""" + + etag: Optional[str] = None + """Resource version for optimistic concurrency control. Clients read this from GetMigration or + ListMigrations and pass it back as the request ``etag`` on update, delete, approve, reject, and + cancel methods.""" + + initiator: Optional[MigrationInitiator] = None + """The user who initiated this migration and their role. Set at creation.""" + + migration_id: Optional[str] = None + """Server-generated unique identifier for this migration.""" + + migration_report: Optional[MigrationReport] = None + """Agent-produced planning report, presented to the user at the review gates.""" + + run_validation: Optional[bool] = None + """When true, Databricks runs the working copy to validate it after the translation review. + Unset/false skips the run, so the migration goes from the translation review straight to + execution review. Whether the translation review waits for the user is controlled by + ``auto_approve``, not this field. Set at creation; cannot be changed afterward.""" + + state: Optional[MigrationState] = None + """Current state of the migration in the state machine. Managed by the server.""" + + state_history: Optional[List[MigrationStateTransition]] = None + """History of state transitions, ordered chronologically. Each entry records the state entered and + the timestamp of the transition.""" + + validation_report: Optional[ValidationReport] = None + """Agent-produced validation results from sandbox execution.""" + + def as_dict(self) -> dict: + """Serializes the Migration into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.asset_type is not None: + body["asset_type"] = self.asset_type.value + if self.asset_uri is not None: + body["asset_uri"] = self.asset_uri + if self.auto_approve is not None: + body["auto_approve"] = self.auto_approve + if self.baking_config: + body["baking_config"] = self.baking_config.as_dict() + if self.baking_report: + body["baking_report"] = self.baking_report.as_dict() + if self.create_time is not None: + body["create_time"] = self.create_time.ToJsonString() + if self.etag is not None: + body["etag"] = self.etag + if self.initiator: + body["initiator"] = self.initiator.as_dict() + if self.migration_id is not None: + body["migration_id"] = self.migration_id + if self.migration_report: + body["migration_report"] = self.migration_report.as_dict() + if self.run_validation is not None: + body["run_validation"] = self.run_validation + if self.state is not None: + body["state"] = self.state.value + if self.state_history: + body["state_history"] = [v.as_dict() for v in self.state_history] + if self.validation_report: + body["validation_report"] = self.validation_report.as_dict() + return body + + def as_shallow_dict(self) -> dict: + """Serializes the Migration into a shallow dictionary of its immediate attributes.""" + body = {} + if self.asset_type is not None: + body["asset_type"] = self.asset_type + if self.asset_uri is not None: + body["asset_uri"] = self.asset_uri + if self.auto_approve is not None: + body["auto_approve"] = self.auto_approve + if self.baking_config: + body["baking_config"] = self.baking_config + if self.baking_report: + body["baking_report"] = self.baking_report + if self.create_time is not None: + body["create_time"] = self.create_time + if self.etag is not None: + body["etag"] = self.etag + if self.initiator: + body["initiator"] = self.initiator + if self.migration_id is not None: + body["migration_id"] = self.migration_id + if self.migration_report: + body["migration_report"] = self.migration_report + if self.run_validation is not None: + body["run_validation"] = self.run_validation + if self.state is not None: + body["state"] = self.state + if self.state_history: + body["state_history"] = self.state_history + if self.validation_report: + body["validation_report"] = self.validation_report + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> Migration: + """Deserializes the Migration from a dictionary.""" + return cls( + asset_type=_enum(d, "asset_type", MigrationAssetType), + asset_uri=d.get("asset_uri", None), + auto_approve=d.get("auto_approve", None), + baking_config=_from_dict(d, "baking_config", BakingConfig), + baking_report=_from_dict(d, "baking_report", BakingReport), + create_time=_timestamp(d, "create_time"), + etag=d.get("etag", None), + initiator=_from_dict(d, "initiator", MigrationInitiator), + migration_id=d.get("migration_id", None), + migration_report=_from_dict(d, "migration_report", MigrationReport), + run_validation=d.get("run_validation", None), + state=_enum(d, "state", MigrationState), + state_history=_repeated_dict(d, "state_history", MigrationStateTransition), + validation_report=_from_dict(d, "validation_report", ValidationReport), + ) + + +@dataclass +class MigrationAsset: + """A working copy (job or notebook) the migration agent created for the migration.""" + + asset_type: Optional[MigrationAssetType] = None + """Type of asset this working copy represents.""" + + id: Optional[str] = None + """Identifier of the asset.""" + + uri: Optional[str] = None + """Link to the asset, e.g. "jobs/12345".""" + + def as_dict(self) -> dict: + """Serializes the MigrationAsset into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.asset_type is not None: + body["asset_type"] = self.asset_type.value + if self.id is not None: + body["id"] = self.id + if self.uri is not None: + body["uri"] = self.uri + return body + + def as_shallow_dict(self) -> dict: + """Serializes the MigrationAsset into a shallow dictionary of its immediate attributes.""" + body = {} + if self.asset_type is not None: + body["asset_type"] = self.asset_type + if self.id is not None: + body["id"] = self.id + if self.uri is not None: + body["uri"] = self.uri + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> MigrationAsset: + """Deserializes the MigrationAsset from a dictionary.""" + return cls(asset_type=_enum(d, "asset_type", MigrationAssetType), id=d.get("id", None), uri=d.get("uri", None)) + + +class MigrationAssetType(Enum): + """Type of a migration asset: the Classic asset being migrated, or an asset the migration agent + created for it. Migrations can only be created for JOB assets.""" + + FILE = "FILE" + INTERACTIVE = "INTERACTIVE" + JOB = "JOB" + + +@dataclass +class MigrationInitiator: + """The initiator of a migration, including their role and user identity.""" + + role: Optional[MigrationInitiatorRole] = None + """The role of the initiator, determining approval authority. Set by the server: the role is + resolved from the caller's identity and permissions on the asset, never taken from the request, + since a caller must not be able to assert its own approval authority.""" + + user_id: Optional[int] = None + """The user ID of the initiator. Set by the server from the request context.""" + + def as_dict(self) -> dict: + """Serializes the MigrationInitiator into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.role is not None: + body["role"] = self.role.value + if self.user_id is not None: + body["user_id"] = self.user_id + return body + + def as_shallow_dict(self) -> dict: + """Serializes the MigrationInitiator into a shallow dictionary of its immediate attributes.""" + body = {} + if self.role is not None: + body["role"] = self.role + if self.user_id is not None: + body["user_id"] = self.user_id + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> MigrationInitiator: + """Deserializes the MigrationInitiator from a dictionary.""" + return cls(role=_enum(d, "role", MigrationInitiatorRole), user_id=_int64(d, "user_id")) + + +class MigrationInitiatorRole(Enum): + """Role of the migration initiator, determining approval authority.""" + + ACCOUNT_ADMIN = "ACCOUNT_ADMIN" + JOB_ADMIN = "JOB_ADMIN" + SYSTEM = "SYSTEM" + WORKSPACE_ADMIN = "WORKSPACE_ADMIN" + + +@dataclass +class MigrationReport: + """Agent-produced planning report, presented to the user at the review gates.""" + + primary_asset: Optional[MigrationAsset] = None + """The primary asset the agent created — the entry point / validation target.""" + + secondary_assets: Optional[List[MigrationAsset]] = None + """Additional assets the agent created (e.g. notebooks or jobs the primary asset invokes).""" + + def as_dict(self) -> dict: + """Serializes the MigrationReport into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.primary_asset: + body["primary_asset"] = self.primary_asset.as_dict() + if self.secondary_assets: + body["secondary_assets"] = [v.as_dict() for v in self.secondary_assets] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the MigrationReport into a shallow dictionary of its immediate attributes.""" + body = {} + if self.primary_asset: + body["primary_asset"] = self.primary_asset + if self.secondary_assets: + body["secondary_assets"] = self.secondary_assets + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> MigrationReport: + """Deserializes the MigrationReport from a dictionary.""" + return cls( + primary_asset=_from_dict(d, "primary_asset", MigrationAsset), + secondary_assets=_repeated_dict(d, "secondary_assets", MigrationAsset), + ) + + +class MigrationState(Enum): + """The state of a migration in its lifecycle.""" + + CANCELLED = "CANCELLED" + FAILED = "FAILED" + MIGRATION_COMPLETE = "MIGRATION_COMPLETE" + MIGRATION_OBSERVATION = "MIGRATION_OBSERVATION" + MIGRATION_PENDING = "MIGRATION_PENDING" + MIGRATION_REVIEW = "MIGRATION_REVIEW" + MIGRATION_RUNNING = "MIGRATION_RUNNING" + NOT_SUPPORTED = "NOT_SUPPORTED" + PLANNING_PENDING = "PLANNING_PENDING" + PLANNING_REVIEW = "PLANNING_REVIEW" + PLANNING_RUNNING = "PLANNING_RUNNING" + REVISION_PENDING = "REVISION_PENDING" + TRANSLATION_CREATING = "TRANSLATION_CREATING" + TRANSLATION_PENDING = "TRANSLATION_PENDING" + TRANSLATION_READY = "TRANSLATION_READY" + VALIDATION_RUNNING = "VALIDATION_RUNNING" + VALIDATION_SCHEDULING = "VALIDATION_SCHEDULING" + + +@dataclass +class MigrationStateTransition: + """A record of a state transition in the migration lifecycle.""" + + internal_comment: Optional[str] = None + """Free-text the service recorded for this transition — e.g. an internally-determined + cancellation or failure reason. Empty for user-driven transitions.""" + + state: Optional[MigrationState] = None + """The state that was entered.""" + + transition_time: Optional[Timestamp] = None + """Timestamp when this state was entered.""" + + user_comment: Optional[str] = None + """Free-text the initiating user supplied at a user-driven transition — reviewer feedback on an + approval, or the reason for a rejection or cancellation. Empty for transitions the service + drives internally.""" + + def as_dict(self) -> dict: + """Serializes the MigrationStateTransition into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.internal_comment is not None: + body["internal_comment"] = self.internal_comment + if self.state is not None: + body["state"] = self.state.value + if self.transition_time is not None: + body["transition_time"] = self.transition_time.ToJsonString() + if self.user_comment is not None: + body["user_comment"] = self.user_comment + return body + + def as_shallow_dict(self) -> dict: + """Serializes the MigrationStateTransition into a shallow dictionary of its immediate attributes.""" + body = {} + if self.internal_comment is not None: + body["internal_comment"] = self.internal_comment + if self.state is not None: + body["state"] = self.state + if self.transition_time is not None: + body["transition_time"] = self.transition_time + if self.user_comment is not None: + body["user_comment"] = self.user_comment + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> MigrationStateTransition: + """Deserializes the MigrationStateTransition from a dictionary.""" + return cls( + internal_comment=d.get("internal_comment", None), + state=_enum(d, "state", MigrationState), + transition_time=_timestamp(d, "transition_time"), + user_comment=d.get("user_comment", None), + ) + + +@dataclass +class ValidationReport: + """Validation results from running the translated workload in the sandbox.""" + + duration_ms: Optional[int] = None + """Duration of the sandbox validation run in milliseconds. Pair with start_time to derive the end + time.""" + + passed: Optional[bool] = None + """Databricks-internal recommendation on whether to proceed with the migration, based on the + validation run. May differ from run_status: a run can succeed while the recommendation is + negative, or vice versa.""" + + run_status: Optional[str] = None + """Status of the validation run itself (the objective run outcome), as reported by the run system, + e.g. "SUCCESS" / "FAILED" / "TIMED_OUT".""" + + run_url: Optional[str] = None + """Link to the validation run (job or notebook).""" + + start_time: Optional[Timestamp] = None + """Timestamp when validation started in the sandbox.""" + + def as_dict(self) -> dict: + """Serializes the ValidationReport into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.duration_ms is not None: + body["duration_ms"] = self.duration_ms + if self.passed is not None: + body["passed"] = self.passed + if self.run_status is not None: + body["run_status"] = self.run_status + if self.run_url is not None: + body["run_url"] = self.run_url + if self.start_time is not None: + body["start_time"] = self.start_time.ToJsonString() + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ValidationReport into a shallow dictionary of its immediate attributes.""" + body = {} + if self.duration_ms is not None: + body["duration_ms"] = self.duration_ms + if self.passed is not None: + body["passed"] = self.passed + if self.run_status is not None: + body["run_status"] = self.run_status + if self.run_url is not None: + body["run_url"] = self.run_url + if self.start_time is not None: + body["start_time"] = self.start_time + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ValidationReport: + """Deserializes the ValidationReport from a dictionary.""" + return cls( + duration_ms=_int64(d, "duration_ms"), + passed=d.get("passed", None), + run_status=d.get("run_status", None), + run_url=d.get("run_url", None), + start_time=_timestamp(d, "start_time"), + ) + + +class ComputeMigrationsAPI: + """Manage migrations of Classic compute workloads to Serverless.""" + + def __init__(self, api_client): + self._api = api_client + + def approve_migration_execution( + self, migration_id: str, *, additional_feedback: Optional[str] = None, etag: Optional[str] = None + ) -> Migration: + """Approves migration execution at the MIGRATION_REVIEW gate, advancing to execution. + + Valid only when the migration is in MIGRATION_REVIEW. If supplied, the ``etag`` must match the current + server value. + + Returns NOT_FOUND if the migration does not exist. Returns FAILED_PRECONDITION if the migration is not + in MIGRATION_REVIEW. Returns ABORTED if the etag does not match the current resource version. + + :param migration_id: str + The ID of the migration to approve for execution. + :param additional_feedback: str (optional) + Optional free-text feedback or instructions the user provides along with the approval, carried into + the next stage of the migration. + :param etag: str (optional) + Optional resource version for optimistic concurrency control. If supplied, it must match the current + server value. + + :returns: :class:`Migration` + """ + + body = {} + if additional_feedback is not None: + body["additional_feedback"] = additional_feedback + if etag is not None: + body["etag"] = etag + headers = { + "Accept": "application/json", + "Content-Type": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do( + "POST", f"/api/compute-migrations/v1/migrations/{migration_id}/approveExecution", body=body, headers=headers + ) + return Migration.from_dict(res) + + def approve_migration_translation( + self, migration_id: str, *, additional_feedback: Optional[str] = None, etag: Optional[str] = None + ) -> Migration: + """Approves the working copy at the TRANSLATION_READY gate. If the migration was created with + ``run_validation``, the working copy is then run to verify it executes successfully before the + migration is presented for execution review; otherwise it advances straight to MIGRATION_REVIEW. + + Valid only when the migration is in TRANSLATION_READY. If supplied, the ``etag`` must match the + current server value. + + Returns NOT_FOUND if the migration does not exist. Returns FAILED_PRECONDITION if the migration is not + in TRANSLATION_READY. Returns ABORTED if the etag does not match the current resource version. + + :param migration_id: str + The ID of the migration to approve. + :param additional_feedback: str (optional) + Optional free-text feedback or instructions the user provides along with the approval, carried into + the next stage of the migration. + :param etag: str (optional) + Optional resource version for optimistic concurrency control. If supplied, it must match the current + server value. + + :returns: :class:`Migration` + """ + + body = {} + if additional_feedback is not None: + body["additional_feedback"] = additional_feedback + if etag is not None: + body["etag"] = etag + headers = { + "Accept": "application/json", + "Content-Type": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do( + "POST", + f"/api/compute-migrations/v1/migrations/{migration_id}/approveTranslation", + body=body, + headers=headers, + ) + return Migration.from_dict(res) + + def cancel_migration( + self, migration_id: str, *, cancellation_reason: Optional[str] = None, etag: Optional[str] = None + ) -> Migration: + """Cancels the migration, transitioning it to CANCELLED. + + Valid from any non-terminal state, unlike the review-gate transitions, which each require a specific + state. Any in-flight agent work is abandoned and the migration stops advancing. The resource is + retained for auditability. + + Returns NOT_FOUND if the migration does not exist. Returns FAILED_PRECONDITION if the migration is + already in a terminal state. Returns ABORTED if the etag does not match the current resource version. + + :param migration_id: str + The ID of the migration to cancel. + :param cancellation_reason: str (optional) + Optional free-text reason for the cancellation. + :param etag: str (optional) + Optional resource version for optimistic concurrency control. If supplied, it must match the current + server value. + + :returns: :class:`Migration` + """ + + body = {} + if cancellation_reason is not None: + body["cancellation_reason"] = cancellation_reason + if etag is not None: + body["etag"] = etag + headers = { + "Accept": "application/json", + "Content-Type": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do( + "POST", f"/api/compute-migrations/v1/migrations/{migration_id}/cancel", body=body, headers=headers + ) + return Migration.from_dict(res) + + def create(self, migration: Migration, *, request_id: Optional[str] = None) -> Migration: + """Creates a new migration for a Classic compute asset. + + The migration is created in PLANNING_PENDING and advances through the state machine as planning work + is scheduled. The caller must have appropriate permissions: workspace admin for infrastructure-scoped + migrations, job owner/admin for individual workload migrations. + + At most one migration may hold an asset. Returns ALREADY_EXISTS if the asset has a migration in + flight, or one that already migrated it successfully; a migration that ended without succeeding, or + was deleted, releases the asset so it can be migrated again. Returns INVALID_PARAMETER_VALUE if the + asset is not eligible for migration. + + :param migration: :class:`Migration` + The migration to create. The server generates ``migration_id`` and sets ``state`` to + PLANNING_PENDING. Required fields: ``asset_type``, ``asset_uri``. + :param request_id: str (optional) + Client-generated request ID for idempotent request deduplication. The server de-duplicates requests + with the same request_id for at least 24 hours, returning the existing migration instead of creating + a duplicate. + + :returns: :class:`Migration` + """ + + if request_id is None or request_id == "": + request_id = str(uuid.uuid4()) + body = migration.as_dict() + query = {} + if request_id is not None: + query["request_id"] = request_id + headers = { + "Accept": "application/json", + "Content-Type": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do("POST", "/api/compute-migrations/v1/migrations", query=query, body=body, headers=headers) + return Migration.from_dict(res) + + def delete(self, migration_id: str, *, etag: Optional[str] = None) -> Migration: + """Deletes a migration (soft-delete). + + Allowed from any state. Deleting a migration marks the record deleted only — it does not roll back a + migration that already ran, and resources created during the migration are not cleaned up. The + migration resource is retained for auditability. To stop an in-flight migration (agent teardown, + transition to CANCELLED), use CancelMigration instead. + + Returns NOT_FOUND if the migration does not exist, or was already deleted: a deleted migration is + hidden from every method (Get, Update, List, and Delete), so it cannot be deleted again. Returns + ABORTED if a supplied nonempty request etag does not match the resource version. + + :param migration_id: str + The ID of the migration to cancel. + :param etag: str (optional) + If nonempty, must match the current resource etag or the request fails with ABORTED. + + :returns: :class:`Migration` + """ + + query = {} + if etag is not None: + query["etag"] = etag + headers = { + "Accept": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do( + "DELETE", f"/api/compute-migrations/v1/migrations/{migration_id}", query=query, headers=headers + ) + return Migration.from_dict(res) + + def get(self, migration_id: str) -> Migration: + """Retrieves a migration by its ID. + + Returns the current state, configuration, and output artifacts of the migration. Use this to poll + migration progress or inspect artifacts after state transitions. Returns NOT_FOUND if the migration + does not exist in the caller's workspace. + + :param migration_id: str + The ID of the migration to retrieve. + + :returns: :class:`Migration` + """ + + headers = { + "Accept": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do("GET", f"/api/compute-migrations/v1/migrations/{migration_id}", headers=headers) + return Migration.from_dict(res) + + def list( + self, + *, + asset_type: Optional[MigrationAssetType] = None, + page_size: Optional[int] = None, + page_token: Optional[str] = None, + states: Optional[List[MigrationState]] = None, + ) -> Iterator[Migration]: + """Lists migrations in the caller's workspace, ordered by migration ID, with cursor-based pagination and + filtering by asset type and state. + + Eventually consistent: results are read from a cache, and pagination is not a point-in-time snapshot, + so a migration created during a List may not appear until a subsequent one. + + :param asset_type: :class:`MigrationAssetType` (optional) + Filter by asset type. + :param page_size: int (optional) + Maximum number of migrations to return. Default: 50, maximum: 1000. + :param page_token: str (optional) + Page token from a previous ListMigrations call. When paginating, all other parameters must match the + call that provided the page token. + :param states: List[:class:`MigrationState`] (optional) + Filter by migration state. Returns migrations in any of the specified states. + + :returns: Iterator over :class:`Migration` + """ + + query = {} + if asset_type is not None: + query["asset_type"] = asset_type.value + if page_size is not None: + query["page_size"] = page_size + if page_token is not None: + query["page_token"] = page_token + if states is not None: + query["states"] = [v.value for v in states] + headers = { + "Accept": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + while True: + json = self._api.do("GET", "/api/compute-migrations/v1/migrations", query=query, headers=headers) + if "migrations" in json: + for v in json["migrations"]: + yield Migration.from_dict(v) + if "next_page_token" not in json or not json["next_page_token"]: + return + query["page_token"] = json["next_page_token"] + + def reject_migration_execution( + self, migration_id: str, rejection_reason: str, *, etag: Optional[str] = None + ) -> Migration: + """Rejects migration execution at the MIGRATION_REVIEW gate, sending the migration to REVISION_PENDING to + loop back for re-planning. + + Valid only when the migration is in MIGRATION_REVIEW. If supplied, the ``etag`` must match the current + server value. + + Returns NOT_FOUND if the migration does not exist. Returns FAILED_PRECONDITION if the migration is not + in MIGRATION_REVIEW. Returns ABORTED if the etag does not match the current resource version. + + :param migration_id: str + The ID of the migration to reject. + :param rejection_reason: str + Required free-text reason for the rejection, fed back into re-planning so the next plan can address + the user's concerns. + :param etag: str (optional) + Optional resource version for optimistic concurrency control. If supplied, it must match the current + server value. + + :returns: :class:`Migration` + """ + + body = {} + if etag is not None: + body["etag"] = etag + if rejection_reason is not None: + body["rejection_reason"] = rejection_reason + headers = { + "Accept": "application/json", + "Content-Type": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do( + "POST", f"/api/compute-migrations/v1/migrations/{migration_id}/rejectExecution", body=body, headers=headers + ) + return Migration.from_dict(res) + + def reject_migration_translation( + self, migration_id: str, rejection_reason: str, *, etag: Optional[str] = None + ) -> Migration: + """Rejects the working copy at the TRANSLATION_READY gate, sending the migration to REVISION_PENDING to + loop back for re-planning. + + Valid only when the migration is in TRANSLATION_READY. If supplied, the ``etag`` must match the + current server value. + + Returns NOT_FOUND if the migration does not exist. Returns FAILED_PRECONDITION if the migration is not + in TRANSLATION_READY. Returns ABORTED if the etag does not match the current resource version. + + :param migration_id: str + The ID of the migration to reject. + :param rejection_reason: str + Required free-text reason for the rejection, fed back into re-planning so the next plan can address + the user's concerns. + :param etag: str (optional) + Optional resource version for optimistic concurrency control. If supplied, it must match the current + server value. + + :returns: :class:`Migration` + """ + + body = {} + if etag is not None: + body["etag"] = etag + if rejection_reason is not None: + body["rejection_reason"] = rejection_reason + headers = { + "Accept": "application/json", + "Content-Type": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do( + "POST", + f"/api/compute-migrations/v1/migrations/{migration_id}/rejectTranslation", + body=body, + headers=headers, + ) + return Migration.from_dict(res) diff --git a/databricks/sdk/service/genieagents.py b/databricks/sdk/service/genieagents.py new file mode 100644 index 000000000..20324dd33 --- /dev/null +++ b/databricks/sdk/service/genieagents.py @@ -0,0 +1,1524 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. +# ruff: noqa: F811, F841 +# F401 is intentionally NOT covered: `make fmt` uses `ruff check --fix-only` +# to strip the fat-import header below; ignoring F401 would defeat that. + +from __future__ import annotations +from dataclasses import dataclass +from enum import Enum +from typing import Dict, List, Any, Iterator, Optional + +from google.protobuf.timestamp_pb2 import Timestamp + +import logging + +from databricks.sdk.service._internal import ( + _enum, + _from_dict, + _repeated_dict, + _timestamp, +) +from databricks.sdk.common.types.fieldmask import FieldMask + + +_LOG = logging.getLogger("databricks.sdk") + + +# all definitions in this file are in alphabetical order + + +@dataclass +class GenieAgent: + """A Genie Agent resource.""" + + display_name: str + """Display name of the Genie Agent. Must be unique within its folder (effective_parent_path): a + create, rename, or move that would collide with an existing Genie Agent's name in the target + folder is rejected with ALREADY_EXISTS. + + The maximum length is 63 characters. Requests with a longer value are rejected with + INVALID_PARAMETER_VALUE.""" + + definition: GenieAgentDefinition + """The definition of the Genie Agent: its sample questions, data sources, instructions, and + benchmarks. Must set at least one field. Returned on get and on list with + ``GENIE_AGENT_VIEW_FULL``. Create and update return a submitted definition unchanged; one the + service would store differently, or with a column config the workspace can't keep (unknown + column, or a display name while those are disabled), is rejected with INVALID_PARAMETER_VALUE.""" + + warehouse_id: str + """Warehouse associated with the Genie Agent.""" + + create_time: Optional[Timestamp] = None + """Time when the Genie Agent was created.""" + + description: Optional[str] = None + """Description of the Genie Agent. + + The maximum length is 4,096 characters. Requests with a longer value are rejected with + INVALID_PARAMETER_VALUE.""" + + effective_parent_path: Optional[str] = None + """The workspace folder that contains the Genie Agent.""" + + etag: Optional[str] = None + """Weak etag of ``definition`` only; metadata changes do not affect it. Send it on update or delete + to detect concurrent definition changes.""" + + name: Optional[str] = None + """The resource name of the Genie Agent, in the form ``agents/{genie_agent_id}``. Ignored on + create.""" + + parent_path: Optional[str] = None + """The workspace folder to create the Genie Agent in. If omitted, the caller's home folder is used. + Returned exactly as submitted; see ``effective_parent_path`` for the current folder. On update, + setting it moves the Genie Agent and clearing it does not. At most 4,096 characters.""" + + update_time: Optional[Timestamp] = None + """Time when the Genie Agent was last modified.""" + + def as_dict(self) -> dict: + """Serializes the GenieAgent into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.create_time is not None: + body["create_time"] = self.create_time.ToJsonString() + if self.definition: + body["definition"] = self.definition.as_dict() + if self.description is not None: + body["description"] = self.description + if self.display_name is not None: + body["display_name"] = self.display_name + if self.effective_parent_path is not None: + body["effective_parent_path"] = self.effective_parent_path + if self.etag is not None: + body["etag"] = self.etag + if self.name is not None: + body["name"] = self.name + if self.parent_path is not None: + body["parent_path"] = self.parent_path + if self.update_time is not None: + body["update_time"] = self.update_time.ToJsonString() + if self.warehouse_id is not None: + body["warehouse_id"] = self.warehouse_id + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieAgent into a shallow dictionary of its immediate attributes.""" + body = {} + if self.create_time is not None: + body["create_time"] = self.create_time + if self.definition: + body["definition"] = self.definition + if self.description is not None: + body["description"] = self.description + if self.display_name is not None: + body["display_name"] = self.display_name + if self.effective_parent_path is not None: + body["effective_parent_path"] = self.effective_parent_path + if self.etag is not None: + body["etag"] = self.etag + if self.name is not None: + body["name"] = self.name + if self.parent_path is not None: + body["parent_path"] = self.parent_path + if self.update_time is not None: + body["update_time"] = self.update_time + if self.warehouse_id is not None: + body["warehouse_id"] = self.warehouse_id + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieAgent: + """Deserializes the GenieAgent from a dictionary.""" + return cls( + create_time=_timestamp(d, "create_time"), + definition=_from_dict(d, "definition", GenieAgentDefinition), + description=d.get("description", None), + display_name=d.get("display_name", None), + effective_parent_path=d.get("effective_parent_path", None), + etag=d.get("etag", None), + name=d.get("name", None), + parent_path=d.get("parent_path", None), + update_time=_timestamp(d, "update_time"), + warehouse_id=d.get("warehouse_id", None), + ) + + +@dataclass +class GenieAgentBenchmarkAnswer: + """The expected answer to a benchmark question.""" + + format: GenieAgentBenchmarkAnswerFormat + """The format of the answer.""" + + content: str + """The answer.""" + + def as_dict(self) -> dict: + """Serializes the GenieAgentBenchmarkAnswer into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.content is not None: + body["content"] = self.content + if self.format is not None: + body["format"] = self.format.value + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieAgentBenchmarkAnswer into a shallow dictionary of its immediate attributes.""" + body = {} + if self.content is not None: + body["content"] = self.content + if self.format is not None: + body["format"] = self.format + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieAgentBenchmarkAnswer: + """Deserializes the GenieAgentBenchmarkAnswer from a dictionary.""" + return cls(content=d.get("content", None), format=_enum(d, "format", GenieAgentBenchmarkAnswerFormat)) + + +class GenieAgentBenchmarkAnswerFormat(Enum): + """The format of a benchmark answer.""" + + GENIE_AGENT_BENCHMARK_ANSWER_FORMAT_SQL = "GENIE_AGENT_BENCHMARK_ANSWER_FORMAT_SQL" + + +@dataclass +class GenieAgentBenchmarkQuestion: + """A benchmark question with its expected answer.""" + + id: str + """Identifier of the question: 32 lowercase hexadecimal characters, unique within the Genie Agent.""" + + question: str + """The question.""" + + answer: Optional[List[GenieAgentBenchmarkAnswer]] = None + """The expected answer. At most one answer is supported; omit it when the question has no expected + answer.""" + + evaluation_note: Optional[str] = None + """Notes on how to evaluate answers to the question.""" + + def as_dict(self) -> dict: + """Serializes the GenieAgentBenchmarkQuestion into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.answer: + body["answer"] = [v.as_dict() for v in self.answer] + if self.evaluation_note is not None: + body["evaluation_note"] = self.evaluation_note + if self.id is not None: + body["id"] = self.id + if self.question is not None: + body["question"] = self.question + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieAgentBenchmarkQuestion into a shallow dictionary of its immediate attributes.""" + body = {} + if self.answer: + body["answer"] = self.answer + if self.evaluation_note is not None: + body["evaluation_note"] = self.evaluation_note + if self.id is not None: + body["id"] = self.id + if self.question is not None: + body["question"] = self.question + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieAgentBenchmarkQuestion: + """Deserializes the GenieAgentBenchmarkQuestion from a dictionary.""" + return cls( + answer=_repeated_dict(d, "answer", GenieAgentBenchmarkAnswer), + evaluation_note=d.get("evaluation_note", None), + id=d.get("id", None), + question=d.get("question", None), + ) + + +@dataclass +class GenieAgentBenchmarks: + """Benchmarks for evaluating a Genie Agent.""" + + questions: Optional[List[GenieAgentBenchmarkQuestion]] = None + """The benchmark questions, sorted by ``id`` in descending order.""" + + def as_dict(self) -> dict: + """Serializes the GenieAgentBenchmarks into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.questions: + body["questions"] = [v.as_dict() for v in self.questions] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieAgentBenchmarks into a shallow dictionary of its immediate attributes.""" + body = {} + if self.questions: + body["questions"] = self.questions + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieAgentBenchmarks: + """Deserializes the GenieAgentBenchmarks from a dictionary.""" + return cls(questions=_repeated_dict(d, "questions", GenieAgentBenchmarkQuestion)) + + +@dataclass +class GenieAgentColumnConfig: + """Settings for one column of a table or metric view.""" + + column: str + """The name of the column. Must exist in the table or metric view.""" + + description: Optional[str] = None + """Description of the column, overriding its Unity Catalog comment.""" + + display_name: Optional[str] = None + """Display name of the column in the Genie Agent. Accepted only where column display names are + enabled for the workspace.""" + + enable_entity_matching: Optional[bool] = None + """Whether the Genie Agent indexes the column's values to match them in questions.""" + + enable_format_assistance: Optional[bool] = None + """Whether the Genie Agent samples example values from the column to learn its format.""" + + exclude: Optional[bool] = None + """Whether the column is hidden from the Genie Agent.""" + + synonyms: Optional[List[str]] = None + """Synonyms for the column name.""" + + def as_dict(self) -> dict: + """Serializes the GenieAgentColumnConfig into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.column is not None: + body["column"] = self.column + if self.description is not None: + body["description"] = self.description + if self.display_name is not None: + body["display_name"] = self.display_name + if self.enable_entity_matching is not None: + body["enable_entity_matching"] = self.enable_entity_matching + if self.enable_format_assistance is not None: + body["enable_format_assistance"] = self.enable_format_assistance + if self.exclude is not None: + body["exclude"] = self.exclude + if self.synonyms: + body["synonyms"] = [v for v in self.synonyms] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieAgentColumnConfig into a shallow dictionary of its immediate attributes.""" + body = {} + if self.column is not None: + body["column"] = self.column + if self.description is not None: + body["description"] = self.description + if self.display_name is not None: + body["display_name"] = self.display_name + if self.enable_entity_matching is not None: + body["enable_entity_matching"] = self.enable_entity_matching + if self.enable_format_assistance is not None: + body["enable_format_assistance"] = self.enable_format_assistance + if self.exclude is not None: + body["exclude"] = self.exclude + if self.synonyms: + body["synonyms"] = self.synonyms + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieAgentColumnConfig: + """Deserializes the GenieAgentColumnConfig from a dictionary.""" + return cls( + column=d.get("column", None), + description=d.get("description", None), + display_name=d.get("display_name", None), + enable_entity_matching=d.get("enable_entity_matching", None), + enable_format_assistance=d.get("enable_format_assistance", None), + exclude=d.get("exclude", None), + synonyms=d.get("synonyms", None), + ) + + +@dataclass +class GenieAgentConfig: + """General settings of a Genie Agent. The warehouse, display name, and description are separate + fields of the Genie Agent.""" + + sample_questions: Optional[List[GenieAgentSampleQuestion]] = None + """Sample questions shown to users of the Genie Agent, sorted by ``id``.""" + + def as_dict(self) -> dict: + """Serializes the GenieAgentConfig into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.sample_questions: + body["sample_questions"] = [v.as_dict() for v in self.sample_questions] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieAgentConfig into a shallow dictionary of its immediate attributes.""" + body = {} + if self.sample_questions: + body["sample_questions"] = self.sample_questions + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieAgentConfig: + """Deserializes the GenieAgentConfig from a dictionary.""" + return cls(sample_questions=_repeated_dict(d, "sample_questions", GenieAgentSampleQuestion)) + + +@dataclass +class GenieAgentDataSources: + """The data a Genie Agent can access.""" + + mcp_connections: Optional[List[GenieAgentMcpConnection]] = None + """Unity Catalog MCP connections, sorted by ``mcp_connection``.""" + + tables: Optional[List[GenieAgentTable]] = None + """Unity Catalog tables, views, and metric views, in the order they were added.""" + + volumes: Optional[List[GenieAgentVolume]] = None + """Unity Catalog volumes, in the order they were added.""" + + def as_dict(self) -> dict: + """Serializes the GenieAgentDataSources into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.mcp_connections: + body["mcp_connections"] = [v.as_dict() for v in self.mcp_connections] + if self.tables: + body["tables"] = [v.as_dict() for v in self.tables] + if self.volumes: + body["volumes"] = [v.as_dict() for v in self.volumes] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieAgentDataSources into a shallow dictionary of its immediate attributes.""" + body = {} + if self.mcp_connections: + body["mcp_connections"] = self.mcp_connections + if self.tables: + body["tables"] = self.tables + if self.volumes: + body["volumes"] = self.volumes + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieAgentDataSources: + """Deserializes the GenieAgentDataSources from a dictionary.""" + return cls( + mcp_connections=_repeated_dict(d, "mcp_connections", GenieAgentMcpConnection), + tables=_repeated_dict(d, "tables", GenieAgentTable), + volumes=_repeated_dict(d, "volumes", GenieAgentVolume), + ) + + +@dataclass +class GenieAgentDefinition: + """The definition of a Genie Agent.""" + + benchmarks: Optional[GenieAgentBenchmarks] = None + """Benchmarks for evaluating the quality of the Genie Agent's answers.""" + + config: Optional[GenieAgentConfig] = None + """General settings of the Genie Agent.""" + + data_sources: Optional[GenieAgentDataSources] = None + """The data the Genie Agent can access.""" + + instructions: Optional[GenieAgentInstructions] = None + """Instructions and SQL that guide how the Genie Agent answers questions.""" + + def as_dict(self) -> dict: + """Serializes the GenieAgentDefinition into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.benchmarks: + body["benchmarks"] = self.benchmarks.as_dict() + if self.config: + body["config"] = self.config.as_dict() + if self.data_sources: + body["data_sources"] = self.data_sources.as_dict() + if self.instructions: + body["instructions"] = self.instructions.as_dict() + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieAgentDefinition into a shallow dictionary of its immediate attributes.""" + body = {} + if self.benchmarks: + body["benchmarks"] = self.benchmarks + if self.config: + body["config"] = self.config + if self.data_sources: + body["data_sources"] = self.data_sources + if self.instructions: + body["instructions"] = self.instructions + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieAgentDefinition: + """Deserializes the GenieAgentDefinition from a dictionary.""" + return cls( + benchmarks=_from_dict(d, "benchmarks", GenieAgentBenchmarks), + config=_from_dict(d, "config", GenieAgentConfig), + data_sources=_from_dict(d, "data_sources", GenieAgentDataSources), + instructions=_from_dict(d, "instructions", GenieAgentInstructions), + ) + + +@dataclass +class GenieAgentExampleQuestionSql: + """An example question paired with the SQL that answers it.""" + + id: str + """Identifier of the example: 32 lowercase hexadecimal characters, unique within the Genie Agent.""" + + question: str + """The question.""" + + sql: str + """The SQL that answers the question.""" + + parameters: Optional[List[GenieAgentSqlParameter]] = None + """Parameters used in the SQL, in the order given.""" + + usage_guidance: Optional[str] = None + """Guidance on when to use this example.""" + + def as_dict(self) -> dict: + """Serializes the GenieAgentExampleQuestionSql into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.id is not None: + body["id"] = self.id + if self.parameters: + body["parameters"] = [v.as_dict() for v in self.parameters] + if self.question is not None: + body["question"] = self.question + if self.sql is not None: + body["sql"] = self.sql + if self.usage_guidance is not None: + body["usage_guidance"] = self.usage_guidance + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieAgentExampleQuestionSql into a shallow dictionary of its immediate attributes.""" + body = {} + if self.id is not None: + body["id"] = self.id + if self.parameters: + body["parameters"] = self.parameters + if self.question is not None: + body["question"] = self.question + if self.sql is not None: + body["sql"] = self.sql + if self.usage_guidance is not None: + body["usage_guidance"] = self.usage_guidance + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieAgentExampleQuestionSql: + """Deserializes the GenieAgentExampleQuestionSql from a dictionary.""" + return cls( + id=d.get("id", None), + parameters=_repeated_dict(d, "parameters", GenieAgentSqlParameter), + question=d.get("question", None), + sql=d.get("sql", None), + usage_guidance=d.get("usage_guidance", None), + ) + + +@dataclass +class GenieAgentInstructions: + """Instructions and SQL that guide how a Genie Agent answers questions.""" + + example_question_sqls: Optional[List[GenieAgentExampleQuestionSql]] = None + """Example questions paired with the SQL that answers them.""" + + join_specs: Optional[List[GenieAgentJoinSpec]] = None + """Predefined joins between tables or metric views.""" + + sql_functions: Optional[List[GenieAgentSqlFunction]] = None + """Unity Catalog SQL functions the generated SQL may call.""" + + sql_snippets: Optional[GenieAgentSqlSnippets] = None + """Reusable SQL filters, expressions, and measures.""" + + text_instructions: Optional[List[GenieAgentTextInstruction]] = None + """General text instructions. At most one text instruction is supported.""" + + def as_dict(self) -> dict: + """Serializes the GenieAgentInstructions into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.example_question_sqls: + body["example_question_sqls"] = [v.as_dict() for v in self.example_question_sqls] + if self.join_specs: + body["join_specs"] = [v.as_dict() for v in self.join_specs] + if self.sql_functions: + body["sql_functions"] = [v.as_dict() for v in self.sql_functions] + if self.sql_snippets: + body["sql_snippets"] = self.sql_snippets.as_dict() + if self.text_instructions: + body["text_instructions"] = [v.as_dict() for v in self.text_instructions] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieAgentInstructions into a shallow dictionary of its immediate attributes.""" + body = {} + if self.example_question_sqls: + body["example_question_sqls"] = self.example_question_sqls + if self.join_specs: + body["join_specs"] = self.join_specs + if self.sql_functions: + body["sql_functions"] = self.sql_functions + if self.sql_snippets: + body["sql_snippets"] = self.sql_snippets + if self.text_instructions: + body["text_instructions"] = self.text_instructions + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieAgentInstructions: + """Deserializes the GenieAgentInstructions from a dictionary.""" + return cls( + example_question_sqls=_repeated_dict(d, "example_question_sqls", GenieAgentExampleQuestionSql), + join_specs=_repeated_dict(d, "join_specs", GenieAgentJoinSpec), + sql_functions=_repeated_dict(d, "sql_functions", GenieAgentSqlFunction), + sql_snippets=_from_dict(d, "sql_snippets", GenieAgentSqlSnippets), + text_instructions=_repeated_dict(d, "text_instructions", GenieAgentTextInstruction), + ) + + +class GenieAgentJoinRelationshipType(Enum): + """How the rows of the two sides of a join relate.""" + + GENIE_AGENT_JOIN_RELATIONSHIP_TYPE_MANY_TO_MANY = "GENIE_AGENT_JOIN_RELATIONSHIP_TYPE_MANY_TO_MANY" + GENIE_AGENT_JOIN_RELATIONSHIP_TYPE_MANY_TO_ONE = "GENIE_AGENT_JOIN_RELATIONSHIP_TYPE_MANY_TO_ONE" + GENIE_AGENT_JOIN_RELATIONSHIP_TYPE_ONE_TO_MANY = "GENIE_AGENT_JOIN_RELATIONSHIP_TYPE_ONE_TO_MANY" + GENIE_AGENT_JOIN_RELATIONSHIP_TYPE_ONE_TO_ONE = "GENIE_AGENT_JOIN_RELATIONSHIP_TYPE_ONE_TO_ONE" + + +@dataclass +class GenieAgentJoinSource: + """One side of a join.""" + + identifier: str + """The full three-level name of the table or metric view.""" + + alias: Optional[str] = None + """The alias the join condition uses for this side.""" + + def as_dict(self) -> dict: + """Serializes the GenieAgentJoinSource into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.alias is not None: + body["alias"] = self.alias + if self.identifier is not None: + body["identifier"] = self.identifier + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieAgentJoinSource into a shallow dictionary of its immediate attributes.""" + body = {} + if self.alias is not None: + body["alias"] = self.alias + if self.identifier is not None: + body["identifier"] = self.identifier + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieAgentJoinSource: + """Deserializes the GenieAgentJoinSource from a dictionary.""" + return cls(alias=d.get("alias", None), identifier=d.get("identifier", None)) + + +@dataclass +class GenieAgentJoinSpec: + """A predefined join between two tables or metric views.""" + + id: str + """Identifier of the join: 32 lowercase hexadecimal characters, unique within the Genie Agent.""" + + left: GenieAgentJoinSource + """The left side of the join.""" + + right: GenieAgentJoinSource + """The right side of the join.""" + + sql: str + """The join condition, for example ``orders.customer_id = customers.id``.""" + + comment: Optional[str] = None + """Description of the join's purpose.""" + + instruction: Optional[str] = None + """Guidance on when to use the join.""" + + relationship_type: Optional[GenieAgentJoinRelationshipType] = None + """How the rows of the two sides relate.""" + + def as_dict(self) -> dict: + """Serializes the GenieAgentJoinSpec into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.comment is not None: + body["comment"] = self.comment + if self.id is not None: + body["id"] = self.id + if self.instruction is not None: + body["instruction"] = self.instruction + if self.left: + body["left"] = self.left.as_dict() + if self.relationship_type is not None: + body["relationship_type"] = self.relationship_type.value + if self.right: + body["right"] = self.right.as_dict() + if self.sql is not None: + body["sql"] = self.sql + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieAgentJoinSpec into a shallow dictionary of its immediate attributes.""" + body = {} + if self.comment is not None: + body["comment"] = self.comment + if self.id is not None: + body["id"] = self.id + if self.instruction is not None: + body["instruction"] = self.instruction + if self.left: + body["left"] = self.left + if self.relationship_type is not None: + body["relationship_type"] = self.relationship_type + if self.right: + body["right"] = self.right + if self.sql is not None: + body["sql"] = self.sql + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieAgentJoinSpec: + """Deserializes the GenieAgentJoinSpec from a dictionary.""" + return cls( + comment=d.get("comment", None), + id=d.get("id", None), + instruction=d.get("instruction", None), + left=_from_dict(d, "left", GenieAgentJoinSource), + relationship_type=_enum(d, "relationship_type", GenieAgentJoinRelationshipType), + right=_from_dict(d, "right", GenieAgentJoinSource), + sql=d.get("sql", None), + ) + + +@dataclass +class GenieAgentMcpConnection: + """A Unity Catalog MCP connection the Genie Agent can use.""" + + mcp_connection: str + """The name of the Unity Catalog connection, for example ``system.ai.google_calendar``.""" + + def as_dict(self) -> dict: + """Serializes the GenieAgentMcpConnection into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.mcp_connection is not None: + body["mcp_connection"] = self.mcp_connection + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieAgentMcpConnection into a shallow dictionary of its immediate attributes.""" + body = {} + if self.mcp_connection is not None: + body["mcp_connection"] = self.mcp_connection + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieAgentMcpConnection: + """Deserializes the GenieAgentMcpConnection from a dictionary.""" + return cls(mcp_connection=d.get("mcp_connection", None)) + + +@dataclass +class GenieAgentSampleQuestion: + """A sample question shown to users of the Genie Agent.""" + + id: str + """Identifier of the sample question: 32 lowercase hexadecimal characters, unique within the Genie + Agent.""" + + question: str + """The text of the question.""" + + def as_dict(self) -> dict: + """Serializes the GenieAgentSampleQuestion into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.id is not None: + body["id"] = self.id + if self.question is not None: + body["question"] = self.question + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieAgentSampleQuestion into a shallow dictionary of its immediate attributes.""" + body = {} + if self.id is not None: + body["id"] = self.id + if self.question is not None: + body["question"] = self.question + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieAgentSampleQuestion: + """Deserializes the GenieAgentSampleQuestion from a dictionary.""" + return cls(id=d.get("id", None), question=d.get("question", None)) + + +@dataclass +class GenieAgentSqlExpression: + """A reusable non-aggregated SQL expression (dimension).""" + + id: str + """Identifier of the expression: 32 lowercase hexadecimal characters, unique within the Genie + Agent.""" + + sql: str + """The SQL of the expression.""" + + alias: Optional[str] = None + """The alias the Genie Agent uses for the expression in generated SQL.""" + + display_name: Optional[str] = None + """Display name of the expression.""" + + instruction: Optional[str] = None + """Guidance on when to use the expression.""" + + synonyms: Optional[List[str]] = None + """Synonyms for the expression.""" + + def as_dict(self) -> dict: + """Serializes the GenieAgentSqlExpression into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.alias is not None: + body["alias"] = self.alias + if self.display_name is not None: + body["display_name"] = self.display_name + if self.id is not None: + body["id"] = self.id + if self.instruction is not None: + body["instruction"] = self.instruction + if self.sql is not None: + body["sql"] = self.sql + if self.synonyms: + body["synonyms"] = [v for v in self.synonyms] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieAgentSqlExpression into a shallow dictionary of its immediate attributes.""" + body = {} + if self.alias is not None: + body["alias"] = self.alias + if self.display_name is not None: + body["display_name"] = self.display_name + if self.id is not None: + body["id"] = self.id + if self.instruction is not None: + body["instruction"] = self.instruction + if self.sql is not None: + body["sql"] = self.sql + if self.synonyms: + body["synonyms"] = self.synonyms + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieAgentSqlExpression: + """Deserializes the GenieAgentSqlExpression from a dictionary.""" + return cls( + alias=d.get("alias", None), + display_name=d.get("display_name", None), + id=d.get("id", None), + instruction=d.get("instruction", None), + sql=d.get("sql", None), + synonyms=d.get("synonyms", None), + ) + + +@dataclass +class GenieAgentSqlFilter: + """A reusable SQL filter.""" + + id: str + """Identifier of the filter: 32 lowercase hexadecimal characters, unique within the Genie Agent.""" + + sql: str + """The SQL of the filter.""" + + display_name: Optional[str] = None + """Display name of the filter.""" + + instruction: Optional[str] = None + """Guidance on when to use the filter.""" + + synonyms: Optional[List[str]] = None + """Synonyms for the filter.""" + + def as_dict(self) -> dict: + """Serializes the GenieAgentSqlFilter into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.display_name is not None: + body["display_name"] = self.display_name + if self.id is not None: + body["id"] = self.id + if self.instruction is not None: + body["instruction"] = self.instruction + if self.sql is not None: + body["sql"] = self.sql + if self.synonyms: + body["synonyms"] = [v for v in self.synonyms] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieAgentSqlFilter into a shallow dictionary of its immediate attributes.""" + body = {} + if self.display_name is not None: + body["display_name"] = self.display_name + if self.id is not None: + body["id"] = self.id + if self.instruction is not None: + body["instruction"] = self.instruction + if self.sql is not None: + body["sql"] = self.sql + if self.synonyms: + body["synonyms"] = self.synonyms + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieAgentSqlFilter: + """Deserializes the GenieAgentSqlFilter from a dictionary.""" + return cls( + display_name=d.get("display_name", None), + id=d.get("id", None), + instruction=d.get("instruction", None), + sql=d.get("sql", None), + synonyms=d.get("synonyms", None), + ) + + +@dataclass +class GenieAgentSqlFunction: + """A Unity Catalog SQL function the generated SQL may call.""" + + id: str + """Identifier of the entry: 32 lowercase hexadecimal characters, unique within the Genie Agent.""" + + identifier: str + """The full three-level name of the function, for example ``catalog.schema.function``.""" + + def as_dict(self) -> dict: + """Serializes the GenieAgentSqlFunction into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.id is not None: + body["id"] = self.id + if self.identifier is not None: + body["identifier"] = self.identifier + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieAgentSqlFunction into a shallow dictionary of its immediate attributes.""" + body = {} + if self.id is not None: + body["id"] = self.id + if self.identifier is not None: + body["identifier"] = self.identifier + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieAgentSqlFunction: + """Deserializes the GenieAgentSqlFunction from a dictionary.""" + return cls(id=d.get("id", None), identifier=d.get("identifier", None)) + + +@dataclass +class GenieAgentSqlMeasure: + """A reusable aggregated SQL measure.""" + + id: str + """Identifier of the measure: 32 lowercase hexadecimal characters, unique within the Genie Agent.""" + + sql: str + """The SQL of the measure.""" + + alias: Optional[str] = None + """The alias the Genie Agent uses for the measure in generated SQL.""" + + display_name: Optional[str] = None + """Display name of the measure.""" + + instruction: Optional[str] = None + """Guidance on when to use the measure.""" + + synonyms: Optional[List[str]] = None + """Synonyms for the measure.""" + + def as_dict(self) -> dict: + """Serializes the GenieAgentSqlMeasure into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.alias is not None: + body["alias"] = self.alias + if self.display_name is not None: + body["display_name"] = self.display_name + if self.id is not None: + body["id"] = self.id + if self.instruction is not None: + body["instruction"] = self.instruction + if self.sql is not None: + body["sql"] = self.sql + if self.synonyms: + body["synonyms"] = [v for v in self.synonyms] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieAgentSqlMeasure into a shallow dictionary of its immediate attributes.""" + body = {} + if self.alias is not None: + body["alias"] = self.alias + if self.display_name is not None: + body["display_name"] = self.display_name + if self.id is not None: + body["id"] = self.id + if self.instruction is not None: + body["instruction"] = self.instruction + if self.sql is not None: + body["sql"] = self.sql + if self.synonyms: + body["synonyms"] = self.synonyms + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieAgentSqlMeasure: + """Deserializes the GenieAgentSqlMeasure from a dictionary.""" + return cls( + alias=d.get("alias", None), + display_name=d.get("display_name", None), + id=d.get("id", None), + instruction=d.get("instruction", None), + sql=d.get("sql", None), + synonyms=d.get("synonyms", None), + ) + + +@dataclass +class GenieAgentSqlParameter: + """A parameter of a parameterized SQL example.""" + + name: str + """The name of the parameter in the SQL.""" + + default_value: Optional[GenieAgentSqlParameterDefaultValue] = None + """The parameter's default value.""" + + description: Optional[str] = None + """Description of the parameter.""" + + type_hint: Optional[str] = None + """The parameter's data type, for example ``STRING`` or ``INTEGER``.""" + + def as_dict(self) -> dict: + """Serializes the GenieAgentSqlParameter into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.default_value: + body["default_value"] = self.default_value.as_dict() + if self.description is not None: + body["description"] = self.description + if self.name is not None: + body["name"] = self.name + if self.type_hint is not None: + body["type_hint"] = self.type_hint + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieAgentSqlParameter into a shallow dictionary of its immediate attributes.""" + body = {} + if self.default_value: + body["default_value"] = self.default_value + if self.description is not None: + body["description"] = self.description + if self.name is not None: + body["name"] = self.name + if self.type_hint is not None: + body["type_hint"] = self.type_hint + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieAgentSqlParameter: + """Deserializes the GenieAgentSqlParameter from a dictionary.""" + return cls( + default_value=_from_dict(d, "default_value", GenieAgentSqlParameterDefaultValue), + description=d.get("description", None), + name=d.get("name", None), + type_hint=d.get("type_hint", None), + ) + + +@dataclass +class GenieAgentSqlParameterDefaultRange: + """A range of default values for a SQL parameter.""" + + max: Optional[str] = None + """The upper end of the range.""" + + min: Optional[str] = None + """The lower end of the range.""" + + def as_dict(self) -> dict: + """Serializes the GenieAgentSqlParameterDefaultRange into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.max is not None: + body["max"] = self.max + if self.min is not None: + body["min"] = self.min + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieAgentSqlParameterDefaultRange into a shallow dictionary of its immediate attributes.""" + body = {} + if self.max is not None: + body["max"] = self.max + if self.min is not None: + body["min"] = self.min + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieAgentSqlParameterDefaultRange: + """Deserializes the GenieAgentSqlParameterDefaultRange from a dictionary.""" + return cls(max=d.get("max", None), min=d.get("min", None)) + + +@dataclass +class GenieAgentSqlParameterDefaultValue: + """The default value of a SQL parameter: a list of values or a range.""" + + value_range: Optional[GenieAgentSqlParameterDefaultRange] = None + """The default range.""" + + values: Optional[List[str]] = None + """The default values. A single default is a one-element list.""" + + def as_dict(self) -> dict: + """Serializes the GenieAgentSqlParameterDefaultValue into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.value_range: + body["value_range"] = self.value_range.as_dict() + if self.values: + body["values"] = [v for v in self.values] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieAgentSqlParameterDefaultValue into a shallow dictionary of its immediate attributes.""" + body = {} + if self.value_range: + body["value_range"] = self.value_range + if self.values: + body["values"] = self.values + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieAgentSqlParameterDefaultValue: + """Deserializes the GenieAgentSqlParameterDefaultValue from a dictionary.""" + return cls( + value_range=_from_dict(d, "value_range", GenieAgentSqlParameterDefaultRange), values=d.get("values", None) + ) + + +@dataclass +class GenieAgentSqlSnippets: + """Reusable SQL snippets.""" + + expressions: Optional[List[GenieAgentSqlExpression]] = None + """Non-aggregated expressions (dimensions), such as ``price * units``.""" + + filters: Optional[List[GenieAgentSqlFilter]] = None + """Filters, such as ``state = 'CA'``.""" + + measures: Optional[List[GenieAgentSqlMeasure]] = None + """Aggregated measures, such as ``SUM(price * units)``.""" + + def as_dict(self) -> dict: + """Serializes the GenieAgentSqlSnippets into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.expressions: + body["expressions"] = [v.as_dict() for v in self.expressions] + if self.filters: + body["filters"] = [v.as_dict() for v in self.filters] + if self.measures: + body["measures"] = [v.as_dict() for v in self.measures] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieAgentSqlSnippets into a shallow dictionary of its immediate attributes.""" + body = {} + if self.expressions: + body["expressions"] = self.expressions + if self.filters: + body["filters"] = self.filters + if self.measures: + body["measures"] = self.measures + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieAgentSqlSnippets: + """Deserializes the GenieAgentSqlSnippets from a dictionary.""" + return cls( + expressions=_repeated_dict(d, "expressions", GenieAgentSqlExpression), + filters=_repeated_dict(d, "filters", GenieAgentSqlFilter), + measures=_repeated_dict(d, "measures", GenieAgentSqlMeasure), + ) + + +@dataclass +class GenieAgentTable: + """A Unity Catalog table, view, or metric view the Genie Agent can access.""" + + identifier: str + """The full three-level name, for example ``catalog.schema.table``.""" + + column_configs: Optional[List[GenieAgentColumnConfig]] = None + """Column settings for the table, sorted by ``column``.""" + + description: Optional[str] = None + """Description of the table.""" + + def as_dict(self) -> dict: + """Serializes the GenieAgentTable into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.column_configs: + body["column_configs"] = [v.as_dict() for v in self.column_configs] + if self.description is not None: + body["description"] = self.description + if self.identifier is not None: + body["identifier"] = self.identifier + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieAgentTable into a shallow dictionary of its immediate attributes.""" + body = {} + if self.column_configs: + body["column_configs"] = self.column_configs + if self.description is not None: + body["description"] = self.description + if self.identifier is not None: + body["identifier"] = self.identifier + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieAgentTable: + """Deserializes the GenieAgentTable from a dictionary.""" + return cls( + column_configs=_repeated_dict(d, "column_configs", GenieAgentColumnConfig), + description=d.get("description", None), + identifier=d.get("identifier", None), + ) + + +@dataclass +class GenieAgentTextInstruction: + """A general text instruction for the Genie Agent.""" + + id: str + """Identifier of the instruction: 32 lowercase hexadecimal characters, unique within the Genie + Agent.""" + + content: str + """The text of the instruction.""" + + def as_dict(self) -> dict: + """Serializes the GenieAgentTextInstruction into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.content is not None: + body["content"] = self.content + if self.id is not None: + body["id"] = self.id + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieAgentTextInstruction into a shallow dictionary of its immediate attributes.""" + body = {} + if self.content is not None: + body["content"] = self.content + if self.id is not None: + body["id"] = self.id + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieAgentTextInstruction: + """Deserializes the GenieAgentTextInstruction from a dictionary.""" + return cls(content=d.get("content", None), id=d.get("id", None)) + + +class GenieAgentView(Enum): + """The fields returned for each Genie Agent in a list response.""" + + GENIE_AGENT_VIEW_BASIC = "GENIE_AGENT_VIEW_BASIC" + GENIE_AGENT_VIEW_FULL = "GENIE_AGENT_VIEW_FULL" + + +@dataclass +class GenieAgentVolume: + """A Unity Catalog volume the Genie Agent can access.""" + + path: str + """The volume path, for example ``/Volumes/catalog/schema/volume``. May name a directory or file + within the volume.""" + + description: Optional[str] = None + """Description of the volume.""" + + def as_dict(self) -> dict: + """Serializes the GenieAgentVolume into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.description is not None: + body["description"] = self.description + if self.path is not None: + body["path"] = self.path + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieAgentVolume into a shallow dictionary of its immediate attributes.""" + body = {} + if self.description is not None: + body["description"] = self.description + if self.path is not None: + body["path"] = self.path + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieAgentVolume: + """Deserializes the GenieAgentVolume from a dictionary.""" + return cls(description=d.get("description", None), path=d.get("path", None)) + + +@dataclass +class ListGenieAgentsResponse: + """Response for a list request: a page of Genie Agents plus the token for the next page.""" + + genie_agents: Optional[List[GenieAgent]] = None + """The Genie Agents, with the fields selected by the request's ``view``.""" + + next_page_token: Optional[str] = None + """Token to get the next page of results.""" + + unreachable: Optional[List[str]] = None + """Resource names, in the form ``agents/{genie_agent_id}``, of Genie Agents in this page that could + not be read and are omitted from ``genie_agents``. Populated only for ``GENIE_AGENT_VIEW_FULL``. + To see why, get the Genie Agent by name; to retry, repeat the request. Other kinds of resource + names may be added to this field in the future.""" + + def as_dict(self) -> dict: + """Serializes the ListGenieAgentsResponse into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.genie_agents: + body["genie_agents"] = [v.as_dict() for v in self.genie_agents] + if self.next_page_token is not None: + body["next_page_token"] = self.next_page_token + if self.unreachable: + body["unreachable"] = [v for v in self.unreachable] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ListGenieAgentsResponse into a shallow dictionary of its immediate attributes.""" + body = {} + if self.genie_agents: + body["genie_agents"] = self.genie_agents + if self.next_page_token is not None: + body["next_page_token"] = self.next_page_token + if self.unreachable: + body["unreachable"] = self.unreachable + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ListGenieAgentsResponse: + """Deserializes the ListGenieAgentsResponse from a dictionary.""" + return cls( + genie_agents=_repeated_dict(d, "genie_agents", GenieAgent), + next_page_token=d.get("next_page_token", None), + unreachable=d.get("unreachable", None), + ) + + +class GenieAgentsAPI: + """Genie provides a no-code experience for business users, powered by AI/BI. Analysts set up Genie Agents + that business users can use to ask questions using natural language. Genie uses data registered to Unity + Catalog and requires at least CAN USE permission on a Pro or Serverless SQL warehouse. Also, Databricks + Assistant must be enabled.""" + + def __init__(self, api_client): + self._api = api_client + + def create_genie_agent(self, genie_agent: GenieAgent, *, genie_agent_id: Optional[str] = None) -> GenieAgent: + """Creates a Genie Agent. + + :param genie_agent: :class:`GenieAgent` + The Genie Agent to create. + :param genie_agent_id: str (optional) + The ID for the new Genie Agent: 32 lowercase hexadecimal characters (a UUID without dashes), not all + zeros. Assigned by the server if omitted. Fails with ALREADY_EXISTS if the ID is in use in the + workspace, including by a Genie Agent in the trash, or PERMISSION_DENIED if the caller cannot view + the Genie Agent using it. + + :returns: :class:`GenieAgent` + """ + + body = genie_agent.as_dict() + query = {} + if genie_agent_id is not None: + query["genie_agent_id"] = genie_agent_id + headers = { + "Accept": "application/json", + "Content-Type": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do("POST", "/api/2.0/genie/agents", query=query, body=body, headers=headers) + return GenieAgent.from_dict(res) + + def delete_genie_agent(self, name: str, *, etag: Optional[str] = None): + """Deletes a Genie Agent, moving it to trash. + + :param name: str + The resource name of the Genie Agent to delete, in the form ``agents/{genie_agent_id}``. + :param etag: str (optional) + Etag from a previous read. On mismatch the delete fails with ABORTED and nothing is deleted. + + + """ + + query = {} + if etag is not None: + query["etag"] = etag + headers = { + "Accept": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + self._api.do("DELETE", f"/api/2.0/genie/{name}", query=query, headers=headers) + + def get_genie_agent(self, name: str) -> GenieAgent: + """Gets a Genie Agent. + + :param name: str + The resource name of the Genie Agent to retrieve, in the form ``agents/{genie_agent_id}``. + + :returns: :class:`GenieAgent` + """ + + headers = { + "Accept": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do("GET", f"/api/2.0/genie/{name}", headers=headers) + return GenieAgent.from_dict(res) + + def list_genie_agents( + self, + *, + page_size: Optional[int] = None, + page_token: Optional[str] = None, + view: Optional[GenieAgentView] = None, + ) -> Iterator[GenieAgent]: + """Lists the Genie Agents in the workspace. Genie Agents created from an integration (for example, from a + dashboard) are not listed. + + :param page_size: int (optional) + Maximum number of Genie Agents to return per page. Defaults to 20 when unset or 0; values above 100 + are coerced to 100. Negative values are rejected with INVALID_PARAMETER_VALUE. With + ``GENIE_AGENT_VIEW_FULL``, a smaller per-page maximum applies. + :param page_token: str (optional) + Pagination token from a previous response's ``next_page_token``. A token that was not returned by + this method is rejected with INVALID_PARAMETER_VALUE. + :param view: :class:`GenieAgentView` (optional) + The fields to return for each Genie Agent. Defaults to ``GENIE_AGENT_VIEW_BASIC``. Use + ``GENIE_AGENT_VIEW_FULL`` only when the agent bodies are needed, since it returns smaller pages. + + :returns: Iterator over :class:`GenieAgent` + """ + + query = {} + if page_size is not None: + query["page_size"] = page_size + if page_token is not None: + query["page_token"] = page_token + if view is not None: + query["view"] = view.value + headers = { + "Accept": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + while True: + json = self._api.do("GET", "/api/2.0/genie/agents", query=query, headers=headers) + if "genie_agents" in json: + for v in json["genie_agents"]: + yield GenieAgent.from_dict(v) + if "next_page_token" not in json or not json["next_page_token"]: + return + query["page_token"] = json["next_page_token"] + + def update_genie_agent( + self, name: str, genie_agent: GenieAgent, update_mask: FieldMask, *, etag: Optional[str] = None + ) -> GenieAgent: + """Updates a Genie Agent. Only the fields named in update_mask are applied. + + :param name: str + The resource name of the Genie Agent, in the form ``agents/{genie_agent_id}``. Ignored on create. + :param genie_agent: :class:`GenieAgent` + The Genie Agent to update. Its ``name`` identifies the resource; fields named in update_mask are + applied. + :param update_mask: FieldMask + The set of fields to update. The mask is required and must name at least one field; an empty or + absent mask is rejected with INVALID_PARAMETER_VALUE, as is a path that is not an updatable field. + + A field named in the mask takes the value from ``genie_agent``, even when it is unset there: an + unset ``description`` is cleared; an unset ``display_name``, ``warehouse_id``, or ``definition`` is + rejected with INVALID_PARAMETER_VALUE; an unset ``parent_path`` is cleared without moving the Genie + Agent. ``definition`` is replaced as a whole; paths within it are not supported. + + The special value ``*`` replaces the whole Genie Agent, as if every updatable field were named. It + must be the only path in the mask. Prefer naming fields explicitly: with ``*``, any field added to + Genie Agent later is also replaced, so a client unaware of it may clear it. + :param etag: str (optional) + Etag from a previous read. On mismatch the update fails with ABORTED (metadata changes may already + be applied); get the agent and retry. May return TEMPORARILY_UNAVAILABLE. + + :returns: :class:`GenieAgent` + """ + + body = genie_agent.as_dict() + query = {} + if etag is not None: + query["etag"] = etag + if update_mask is not None: + query["update_mask"] = update_mask.ToJsonString() + headers = { + "Accept": "application/json", + "Content-Type": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do("PATCH", f"/api/2.0/genie/{name}", query=query, body=body, headers=headers) + return GenieAgent.from_dict(res) diff --git a/databricks/sdk/service/iamv2.py b/databricks/sdk/service/iamv2.py index 211272f3b..689622f99 100644 --- a/databricks/sdk/service/iamv2.py +++ b/databricks/sdk/service/iamv2.py @@ -46,6 +46,19 @@ class DirectGroupMember: membership_source: Optional[GroupMembershipSource] = None """The source of group membership (internal or from identity provider).""" + name: Optional[str] = None + """The resource name of the group membership. The format depends on the API that returned it: + + - Account-scoped: accounts/{account_id}/groups/{group_id}/direct-members/{principal_id} + - Workspace-scoped: groups/{group_id}/direct-members/{principal_id}""" + + parent: Optional[str] = None + """The resource name of the group that contains the membership. The format depends on the API that + returned it: + + - Account-scoped: accounts/{account_id}/groups/{group_id} + - Workspace-scoped: groups/{group_id}""" + principal_type: Optional[PrincipalType] = None """The type of the principal (user/service principal/group).""" @@ -60,6 +73,10 @@ def as_dict(self) -> dict: body["group_id"] = self.group_id if self.membership_source is not None: body["membership_source"] = self.membership_source.value + if self.name is not None: + body["name"] = self.name + if self.parent is not None: + body["parent"] = self.parent if self.principal_id is not None: body["principal_id"] = self.principal_id if self.principal_type is not None: @@ -77,6 +94,10 @@ def as_shallow_dict(self) -> dict: body["group_id"] = self.group_id if self.membership_source is not None: body["membership_source"] = self.membership_source + if self.name is not None: + body["name"] = self.name + if self.parent is not None: + body["parent"] = self.parent if self.principal_id is not None: body["principal_id"] = self.principal_id if self.principal_type is not None: @@ -91,6 +112,8 @@ def from_dict(cls, d: Dict[str, Any]) -> DirectGroupMember: external_id=d.get("external_id", None), group_id=_int64(d, "group_id"), membership_source=_enum(d, "membership_source", GroupMembershipSource), + name=d.get("name", None), + parent=d.get("parent", None), principal_id=_int64(d, "principal_id"), principal_type=_enum(d, "principal_type", PrincipalType), ) @@ -426,6 +449,16 @@ class Group: group_name: Optional[str] = None """Display name of the group.""" + name: Optional[str] = None + """The resource name of the group. The format depends on the API that returned it: + + - Account-scoped: accounts/{account_id}/groups/{group_id} + - Workspace-scoped: groups/{group_id}""" + + parent: Optional[str] = None + """The resource name of the account that contains the group, ``accounts/{account_id}``. Both the + account-scoped and workspace-scoped APIs return it.""" + def as_dict(self) -> dict: """Serializes the Group into a dictionary suitable for use as a JSON request body.""" body = {} @@ -437,6 +470,10 @@ def as_dict(self) -> dict: body["group_id"] = self.group_id if self.group_name is not None: body["group_name"] = self.group_name + if self.name is not None: + body["name"] = self.name + if self.parent is not None: + body["parent"] = self.parent return body def as_shallow_dict(self) -> dict: @@ -450,6 +487,10 @@ def as_shallow_dict(self) -> dict: body["group_id"] = self.group_id if self.group_name is not None: body["group_name"] = self.group_name + if self.name is not None: + body["name"] = self.name + if self.parent is not None: + body["parent"] = self.parent return body @classmethod @@ -460,6 +501,8 @@ def from_dict(cls, d: Dict[str, Any]) -> Group: external_id=d.get("external_id", None), group_id=d.get("group_id", None), group_name=d.get("group_name", None), + name=d.get("name", None), + parent=d.get("parent", None), ) @@ -994,6 +1037,16 @@ class ServicePrincipal: external_id: Optional[str] = None """ExternalId of the service principal in the customer's IdP.""" + name: Optional[str] = None + """The resource name of the service principal. The format depends on the API that returned it: + + - Account-scoped: accounts/{account_id}/service-principals/{service_principal_id} + - Workspace-scoped: service-principals/{service_principal_id}""" + + parent: Optional[str] = None + """The resource name of the account that contains the service principal, ``accounts/{account_id}``. + Both the account-scoped and workspace-scoped APIs return it.""" + service_principal_id: Optional[str] = None """Internal service principal ID of the service principal in Databricks.""" @@ -1010,6 +1063,10 @@ def as_dict(self) -> dict: body["display_name"] = self.display_name if self.external_id is not None: body["external_id"] = self.external_id + if self.name is not None: + body["name"] = self.name + if self.parent is not None: + body["parent"] = self.parent if self.service_principal_id is not None: body["service_principal_id"] = self.service_principal_id return body @@ -1027,6 +1084,10 @@ def as_shallow_dict(self) -> dict: body["display_name"] = self.display_name if self.external_id is not None: body["external_id"] = self.external_id + if self.name is not None: + body["name"] = self.name + if self.parent is not None: + body["parent"] = self.parent if self.service_principal_id is not None: body["service_principal_id"] = self.service_principal_id return body @@ -1040,6 +1101,8 @@ def from_dict(cls, d: Dict[str, Any]) -> ServicePrincipal: application_id=d.get("application_id", None), display_name=d.get("display_name", None), external_id=d.get("external_id", None), + name=d.get("name", None), + parent=d.get("parent", None), service_principal_id=d.get("service_principal_id", None), ) @@ -1064,6 +1127,14 @@ class TransitiveParentGroup: group_id: Optional[str] = None """Internal group ID of the group in Databricks.""" + name: Optional[str] = None + """The resource name of the transitive parent group. The format depends on the API that returned + it: + + - Account-scoped: + accounts/{account_id}/principals/{principal_id}/transitive-parent-groups/{group_id} + - Workspace-scoped: principals/{principal_id}/transitive-parent-groups/{group_id}""" + def as_dict(self) -> dict: """Serializes the TransitiveParentGroup into a dictionary suitable for use as a JSON request body.""" body = {} @@ -1073,6 +1144,8 @@ def as_dict(self) -> dict: body["external_id"] = self.external_id if self.group_id is not None: body["group_id"] = self.group_id + if self.name is not None: + body["name"] = self.name return body def as_shallow_dict(self) -> dict: @@ -1084,6 +1157,8 @@ def as_shallow_dict(self) -> dict: body["external_id"] = self.external_id if self.group_id is not None: body["group_id"] = self.group_id + if self.name is not None: + body["name"] = self.name return body @classmethod @@ -1093,6 +1168,7 @@ def from_dict(cls, d: Dict[str, Any]) -> TransitiveParentGroup: account_id=d.get("account_id", None), external_id=d.get("external_id", None), group_id=d.get("group_id", None), + name=d.get("name", None), ) @@ -1114,6 +1190,16 @@ class User: external_id: Optional[str] = None """ExternalId of the user in the customer's IdP.""" + name: Optional[str] = None + """The resource name of the user. The format depends on the API that returned it: + + - Account-scoped: accounts/{account_id}/users/{user_id} + - Workspace-scoped: users/{user_id}""" + + parent: Optional[str] = None + """The resource name of the account that contains the user, ``accounts/{account_id}``. Both the + account-scoped and workspace-scoped APIs return it.""" + user_id: Optional[str] = None """Internal userId of the user in Databricks.""" @@ -1128,6 +1214,10 @@ def as_dict(self) -> dict: body["external_id"] = self.external_id if self.full_name: body["full_name"] = self.full_name.as_dict() + if self.name is not None: + body["name"] = self.name + if self.parent is not None: + body["parent"] = self.parent if self.user_id is not None: body["user_id"] = self.user_id if self.username is not None: @@ -1145,6 +1235,10 @@ def as_shallow_dict(self) -> dict: body["external_id"] = self.external_id if self.full_name: body["full_name"] = self.full_name + if self.name is not None: + body["name"] = self.name + if self.parent is not None: + body["parent"] = self.parent if self.user_id is not None: body["user_id"] = self.user_id if self.username is not None: @@ -1159,6 +1253,8 @@ def from_dict(cls, d: Dict[str, Any]) -> User: account_user_status=_enum(d, "account_user_status", State), external_id=d.get("external_id", None), full_name=_from_dict(d, "full_name", UserFullName), + name=d.get("name", None), + parent=d.get("parent", None), user_id=d.get("user_id", None), username=d.get("username", None), ) @@ -1203,6 +1299,14 @@ class WorkspaceAccessDetail: account_id: Optional[str] = None """The account ID parent of the workspace where the principal has access.""" + name: Optional[str] = None + """The resource name of the workspace access detail. The format depends on the API that returned + it: + + - Account-scoped: + accounts/{account_id}/workspaces/{workspace_id}/workspace-access-details/{principal_id} + - Workspace-scoped: workspace-access-details/{principal_id}""" + permissions: Optional[List[WorkspacePermission]] = None """The permissions granted to the principal in the workspace.""" @@ -1224,6 +1328,8 @@ def as_dict(self) -> dict: body["access_type"] = self.access_type.value if self.account_id is not None: body["account_id"] = self.account_id + if self.name is not None: + body["name"] = self.name if self.permissions: body["permissions"] = [v.value for v in self.permissions] if self.principal_id is not None: @@ -1243,6 +1349,8 @@ def as_shallow_dict(self) -> dict: body["access_type"] = self.access_type if self.account_id is not None: body["account_id"] = self.account_id + if self.name is not None: + body["name"] = self.name if self.permissions: body["permissions"] = self.permissions if self.principal_id is not None: @@ -1261,6 +1369,7 @@ def from_dict(cls, d: Dict[str, Any]) -> WorkspaceAccessDetail: return cls( access_type=_enum(d, "access_type", WorkspaceAccessDetailAccessType), account_id=d.get("account_id", None), + name=d.get("name", None), permissions=_repeated_enum(d, "permissions", WorkspacePermission), principal_id=_int64(d, "principal_id"), principal_type=_enum(d, "principal_type", PrincipalType), @@ -1306,7 +1415,9 @@ class WorkspaceAssignment: permission-assignment APIs, and is intended for account and workspace admins.""" principal_id: int - """The internal ID of the principal (user/sp/group) in Databricks.""" + """The internal ID of the principal (user/sp/group) in Databricks. Set when the assignment is + created and cannot be changed afterwards; to assign a different principal, create a new + assignment.""" account_id: Optional[str] = None """The account ID parent of the workspace where the principal is assigned""" @@ -1321,6 +1432,18 @@ class WorkspaceAssignment: principal also holds through a group. List responses leave this field empty. Get a single principal to read its entitlements.""" + name: Optional[str] = None + """The resource name of the workspace assignment. The format depends on the API that returned it: + + - Account-scoped: + accounts/{account_id}/workspaces/{workspace_id}/workspace-assignments/{principal_id} + - Workspace-scoped: workspace-assignments/{principal_id}""" + + parent: Optional[str] = None + """The resource name of the workspace that contains the assignment, + ``accounts/{account_id}/workspaces/{workspace_id}``. Both the account-scoped and + workspace-scoped APIs return it.""" + principal_type: Optional[PrincipalType] = None """The type of the principal (user/service principal/group) that is assigned.""" @@ -1336,6 +1459,10 @@ def as_dict(self) -> dict: body["effective_entitlements"] = [v.value for v in self.effective_entitlements] if self.entitlements: body["entitlements"] = [v.value for v in self.entitlements] + if self.name is not None: + body["name"] = self.name + if self.parent is not None: + body["parent"] = self.parent if self.principal_id is not None: body["principal_id"] = self.principal_id if self.principal_type is not None: @@ -1353,6 +1480,10 @@ def as_shallow_dict(self) -> dict: body["effective_entitlements"] = self.effective_entitlements if self.entitlements: body["entitlements"] = self.entitlements + if self.name is not None: + body["name"] = self.name + if self.parent is not None: + body["parent"] = self.parent if self.principal_id is not None: body["principal_id"] = self.principal_id if self.principal_type is not None: @@ -1368,6 +1499,8 @@ def from_dict(cls, d: Dict[str, Any]) -> WorkspaceAssignment: account_id=d.get("account_id", None), effective_entitlements=_repeated_enum(d, "effective_entitlements", Entitlement), entitlements=_repeated_enum(d, "entitlements", Entitlement), + name=d.get("name", None), + parent=d.get("parent", None), principal_id=_int64(d, "principal_id"), principal_type=_enum(d, "principal_type", PrincipalType), workspace_id=_int64(d, "workspace_id"), @@ -1471,6 +1604,10 @@ class WorkspaceIdentityDetail: assignment_type: Optional[WorkspaceIdentityDetailAssignmentType] = None """The type of assignment the principal has to the workspace (direct or indirect).""" + name: Optional[str] = None + """The resource name of the workspace identity detail. Format: + principals/{principal_id}/workspace-identity-detail""" + principal_id: Optional[int] = None """The internal ID of the principal (user/sp/group) in Databricks.""" @@ -1485,6 +1622,8 @@ def as_dict(self) -> dict: body = {} if self.assignment_type is not None: body["assignment_type"] = self.assignment_type.value + if self.name is not None: + body["name"] = self.name if self.principal_id is not None: body["principal_id"] = self.principal_id if self.principal_type is not None: @@ -1498,6 +1637,8 @@ def as_shallow_dict(self) -> dict: body = {} if self.assignment_type is not None: body["assignment_type"] = self.assignment_type + if self.name is not None: + body["name"] = self.name if self.principal_id is not None: body["principal_id"] = self.principal_id if self.principal_type is not None: @@ -1511,6 +1652,7 @@ def from_dict(cls, d: Dict[str, Any]) -> WorkspaceIdentityDetail: """Deserializes the WorkspaceIdentityDetail from a dictionary.""" return cls( assignment_type=_enum(d, "assignment_type", WorkspaceIdentityDetailAssignmentType), + name=d.get("name", None), principal_id=_int64(d, "principal_id"), principal_type=_enum(d, "principal_type", PrincipalType), workspace_identity_status=_enum(d, "workspace_identity_status", State), @@ -1537,14 +1679,15 @@ class AccountIamV2API: def __init__(self, api_client): self._api = api_client - def create_direct_group_member(self, group_id: int, direct_group_member: DirectGroupMember) -> DirectGroupMember: + def create_direct_group_member(self, parent: str, direct_group_member: DirectGroupMember) -> DirectGroupMember: """Creates a group membership (assigns a principal to a group). Authorization: the caller must be an account admin or a manager of the group (holds the ``roles/group.manager`` role on it). - :param group_id: int - Required. Internal ID of the group in Databricks. + :param parent: str + The parent resource name of the group to add the member to. Format: + accounts/{account_id}/groups/{group_id} :param direct_group_member: :class:`DirectGroupMember` Required. The direct group member to be added to the group. @@ -1558,15 +1701,10 @@ def create_direct_group_member(self, group_id: int, direct_group_member: DirectG "Content-Type": "application/json", } - res = self._api.do( - "POST", - f"/api/2.0/identity/accounts/{self._api.account_id}/groups/{group_id}/direct-members", - body=body, - headers=headers, - ) + res = self._api.do("POST", f"/api/2.0/identity/{parent}/direct-members", body=body, headers=headers) return DirectGroupMember.from_dict(res) - def create_group(self, group: Group) -> Group: + def create_group(self, parent: str, group: Group) -> Group: """Creates a local group in the Databricks account and returns the created group. A local group is one that is not synced from the customer's identity provider, and can be created whether or not Account Identity Management (AIM) is enabled. @@ -1575,6 +1713,8 @@ def create_group(self, group: Group) -> Group: identity provider, resolve it by its external ID with ResolveGroup; to read an existing external identity, use the ExternalGroup resource. + :param parent: str + The parent resource name of the account in which to create the group. Format: accounts/{account_id} :param group: :class:`Group` Required. Group to be created in Databricks @@ -1588,9 +1728,7 @@ def create_group(self, group: Group) -> Group: "Content-Type": "application/json", } - res = self._api.do( - "POST", f"/api/2.0/identity/accounts/{self._api.account_id}/groups", body=body, headers=headers - ) + res = self._api.do("POST", f"/api/2.0/identity/{parent}/groups", body=body, headers=headers) return Group.from_dict(res) def create_identity_visibility_filter( @@ -1634,7 +1772,7 @@ def create_identity_visibility_filter( ) return IdentityVisibilityFilter.from_dict(res) - def create_service_principal(self, service_principal: ServicePrincipal) -> ServicePrincipal: + def create_service_principal(self, parent: str, service_principal: ServicePrincipal) -> ServicePrincipal: """Creates a local service principal in the Databricks account and returns the created service principal. A local service principal is one that is not synced from the customer's identity provider, and can be created whether or not Account Identity Management (AIM) is enabled. @@ -1643,6 +1781,9 @@ def create_service_principal(self, service_principal: ServicePrincipal) -> Servi identity provider, resolve it by its external ID with ResolveServicePrincipal; to read an existing external identity, use the ExternalServicePrincipal resource. + :param parent: str + The parent resource name of the account in which to create the service principal. Format: + accounts/{account_id} :param service_principal: :class:`ServicePrincipal` Required. Service principal to be created in Databricks @@ -1656,12 +1797,10 @@ def create_service_principal(self, service_principal: ServicePrincipal) -> Servi "Content-Type": "application/json", } - res = self._api.do( - "POST", f"/api/2.0/identity/accounts/{self._api.account_id}/service-principals", body=body, headers=headers - ) + res = self._api.do("POST", f"/api/2.0/identity/{parent}/service-principals", body=body, headers=headers) return ServicePrincipal.from_dict(res) - def create_user(self, user: User) -> User: + def create_user(self, parent: str, user: User) -> User: """Creates a local user in the Databricks account and returns the created user. A local user is one that is not synced from the customer's identity provider, and can be created whether or not Account Identity Management (AIM) is enabled. @@ -1670,6 +1809,8 @@ def create_user(self, user: User) -> User: identity provider, resolve it by its external ID with ResolveUser; to read an existing external identity, use the ExternalUser resource. + :param parent: str + The parent resource name of the account in which to create the user. Format: accounts/{account_id} :param user: :class:`User` Required. User to be created in Databricks @@ -1683,21 +1824,20 @@ def create_user(self, user: User) -> User: "Content-Type": "application/json", } - res = self._api.do( - "POST", f"/api/2.0/identity/accounts/{self._api.account_id}/users", body=body, headers=headers - ) + res = self._api.do("POST", f"/api/2.0/identity/{parent}/users", body=body, headers=headers) return User.from_dict(res) def create_workspace_assignment( - self, workspace_id: int, workspace_assignment: WorkspaceAssignment + self, parent: str, workspace_assignment: WorkspaceAssignment ) -> WorkspaceAssignment: """Creates a workspace assignment for a principal. Entitlements are granted one at a time rather than atomically. If the request fails partway through, the principal stays assigned to the workspace with only some of the requested entitlements. Get the assignment afterwards to confirm which entitlements were granted. - :param workspace_id: int - Required. The workspace ID for which the workspace assignment is being created. + :param parent: str + The parent resource name of the workspace to assign the principal to. Format: + accounts/{account_id}/workspaces/{workspace_id} :param workspace_assignment: :class:`WorkspaceAssignment` Required. Workspace assignment to be created in Databricks. @@ -1711,12 +1851,7 @@ def create_workspace_assignment( "Content-Type": "application/json", } - res = self._api.do( - "POST", - f"/api/2.0/identity/accounts/{self._api.account_id}/workspaces/{workspace_id}/workspace-assignments", - body=body, - headers=headers, - ) + res = self._api.do("POST", f"/api/2.0/identity/{parent}/workspace-assignments", body=body, headers=headers) return WorkspaceAssignment.from_dict(res) def create_workspace_assignment_detail( @@ -1747,16 +1882,15 @@ def create_workspace_assignment_detail( ) return WorkspaceAssignmentDetail.from_dict(res) - def delete_direct_group_member(self, group_id: int, principal_id: int): + def delete_direct_group_member(self, name: str): """Deletes a group membership (unassigns a principal from a group). Authorization: the caller must be an account admin or a manager of the group (holds the ``roles/group.manager`` role on it). - :param group_id: int - Required. Internal ID of the group in Databricks. - :param principal_id: int - Required. Internal ID of the principal to be unassigned from the group. + :param name: str + The resource name of the group membership. Format: + accounts/{account_id}/groups/{group_id}/direct-members/{principal_id} """ @@ -1765,20 +1899,16 @@ def delete_direct_group_member(self, group_id: int, principal_id: int): "Accept": "application/json", } - self._api.do( - "DELETE", - f"/api/2.0/identity/accounts/{self._api.account_id}/groups/{group_id}/direct-members/{principal_id}", - headers=headers, - ) + self._api.do("DELETE", f"/api/2.0/identity/{name}", headers=headers) - def delete_group(self, group_id: str): + def delete_group(self, name: str): """Deletes a group from the Databricks account by its internal ID. Authorization: the caller must be an account admin or a manager of the group (holds the ``roles/group.manager`` role on it). - :param group_id: str - Required. Internal ID of the group in Databricks. + :param name: str + The resource name of the group. Format: accounts/{account_id}/groups/{group_id} """ @@ -1787,7 +1917,7 @@ def delete_group(self, group_id: str): "Accept": "application/json", } - self._api.do("DELETE", f"/api/2.0/identity/accounts/{self._api.account_id}/groups/{group_id}", headers=headers) + self._api.do("DELETE", f"/api/2.0/identity/{name}", headers=headers) def delete_identity_visibility_filter(self, name: str): """Deletes a group visibility filter by its resource name. Requires account admin privileges. Deleting a @@ -1809,11 +1939,12 @@ def delete_identity_visibility_filter(self, name: str): self._api.do("DELETE", f"/api/2.0/identity/{name}", headers=headers) - def delete_service_principal(self, service_principal_id: str): + def delete_service_principal(self, name: str): """Deletes a service principal from the Databricks account by its internal ID. - :param service_principal_id: str - Required. Internal ID of the service principal in Databricks. + :param name: str + The resource name of the service principal. Format: + accounts/{account_id}/service-principals/{service_principal_id} """ @@ -1822,17 +1953,13 @@ def delete_service_principal(self, service_principal_id: str): "Accept": "application/json", } - self._api.do( - "DELETE", - f"/api/2.0/identity/accounts/{self._api.account_id}/service-principals/{service_principal_id}", - headers=headers, - ) + self._api.do("DELETE", f"/api/2.0/identity/{name}", headers=headers) - def delete_user(self, user_id: str): + def delete_user(self, name: str): """Deletes a user from the Databricks account by its internal ID. - :param user_id: str - Required. Internal ID of the user in Databricks. + :param name: str + The resource name of the user. Format: accounts/{account_id}/users/{user_id} """ @@ -1841,17 +1968,16 @@ def delete_user(self, user_id: str): "Accept": "application/json", } - self._api.do("DELETE", f"/api/2.0/identity/accounts/{self._api.account_id}/users/{user_id}", headers=headers) + self._api.do("DELETE", f"/api/2.0/identity/{name}", headers=headers) - def delete_workspace_assignment(self, workspace_id: int, principal_id: int): + def delete_workspace_assignment(self, name: str): """Deletes a workspace assignment for a principal, revoking all of its entitlements. Entitlements are revoked one at a time rather than atomically. If the request fails partway through, the principal stays assigned with some of its original entitlements. Retrying is safe. - :param workspace_id: int - The workspace ID where the principal has access. - :param principal_id: int - Required. ID of the principal in Databricks to delete workspace assignment for. + :param name: str + The resource name of the workspace assignment. Format: + accounts/{account_id}/workspaces/{workspace_id}/workspace-assignments/{principal_id} """ @@ -1860,11 +1986,7 @@ def delete_workspace_assignment(self, workspace_id: int, principal_id: int): "Accept": "application/json", } - self._api.do( - "DELETE", - f"/api/2.0/identity/accounts/{self._api.account_id}/workspaces/{workspace_id}/workspace-assignments/{principal_id}", - headers=headers, - ) + self._api.do("DELETE", f"/api/2.0/identity/{name}", headers=headers) def delete_workspace_assignment_detail(self, workspace_id: int, principal_id: int): """Deletes a workspace assignment detail for a principal, revoking all of its entitlements. Entitlements @@ -1889,13 +2011,12 @@ def delete_workspace_assignment_detail(self, workspace_id: int, principal_id: in headers=headers, ) - def get_direct_group_member(self, group_id: int, principal_id: int) -> DirectGroupMember: + def get_direct_group_member(self, name: str) -> DirectGroupMember: """Gets a provisioned direct member of a group. - :param group_id: int - Required. Internal ID of the group in Databricks. - :param principal_id: int - Required. Internal ID of the principal belonging to the group in Databricks. + :param name: str + The resource name of the group membership. Format: + accounts/{account_id}/groups/{group_id}/direct-members/{principal_id} :returns: :class:`DirectGroupMember` """ @@ -1904,11 +2025,7 @@ def get_direct_group_member(self, group_id: int, principal_id: int) -> DirectGro "Accept": "application/json", } - res = self._api.do( - "GET", - f"/api/2.0/identity/accounts/{self._api.account_id}/groups/{group_id}/direct-members/{principal_id}", - headers=headers, - ) + res = self._api.do("GET", f"/api/2.0/identity/{name}", headers=headers) return DirectGroupMember.from_dict(res) def get_external_group(self, name: str) -> ExternalGroup: @@ -1968,11 +2085,11 @@ def get_external_user(self, name: str) -> ExternalUser: res = self._api.do("GET", f"/api/2.0/identity/{name}", headers=headers) return ExternalUser.from_dict(res) - def get_group(self, group_id: str) -> Group: + def get_group(self, name: str) -> Group: """Fetches a group from the Databricks account by its internal ID. - :param group_id: str - Required. Internal ID of the group in Databricks. + :param name: str + The resource name of the group. Format: accounts/{account_id}/groups/{group_id} :returns: :class:`Group` """ @@ -1981,9 +2098,7 @@ def get_group(self, group_id: str) -> Group: "Accept": "application/json", } - res = self._api.do( - "GET", f"/api/2.0/identity/accounts/{self._api.account_id}/groups/{group_id}", headers=headers - ) + res = self._api.do("GET", f"/api/2.0/identity/{name}", headers=headers) return Group.from_dict(res) def get_identity_visibility_filter(self, name: str) -> IdentityVisibilityFilter: @@ -2004,11 +2119,12 @@ def get_identity_visibility_filter(self, name: str) -> IdentityVisibilityFilter: res = self._api.do("GET", f"/api/2.0/identity/{name}", headers=headers) return IdentityVisibilityFilter.from_dict(res) - def get_service_principal(self, service_principal_id: str) -> ServicePrincipal: + def get_service_principal(self, name: str) -> ServicePrincipal: """Fetches a service principal from the Databricks account by its internal ID. - :param service_principal_id: str - Required. Internal ID of the service principal in Databricks. + :param name: str + The resource name of the service principal. Format: + accounts/{account_id}/service-principals/{service_principal_id} :returns: :class:`ServicePrincipal` """ @@ -2017,18 +2133,14 @@ def get_service_principal(self, service_principal_id: str) -> ServicePrincipal: "Accept": "application/json", } - res = self._api.do( - "GET", - f"/api/2.0/identity/accounts/{self._api.account_id}/service-principals/{service_principal_id}", - headers=headers, - ) + res = self._api.do("GET", f"/api/2.0/identity/{name}", headers=headers) return ServicePrincipal.from_dict(res) - def get_user(self, user_id: str) -> User: + def get_user(self, name: str) -> User: """Fetches a user from the Databricks account by its internal ID. - :param user_id: str - Required. Internal ID of the user in Databricks. + :param name: str + The resource name of the user. Format: accounts/{account_id}/users/{user_id} :returns: :class:`User` """ @@ -2037,11 +2149,11 @@ def get_user(self, user_id: str) -> User: "Accept": "application/json", } - res = self._api.do("GET", f"/api/2.0/identity/accounts/{self._api.account_id}/users/{user_id}", headers=headers) + res = self._api.do("GET", f"/api/2.0/identity/{name}", headers=headers) return User.from_dict(res) def get_workspace_access_detail( - self, workspace_id: int, principal_id: int, *, view: Optional[WorkspaceAccessDetailView] = None + self, name: str, *, view: Optional[WorkspaceAccessDetailView] = None ) -> WorkspaceAccessDetail: """Returns the access details for a principal in a workspace. Allows for checking access details for any provisioned principal (user, service principal, or group) in a workspace. @@ -2050,11 +2162,9 @@ def get_workspace_access_detail( IdP or added explicitly to Databricks via SCIM/UI. Allows for passing in a "view" parameter to control what fields are returned (BASIC by default or FULL). - :param workspace_id: int - Required. The workspace ID for which the access details are being requested. - :param principal_id: int - Required. The internal ID of the principal (user/sp/group) for which the access details are being - requested. + :param name: str + The resource name of the workspace access detail. Format: + accounts/{account_id}/workspaces/{workspace_id}/workspace-access-details/{principal_id} :param view: :class:`WorkspaceAccessDetailView` (optional) Controls what fields are returned. @@ -2068,22 +2178,15 @@ def get_workspace_access_detail( "Accept": "application/json", } - res = self._api.do( - "GET", - f"/api/2.0/identity/accounts/{self._api.account_id}/workspaces/{workspace_id}/workspace-access-details/{principal_id}", - query=query, - headers=headers, - ) + res = self._api.do("GET", f"/api/2.0/identity/{name}", query=query, headers=headers) return WorkspaceAccessDetail.from_dict(res) - def get_workspace_assignment(self, workspace_id: int, principal_id: int) -> WorkspaceAssignment: + def get_workspace_assignment(self, name: str) -> WorkspaceAssignment: """Returns the assignment for a principal in a workspace. - :param workspace_id: int - Required. The workspace ID for which the assignment is being requested. - :param principal_id: int - Required. The internal ID of the principal (user/sp/group) for which the assignment is being - requested. + :param name: str + The resource name of the workspace assignment. Format: + accounts/{account_id}/workspaces/{workspace_id}/workspace-assignments/{principal_id} :returns: :class:`WorkspaceAssignment` """ @@ -2092,11 +2195,7 @@ def get_workspace_assignment(self, workspace_id: int, principal_id: int) -> Work "Accept": "application/json", } - res = self._api.do( - "GET", - f"/api/2.0/identity/accounts/{self._api.account_id}/workspaces/{workspace_id}/workspace-assignments/{principal_id}", - headers=headers, - ) + res = self._api.do("GET", f"/api/2.0/identity/{name}", headers=headers) return WorkspaceAssignment.from_dict(res) def get_workspace_assignment_detail(self, workspace_id: int, principal_id: int) -> WorkspaceAssignmentDetail: @@ -2123,13 +2222,14 @@ def get_workspace_assignment_detail(self, workspace_id: int, principal_id: int) return WorkspaceAssignmentDetail.from_dict(res) def list_direct_group_members( - self, group_id: int, *, page_size: Optional[int] = None, page_token: Optional[str] = None + self, parent: str, *, page_size: Optional[int] = None, page_token: Optional[str] = None ) -> Iterator[DirectGroupMember]: """Lists provisioned direct members of a group with their membership source (internal or from identity provider). - :param group_id: int - Required. Internal ID of the group in Databricks whose direct members are being listed. + :param parent: str + The parent resource name of the group whose direct members to list. Format: + accounts/{account_id}/groups/{group_id} :param page_size: int (optional) The maximum number of members to return. The service may return fewer than this value. If not provided, defaults to 1000, which is also the maximum allowed. Requests for more than the maximum @@ -2151,12 +2251,7 @@ def list_direct_group_members( } while True: - json = self._api.do( - "GET", - f"/api/2.0/identity/accounts/{self._api.account_id}/groups/{group_id}/direct-members", - query=query, - headers=headers, - ) + json = self._api.do("GET", f"/api/2.0/identity/{parent}/direct-members", query=query, headers=headers) if "direct_group_members" in json: for v in json["direct_group_members"]: yield DirectGroupMember.from_dict(v) @@ -2165,11 +2260,18 @@ def list_direct_group_members( query["page_token"] = json["next_page_token"] def list_groups( - self, *, filter: Optional[str] = None, page_size: Optional[int] = None, page_token: Optional[str] = None + self, + parent: str, + *, + filter: Optional[str] = None, + page_size: Optional[int] = None, + page_token: Optional[str] = None, ) -> Iterator[Group]: """Lists the groups in the Databricks account, returning one page per call. Supports filtering by group name or external ID. + :param parent: str + The parent resource name of the account whose groups to list. Format: accounts/{account_id} :param filter: str (optional) Optional. Allows filtering groups by group name or external id. :param page_size: int (optional) @@ -2195,9 +2297,7 @@ def list_groups( } while True: - json = self._api.do( - "GET", f"/api/2.0/identity/accounts/{self._api.account_id}/groups", query=query, headers=headers - ) + json = self._api.do("GET", f"/api/2.0/identity/{parent}/groups", query=query, headers=headers) if "groups" in json: for v in json["groups"]: yield Group.from_dict(v) @@ -2246,11 +2346,19 @@ def list_identity_visibility_filters( query["page_token"] = json["next_page_token"] def list_service_principals( - self, *, filter: Optional[str] = None, page_size: Optional[int] = None, page_token: Optional[str] = None + self, + parent: str, + *, + filter: Optional[str] = None, + page_size: Optional[int] = None, + page_token: Optional[str] = None, ) -> Iterator[ServicePrincipal]: """Lists the service principals in the Databricks account, returning one page per call. Supports filtering by application ID or external ID. + :param parent: str + The parent resource name of the account whose service principals to list. Format: + accounts/{account_id} :param filter: str (optional) Optional. Allows filtering service principals by application id or external id. :param page_size: int (optional) @@ -2276,12 +2384,7 @@ def list_service_principals( } while True: - json = self._api.do( - "GET", - f"/api/2.0/identity/accounts/{self._api.account_id}/service-principals", - query=query, - headers=headers, - ) + json = self._api.do("GET", f"/api/2.0/identity/{parent}/service-principals", query=query, headers=headers) if "service_principals" in json: for v in json["service_principals"]: yield ServicePrincipal.from_dict(v) @@ -2290,13 +2393,13 @@ def list_service_principals( query["page_token"] = json["next_page_token"] def list_transitive_parent_groups( - self, principal_id: int, *, page_size: Optional[int] = None, page_token: Optional[str] = None + self, parent: str, *, page_size: Optional[int] = None, page_token: Optional[str] = None ) -> ListTransitiveParentGroupsResponse: """Lists all transitive parent groups of a principal. - :param principal_id: int - Required. Internal ID of the principal in Databricks whose transitive parent groups are being - listed. + :param parent: str + The parent resource name of the principal whose transitive parent groups to list. Format: + accounts/{account_id}/principals/{principal_id} :param page_size: int (optional) The maximum number of parent groups to return. The service may return fewer than this value. If not provided, defaults to 1000, which is also the maximum allowed. Requests for more than the maximum @@ -2317,20 +2420,22 @@ def list_transitive_parent_groups( "Accept": "application/json", } - res = self._api.do( - "GET", - f"/api/2.0/identity/accounts/{self._api.account_id}/principals/{principal_id}/transitive-parent-groups", - query=query, - headers=headers, - ) + res = self._api.do("GET", f"/api/2.0/identity/{parent}/transitive-parent-groups", query=query, headers=headers) return ListTransitiveParentGroupsResponse.from_dict(res) def list_users( - self, *, filter: Optional[str] = None, page_size: Optional[int] = None, page_token: Optional[str] = None + self, + parent: str, + *, + filter: Optional[str] = None, + page_size: Optional[int] = None, + page_token: Optional[str] = None, ) -> Iterator[User]: """Lists the users in the Databricks account, returning one page per call. Supports filtering by username or external ID. + :param parent: str + The parent resource name of the account whose users to list. Format: accounts/{account_id} :param filter: str (optional) Optional. Allows filtering users by username or external id. :param page_size: int (optional) @@ -2355,9 +2460,7 @@ def list_users( } while True: - json = self._api.do( - "GET", f"/api/2.0/identity/accounts/{self._api.account_id}/users", query=query, headers=headers - ) + json = self._api.do("GET", f"/api/2.0/identity/{parent}/users", query=query, headers=headers) if "users" in json: for v in json["users"]: yield User.from_dict(v) @@ -2409,14 +2512,15 @@ def list_workspace_assignment_details( query["page_token"] = json["next_page_token"] def list_workspace_assignments( - self, workspace_id: int, *, page_size: Optional[int] = None, page_token: Optional[str] = None + self, parent: str, *, page_size: Optional[int] = None, page_token: Optional[str] = None ) -> Iterator[WorkspaceAssignment]: """Lists workspace assignments for a workspace. The response omits the per-principal entitlement fields (``entitlements`` and ``effective_entitlements``). To read the entitlements for a single principal, get that principal's assignment. - :param workspace_id: int - Required. The workspace ID for which the workspace assignments are being fetched. + :param parent: str + The parent resource name of the workspace whose assignments to list. Format: + accounts/{account_id}/workspaces/{workspace_id} :param page_size: int (optional) The maximum number of workspace assignments to return. The service may return fewer than this value. If not provided, defaults to 1000, which is also the maximum allowed. Requests for more than the @@ -2439,10 +2543,7 @@ def list_workspace_assignments( while True: json = self._api.do( - "GET", - f"/api/2.0/identity/accounts/{self._api.account_id}/workspaces/{workspace_id}/workspace-assignments", - query=query, - headers=headers, + "GET", f"/api/2.0/identity/{parent}/workspace-assignments", query=query, headers=headers ) if "workspace_assignments" in json: for v in json["workspace_assignments"]: @@ -2532,7 +2633,7 @@ def resolve_user(self, external_id: str) -> ResolveUserResponse: ) return ResolveUserResponse.from_dict(res) - def update_group(self, group_id: str, group: Group, update_mask: str) -> Group: + def update_group(self, name: str, group: Group, update_mask: str) -> Group: """Updates an existing group in the Databricks account. Only the fields named in the update mask are modified. Returns the updated Group resource. @@ -2544,8 +2645,11 @@ def update_group(self, group_id: str, group: Group, update_mask: str) -> Group: Authorization: the caller must be an account admin or a manager of the group (holds the ``roles/group.manager`` role on it). - :param group_id: str - Required. Internal ID of the group in Databricks. + :param name: str + The resource name of the group. The format depends on the API that returned it: + + - Account-scoped: accounts/{account_id}/groups/{group_id} + - Workspace-scoped: groups/{group_id} :param group: :class:`Group` Required. Group to be updated in Databricks :param update_mask: str @@ -2563,28 +2667,26 @@ def update_group(self, group_id: str, group: Group, update_mask: str) -> Group: "Content-Type": "application/json", } - res = self._api.do( - "PATCH", - f"/api/2.0/identity/accounts/{self._api.account_id}/groups/{group_id}", - query=query, - body=body, - headers=headers, - ) + res = self._api.do("PATCH", f"/api/2.0/identity/{name}", query=query, body=body, headers=headers) return Group.from_dict(res) def update_service_principal( - self, service_principal_id: str, service_principal: ServicePrincipal, update_mask: str + self, name: str, service_principal: ServicePrincipal, update_mask: str ) -> ServicePrincipal: """Updates an existing service principal in the Databricks account. Only the fields named in the update mask are modified. Returns the updated ServicePrincipal resource. - For locally managed service principals, is_agent must be set to true or false when updated. + To change is_agent on a locally managed service principal, set it to true or false. Naming is_agent in + the update mask without a value leaves it unchanged. When AIM is enabled and the service principal is an external identity (its external_id is set), only external_id can be updated; its other fields are sourced from your identity provider. - :param service_principal_id: str - Required. Internal ID of the service principal in Databricks. + :param name: str + The resource name of the service principal. The format depends on the API that returned it: + + - Account-scoped: accounts/{account_id}/service-principals/{service_principal_id} + - Workspace-scoped: service-principals/{service_principal_id} :param service_principal: :class:`ServicePrincipal` Required. Service Principal to be updated in Databricks :param update_mask: str @@ -2602,16 +2704,10 @@ def update_service_principal( "Content-Type": "application/json", } - res = self._api.do( - "PATCH", - f"/api/2.0/identity/accounts/{self._api.account_id}/service-principals/{service_principal_id}", - query=query, - body=body, - headers=headers, - ) + res = self._api.do("PATCH", f"/api/2.0/identity/{name}", query=query, body=body, headers=headers) return ServicePrincipal.from_dict(res) - def update_user(self, user_id: str, user: User, update_mask: str) -> User: + def update_user(self, name: str, user: User, update_mask: str) -> User: """Updates an existing user in the Databricks account and returns the updated user. Only the fields named in the update mask are modified. The updatable fields are fullName.givenName, fullName.familyName, status, and externalId. @@ -2619,8 +2715,11 @@ def update_user(self, user_id: str, user: User, update_mask: str) -> User: When AIM is enabled and the user is an external identity (its external_id is set), only external_id can be updated; its other fields are sourced from your identity provider. - :param user_id: str - Required. Internal ID of the user in Databricks. + :param name: str + The resource name of the user. The format depends on the API that returned it: + + - Account-scoped: accounts/{account_id}/users/{user_id} + - Workspace-scoped: users/{user_id} :param user: :class:`User` Required. User to be updated in Databricks :param update_mask: str @@ -2638,26 +2737,22 @@ def update_user(self, user_id: str, user: User, update_mask: str) -> User: "Content-Type": "application/json", } - res = self._api.do( - "PATCH", - f"/api/2.0/identity/accounts/{self._api.account_id}/users/{user_id}", - query=query, - body=body, - headers=headers, - ) + res = self._api.do("PATCH", f"/api/2.0/identity/{name}", query=query, body=body, headers=headers) return User.from_dict(res) def update_workspace_assignment( - self, workspace_id: int, principal_id: int, workspace_assignment: WorkspaceAssignment, update_mask: FieldMask + self, name: str, workspace_assignment: WorkspaceAssignment, update_mask: FieldMask ) -> WorkspaceAssignment: """Updates the entitlements of a directly assigned principal in a workspace. Changes are applied one at a time rather than atomically. If the request fails partway through, only some of the requested changes take effect. Get the assignment afterwards to confirm the final state. - :param workspace_id: int - Required. The workspace ID for which the workspace assignment is being updated. - :param principal_id: int - Required. ID of the principal in Databricks. + :param name: str + The resource name of the workspace assignment. The format depends on the API that returned it: + + - Account-scoped: + accounts/{account_id}/workspaces/{workspace_id}/workspace-assignments/{principal_id} + - Workspace-scoped: workspace-assignments/{principal_id} :param workspace_assignment: :class:`WorkspaceAssignment` Required. Workspace assignment to be updated in Databricks. :param update_mask: FieldMask @@ -2675,13 +2770,7 @@ def update_workspace_assignment( "Content-Type": "application/json", } - res = self._api.do( - "PATCH", - f"/api/2.0/identity/accounts/{self._api.account_id}/workspaces/{workspace_id}/workspace-assignments/{principal_id}", - query=query, - body=body, - headers=headers, - ) + res = self._api.do("PATCH", f"/api/2.0/identity/{name}", query=query, body=body, headers=headers) return WorkspaceAssignment.from_dict(res) def update_workspace_assignment_detail( @@ -2733,12 +2822,12 @@ def __init__(self, api_client): self._api = api_client def create_direct_group_member_proxy( - self, group_id: int, direct_group_member: DirectGroupMember + self, parent: str, direct_group_member: DirectGroupMember ) -> DirectGroupMember: """Creates a group membership (assigns a principal to a group). - :param group_id: int - Required. Internal ID of the group in Databricks. + :param parent: str + The parent resource name of the group to add the member to. Format: groups/{group_id} :param direct_group_member: :class:`DirectGroupMember` Required. The group membership to create. @@ -2756,10 +2845,10 @@ def create_direct_group_member_proxy( if cfg.workspace_id: headers["X-Databricks-Workspace-Id"] = cfg.workspace_id - res = self._api.do("POST", f"/api/2.0/identity/groups/{group_id}/direct-members", body=body, headers=headers) + res = self._api.do("POST", f"/api/2.0/identity/{parent}/direct-members", body=body, headers=headers) return DirectGroupMember.from_dict(res) - def create_group_proxy(self, group: Group) -> Group: + def create_group_proxy(self, group: Group, *, parent: Optional[str] = None) -> Group: """Creates a local group in the Databricks account that parents the calling workspace and returns the created group. A local group is one that is not synced from the customer's identity provider, and can be created whether or not Account Identity Management (AIM) is enabled. @@ -2770,12 +2859,17 @@ def create_group_proxy(self, group: Group) -> Group: :param group: :class:`Group` Required. Group to be created in Databricks + :param parent: str (optional) + The resource name of the account to create the group in, ``accounts/{account_id}``. Optional: + defaults to the account of the calling workspace, and any other account is rejected. :returns: :class:`Group` """ body = group.as_dict() query = {} + if parent is not None: + query["parent"] = parent headers = { "Accept": "application/json", "Content-Type": "application/json", @@ -2785,10 +2879,12 @@ def create_group_proxy(self, group: Group) -> Group: if cfg.workspace_id: headers["X-Databricks-Workspace-Id"] = cfg.workspace_id - res = self._api.do("POST", "/api/2.0/identity/groups", body=body, headers=headers) + res = self._api.do("POST", "/api/2.0/identity/groups", query=query, body=body, headers=headers) return Group.from_dict(res) - def create_service_principal_proxy(self, service_principal: ServicePrincipal) -> ServicePrincipal: + def create_service_principal_proxy( + self, service_principal: ServicePrincipal, *, parent: Optional[str] = None + ) -> ServicePrincipal: """Creates a local service principal in the Databricks account that parents the calling workspace and returns the created service principal. A local service principal is one that is not synced from the customer's identity provider, and can be created whether or not Account Identity Management (AIM) is @@ -2800,12 +2896,17 @@ def create_service_principal_proxy(self, service_principal: ServicePrincipal) -> :param service_principal: :class:`ServicePrincipal` Required. Service principal to be created in Databricks + :param parent: str (optional) + The resource name of the account to create the service principal in, ``accounts/{account_id}``. + Optional: defaults to the account of the calling workspace, and any other account is rejected. :returns: :class:`ServicePrincipal` """ body = service_principal.as_dict() query = {} + if parent is not None: + query["parent"] = parent headers = { "Accept": "application/json", "Content-Type": "application/json", @@ -2815,10 +2916,10 @@ def create_service_principal_proxy(self, service_principal: ServicePrincipal) -> if cfg.workspace_id: headers["X-Databricks-Workspace-Id"] = cfg.workspace_id - res = self._api.do("POST", "/api/2.0/identity/service-principals", body=body, headers=headers) + res = self._api.do("POST", "/api/2.0/identity/service-principals", query=query, body=body, headers=headers) return ServicePrincipal.from_dict(res) - def create_user_proxy(self, user: User) -> User: + def create_user_proxy(self, user: User, *, parent: Optional[str] = None) -> User: """Creates a local user in the Databricks account that parents the calling workspace and returns the created user. A local user is one that is not synced from the customer's identity provider, and can be created whether or not Account Identity Management (AIM) is enabled. @@ -2829,12 +2930,17 @@ def create_user_proxy(self, user: User) -> User: :param user: :class:`User` Required. User to be created in Databricks + :param parent: str (optional) + The resource name of the account to create the user in, ``accounts/{account_id}``. Optional: + defaults to the account of the calling workspace, and any other account is rejected. :returns: :class:`User` """ body = user.as_dict() query = {} + if parent is not None: + query["parent"] = parent headers = { "Accept": "application/json", "Content-Type": "application/json", @@ -2844,7 +2950,7 @@ def create_user_proxy(self, user: User) -> User: if cfg.workspace_id: headers["X-Databricks-Workspace-Id"] = cfg.workspace_id - res = self._api.do("POST", "/api/2.0/identity/users", body=body, headers=headers) + res = self._api.do("POST", "/api/2.0/identity/users", query=query, body=body, headers=headers) return User.from_dict(res) def create_workspace_assignment_detail_proxy( @@ -2872,7 +2978,9 @@ def create_workspace_assignment_detail_proxy( res = self._api.do("POST", "/api/2.0/identity/workspace-assignment-details", body=body, headers=headers) return WorkspaceAssignmentDetail.from_dict(res) - def create_workspace_assignment_proxy(self, workspace_assignment: WorkspaceAssignment) -> WorkspaceAssignment: + def create_workspace_assignment_proxy( + self, workspace_assignment: WorkspaceAssignment, *, parent: Optional[str] = None + ) -> WorkspaceAssignment: """Creates a workspace assignment for a principal in the calling workspace. Entitlements are granted one at a time rather than atomically. If the request fails partway through, the principal stays assigned to the workspace with only some of the requested entitlements. Get the assignment afterwards to @@ -2880,12 +2988,18 @@ def create_workspace_assignment_proxy(self, workspace_assignment: WorkspaceAssig :param workspace_assignment: :class:`WorkspaceAssignment` Required. Workspace assignment to be created in Databricks. + :param parent: str (optional) + The resource name of the workspace to assign the principal to, + ``accounts/{account_id}/workspaces/{workspace_id}``. Optional: defaults to the calling workspace, + and any other workspace is rejected. :returns: :class:`WorkspaceAssignment` """ body = workspace_assignment.as_dict() query = {} + if parent is not None: + query["parent"] = parent headers = { "Accept": "application/json", "Content-Type": "application/json", @@ -2895,16 +3009,14 @@ def create_workspace_assignment_proxy(self, workspace_assignment: WorkspaceAssig if cfg.workspace_id: headers["X-Databricks-Workspace-Id"] = cfg.workspace_id - res = self._api.do("POST", "/api/2.0/identity/workspace-assignments", body=body, headers=headers) + res = self._api.do("POST", "/api/2.0/identity/workspace-assignments", query=query, body=body, headers=headers) return WorkspaceAssignment.from_dict(res) - def delete_direct_group_member_proxy(self, group_id: int, principal_id: int): + def delete_direct_group_member_proxy(self, name: str): """Deletes a group membership (unassigns a principal from a group). - :param group_id: int - Required. Internal ID of the group in Databricks. - :param principal_id: int - Required. Internal ID of the principal to be unassigned from the group. + :param name: str + The resource name of the group membership. Format: groups/{group_id}/direct-members/{principal_id} """ @@ -2917,13 +3029,13 @@ def delete_direct_group_member_proxy(self, group_id: int, principal_id: int): if cfg.workspace_id: headers["X-Databricks-Workspace-Id"] = cfg.workspace_id - self._api.do("DELETE", f"/api/2.0/identity/groups/{group_id}/direct-members/{principal_id}", headers=headers) + self._api.do("DELETE", f"/api/2.0/identity/{name}", headers=headers) - def delete_group_proxy(self, group_id: str): + def delete_group_proxy(self, name: str): """Deletes a group by its internal ID from the Databricks account that parents the calling workspace. - :param group_id: str - Required. Internal ID of the group in Databricks. + :param name: str + The resource name of the group. Format: groups/{group_id} """ @@ -2936,14 +3048,14 @@ def delete_group_proxy(self, group_id: str): if cfg.workspace_id: headers["X-Databricks-Workspace-Id"] = cfg.workspace_id - self._api.do("DELETE", f"/api/2.0/identity/groups/{group_id}", headers=headers) + self._api.do("DELETE", f"/api/2.0/identity/{name}", headers=headers) - def delete_service_principal_proxy(self, service_principal_id: str): + def delete_service_principal_proxy(self, name: str): """Deletes a service principal by its internal ID from the Databricks account that parents the calling workspace. - :param service_principal_id: str - Required. Internal ID of the service principal in Databricks. + :param name: str + The resource name of the service principal. Format: service-principals/{service_principal_id} """ @@ -2956,13 +3068,13 @@ def delete_service_principal_proxy(self, service_principal_id: str): if cfg.workspace_id: headers["X-Databricks-Workspace-Id"] = cfg.workspace_id - self._api.do("DELETE", f"/api/2.0/identity/service-principals/{service_principal_id}", headers=headers) + self._api.do("DELETE", f"/api/2.0/identity/{name}", headers=headers) - def delete_user_proxy(self, user_id: str): + def delete_user_proxy(self, name: str): """Deletes a user by its internal ID from the Databricks account that parents the calling workspace. - :param user_id: str - Required. Internal ID of the user in Databricks. + :param name: str + The resource name of the user. Format: users/{user_id} """ @@ -2975,7 +3087,7 @@ def delete_user_proxy(self, user_id: str): if cfg.workspace_id: headers["X-Databricks-Workspace-Id"] = cfg.workspace_id - self._api.do("DELETE", f"/api/2.0/identity/users/{user_id}", headers=headers) + self._api.do("DELETE", f"/api/2.0/identity/{name}", headers=headers) def delete_workspace_assignment_detail_proxy(self, principal_id: int): """Deletes a workspace assignment detail for a principal in the calling workspace, revoking all of its @@ -2999,14 +3111,14 @@ def delete_workspace_assignment_detail_proxy(self, principal_id: int): self._api.do("DELETE", f"/api/2.0/identity/workspace-assignment-details/{principal_id}", headers=headers) - def delete_workspace_assignment_proxy(self, principal_id: int): + def delete_workspace_assignment_proxy(self, name: str): """Deletes a workspace assignment for a principal in the calling workspace, revoking all of its entitlements. Entitlements are revoked one at a time rather than atomically. If the request fails partway through, the principal stays assigned with some of its original entitlements. Retrying is safe. - :param principal_id: int - Required. ID of the principal in Databricks to delete workspace assignment for. + :param name: str + The resource name of the workspace assignment. Format: workspace-assignments/{principal_id} """ @@ -3019,15 +3131,13 @@ def delete_workspace_assignment_proxy(self, principal_id: int): if cfg.workspace_id: headers["X-Databricks-Workspace-Id"] = cfg.workspace_id - self._api.do("DELETE", f"/api/2.0/identity/workspace-assignments/{principal_id}", headers=headers) + self._api.do("DELETE", f"/api/2.0/identity/{name}", headers=headers) - def get_direct_group_member_proxy(self, group_id: int, principal_id: int) -> DirectGroupMember: + def get_direct_group_member_proxy(self, name: str) -> DirectGroupMember: """Gets a provisioned direct member of a group. - :param group_id: int - Required. Internal ID of the group in Databricks. - :param principal_id: int - Required. Internal ID of the principal belonging to the group in Databricks. + :param name: str + The resource name of the group membership. Format: groups/{group_id}/direct-members/{principal_id} :returns: :class:`DirectGroupMember` """ @@ -3040,7 +3150,7 @@ def get_direct_group_member_proxy(self, group_id: int, principal_id: int) -> Dir if cfg.workspace_id: headers["X-Databricks-Workspace-Id"] = cfg.workspace_id - res = self._api.do("GET", f"/api/2.0/identity/groups/{group_id}/direct-members/{principal_id}", headers=headers) + res = self._api.do("GET", f"/api/2.0/identity/{name}", headers=headers) return DirectGroupMember.from_dict(res) def get_external_group_proxy(self, name: str) -> ExternalGroup: @@ -3112,11 +3222,11 @@ def get_external_user_proxy(self, name: str) -> ExternalUser: res = self._api.do("GET", f"/api/2.0/identity/{name}", headers=headers) return ExternalUser.from_dict(res) - def get_group_proxy(self, group_id: str) -> Group: + def get_group_proxy(self, name: str) -> Group: """Fetches a group by its internal ID from the Databricks account that parents the calling workspace. - :param group_id: str - Required. Internal ID of the group in Databricks. + :param name: str + The resource name of the group. Format: groups/{group_id} :returns: :class:`Group` """ @@ -3129,15 +3239,15 @@ def get_group_proxy(self, group_id: str) -> Group: if cfg.workspace_id: headers["X-Databricks-Workspace-Id"] = cfg.workspace_id - res = self._api.do("GET", f"/api/2.0/identity/groups/{group_id}", headers=headers) + res = self._api.do("GET", f"/api/2.0/identity/{name}", headers=headers) return Group.from_dict(res) - def get_service_principal_proxy(self, service_principal_id: str) -> ServicePrincipal: + def get_service_principal_proxy(self, name: str) -> ServicePrincipal: """Fetches a service principal by its internal ID from the Databricks account that parents the calling workspace. - :param service_principal_id: str - Required. Internal ID of the service principal in Databricks. + :param name: str + The resource name of the service principal. Format: service-principals/{service_principal_id} :returns: :class:`ServicePrincipal` """ @@ -3150,14 +3260,14 @@ def get_service_principal_proxy(self, service_principal_id: str) -> ServicePrinc if cfg.workspace_id: headers["X-Databricks-Workspace-Id"] = cfg.workspace_id - res = self._api.do("GET", f"/api/2.0/identity/service-principals/{service_principal_id}", headers=headers) + res = self._api.do("GET", f"/api/2.0/identity/{name}", headers=headers) return ServicePrincipal.from_dict(res) - def get_user_proxy(self, user_id: str) -> User: + def get_user_proxy(self, name: str) -> User: """Fetches a user by its internal ID from the Databricks account that parents the calling workspace. - :param user_id: str - Required. Internal ID of the user in Databricks. + :param name: str + The resource name of the user. Format: users/{user_id} :returns: :class:`User` """ @@ -3170,11 +3280,11 @@ def get_user_proxy(self, user_id: str) -> User: if cfg.workspace_id: headers["X-Databricks-Workspace-Id"] = cfg.workspace_id - res = self._api.do("GET", f"/api/2.0/identity/users/{user_id}", headers=headers) + res = self._api.do("GET", f"/api/2.0/identity/{name}", headers=headers) return User.from_dict(res) def get_workspace_access_detail_local( - self, principal_id: int, *, view: Optional[WorkspaceAccessDetailView] = None + self, name: str, *, view: Optional[WorkspaceAccessDetailView] = None ) -> WorkspaceAccessDetail: """Returns the access details for a principal in the current workspace. Allows for checking access details for any provisioned principal (user, service principal, or group) in the current workspace. @@ -3183,9 +3293,8 @@ def get_workspace_access_detail_local( IdP or added explicitly to Databricks via SCIM/UI. Allows for passing in a "view" parameter to control what fields are returned (BASIC by default or FULL). - :param principal_id: int - Required. The internal ID of the principal (user/sp/group) for which the access details are being - requested. + :param name: str + The resource name of the workspace access detail. Format: workspace-access-details/{principal_id} :param view: :class:`WorkspaceAccessDetailView` (optional) Controls what fields are returned. @@ -3203,9 +3312,7 @@ def get_workspace_access_detail_local( if cfg.workspace_id: headers["X-Databricks-Workspace-Id"] = cfg.workspace_id - res = self._api.do( - "GET", f"/api/2.0/identity/workspace-access-details/{principal_id}", query=query, headers=headers - ) + res = self._api.do("GET", f"/api/2.0/identity/{name}", query=query, headers=headers) return WorkspaceAccessDetail.from_dict(res) def get_workspace_assignment_detail_proxy(self, principal_id: int) -> WorkspaceAssignmentDetail: @@ -3229,12 +3336,11 @@ def get_workspace_assignment_detail_proxy(self, principal_id: int) -> WorkspaceA res = self._api.do("GET", f"/api/2.0/identity/workspace-assignment-details/{principal_id}", headers=headers) return WorkspaceAssignmentDetail.from_dict(res) - def get_workspace_assignment_proxy(self, principal_id: int) -> WorkspaceAssignment: + def get_workspace_assignment_proxy(self, name: str) -> WorkspaceAssignment: """Returns the assignment for a principal in the calling workspace. - :param principal_id: int - Required. The internal ID of the principal (user/sp/group) for which the assignment is being - requested. + :param name: str + The resource name of the workspace assignment. Format: workspace-assignments/{principal_id} :returns: :class:`WorkspaceAssignment` """ @@ -3247,15 +3353,15 @@ def get_workspace_assignment_proxy(self, principal_id: int) -> WorkspaceAssignme if cfg.workspace_id: headers["X-Databricks-Workspace-Id"] = cfg.workspace_id - res = self._api.do("GET", f"/api/2.0/identity/workspace-assignments/{principal_id}", headers=headers) + res = self._api.do("GET", f"/api/2.0/identity/{name}", headers=headers) return WorkspaceAssignment.from_dict(res) - def get_workspace_identity_detail(self, principal_id: int) -> WorkspaceIdentityDetail: + def get_workspace_identity_detail(self, name: str) -> WorkspaceIdentityDetail: """Returns the identity details for a principal in a workspace. - :param principal_id: int - Required. The internal ID of the principal (user/sp/group) for which the identity details are being - requested. + :param name: str + The resource name of the workspace identity detail. Format: + principals/{principal_id}/workspace-identity-detail :returns: :class:`WorkspaceIdentityDetail` """ @@ -3268,17 +3374,17 @@ def get_workspace_identity_detail(self, principal_id: int) -> WorkspaceIdentityD if cfg.workspace_id: headers["X-Databricks-Workspace-Id"] = cfg.workspace_id - res = self._api.do("GET", f"/api/2.0/identity/workspace-identity-details/{principal_id}", headers=headers) + res = self._api.do("GET", f"/api/2.0/identity/{name}", headers=headers) return WorkspaceIdentityDetail.from_dict(res) def list_direct_group_members_proxy( - self, group_id: int, *, page_size: Optional[int] = None, page_token: Optional[str] = None + self, parent: str, *, page_size: Optional[int] = None, page_token: Optional[str] = None ) -> Iterator[DirectGroupMember]: """Lists provisioned direct members of a group with their membership source (internal or from identity provider). - :param group_id: int - Required. Internal ID of the group in Databricks whose direct members are being listed. + :param parent: str + The parent resource name of the group whose direct members to list. Format: groups/{group_id} :param page_size: int (optional) The maximum number of members to return. The service may return fewer than this value. If not provided, defaults to 1000, which is also the maximum allowed. Requests for more than the maximum @@ -3303,9 +3409,7 @@ def list_direct_group_members_proxy( headers["X-Databricks-Workspace-Id"] = cfg.workspace_id while True: - json = self._api.do( - "GET", f"/api/2.0/identity/groups/{group_id}/direct-members", query=query, headers=headers - ) + json = self._api.do("GET", f"/api/2.0/identity/{parent}/direct-members", query=query, headers=headers) if "direct_group_members" in json: for v in json["direct_group_members"]: yield DirectGroupMember.from_dict(v) @@ -3314,7 +3418,12 @@ def list_direct_group_members_proxy( query["page_token"] = json["next_page_token"] def list_groups_proxy( - self, *, filter: Optional[str] = None, page_size: Optional[int] = None, page_token: Optional[str] = None + self, + *, + filter: Optional[str] = None, + page_size: Optional[int] = None, + page_token: Optional[str] = None, + parent: Optional[str] = None, ) -> Iterator[Group]: """Lists the groups in the Databricks account that parents the calling workspace, returning one page per call. Supports filtering by group name or external ID. @@ -3328,6 +3437,9 @@ def list_groups_proxy( :param page_token: str (optional) A page token, received from a previous ListGroups call. Provide this to retrieve the subsequent page. + :param parent: str (optional) + The resource name of the account to list groups in, ``accounts/{account_id}``. Optional: defaults to + the account of the calling workspace, and any other account is rejected. :returns: Iterator over :class:`Group` """ @@ -3339,6 +3451,8 @@ def list_groups_proxy( query["page_size"] = page_size if page_token is not None: query["page_token"] = page_token + if parent is not None: + query["parent"] = parent headers = { "Accept": "application/json", } @@ -3357,7 +3471,12 @@ def list_groups_proxy( query["page_token"] = json["next_page_token"] def list_service_principals_proxy( - self, *, filter: Optional[str] = None, page_size: Optional[int] = None, page_token: Optional[str] = None + self, + *, + filter: Optional[str] = None, + page_size: Optional[int] = None, + page_token: Optional[str] = None, + parent: Optional[str] = None, ) -> Iterator[ServicePrincipal]: """Lists the service principals in the Databricks account that parents the calling workspace, returning one page per call. Supports filtering by application ID or external ID. @@ -3371,6 +3490,9 @@ def list_service_principals_proxy( :param page_token: str (optional) A page token, received from a previous ListServicePrincipals call. Provide this to retrieve the subsequent page. + :param parent: str (optional) + The resource name of the account to list service principals in, ``accounts/{account_id}``. Optional: + defaults to the account of the calling workspace, and any other account is rejected. :returns: Iterator over :class:`ServicePrincipal` """ @@ -3382,6 +3504,8 @@ def list_service_principals_proxy( query["page_size"] = page_size if page_token is not None: query["page_token"] = page_token + if parent is not None: + query["parent"] = parent headers = { "Accept": "application/json", } @@ -3400,13 +3524,13 @@ def list_service_principals_proxy( query["page_token"] = json["next_page_token"] def list_transitive_parent_groups_proxy( - self, principal_id: int, *, page_size: Optional[int] = None, page_token: Optional[str] = None + self, parent: str, *, page_size: Optional[int] = None, page_token: Optional[str] = None ) -> ListTransitiveParentGroupsResponse: """Lists all transitive parent groups of a principal. - :param principal_id: int - Required. Internal ID of the principal in Databricks whose transitive parent groups are being - listed. + :param parent: str + The parent resource name of the principal whose transitive parent groups to list. Format: + principals/{principal_id} :param page_size: int (optional) The maximum number of parent groups to return. The service may return fewer than this value. If not provided, defaults to 1000, which is also the maximum allowed. Requests for more than the maximum @@ -3431,13 +3555,16 @@ def list_transitive_parent_groups_proxy( if cfg.workspace_id: headers["X-Databricks-Workspace-Id"] = cfg.workspace_id - res = self._api.do( - "GET", f"/api/2.0/identity/principals/{principal_id}/transitive-parent-groups", query=query, headers=headers - ) + res = self._api.do("GET", f"/api/2.0/identity/{parent}/transitive-parent-groups", query=query, headers=headers) return ListTransitiveParentGroupsResponse.from_dict(res) def list_users_proxy( - self, *, filter: Optional[str] = None, page_size: Optional[int] = None, page_token: Optional[str] = None + self, + *, + filter: Optional[str] = None, + page_size: Optional[int] = None, + page_token: Optional[str] = None, + parent: Optional[str] = None, ) -> Iterator[User]: """Lists the users in the Databricks account that parents the calling workspace, returning one page per call. Supports filtering by username or external ID. @@ -3450,6 +3577,9 @@ def list_users_proxy( are clamped to 1000. :param page_token: str (optional) A page token, received from a previous ListUsers call. Provide this to retrieve the subsequent page. + :param parent: str (optional) + The resource name of the account to list users in, ``accounts/{account_id}``. Optional: defaults to + the account of the calling workspace, and any other account is rejected. :returns: Iterator over :class:`User` """ @@ -3461,6 +3591,8 @@ def list_users_proxy( query["page_size"] = page_size if page_token is not None: query["page_token"] = page_token + if parent is not None: + query["parent"] = parent headers = { "Accept": "application/json", } @@ -3518,7 +3650,7 @@ def list_workspace_assignment_details_proxy( query["page_token"] = json["next_page_token"] def list_workspace_assignments_proxy( - self, *, page_size: Optional[int] = None, page_token: Optional[str] = None + self, *, page_size: Optional[int] = None, page_token: Optional[str] = None, parent: Optional[str] = None ) -> Iterator[WorkspaceAssignment]: """Lists workspace assignments for the calling workspace. The response omits the per-principal entitlement fields (``entitlements`` and ``effective_entitlements``). To read the entitlements for a @@ -3530,6 +3662,10 @@ def list_workspace_assignments_proxy( maximum are clamped to 1000. :param page_token: str (optional) A page token from a previous list call. Provide this to retrieve the subsequent page. + :param parent: str (optional) + The resource name of the workspace to list assignments of, + ``accounts/{account_id}/workspaces/{workspace_id}``. Optional: defaults to the calling workspace, + and any other workspace is rejected. :returns: Iterator over :class:`WorkspaceAssignment` """ @@ -3539,6 +3675,8 @@ def list_workspace_assignments_proxy( query["page_size"] = page_size if page_token is not None: query["page_token"] = page_token + if parent is not None: + query["parent"] = parent headers = { "Accept": "application/json", } @@ -3636,7 +3774,7 @@ def resolve_user_proxy(self, external_id: str) -> ResolveUserResponse: res = self._api.do("POST", "/api/2.0/identity/users/resolve-by-external-id", body=body, headers=headers) return ResolveUserResponse.from_dict(res) - def update_group_proxy(self, group_id: str, group: Group, update_mask: str) -> Group: + def update_group_proxy(self, name: str, group: Group, update_mask: str) -> Group: """Updates an existing group in the Databricks account that parents the calling workspace. Only the fields named in the update mask are modified. Returns the updated Group resource. @@ -3645,8 +3783,11 @@ def update_group_proxy(self, group_id: str, group: Group, update_mask: str) -> G visibility filters, a new nonempty external_id must identify a visible group. Clearing external_id is allowed. - :param group_id: str - Required. Internal ID of the group in Databricks. + :param name: str + The resource name of the group. The format depends on the API that returned it: + + - Account-scoped: accounts/{account_id}/groups/{group_id} + - Workspace-scoped: groups/{group_id} :param group: :class:`Group` Required. Group to be updated in Databricks :param update_mask: str @@ -3668,22 +3809,26 @@ def update_group_proxy(self, group_id: str, group: Group, update_mask: str) -> G if cfg.workspace_id: headers["X-Databricks-Workspace-Id"] = cfg.workspace_id - res = self._api.do("PATCH", f"/api/2.0/identity/groups/{group_id}", query=query, body=body, headers=headers) + res = self._api.do("PATCH", f"/api/2.0/identity/{name}", query=query, body=body, headers=headers) return Group.from_dict(res) def update_service_principal_proxy( - self, service_principal_id: str, service_principal: ServicePrincipal, update_mask: str + self, name: str, service_principal: ServicePrincipal, update_mask: str ) -> ServicePrincipal: """Updates an existing service principal in the Databricks account that parents the calling workspace. Only the fields named in the update mask are modified. Returns the updated ServicePrincipal resource. - For locally managed service principals, is_agent must be set to true or false when updated. + To change is_agent on a locally managed service principal, set it to true or false. Naming is_agent in + the update mask without a value leaves it unchanged. When AIM is enabled and the service principal is an external identity (its external_id is set), only external_id can be updated; its other fields are sourced from your identity provider. - :param service_principal_id: str - Required. Internal ID of the service principal in Databricks. + :param name: str + The resource name of the service principal. The format depends on the API that returned it: + + - Account-scoped: accounts/{account_id}/service-principals/{service_principal_id} + - Workspace-scoped: service-principals/{service_principal_id} :param service_principal: :class:`ServicePrincipal` Required. Service principal to be updated in Databricks :param update_mask: str @@ -3705,16 +3850,10 @@ def update_service_principal_proxy( if cfg.workspace_id: headers["X-Databricks-Workspace-Id"] = cfg.workspace_id - res = self._api.do( - "PATCH", - f"/api/2.0/identity/service-principals/{service_principal_id}", - query=query, - body=body, - headers=headers, - ) + res = self._api.do("PATCH", f"/api/2.0/identity/{name}", query=query, body=body, headers=headers) return ServicePrincipal.from_dict(res) - def update_user_proxy(self, user_id: str, user: User, update_mask: str) -> User: + def update_user_proxy(self, name: str, user: User, update_mask: str) -> User: """Updates an existing user in the Databricks account that parents the calling workspace and returns the updated user. Only the fields named in the update mask are modified. The updatable fields are fullName.givenName, fullName.familyName, status, and externalId. @@ -3722,8 +3861,11 @@ def update_user_proxy(self, user_id: str, user: User, update_mask: str) -> User: When AIM is enabled and the user is an external identity (its external_id is set), only external_id can be updated; its other fields are sourced from your identity provider. - :param user_id: str - Required. Internal ID of the user in Databricks. + :param name: str + The resource name of the user. The format depends on the API that returned it: + + - Account-scoped: accounts/{account_id}/users/{user_id} + - Workspace-scoped: users/{user_id} :param user: :class:`User` Required. User to be updated in Databricks :param update_mask: str @@ -3745,7 +3887,7 @@ def update_user_proxy(self, user_id: str, user: User, update_mask: str) -> User: if cfg.workspace_id: headers["X-Databricks-Workspace-Id"] = cfg.workspace_id - res = self._api.do("PATCH", f"/api/2.0/identity/users/{user_id}", query=query, body=body, headers=headers) + res = self._api.do("PATCH", f"/api/2.0/identity/{name}", query=query, body=body, headers=headers) return User.from_dict(res) def update_workspace_assignment_detail_proxy( @@ -3788,14 +3930,18 @@ def update_workspace_assignment_detail_proxy( return WorkspaceAssignmentDetail.from_dict(res) def update_workspace_assignment_proxy( - self, principal_id: int, workspace_assignment: WorkspaceAssignment, update_mask: FieldMask + self, name: str, workspace_assignment: WorkspaceAssignment, update_mask: FieldMask ) -> WorkspaceAssignment: """Updates the entitlements of a directly assigned principal in the calling workspace. Changes are applied one at a time rather than atomically. If the request fails partway through, only some of the requested changes take effect. Get the assignment afterwards to confirm the final state. - :param principal_id: int - Required. ID of the principal in Databricks. + :param name: str + The resource name of the workspace assignment. The format depends on the API that returned it: + + - Account-scoped: + accounts/{account_id}/workspaces/{workspace_id}/workspace-assignments/{principal_id} + - Workspace-scoped: workspace-assignments/{principal_id} :param workspace_assignment: :class:`WorkspaceAssignment` Required. Workspace assignment to be updated in Databricks. :param update_mask: FieldMask @@ -3817,18 +3963,17 @@ def update_workspace_assignment_proxy( if cfg.workspace_id: headers["X-Databricks-Workspace-Id"] = cfg.workspace_id - res = self._api.do( - "PATCH", f"/api/2.0/identity/workspace-assignments/{principal_id}", query=query, body=body, headers=headers - ) + res = self._api.do("PATCH", f"/api/2.0/identity/{name}", query=query, body=body, headers=headers) return WorkspaceAssignment.from_dict(res) def update_workspace_identity_detail( - self, principal_id: int, workspace_identity_detail: WorkspaceIdentityDetail, update_mask: FieldMask + self, name: str, workspace_identity_detail: WorkspaceIdentityDetail, update_mask: FieldMask ) -> WorkspaceIdentityDetail: """Updates a workspace identity detail for a principal. - :param principal_id: int - Required. ID of the principal in Databricks. + :param name: str + The resource name of the workspace identity detail. Format: + principals/{principal_id}/workspace-identity-detail :param workspace_identity_detail: :class:`WorkspaceIdentityDetail` Required. Workspace identity detail to be updated in Databricks. :param update_mask: FieldMask @@ -3850,11 +3995,5 @@ def update_workspace_identity_detail( if cfg.workspace_id: headers["X-Databricks-Workspace-Id"] = cfg.workspace_id - res = self._api.do( - "PATCH", - f"/api/2.0/identity/workspace-identity-details/{principal_id}", - query=query, - body=body, - headers=headers, - ) + res = self._api.do("PATCH", f"/api/2.0/identity/{name}", query=query, body=body, headers=headers) return WorkspaceIdentityDetail.from_dict(res) diff --git a/databricks/sdk/service/jobs.py b/databricks/sdk/service/jobs.py index 7dd7966c8..c0ee023cc 100644 --- a/databricks/sdk/service/jobs.py +++ b/databricks/sdk/service/jobs.py @@ -9,6 +9,8 @@ from enum import Enum from typing import Dict, List, Any, Iterator, Callable, Optional +from google.protobuf.duration_pb2 import Duration +from google.protobuf.timestamp_pb2 import Timestamp import time import random @@ -16,11 +18,14 @@ from ..errors import OperationFailed from databricks.sdk.service._internal import ( + _duration, _enum, _from_dict, _int64, _repeated_dict, + _repeated_enum, _repeated_int64, + _timestamp, Wait, ) @@ -34,6 +39,423 @@ # all definitions in this file are in alphabetical order +@dataclass +class AiClassifyFunction: + """Configuration for a task that runs the ``ai_classify`` function. + + The label set comes from exactly one of two places: ``inline_labels``, which holds up to 10 + labels, or a Unity Catalog table named by ``label_table``. Use ``label_table`` for a larger + label set.""" + + version: str + """The ai_classify function version to invoke (for example, ``2.1``).""" + + enable_confidence_scores: Optional[bool] = None + """When true, includes a per-label confidence score in the response.""" + + enable_rationales: Optional[bool] = None + """When true, includes a rationale explaining each classification in the response.""" + + inline_labels: Optional[AiClassifyFunctionInlineLabels] = None + """Labels listed inline, for a small label set.""" + + instructions: Optional[str] = None + """Global description of the task and domain, passed to the model to improve classification + quality.""" + + multilabel: Optional[bool] = None + """When true, more than one label can be returned for each input.""" + + def as_dict(self) -> dict: + """Serializes the AiClassifyFunction into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.enable_confidence_scores is not None: + body["enable_confidence_scores"] = self.enable_confidence_scores + if self.enable_rationales is not None: + body["enable_rationales"] = self.enable_rationales + if self.inline_labels: + body["inline_labels"] = self.inline_labels.as_dict() + if self.instructions is not None: + body["instructions"] = self.instructions + if self.multilabel is not None: + body["multilabel"] = self.multilabel + if self.version is not None: + body["version"] = self.version + return body + + def as_shallow_dict(self) -> dict: + """Serializes the AiClassifyFunction into a shallow dictionary of its immediate attributes.""" + body = {} + if self.enable_confidence_scores is not None: + body["enable_confidence_scores"] = self.enable_confidence_scores + if self.enable_rationales is not None: + body["enable_rationales"] = self.enable_rationales + if self.inline_labels: + body["inline_labels"] = self.inline_labels + if self.instructions is not None: + body["instructions"] = self.instructions + if self.multilabel is not None: + body["multilabel"] = self.multilabel + if self.version is not None: + body["version"] = self.version + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> AiClassifyFunction: + """Deserializes the AiClassifyFunction from a dictionary.""" + return cls( + enable_confidence_scores=d.get("enable_confidence_scores", None), + enable_rationales=d.get("enable_rationales", None), + inline_labels=_from_dict(d, "inline_labels", AiClassifyFunctionInlineLabels), + instructions=d.get("instructions", None), + multilabel=d.get("multilabel", None), + version=d.get("version", None), + ) + + +@dataclass +class AiClassifyFunctionInlineLabels: + """A label set listed inline.""" + + labels: List[AiClassifyFunctionLabel] + """The labels an input can be classified as. At least two labels are required unless ``multilabel`` + is true, which also accepts a single label.""" + + def as_dict(self) -> dict: + """Serializes the AiClassifyFunctionInlineLabels into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.labels: + body["labels"] = [v.as_dict() for v in self.labels] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the AiClassifyFunctionInlineLabels into a shallow dictionary of its immediate attributes.""" + body = {} + if self.labels: + body["labels"] = self.labels + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> AiClassifyFunctionInlineLabels: + """Deserializes the AiClassifyFunctionInlineLabels from a dictionary.""" + return cls(labels=_repeated_dict(d, "labels", AiClassifyFunctionLabel)) + + +@dataclass +class AiClassifyFunctionLabel: + """One label the input can be classified as.""" + + name: str + """The name returned when an input is classified as this label.""" + + description: Optional[str] = None + """Guidance describing what belongs under this label.""" + + def as_dict(self) -> dict: + """Serializes the AiClassifyFunctionLabel into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.description is not None: + body["description"] = self.description + if self.name is not None: + body["name"] = self.name + return body + + def as_shallow_dict(self) -> dict: + """Serializes the AiClassifyFunctionLabel into a shallow dictionary of its immediate attributes.""" + body = {} + if self.description is not None: + body["description"] = self.description + if self.name is not None: + body["name"] = self.name + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> AiClassifyFunctionLabel: + """Deserializes the AiClassifyFunctionLabel from a dictionary.""" + return cls(description=d.get("description", None), name=d.get("name", None)) + + +@dataclass +class AiExtractFunction: + """Configuration for a task that runs the ``ai_extract`` function. + + The fields to extract come from exactly one of two places: an existing AI Extract agent, named + by ``agent_id``, or the inline ``schema`` and ``instructions``. Set ``agent_id`` or the inline + fields, not both.""" + + version: str + """The ai_extract function version to invoke (for example, ``2.1``).""" + + enable_citations: Optional[bool] = None + """When true, includes citation metadata locating each extracted value in the source.""" + + enable_confidence_scores: Optional[bool] = None + """When true, includes a per-field confidence score in the response.""" + + instructions: Optional[str] = None + """Global description of the task and domain, passed to the model to improve extraction quality.""" + + mode: Optional[AiExtractFunctionMode] = None + """The extraction mode. Defaults to standard extraction when unspecified.""" + + schema: Optional[str] = None + """The extraction schema defining the fields to extract, as a JSON string. Either an array of field + names (e.g. ["company", "valuation"]) or an object mapping each field name to a type and an + optional description. Supported types are string, integer, number, boolean, enum, object, and + array; all fields are optional.""" + + def as_dict(self) -> dict: + """Serializes the AiExtractFunction into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.enable_citations is not None: + body["enable_citations"] = self.enable_citations + if self.enable_confidence_scores is not None: + body["enable_confidence_scores"] = self.enable_confidence_scores + if self.instructions is not None: + body["instructions"] = self.instructions + if self.mode is not None: + body["mode"] = self.mode.value + if self.schema is not None: + body["schema"] = self.schema + if self.version is not None: + body["version"] = self.version + return body + + def as_shallow_dict(self) -> dict: + """Serializes the AiExtractFunction into a shallow dictionary of its immediate attributes.""" + body = {} + if self.enable_citations is not None: + body["enable_citations"] = self.enable_citations + if self.enable_confidence_scores is not None: + body["enable_confidence_scores"] = self.enable_confidence_scores + if self.instructions is not None: + body["instructions"] = self.instructions + if self.mode is not None: + body["mode"] = self.mode + if self.schema is not None: + body["schema"] = self.schema + if self.version is not None: + body["version"] = self.version + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> AiExtractFunction: + """Deserializes the AiExtractFunction from a dictionary.""" + return cls( + enable_citations=d.get("enable_citations", None), + enable_confidence_scores=d.get("enable_confidence_scores", None), + instructions=d.get("instructions", None), + mode=_enum(d, "mode", AiExtractFunctionMode), + schema=d.get("schema", None), + version=d.get("version", None), + ) + + +class AiExtractFunctionMode(Enum): + """The extraction mode.""" + + MODE_PRECISION = "MODE_PRECISION" + + +@dataclass +class AiFunctionTask: + """A task that applies an AI function to a set of input documents or rows.""" + + output: QualifiedName + """Output destination table. When ``table`` is unset the task generates a stable name from the + function and task identity so a downstream task can reference it.""" + + ai_classify: Optional[AiClassifyFunction] = None + """Classifies each input document or row into one or more labels.""" + + ai_extract: Optional[AiExtractFunction] = None + """Extracts structured fields from the input documents or rows.""" + + ai_parse_document: Optional[AiParseDocumentFunction] = None + """Parses the input documents into structured text.""" + + delta_table: Optional[DeltaTableSource] = None + """A Unity Catalog Delta table. The function is applied to each row.""" + + volume: Optional[VolumeSource] = None + """A Unity Catalog volume. The function is applied to each file.""" + + def as_dict(self) -> dict: + """Serializes the AiFunctionTask into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.ai_classify: + body["ai_classify"] = self.ai_classify.as_dict() + if self.ai_extract: + body["ai_extract"] = self.ai_extract.as_dict() + if self.ai_parse_document: + body["ai_parse_document"] = self.ai_parse_document.as_dict() + if self.delta_table: + body["delta_table"] = self.delta_table.as_dict() + if self.output: + body["output"] = self.output.as_dict() + if self.volume: + body["volume"] = self.volume.as_dict() + return body + + def as_shallow_dict(self) -> dict: + """Serializes the AiFunctionTask into a shallow dictionary of its immediate attributes.""" + body = {} + if self.ai_classify: + body["ai_classify"] = self.ai_classify + if self.ai_extract: + body["ai_extract"] = self.ai_extract + if self.ai_parse_document: + body["ai_parse_document"] = self.ai_parse_document + if self.delta_table: + body["delta_table"] = self.delta_table + if self.output: + body["output"] = self.output + if self.volume: + body["volume"] = self.volume + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> AiFunctionTask: + """Deserializes the AiFunctionTask from a dictionary.""" + return cls( + ai_classify=_from_dict(d, "ai_classify", AiClassifyFunction), + ai_extract=_from_dict(d, "ai_extract", AiExtractFunction), + ai_parse_document=_from_dict(d, "ai_parse_document", AiParseDocumentFunction), + delta_table=_from_dict(d, "delta_table", DeltaTableSource), + output=_from_dict(d, "output", QualifiedName), + volume=_from_dict(d, "volume", VolumeSource), + ) + + +@dataclass +class AiFunctionTaskOutput: + """The result of an AI function task run. Also reported mid-run, so the counts can be partial.""" + + estimated_time_remaining: Optional[Duration] = None + """Coarse estimate of the time left to run the function over the inputs, from the throughput so + far. Covers that work only, not the merge into the output table that follows it.""" + + failed_row_count: Optional[int] = None + """Inputs permanently failed so far. Each lands in the output table as an error row, and does not + on its own fail the run.""" + + output_table: Optional[QualifiedName] = None + """The table the run writes its results to. Reported even when ``output.table`` is unset and the + name is generated, so a client can read the table.""" + + succeeded_row_count: Optional[int] = None + """Inputs with a result so far. Reported per batch, so it can lag mid-run.""" + + total_row_count: Optional[int] = None + """Inputs to process: rows for a Delta table source, files for a volume source. Stays well inside + 2^53, so a JavaScript client can hold it exactly.""" + + def as_dict(self) -> dict: + """Serializes the AiFunctionTaskOutput into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.estimated_time_remaining is not None: + body["estimated_time_remaining"] = self.estimated_time_remaining.ToJsonString() + if self.failed_row_count is not None: + body["failed_row_count"] = self.failed_row_count + if self.output_table: + body["output_table"] = self.output_table.as_dict() + if self.succeeded_row_count is not None: + body["succeeded_row_count"] = self.succeeded_row_count + if self.total_row_count is not None: + body["total_row_count"] = self.total_row_count + return body + + def as_shallow_dict(self) -> dict: + """Serializes the AiFunctionTaskOutput into a shallow dictionary of its immediate attributes.""" + body = {} + if self.estimated_time_remaining is not None: + body["estimated_time_remaining"] = self.estimated_time_remaining + if self.failed_row_count is not None: + body["failed_row_count"] = self.failed_row_count + if self.output_table: + body["output_table"] = self.output_table + if self.succeeded_row_count is not None: + body["succeeded_row_count"] = self.succeeded_row_count + if self.total_row_count is not None: + body["total_row_count"] = self.total_row_count + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> AiFunctionTaskOutput: + """Deserializes the AiFunctionTaskOutput from a dictionary.""" + return cls( + estimated_time_remaining=_duration(d, "estimated_time_remaining"), + failed_row_count=_int64(d, "failed_row_count"), + output_table=_from_dict(d, "output_table", QualifiedName), + succeeded_row_count=_int64(d, "succeeded_row_count"), + total_row_count=_int64(d, "total_row_count"), + ) + + +@dataclass +class AiParseDocumentFunction: + """Configuration for a task that runs the ``ai_parse_document`` function.""" + + version: str + """The ai_parse_document output schema version to invoke (for example, ``2.0``).""" + + description_element_types: Optional[List[AiParseDocumentFunctionDescriptionElementType]] = None + """Element types to generate AI descriptions for. Leave empty to generate no descriptions; use ALL + for every supported type, or list specific types to restrict generation to only those.""" + + image_output_path: Optional[str] = None + """Unity Catalog volume path where rendered page and element images are written.""" + + page_range: Optional[str] = None + """Pages to parse (1-indexed), as a comma-separated list of page numbers or ranges (e.g. + "1,3,5-10").""" + + def as_dict(self) -> dict: + """Serializes the AiParseDocumentFunction into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.description_element_types: + body["description_element_types"] = [v.value for v in self.description_element_types] + if self.image_output_path is not None: + body["image_output_path"] = self.image_output_path + if self.page_range is not None: + body["page_range"] = self.page_range + if self.version is not None: + body["version"] = self.version + return body + + def as_shallow_dict(self) -> dict: + """Serializes the AiParseDocumentFunction into a shallow dictionary of its immediate attributes.""" + body = {} + if self.description_element_types: + body["description_element_types"] = self.description_element_types + if self.image_output_path is not None: + body["image_output_path"] = self.image_output_path + if self.page_range is not None: + body["page_range"] = self.page_range + if self.version is not None: + body["version"] = self.version + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> AiParseDocumentFunction: + """Deserializes the AiParseDocumentFunction from a dictionary.""" + return cls( + description_element_types=_repeated_enum( + d, "description_element_types", AiParseDocumentFunctionDescriptionElementType + ), + image_output_path=d.get("image_output_path", None), + page_range=d.get("page_range", None), + version=d.get("version", None), + ) + + +class AiParseDocumentFunctionDescriptionElementType(Enum): + """An element type that AI descriptions can be generated for.""" + + DESCRIPTION_ELEMENT_TYPE_ALL = "DESCRIPTION_ELEMENT_TYPE_ALL" + DESCRIPTION_ELEMENT_TYPE_FIGURE = "DESCRIPTION_ELEMENT_TYPE_FIGURE" + + @dataclass class AiRuntimeTask: """AiRuntimeTask: multi-node GPU compute task definition for Databricks AI Runtime workloads. @@ -1105,6 +1527,67 @@ def from_dict(cls, d: Dict[str, Any]) -> ClusterSpec: ) +@dataclass +class CompletedJob: + """A single upstream completion that contributed to triggering the run, captured at trigger time""" + + completion_time: Optional[Timestamp] = None + """The completion time of the upstream run""" + + job_id: Optional[str] = None + """The ID of the upstream job""" + + run_id: Optional[str] = None + """The ID of the upstream run whose completion contributed to the trigger""" + + termination_code: Optional[TerminationCodeCode] = None + """The termination code of the upstream run""" + + workspace_id: Optional[str] = None + """The ID of the workspace containing the upstream job""" + + def as_dict(self) -> dict: + """Serializes the CompletedJob into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.completion_time is not None: + body["completion_time"] = self.completion_time.ToJsonString() + if self.job_id is not None: + body["job_id"] = self.job_id + if self.run_id is not None: + body["run_id"] = self.run_id + if self.termination_code is not None: + body["termination_code"] = self.termination_code.value + if self.workspace_id is not None: + body["workspace_id"] = self.workspace_id + return body + + def as_shallow_dict(self) -> dict: + """Serializes the CompletedJob into a shallow dictionary of its immediate attributes.""" + body = {} + if self.completion_time is not None: + body["completion_time"] = self.completion_time + if self.job_id is not None: + body["job_id"] = self.job_id + if self.run_id is not None: + body["run_id"] = self.run_id + if self.termination_code is not None: + body["termination_code"] = self.termination_code + if self.workspace_id is not None: + body["workspace_id"] = self.workspace_id + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> CompletedJob: + """Deserializes the CompletedJob from a dictionary.""" + return cls( + completion_time=_timestamp(d, "completion_time"), + job_id=d.get("job_id", None), + run_id=d.get("run_id", None), + termination_code=_enum(d, "termination_code", TerminationCodeCode), + workspace_id=d.get("workspace_id", None), + ) + + @dataclass class Compute: hardware_accelerator: Optional[compute.HardwareAcceleratorType] = None @@ -2116,6 +2599,55 @@ def from_dict(cls, d: Dict[str, Any]) -> DbtTask: ) +@dataclass +class DeltaTableSource: + """A Unity Catalog Delta table input source. The function is applied to each row.""" + + table_name: str + """Fully qualified three-part name of the source table (catalog.schema.table).""" + + input_column: str + """Name of the column holding the content the function is applied to.""" + + merge_key_columns: Optional[List[str]] = None + """Columns that together identify a row across incremental runs, so results can be matched back to + their source rows. When unset, the task derives an implicit key from Delta row tracking. If row + tracking is also disabled there is no stable key, so each run reprocesses the entire source (a + full refresh) rather than only the changed rows; set a key column or enable row tracking to keep + runs incremental.""" + + def as_dict(self) -> dict: + """Serializes the DeltaTableSource into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.input_column is not None: + body["input_column"] = self.input_column + if self.merge_key_columns: + body["merge_key_columns"] = [v for v in self.merge_key_columns] + if self.table_name is not None: + body["table_name"] = self.table_name + return body + + def as_shallow_dict(self) -> dict: + """Serializes the DeltaTableSource into a shallow dictionary of its immediate attributes.""" + body = {} + if self.input_column is not None: + body["input_column"] = self.input_column + if self.merge_key_columns: + body["merge_key_columns"] = self.merge_key_columns + if self.table_name is not None: + body["table_name"] = self.table_name + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> DeltaTableSource: + """Deserializes the DeltaTableSource from a dictionary.""" + return cls( + input_column=d.get("input_column", None), + merge_key_columns=d.get("merge_key_columns", None), + table_name=d.get("table_name", None), + ) + + @dataclass class DeploymentSpec: """DeploymentSpec: configuration for one deployment within an AiRuntimeTask. Each entry in @@ -3113,6 +3645,11 @@ class JobCluster: new_cluster: Optional[compute.ClusterSpec] = None """If new_cluster, a description of a cluster that is created for each task.""" + serverless_compute: Optional[JobClusterServerlessCompute] = None + """Serverless compute settings for this cluster. At most one job cluster per job may set this + field; the rate limit defined on the referenced policy applies across all tasks bound to this + cluster.""" + serverless_compute_id: Optional[str] = None """The ID of the serverless compute object to bind this cluster to. At most one JobCluster per job may set this field; the rate limit defined on the referenced serverless compute applies across @@ -3125,6 +3662,8 @@ def as_dict(self) -> dict: body["job_cluster_key"] = self.job_cluster_key if self.new_cluster: body["new_cluster"] = self.new_cluster.as_dict() + if self.serverless_compute: + body["serverless_compute"] = self.serverless_compute.as_dict() if self.serverless_compute_id is not None: body["serverless_compute_id"] = self.serverless_compute_id return body @@ -3136,6 +3675,8 @@ def as_shallow_dict(self) -> dict: body["job_cluster_key"] = self.job_cluster_key if self.new_cluster: body["new_cluster"] = self.new_cluster + if self.serverless_compute: + body["serverless_compute"] = self.serverless_compute if self.serverless_compute_id is not None: body["serverless_compute_id"] = self.serverless_compute_id return body @@ -3146,10 +3687,262 @@ def from_dict(cls, d: Dict[str, Any]) -> JobCluster: return cls( job_cluster_key=d.get("job_cluster_key", None), new_cluster=_from_dict(d, "new_cluster", compute.ClusterSpec), + serverless_compute=_from_dict(d, "serverless_compute", JobClusterServerlessCompute), serverless_compute_id=d.get("serverless_compute_id", None), ) +@dataclass +class JobClusterServerlessCompute: + """Serverless compute settings for a job cluster.""" + + policy_id: Optional[str] = None + """The ID of the serverless compute policy to apply to this cluster.""" + + def as_dict(self) -> dict: + """Serializes the JobClusterServerlessCompute into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.policy_id is not None: + body["policy_id"] = self.policy_id + return body + + def as_shallow_dict(self) -> dict: + """Serializes the JobClusterServerlessCompute into a shallow dictionary of its immediate attributes.""" + body = {} + if self.policy_id is not None: + body["policy_id"] = self.policy_id + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> JobClusterServerlessCompute: + """Deserializes the JobClusterServerlessCompute from a dictionary.""" + return cls(policy_id=d.get("policy_id", None)) + + +@dataclass +class JobCompletionTriggerConfiguration: + """A trigger that starts a run after every upstream job completes with a configured status.""" + + jobs: List[JobCompletionTriggerConfigurationMonitoredJob] + """The upstream jobs to monitor. Each job must be in a workspace attached to the same Unity Catalog + metastore as the downstream job. A maximum of 10 jobs can be monitored.""" + + if_status: Optional[JobCompletionTriggerConfigurationIfStatus] = None + """The status that all upstream jobs must have before this trigger fires. Defaults to SUCCESS.""" + + def as_dict(self) -> dict: + """Serializes the JobCompletionTriggerConfiguration into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.if_status is not None: + body["if_status"] = self.if_status.value + if self.jobs: + body["jobs"] = [v.as_dict() for v in self.jobs] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the JobCompletionTriggerConfiguration into a shallow dictionary of its immediate attributes.""" + body = {} + if self.if_status is not None: + body["if_status"] = self.if_status + if self.jobs: + body["jobs"] = self.jobs + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> JobCompletionTriggerConfiguration: + """Deserializes the JobCompletionTriggerConfiguration from a dictionary.""" + return cls( + if_status=_enum(d, "if_status", JobCompletionTriggerConfigurationIfStatus), + jobs=_repeated_dict(d, "jobs", JobCompletionTriggerConfigurationMonitoredJob), + ) + + +class JobCompletionTriggerConfigurationIfStatus(Enum): + """The status that all upstream jobs must have before a job completion trigger fires.""" + + DONE = "DONE" + FAILED = "FAILED" + SUCCESS = "SUCCESS" + + +@dataclass +class JobCompletionTriggerConfigurationMonitoredJob: + """An upstream job to monitor.""" + + workspace_id: str + """The ID of the workspace containing the upstream job. This can differ from the caller's workspace + and downstream services enforce correct entitlements and CAN_VIEW permissions for the user + job + combination. SDR: https://databricks.atlassian.net/browse/SDR-5267""" + + job_id: str + """The ID of the upstream job.""" + + def as_dict(self) -> dict: + """Serializes the JobCompletionTriggerConfigurationMonitoredJob into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.job_id is not None: + body["job_id"] = self.job_id + if self.workspace_id is not None: + body["workspace_id"] = self.workspace_id + return body + + def as_shallow_dict(self) -> dict: + """Serializes the JobCompletionTriggerConfigurationMonitoredJob into a shallow dictionary of its immediate attributes.""" + body = {} + if self.job_id is not None: + body["job_id"] = self.job_id + if self.workspace_id is not None: + body["workspace_id"] = self.workspace_id + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> JobCompletionTriggerConfigurationMonitoredJob: + """Deserializes the JobCompletionTriggerConfigurationMonitoredJob from a dictionary.""" + return cls(job_id=d.get("job_id", None), workspace_id=d.get("workspace_id", None)) + + +@dataclass +class JobCompletionTriggerRunInfoDetails: + """Contains the immutable set of upstream completions that triggered the run, in configured jobs + order""" + + completed_jobs: Optional[List[CompletedJob]] = None + + def as_dict(self) -> dict: + """Serializes the JobCompletionTriggerRunInfoDetails into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.completed_jobs: + body["completed_jobs"] = [v.as_dict() for v in self.completed_jobs] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the JobCompletionTriggerRunInfoDetails into a shallow dictionary of its immediate attributes.""" + body = {} + if self.completed_jobs: + body["completed_jobs"] = self.completed_jobs + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> JobCompletionTriggerRunInfoDetails: + """Deserializes the JobCompletionTriggerRunInfoDetails from a dictionary.""" + return cls(completed_jobs=_repeated_dict(d, "completed_jobs", CompletedJob)) + + +@dataclass +class JobCompletionTriggerState: + """Runtime state for a job completion trigger.""" + + monitored_job_states: Optional[List[JobCompletionTriggerStateMonitoredJobState]] = None + """State for each configured upstream job, in the same order as ``jobs``.""" + + def as_dict(self) -> dict: + """Serializes the JobCompletionTriggerState into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.monitored_job_states: + body["monitored_job_states"] = [v.as_dict() for v in self.monitored_job_states] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the JobCompletionTriggerState into a shallow dictionary of its immediate attributes.""" + body = {} + if self.monitored_job_states: + body["monitored_job_states"] = self.monitored_job_states + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> JobCompletionTriggerState: + """Deserializes the JobCompletionTriggerState from a dictionary.""" + return cls( + monitored_job_states=_repeated_dict(d, "monitored_job_states", JobCompletionTriggerStateMonitoredJobState) + ) + + +@dataclass +class JobCompletionTriggerStateMonitoredJobState: + has_seen_completion: Optional[bool] = None + """Whether this job has contributed a matching completion toward the next trigger fire.""" + + is_deleted: Optional[bool] = None + """Whether the upstream job was missing or deleted when last observed by this trigger. Omitted when + deletion status is unknown. This observation may become stale while the trigger is paused or + evaluation cannot access the upstream job.""" + + job_id: Optional[str] = None + """The ID of the upstream job.""" + + last_seen_completion_time: Optional[Timestamp] = None + """The completion time of the latest matching upstream run observed by the trigger.""" + + last_seen_result: Optional[TerminationCodeCode] = None + """The termination code of the latest matching upstream run observed by the trigger.""" + + last_seen_run_id: Optional[str] = None + """The ID of the latest matching upstream run observed by the trigger.""" + + latest_non_matching_result: Optional[TerminationCodeCode] = None + """The termination code of the latest upstream run if it does not satisfy ``if_status``. Unset when + the latest run matches or no run has been observed.""" + + workspace_id: Optional[str] = None + """The ID of the workspace containing the upstream job.""" + + def as_dict(self) -> dict: + """Serializes the JobCompletionTriggerStateMonitoredJobState into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.has_seen_completion is not None: + body["has_seen_completion"] = self.has_seen_completion + if self.is_deleted is not None: + body["is_deleted"] = self.is_deleted + if self.job_id is not None: + body["job_id"] = self.job_id + if self.last_seen_completion_time is not None: + body["last_seen_completion_time"] = self.last_seen_completion_time.ToJsonString() + if self.last_seen_result is not None: + body["last_seen_result"] = self.last_seen_result.value + if self.last_seen_run_id is not None: + body["last_seen_run_id"] = self.last_seen_run_id + if self.latest_non_matching_result is not None: + body["latest_non_matching_result"] = self.latest_non_matching_result.value + if self.workspace_id is not None: + body["workspace_id"] = self.workspace_id + return body + + def as_shallow_dict(self) -> dict: + """Serializes the JobCompletionTriggerStateMonitoredJobState into a shallow dictionary of its immediate attributes.""" + body = {} + if self.has_seen_completion is not None: + body["has_seen_completion"] = self.has_seen_completion + if self.is_deleted is not None: + body["is_deleted"] = self.is_deleted + if self.job_id is not None: + body["job_id"] = self.job_id + if self.last_seen_completion_time is not None: + body["last_seen_completion_time"] = self.last_seen_completion_time + if self.last_seen_result is not None: + body["last_seen_result"] = self.last_seen_result + if self.last_seen_run_id is not None: + body["last_seen_run_id"] = self.last_seen_run_id + if self.latest_non_matching_result is not None: + body["latest_non_matching_result"] = self.latest_non_matching_result + if self.workspace_id is not None: + body["workspace_id"] = self.workspace_id + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> JobCompletionTriggerStateMonitoredJobState: + """Deserializes the JobCompletionTriggerStateMonitoredJobState from a dictionary.""" + return cls( + has_seen_completion=d.get("has_seen_completion", None), + is_deleted=d.get("is_deleted", None), + job_id=d.get("job_id", None), + last_seen_completion_time=_timestamp(d, "last_seen_completion_time"), + last_seen_result=_enum(d, "last_seen_result", TerminationCodeCode), + last_seen_run_id=d.get("last_seen_run_id", None), + latest_non_matching_result=_enum(d, "latest_non_matching_result", TerminationCodeCode), + workspace_id=d.get("workspace_id", None), + ) + + @dataclass class JobCompliance: job_id: int @@ -3825,10 +4618,11 @@ class JobSettings: serverless compute or a compute resource that uses Environments mode. For notebook tasks that use serverless compute or a compute resource that uses Environments - mode, if the environment_key is not specified, the notebook environment will be used if present. - If a jobs environment is specified, it will override the notebook environment. For other tasks - that use serverless compute or a compute resource that uses Environments mode, the task - environment is required to be specified using environment_key in the task settings.""" + mode, if the ``environment_key`` is not specified, the notebook environment will be used if + present. If a jobs environment is specified, it will override the notebook environment. + + For other tasks that use serverless compute or a compute resource that uses Environments mode, + the task environment is required to be specified using ``environment_key`` in the task settings.""" format: Optional[Format] = None """Used to tell what is the format of the job. This field is ignored in Create/Update/Reset calls. @@ -4719,6 +5513,8 @@ class PerTriggerState: file_arrival: Optional[FileArrivalTriggerState] = None + job_completion: Optional[JobCompletionTriggerState] = None + model: Optional[ModelTriggerState] = None pause_status: Optional[PauseStatus] = None @@ -4740,6 +5536,8 @@ def as_dict(self) -> dict: body["continuous"] = self.continuous.as_dict() if self.file_arrival: body["file_arrival"] = self.file_arrival.as_dict() + if self.job_completion: + body["job_completion"] = self.job_completion.as_dict() if self.model: body["model"] = self.model.as_dict() if self.pause_status is not None: @@ -4761,6 +5559,8 @@ def as_shallow_dict(self) -> dict: body["continuous"] = self.continuous if self.file_arrival: body["file_arrival"] = self.file_arrival + if self.job_completion: + body["job_completion"] = self.job_completion if self.model: body["model"] = self.model if self.pause_status is not None: @@ -4781,6 +5581,7 @@ def from_dict(cls, d: Dict[str, Any]) -> PerTriggerState: return cls( continuous=_from_dict(d, "continuous", ContinuousTriggerState), file_arrival=_from_dict(d, "file_arrival", FileArrivalTriggerState), + job_completion=_from_dict(d, "job_completion", JobCompletionTriggerState), model=_from_dict(d, "model", ModelTriggerState), pause_status=_enum(d, "pause_status", PauseStatus), periodic=_from_dict(d, "periodic", PeriodicTriggerState), @@ -5285,6 +6086,47 @@ def from_dict(cls, d: Dict[str, Any]) -> PythonWheelTask: ) +@dataclass +class QualifiedName: + """A three-level Unity Catalog name (catalog.schema.table); any level may be unset.""" + + catalog: Optional[str] = None + """Unity Catalog catalog.""" + + schema: Optional[str] = None + """Unity Catalog schema, within the catalog.""" + + table: Optional[str] = None + """Table, within the catalog and schema.""" + + def as_dict(self) -> dict: + """Serializes the QualifiedName into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.catalog is not None: + body["catalog"] = self.catalog + if self.schema is not None: + body["schema"] = self.schema + if self.table is not None: + body["table"] = self.table + return body + + def as_shallow_dict(self) -> dict: + """Serializes the QualifiedName into a shallow dictionary of its immediate attributes.""" + body = {} + if self.catalog is not None: + body["catalog"] = self.catalog + if self.schema is not None: + body["schema"] = self.schema + if self.table is not None: + body["table"] = self.table + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> QualifiedName: + """Deserializes the QualifiedName from a dictionary.""" + return cls(catalog=d.get("catalog", None), schema=d.get("schema", None), table=d.get("table", None)) + + @dataclass class QueueDetails: code: Optional[QueueDetailsCodeCode] = None @@ -5696,6 +6538,10 @@ def from_dict(cls, d: Dict[str, Any]) -> ResolvedStringParamsValues: @dataclass class ResolvedValues: + ai_function_task: Optional[ResolvedValuesAiFunctionTaskResolvedValues] = None + """Resolved values for an AI Function task: the Delta table input with upstream-task references + replaced by concrete values before the run.""" + ai_runtime_task: Optional[ResolvedValuesAiRuntimeTaskResolvedValues] = None """Resolved values for an AI Runtime task — env_vars with ``{{tasks..values.}}`` references substituted to concrete values before submission to the training service.""" @@ -5725,6 +6571,8 @@ class ResolvedValues: def as_dict(self) -> dict: """Serializes the ResolvedValues into a dictionary suitable for use as a JSON request body.""" body = {} + if self.ai_function_task: + body["ai_function_task"] = self.ai_function_task.as_dict() if self.ai_runtime_task: body["ai_runtime_task"] = self.ai_runtime_task.as_dict() if self.condition_task: @@ -5754,6 +6602,8 @@ def as_dict(self) -> dict: def as_shallow_dict(self) -> dict: """Serializes the ResolvedValues into a shallow dictionary of its immediate attributes.""" body = {} + if self.ai_function_task: + body["ai_function_task"] = self.ai_function_task if self.ai_runtime_task: body["ai_runtime_task"] = self.ai_runtime_task if self.condition_task: @@ -5784,6 +6634,7 @@ def as_shallow_dict(self) -> dict: def from_dict(cls, d: Dict[str, Any]) -> ResolvedValues: """Deserializes the ResolvedValues from a dictionary.""" return cls( + ai_function_task=_from_dict(d, "ai_function_task", ResolvedValuesAiFunctionTaskResolvedValues), ai_runtime_task=_from_dict(d, "ai_runtime_task", ResolvedValuesAiRuntimeTaskResolvedValues), condition_task=_from_dict(d, "condition_task", ResolvedConditionTaskValues), dbt_task=_from_dict(d, "dbt_task", ResolvedDbtTaskValues), @@ -5799,6 +6650,26 @@ def from_dict(cls, d: Dict[str, Any]) -> ResolvedValues: ) +@dataclass +class ResolvedValuesAiFunctionTaskResolvedValues: + """Resolved values for an AI Function task's Delta table input.""" + + def as_dict(self) -> dict: + """Serializes the ResolvedValuesAiFunctionTaskResolvedValues into a dictionary suitable for use as a JSON request body.""" + body = {} + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ResolvedValuesAiFunctionTaskResolvedValues into a shallow dictionary of its immediate attributes.""" + body = {} + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ResolvedValuesAiFunctionTaskResolvedValues: + """Deserializes the ResolvedValuesAiFunctionTaskResolvedValues from a dictionary.""" + return cls() + + @dataclass class ResolvedValuesAiRuntimeTaskResolvedValues: """Resolved values for an AiRuntimeTask after dynamic-value substitution, so Jobs can expand @@ -6610,6 +7481,9 @@ def from_dict(cls, d: Dict[str, Any]) -> RunNowResponse: class RunOutput: """Run output was retrieved successfully.""" + ai_function_task_output: Optional[AiFunctionTaskOutput] = None + """The output of an AI function task, if available""" + ai_runtime_task_output: Optional[AiRuntimeTaskOutput] = None """The output of an AiRuntimeTask, if available — MLflow identifiers, artifact paths, and per-replica allocated compute. Run lifecycle / termination status lives on the surrounding @@ -6673,6 +7547,8 @@ class RunOutput: def as_dict(self) -> dict: """Serializes the RunOutput into a dictionary suitable for use as a JSON request body.""" body = {} + if self.ai_function_task_output: + body["ai_function_task_output"] = self.ai_function_task_output.as_dict() if self.ai_runtime_task_output: body["ai_runtime_task_output"] = self.ai_runtime_task_output.as_dict() if self.alert_output: @@ -6710,6 +7586,8 @@ def as_dict(self) -> dict: def as_shallow_dict(self) -> dict: """Serializes the RunOutput into a shallow dictionary of its immediate attributes.""" body = {} + if self.ai_function_task_output: + body["ai_function_task_output"] = self.ai_function_task_output if self.ai_runtime_task_output: body["ai_runtime_task_output"] = self.ai_runtime_task_output if self.alert_output: @@ -6748,6 +7626,7 @@ def as_shallow_dict(self) -> dict: def from_dict(cls, d: Dict[str, Any]) -> RunOutput: """Deserializes the RunOutput from a dictionary.""" return cls( + ai_function_task_output=_from_dict(d, "ai_function_task_output", AiFunctionTaskOutput), ai_runtime_task_output=_from_dict(d, "ai_runtime_task_output", AiRuntimeTaskOutput), alert_output=_from_dict(d, "alert_output", AlertTaskOutput), clean_rooms_notebook_output=_from_dict( @@ -7056,6 +7935,9 @@ class RunTask: field is required and must be unique within its parent job. On Update or Reset, this field is used to reference the tasks to be updated or reset.""" + ai_function_task: Optional[AiFunctionTask] = None + """Applies an AI function to a set of input documents or rows.""" + ai_runtime_task: Optional[AiRuntimeTask] = None """The task runs a multi-gpu compute workload on Databricks AI Runtime. Specify the accelerator type and count, the command to run, and where the workload's code and MLflow output are stored.""" @@ -7288,6 +8170,8 @@ class RunTask: def as_dict(self) -> dict: """Serializes the RunTask into a dictionary suitable for use as a JSON request body.""" body = {} + if self.ai_function_task: + body["ai_function_task"] = self.ai_function_task.as_dict() if self.ai_runtime_task: body["ai_runtime_task"] = self.ai_runtime_task.as_dict() if self.alert_task: @@ -7405,6 +8289,8 @@ def as_dict(self) -> dict: def as_shallow_dict(self) -> dict: """Serializes the RunTask into a shallow dictionary of its immediate attributes.""" body = {} + if self.ai_function_task: + body["ai_function_task"] = self.ai_function_task if self.ai_runtime_task: body["ai_runtime_task"] = self.ai_runtime_task if self.alert_task: @@ -7523,6 +8409,7 @@ def as_shallow_dict(self) -> dict: def from_dict(cls, d: Dict[str, Any]) -> RunTask: """Deserializes the RunTask from a dictionary.""" return cls( + ai_function_task=_from_dict(d, "ai_function_task", AiFunctionTask), ai_runtime_task=_from_dict(d, "ai_runtime_task", AiRuntimeTask), alert_task=_from_dict(d, "alert_task", AlertTask), attempt_number=d.get("attempt_number", None), @@ -8775,6 +9662,9 @@ class SubmitTask: field is required and must be unique within its parent job. On Update or Reset, this field is used to reference the tasks to be updated or reset.""" + ai_function_task: Optional[AiFunctionTask] = None + """Applies an AI function to a set of input documents or rows.""" + ai_runtime_task: Optional[AiRuntimeTask] = None """The task runs a multi-gpu compute workload on Databricks AI Runtime. Specify the accelerator type and count, the command to run, and where the workload's code and MLflow output are stored.""" @@ -8926,6 +9816,8 @@ class SubmitTask: def as_dict(self) -> dict: """Serializes the SubmitTask into a dictionary suitable for use as a JSON request body.""" body = {} + if self.ai_function_task: + body["ai_function_task"] = self.ai_function_task.as_dict() if self.ai_runtime_task: body["ai_runtime_task"] = self.ai_runtime_task.as_dict() if self.alert_task: @@ -9011,6 +9903,8 @@ def as_dict(self) -> dict: def as_shallow_dict(self) -> dict: """Serializes the SubmitTask into a shallow dictionary of its immediate attributes.""" body = {} + if self.ai_function_task: + body["ai_function_task"] = self.ai_function_task if self.ai_runtime_task: body["ai_runtime_task"] = self.ai_runtime_task if self.alert_task: @@ -9097,6 +9991,7 @@ def as_shallow_dict(self) -> dict: def from_dict(cls, d: Dict[str, Any]) -> SubmitTask: """Deserializes the SubmitTask from a dictionary.""" return cls( + ai_function_task=_from_dict(d, "ai_function_task", AiFunctionTask), ai_runtime_task=_from_dict(d, "ai_runtime_task", AiRuntimeTask), alert_task=_from_dict(d, "alert_task", AlertTask), clean_rooms_notebook_task=_from_dict(d, "clean_rooms_notebook_task", CleanRoomsNotebookTask), @@ -9546,6 +10441,9 @@ class Task: field is required and must be unique within its parent job. On Update or Reset, this field is used to reference the tasks to be updated or reset.""" + ai_function_task: Optional[AiFunctionTask] = None + """Applies an AI function to a set of input documents or rows.""" + ai_runtime_task: Optional[AiRuntimeTask] = None """The task runs a multi-gpu compute workload on Databricks AI Runtime. Specify the accelerator type and count, the command to run, and where the workload's code and MLflow output are stored.""" @@ -9705,6 +10603,8 @@ class Task: def as_dict(self) -> dict: """Serializes the Task into a dictionary suitable for use as a JSON request body.""" body = {} + if self.ai_function_task: + body["ai_function_task"] = self.ai_function_task.as_dict() if self.ai_runtime_task: body["ai_runtime_task"] = self.ai_runtime_task.as_dict() if self.alert_task: @@ -9792,6 +10692,8 @@ def as_dict(self) -> dict: def as_shallow_dict(self) -> dict: """Serializes the Task into a shallow dictionary of its immediate attributes.""" body = {} + if self.ai_function_task: + body["ai_function_task"] = self.ai_function_task if self.ai_runtime_task: body["ai_runtime_task"] = self.ai_runtime_task if self.alert_task: @@ -9880,6 +10782,7 @@ def as_shallow_dict(self) -> dict: def from_dict(cls, d: Dict[str, Any]) -> Task: """Deserializes the Task from a dictionary.""" return cls( + ai_function_task=_from_dict(d, "ai_function_task", AiFunctionTask), ai_runtime_task=_from_dict(d, "ai_runtime_task", AiRuntimeTask), alert_task=_from_dict(d, "alert_task", AlertTask), clean_rooms_notebook_task=_from_dict(d, "clean_rooms_notebook_task", CleanRoomsNotebookTask), @@ -10273,6 +11176,9 @@ class TriggerConfiguration: file_arrival: Optional[FileArrivalTriggerConfiguration] = None """File arrival trigger configuration.""" + job_completion: Optional[JobCompletionTriggerConfiguration] = None + """Job completion trigger configuration.""" + model: Optional[ModelTriggerConfiguration] = None """Model trigger configuration.""" @@ -10300,6 +11206,8 @@ def as_dict(self) -> dict: body["continuous"] = self.continuous.as_dict() if self.file_arrival: body["file_arrival"] = self.file_arrival.as_dict() + if self.job_completion: + body["job_completion"] = self.job_completion.as_dict() if self.model: body["model"] = self.model.as_dict() if self.pause_status is not None: @@ -10321,6 +11229,8 @@ def as_shallow_dict(self) -> dict: body["continuous"] = self.continuous if self.file_arrival: body["file_arrival"] = self.file_arrival + if self.job_completion: + body["job_completion"] = self.job_completion if self.model: body["model"] = self.model if self.pause_status is not None: @@ -10341,6 +11251,7 @@ def from_dict(cls, d: Dict[str, Any]) -> TriggerConfiguration: return cls( continuous=_from_dict(d, "continuous", ContinuousTriggerConfiguration), file_arrival=_from_dict(d, "file_arrival", FileArrivalTriggerConfiguration), + job_completion=_from_dict(d, "job_completion", JobCompletionTriggerConfiguration), model=_from_dict(d, "model", ModelTriggerConfiguration), pause_status=_enum(d, "pause_status", PauseStatus), periodic=_from_dict(d, "periodic", PeriodicTriggerConfiguration), @@ -10474,6 +11385,10 @@ def from_dict(cls, d: Dict[str, Any]) -> TriggerHistory: class TriggerInfo: """Additional details about what triggered the run""" + job_completion: Optional[JobCompletionTriggerRunInfoDetails] = None + """Additional details for job runs triggered by a job completion trigger (trigger: + "JOB_COMPLETION")""" + run_id: Optional[int] = None """The run id of the Run Job task run""" @@ -10483,6 +11398,8 @@ class TriggerInfo: def as_dict(self) -> dict: """Serializes the TriggerInfo into a dictionary suitable for use as a JSON request body.""" body = {} + if self.job_completion: + body["job_completion"] = self.job_completion.as_dict() if self.run_id is not None: body["run_id"] = self.run_id if self.sql_condition: @@ -10492,6 +11409,8 @@ def as_dict(self) -> dict: def as_shallow_dict(self) -> dict: """Serializes the TriggerInfo into a shallow dictionary of its immediate attributes.""" body = {} + if self.job_completion: + body["job_completion"] = self.job_completion if self.run_id is not None: body["run_id"] = self.run_id if self.sql_condition: @@ -10501,7 +11420,11 @@ def as_shallow_dict(self) -> dict: @classmethod def from_dict(cls, d: Dict[str, Any]) -> TriggerInfo: """Deserializes the TriggerInfo from a dictionary.""" - return cls(run_id=_int64(d, "run_id"), sql_condition=_from_dict(d, "sql_condition", SqlConditionRunInfoDetails)) + return cls( + job_completion=_from_dict(d, "job_completion", JobCompletionTriggerRunInfoDetails), + run_id=_int64(d, "run_id"), + sql_condition=_from_dict(d, "sql_condition", SqlConditionRunInfoDetails), + ) @dataclass @@ -10509,6 +11432,9 @@ class TriggerSettings: file_arrival: Optional[FileArrivalTriggerConfiguration] = None """File arrival trigger settings.""" + job_completion: Optional[JobCompletionTriggerConfiguration] = None + """Job completion trigger settings.""" + model: Optional[ModelTriggerConfiguration] = None pause_status: Optional[PauseStatus] = None @@ -10528,6 +11454,8 @@ def as_dict(self) -> dict: body = {} if self.file_arrival: body["file_arrival"] = self.file_arrival.as_dict() + if self.job_completion: + body["job_completion"] = self.job_completion.as_dict() if self.model: body["model"] = self.model.as_dict() if self.pause_status is not None: @@ -10545,6 +11473,8 @@ def as_shallow_dict(self) -> dict: body = {} if self.file_arrival: body["file_arrival"] = self.file_arrival + if self.job_completion: + body["job_completion"] = self.job_completion if self.model: body["model"] = self.model if self.pause_status is not None: @@ -10562,6 +11492,7 @@ def from_dict(cls, d: Dict[str, Any]) -> TriggerSettings: """Deserializes the TriggerSettings from a dictionary.""" return cls( file_arrival=_from_dict(d, "file_arrival", FileArrivalTriggerConfiguration), + job_completion=_from_dict(d, "job_completion", JobCompletionTriggerConfiguration), model=_from_dict(d, "model", ModelTriggerConfiguration), pause_status=_enum(d, "pause_status", PauseStatus), periodic=_from_dict(d, "periodic", PeriodicTriggerConfiguration), @@ -10574,6 +11505,8 @@ def from_dict(cls, d: Dict[str, Any]) -> TriggerSettings: class TriggerStateProto: file_arrival: Optional[FileArrivalTriggerState] = None + job_completion: Optional[JobCompletionTriggerState] = None + pause_status: Optional[PauseStatus] = None """Whether this trigger is paused or not. For continuous schedules, it can differ from the configured pause_status whenever a paused continuous job is kickstarted by an operation other @@ -10589,6 +11522,8 @@ def as_dict(self) -> dict: body = {} if self.file_arrival: body["file_arrival"] = self.file_arrival.as_dict() + if self.job_completion: + body["job_completion"] = self.job_completion.as_dict() if self.pause_status is not None: body["pause_status"] = self.pause_status.value if self.sql_condition: @@ -10602,6 +11537,8 @@ def as_shallow_dict(self) -> dict: body = {} if self.file_arrival: body["file_arrival"] = self.file_arrival + if self.job_completion: + body["job_completion"] = self.job_completion if self.pause_status is not None: body["pause_status"] = self.pause_status if self.sql_condition: @@ -10615,6 +11552,7 @@ def from_dict(cls, d: Dict[str, Any]) -> TriggerStateProto: """Deserializes the TriggerStateProto from a dictionary.""" return cls( file_arrival=_from_dict(d, "file_arrival", FileArrivalTriggerState), + job_completion=_from_dict(d, "job_completion", JobCompletionTriggerState), pause_status=_enum(d, "pause_status", PauseStatus), sql_condition=_from_dict(d, "sql_condition", SqlConditionState), table=_from_dict(d, "table", TableTriggerState), @@ -10641,6 +11579,7 @@ class TriggerType(Enum): CONTINUOUS = "CONTINUOUS" CONTINUOUS_RESTART = "CONTINUOUS_RESTART" FILE_ARRIVAL = "FILE_ARRIVAL" + JOB_COMPLETION = "JOB_COMPLETION" ONE_TIME = "ONE_TIME" PERIODIC = "PERIODIC" RETRY = "RETRY" @@ -10706,6 +11645,34 @@ class ViewsToExport(Enum): DASHBOARDS = "DASHBOARDS" +@dataclass +class VolumeSource: + """A Unity Catalog volume input source. The function is applied to each file.""" + + volume_path: str + """Unity Catalog volume path to the input files (e.g. + ``/Volumes////...``).""" + + def as_dict(self) -> dict: + """Serializes the VolumeSource into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.volume_path is not None: + body["volume_path"] = self.volume_path + return body + + def as_shallow_dict(self) -> dict: + """Serializes the VolumeSource into a shallow dictionary of its immediate attributes.""" + body = {} + if self.volume_path is not None: + body["volume_path"] = self.volume_path + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> VolumeSource: + """Deserializes the VolumeSource from a dictionary.""" + return cls(volume_path=d.get("volume_path", None)) + + @dataclass class Webhook: id: str @@ -11031,10 +11998,11 @@ def create( serverless compute or a compute resource that uses Environments mode. For notebook tasks that use serverless compute or a compute resource that uses Environments mode, if - the environment_key is not specified, the notebook environment will be used if present. If a jobs - environment is specified, it will override the notebook environment. For other tasks that use - serverless compute or a compute resource that uses Environments mode, the task environment is - required to be specified using environment_key in the task settings. + the ``environment_key`` is not specified, the notebook environment will be used if present. If a + jobs environment is specified, it will override the notebook environment. + + For other tasks that use serverless compute or a compute resource that uses Environments mode, the + task environment is required to be specified using ``environment_key`` in the task settings. :param format: :class:`Format` (optional) Used to tell what is the format of the job. This field is ignored in Create/Update/Reset calls. When using the Jobs API 2.1 this value is always set to ``"MULTI_TASK"``. diff --git a/databricks/sdk/service/ml.py b/databricks/sdk/service/ml.py index 8f5015786..45331790c 100644 --- a/databricks/sdk/service/ml.py +++ b/databricks/sdk/service/ml.py @@ -559,7 +559,9 @@ class BackfillOperationMetadata: """Progress and configuration for a backfill.""" backfill_ranges: Optional[List[BackfillRange]] = None - """Output ranges targeted by the backfill.""" + """The requested time ranges, each aligned outward to the boundaries at which the features emit + values. Ranges that overlap after alignment are combined into one. The backfill stops at the + most recent materialized values, so a range may extend past the values that it recomputed.""" feature_full_names: Optional[List[str]] = None """Full names of the features targeted by the backfill.""" @@ -611,10 +613,14 @@ class BackfillOperationMetadataState(Enum): @dataclass class BackfillRange: - """A time range for a backfill.""" + """A time range of source data for a backfill. The range is aligned outward to the boundaries at + which the feature emits values: every slide_duration for a sliding window or every + window_duration for a tumbling window, shifted by offset. Feature values up to one + window_duration after the aligned end may change. The backfill stops at the most recent + materialized values and leaves later values to materialization.""" end_time: Optional[Timestamp] = None - """End of the backfill range, exclusive. If unset, defaults to the current time.""" + """End of the backfill range, exclusive. If unset, defaults to the most recent materialized values.""" start_time: Optional[Timestamp] = None """Start of the backfill range, inclusive. If unset, defaults to the earliest source timestamp of @@ -5386,7 +5392,8 @@ def from_dict(cls, d: Dict[str, Any]) -> MaterializationFailureNotification: @dataclass class MaterializationNotifications: - """Notifications for the jobs and pipelines that materialize a feature, one field per trigger.""" + """Notifications for the jobs and pipelines that materialize the features in a set, one field per + trigger.""" on_failure: Optional[MaterializationFailureNotification] = None """Who to notify when a run fails, and on which attempts.""" @@ -5696,16 +5703,14 @@ class MaterializedFeature: operations/{operation_id}.""" materialization_set_id: Optional[str] = None - """The materialization set this materialized feature was created for, through DABs or direct-engine - creation. Unset for a materialized feature that was not created as part of a set.""" + """The ID of the materialization set containing this materialized feature. Unset for a materialized + feature that is not part of a set.""" materialized_feature_id: Optional[str] = None """Server-assigned unique identifier for the materialized feature.""" notifications: Optional[MaterializationNotifications] = None - """Notification configuration around the materialization job or pipeline lifecycle. They are - applied to every job and pipeline that materializes this feature. Features which are - materialized in the same pipeline will share the same notification.""" + """Notification configuration for the jobs and pipelines that materialize this feature.""" offline_store_config: Optional[OfflineStoreConfig] = None """Destination for writing feature values to an offline Delta table. The resulting table is @@ -5720,8 +5725,7 @@ class MaterializedFeature: features.""" pipeline_schedule_state: Optional[MaterializedFeaturePipelineScheduleState] = None - """The schedule state of the materialization pipeline. Hidden from GraphQL: being deprecated, so - not exposed to Catalog Explorer.""" + """Deprecated: This legacy pipeline setting does not control batch job scheduling.""" streaming_mode: Optional[StreamingMode] = None """The Structured Streaming trigger mode used for materialization. Real-time mode (RTM) targets @@ -11881,9 +11885,11 @@ def backfill_features( """Backfill features. :param feature_full_names: List[str] - Full names of the features to backfill. + Full names of the features to backfill. All features must share the same interval between values and + the same offset. The interval is the slide_duration of a sliding window or the window_duration of a + tumbling window. :param backfill_ranges: List[:class:`BackfillRange`] - Output ranges to backfill. + Time ranges of source data to backfill. Ranges must not overlap. :param budget_policy_id: str (optional) The budget policy ID, in UUID format, used to attribute the serverless compute cost of this backfill. If not specified, a default budget policy may be applied. @@ -12715,15 +12721,14 @@ def update_materialization_set( def update_materialized_feature( self, materialized_feature_id: str, materialized_feature: MaterializedFeature, update_mask: str ) -> MaterializedFeature: - """Update a materialized feature (pause/resume). + """Updates a materialized feature. :param materialized_feature_id: str Server-assigned unique identifier for the materialized feature. :param materialized_feature: :class:`MaterializedFeature` The materialized feature to update. :param update_mask: str - Provide the materialization feature fields which should be updated. Currently, only the - pipeline_state field can be updated. + The fields to update. Only ``pipeline_schedule_state`` is supported. :returns: :class:`MaterializedFeature` """ diff --git a/databricks/sdk/service/postgres.py b/databricks/sdk/service/postgres.py index 09d90abbf..24039e8dd 100644 --- a/databricks/sdk/service/postgres.py +++ b/databricks/sdk/service/postgres.py @@ -3270,14 +3270,22 @@ class RoleRoleSpec: - For the regular postgres roles, authentication based on postgres passwords is used. NOTE: for the Databricks identity type GROUP, LAKEBASE_OAUTH_V1 is the default auth method - (group can login as well).""" + (group can login as well). + + This field is set when the role is created and cannot be changed by UpdateRole. If provided in + an UpdateRole request body, its value is ignored and the existing value is retained. Do not + include this field as an explicit path in update_mask.""" identity_type: Optional[RoleIdentityType] = None """The type of role. When specifying a managed-identity, the chosen role_id must be a valid: - application ID for SERVICE_PRINCIPAL - user email for USER - - group name for GROUP""" + - group name for GROUP + + This field is set when the role is created and cannot be changed by UpdateRole. If provided in + an UpdateRole request body, its value is ignored and the existing value is retained. Do not + include this field as an explicit path in update_mask.""" membership_roles: Optional[List[RoleMembershipRole]] = None """An enum value for a standard role that this role is a member of.""" @@ -3295,7 +3303,11 @@ class RoleRoleSpec: 1. user email for IdentityType.USER 2. app ID for IdentityType.SERVICE_PRINCIPAL - 3. group name for IdentityType.GROUP""" + 3. group name for IdentityType.GROUP + + This field is set when the role is created and cannot be changed by UpdateRole. If provided in + an UpdateRole request body, its value is ignored and the existing value is retained. Do not + include this field as an explicit path in update_mask.""" def as_dict(self) -> dict: """Serializes the RoleRoleSpec into a dictionary suitable for use as a JSON request body.""" @@ -4632,12 +4644,16 @@ def create_role( :param role: :class:`Role` The desired specification of a Role. :param replace_existing: bool (optional) - If true, update the role if it already exists instead of returning an error. - - When the role already exists, the provided ``role`` spec fully replaces the existing one: - ``membership_roles`` is overwritten, not merged. Leaving ``membership_roles`` empty clears all of - the role's existing memberships, including ``DATABRICKS_SUPERUSER``. Always send the complete - desired list of memberships when using this field. + If true, update the role if it already exists instead of returning an error. This requires + ``role_id`` to be set. Without ``role_id``, the request follows the normal create path and fails if + the role name already exists. + + When the role already exists, only mutable fields are updated. Attributes omitted from the request + retain their current values. Set an attribute explicitly to false to disable it. ``identity_type``, + ``auth_method``, and ``postgres_role`` retain their existing values; supplied values for these + fields are ignored. ``membership_roles`` is overwritten, not merged. Leaving ``membership_roles`` + empty clears all of the role's existing memberships, including ``DATABRICKS_SUPERUSER``. Always send + the complete desired list of memberships when using this field. :param role_id: str (optional) The ID to use for the Role, which will become the final component of the role's resource name. This ID becomes the role in Postgres. @@ -5851,7 +5867,12 @@ def update_project(self, name: str, project: Project, update_mask: FieldMask) -> return UpdateProjectOperation(self, operation) def update_role(self, name: str, role: Role, update_mask: FieldMask) -> UpdateRoleOperation: - """Update a role for a branch. + """Updates the mutable fields of a Postgres role for a branch. + + The fields ``spec.identity_type``, ``spec.auth_method``, and ``spec.postgres_role`` are fixed when the + role is created. UpdateRole ignores values supplied for these fields in the request body and retains + their existing values. Explicitly naming one of these fields in ``update_mask`` returns an + INVALID_PARAMETER_VALUE error. :param name: str Output only. The full resource path of the role. Format: @@ -5862,7 +5883,17 @@ def update_role(self, name: str, role: Role, update_mask: FieldMask) -> UpdateRo The role's ``name`` field is used to identify the role to update. Format: projects/{project_id}/branches/{branch_id}/roles/{role_id} :param update_mask: FieldMask - The list of fields to update. + The list of fields to update. Supported fields are ``spec.membership_roles``, + ``spec.attributes.createdb``, ``spec.attributes.createrole``, and ``spec.attributes.bypassrls``. + + ``spec.identity_type``, ``spec.auth_method``, and ``spec.postgres_role`` cannot be changed. Values + supplied for these fields in the request body are ignored. Including one of these fields as an + explicit path in update_mask returns an INVALID_PARAMETER_VALUE error. When ``spec`` is the only + path in the mask, only the supplied mutable attributes are updated. ``*`` by itself, or ``spec`` + combined with other supported paths, selects all supported mutable fields. For these masks, all + three attributes (``createdb``, ``createrole``, and ``bypassrls``) must be supplied; omitting any + returns an INVALID_PARAMETER_VALUE error. Each of these mask forms replaces ``membership_roles``. + Omitting it or supplying an empty list clears all memberships, including ``DATABRICKS_SUPERUSER``. :returns: :class:`Operation` """ diff --git a/databricks/sdk/service/sandbox.py b/databricks/sdk/service/sandbox.py index 1d4c4472f..8019d672c 100644 --- a/databricks/sdk/service/sandbox.py +++ b/databricks/sdk/service/sandbox.py @@ -6,7 +6,7 @@ from __future__ import annotations from dataclasses import dataclass from enum import Enum -from typing import Dict, List, Any, Iterator, Optional +from typing import Dict, List, Any, Iterator, Optional, BinaryIO from google.protobuf.duration_pb2 import Duration from google.protobuf.timestamp_pb2 import Timestamp @@ -165,6 +165,80 @@ def from_dict(cls, d: Dict[str, Any]) -> DeletePathResponse: return cls() +@dataclass +class DownloadFileResponse: + accept_ranges: Optional[str] = None + """The supported range unit: bytes.""" + + content_length: Optional[int] = None + """The length of the response body in bytes.""" + + content_range: Optional[str] = None + """The returned byte range and full file length, for a partial response.""" + + content_type: Optional[str] = None + """The media type of the response body: application/octet-stream.""" + + contents: Optional[BinaryIO] = None + + etag: Optional[str] = None + """The entity tag used for conditional requests.""" + + last_modified: Optional[str] = None + """The modification time as an HTTP-date, when representable.""" + + def as_dict(self) -> dict: + """Serializes the DownloadFileResponse into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.accept_ranges is not None: + body["Accept-Ranges"] = self.accept_ranges + if self.content_length is not None: + body["Content-Length"] = self.content_length + if self.content_range is not None: + body["Content-Range"] = self.content_range + if self.content_type is not None: + body["Content-Type"] = self.content_type + if self.contents: + body["contents"] = self.contents + if self.etag is not None: + body["etag"] = self.etag + if self.last_modified is not None: + body["Last-Modified"] = self.last_modified + return body + + def as_shallow_dict(self) -> dict: + """Serializes the DownloadFileResponse into a shallow dictionary of its immediate attributes.""" + body = {} + if self.accept_ranges is not None: + body["Accept-Ranges"] = self.accept_ranges + if self.content_length is not None: + body["Content-Length"] = self.content_length + if self.content_range is not None: + body["Content-Range"] = self.content_range + if self.content_type is not None: + body["Content-Type"] = self.content_type + if self.contents: + body["contents"] = self.contents + if self.etag is not None: + body["etag"] = self.etag + if self.last_modified is not None: + body["Last-Modified"] = self.last_modified + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> DownloadFileResponse: + """Deserializes the DownloadFileResponse from a dictionary.""" + return cls( + accept_ranges=d.get("Accept-Ranges", None), + content_length=_int64(d, "Content-Length"), + content_range=d.get("Content-Range", None), + content_type=d.get("Content-Type", None), + contents=d.get("contents", None), + etag=d.get("etag", None), + last_modified=d.get("Last-Modified", None), + ) + + @dataclass class EnvironmentSpec: image_uri: Optional[str] = None @@ -810,6 +884,99 @@ def delete_sandbox(self, name: str): self._api.do("DELETE", f"/api/2.0/{name}", headers=headers) + def download_file( + self, + parent: str, + path: str, + *, + byte_range: Optional[str] = None, + if_match: Optional[str] = None, + if_modified_since: Optional[str] = None, + if_none_match: Optional[str] = None, + if_range: Optional[str] = None, + if_unmodified_since: Optional[str] = None, + ) -> DownloadFileResponse: + """Downloads a regular file from a sandbox as an application/octet-stream body. The sandbox starts + automatically if it is stopped. Filesystem access uses the sandbox user's permissions and follows + symbolic links. + + A full download returns 200 OK. A satisfiable single byte range returns 206 Partial Content. Multiple + byte ranges are unsupported and return 400. Invalid ranges are ignored; an unsatisfiable range returns + 416 with Content-Range describing the file length. Range is ignored for empty files. + + Conditional requests can return 304 Not Modified or 412 Precondition Failed. If-Range selects a + partial response only when its strong entity tag matches; otherwise, the entire file is returned. + HTTP-date validators in If-Range cause a full download. Reads are not snapshots: concurrent in-place + writes can affect the returned bytes. HEAD is unsupported. + + Range and conditional options may be supplied as query parameters or HTTP headers. A parameter and one + header may both be supplied when their values match exactly. Conflicting values or repeated parameters + return 400. + + :param parent: str + The sandbox resource name, in the form sandboxes/{sandbox_id}. + :param path: str + The absolute UTF-8 path of the file inside the sandbox. + :param byte_range: str (optional) + A single byte range, such as bytes=0-1023 or bytes=-1024. Equivalent to the Range header. + :param if_match: str (optional) + Entity tags or *. A mismatch returns 412 Precondition Failed. Equivalent to the If-Match header. + :param if_modified_since: str (optional) + An HTTP-date. Returns 304 if the file has not changed since this time. Ignored when If-None-Match is + present. Equivalent to the If-Modified-Since header. + :param if_none_match: str (optional) + Entity tags or *. A match returns 304 Not Modified. Equivalent to the If-None-Match header. + :param if_range: str (optional) + A strong entity tag. A mismatch, weak entity tag, or HTTP-date causes a full download. Equivalent to + the If-Range header. + :param if_unmodified_since: str (optional) + An HTTP-date. Returns 412 if the file has changed since this time. Ignored when If-Match is present. + Equivalent to the If-Unmodified-Since header. + + :returns: :class:`DownloadFileResponse` + """ + + query = {} + if byte_range is not None: + query["byte_range"] = byte_range + if if_match is not None: + query["if_match"] = if_match + if if_modified_since is not None: + query["if_modified_since"] = if_modified_since + if if_none_match is not None: + query["if_none_match"] = if_none_match + if if_range is not None: + query["if_range"] = if_range + if if_unmodified_since is not None: + query["if_unmodified_since"] = if_unmodified_since + if path is not None: + query["path"] = path + headers = { + "Accept": "application/octet-stream", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + response_headers = [ + "Accept-Ranges", + "Content-Length", + "Content-Range", + "Content-Type", + "etag", + "Last-Modified", + ] + res = self._api.do( + "GET", + f"/api/2.0/sandbox-files/{parent}/content", + query=query, + headers=headers, + response_headers=response_headers, + raw=True, + ) + return DownloadFileResponse.from_dict(res) + def execute_command_sync( self, name: str, @@ -1119,3 +1286,38 @@ def update_sandbox(self, name: str, sandbox: Sandbox, update_mask: FieldMask) -> res = self._api.do("PATCH", f"/api/2.0/{name}", query=query, body=body, headers=headers) return Sandbox.from_dict(res) + + def upload_file(self, parent: str, contents: BinaryIO, path: str): + """Uploads an application/octet-stream body to a sandbox. The sandbox starts automatically if it is + stopped. Filesystem access uses the sandbox user's permissions and follows symbolic links. + + Creates the file or replaces its entire contents, creating parent directories if missing. The + completed upload becomes visible atomically. A failure before commit leaves an existing destination + unchanged; a lost response after commit does not undo the upload. Completion does not guarantee crash + durability. + + Returns 201 Created for a new file or 204 No Content for a replacement. Partial and conditional + uploads are unsupported: Content-Range, If-Match, If-None-Match, and If-Unmodified-Since return 400. + + :param parent: str + The sandbox resource name, in the form sandboxes/{sandbox_id}. + :param contents: BinaryIO + The complete file contents. An empty body creates an empty file. + :param path: str + The absolute UTF-8 destination path inside the sandbox. + + + """ + + query = {} + if path is not None: + query["path"] = path + headers = { + "Content-Type": "application/octet-stream", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + self._api.do("PUT", f"/api/2.0/sandbox-files/{parent}/content", query=query, headers=headers, data=contents) diff --git a/databricks/sdk/service/settings.py b/databricks/sdk/service/settings.py index 30f37bacc..cc9e828a3 100644 --- a/databricks/sdk/service/settings.py +++ b/databricks/sdk/service/settings.py @@ -5516,6 +5516,10 @@ class NetworkConnectivityConfiguration: account_id: Optional[str] = None """Your Databricks account ID. You can find your account ID in your Databricks accounts console.""" + billing_workspace_id: Optional[int] = None + """The workspace to which charges for this network connectivity configuration are attributed. The + workspace must be active and attached to this network connectivity configuration when set.""" + creation_time: Optional[int] = None """Time in epoch milliseconds when this object was created.""" @@ -5543,6 +5547,8 @@ def as_dict(self) -> dict: body = {} if self.account_id is not None: body["account_id"] = self.account_id + if self.billing_workspace_id is not None: + body["billing_workspace_id"] = self.billing_workspace_id if self.creation_time is not None: body["creation_time"] = self.creation_time if self.egress_config: @@ -5562,6 +5568,8 @@ def as_shallow_dict(self) -> dict: body = {} if self.account_id is not None: body["account_id"] = self.account_id + if self.billing_workspace_id is not None: + body["billing_workspace_id"] = self.billing_workspace_id if self.creation_time is not None: body["creation_time"] = self.creation_time if self.egress_config: @@ -5581,6 +5589,7 @@ def from_dict(cls, d: Dict[str, Any]) -> NetworkConnectivityConfiguration: """Deserializes the NetworkConnectivityConfiguration from a dictionary.""" return cls( account_id=d.get("account_id", None), + billing_workspace_id=_int64(d, "billing_workspace_id"), creation_time=_int64(d, "creation_time"), egress_config=_from_dict(d, "egress_config", NccEgressConfig), name=d.get("name", None), @@ -9336,6 +9345,42 @@ def list_private_endpoint_rules( return query["page_token"] = json["next_page_token"] + def update_network_connectivity_config( + self, + network_connectivity_config_id: str, + network_connectivity_config: NetworkConnectivityConfiguration, + update_mask: FieldMask, + ) -> NetworkConnectivityConfiguration: + """Updates the name or billing workspace for a network connectivity configuration. + + :param network_connectivity_config_id: str + Your Network Connectivity Configuration ID. + :param network_connectivity_config: :class:`NetworkConnectivityConfiguration` + The network connectivity configuration fields to update. + :param update_mask: FieldMask + Must contain at least one of ``name`` or ``billing_workspace_id``. + + :returns: :class:`NetworkConnectivityConfiguration` + """ + + body = network_connectivity_config.as_dict() + query = {} + if update_mask is not None: + query["update_mask"] = update_mask.ToJsonString() + headers = { + "Accept": "application/json", + "Content-Type": "application/json", + } + + res = self._api.do( + "PATCH", + f"/api/2.0/accounts/{self._api.account_id}/network-connectivity-configs/{network_connectivity_config_id}", + query=query, + body=body, + headers=headers, + ) + return NetworkConnectivityConfiguration.from_dict(res) + def update_private_endpoint_rule( self, network_connectivity_config_id: str, diff --git a/databricks/sdk/service/sql.py b/databricks/sdk/service/sql.py index bc49f4cf9..6ec3efb52 100644 --- a/databricks/sdk/service/sql.py +++ b/databricks/sdk/service/sql.py @@ -2587,8 +2587,9 @@ class EndpointInfo: """the jdbc connection string for this warehouse""" max_dbu: Optional[int] = None - """Maximum compute capacity of a real-time SQL warehouse, in DBU per hour. Must be a multiple of 3 - between 3 and 960. The effective maximum cannot be lower than the warehouse's minimum capacity.""" + """Maximum compute capacity of a real-time SQL warehouse, in DBU per hour. Must be at least the + minimum DBU for the selected warehouse size and no greater than the configured DBU limit. Actual + capacity may be below this ceiling.""" max_num_clusters: Optional[int] = None """Maximum number of clusters that the autoscaler will create to handle concurrent queries. @@ -3269,8 +3270,9 @@ class GetWarehouseResponse: """the jdbc connection string for this warehouse""" max_dbu: Optional[int] = None - """Maximum compute capacity of a real-time SQL warehouse, in DBU per hour. Must be a multiple of 3 - between 3 and 960. The effective maximum cannot be lower than the warehouse's minimum capacity.""" + """Maximum compute capacity of a real-time SQL warehouse, in DBU per hour. Must be at least the + minimum DBU for the selected warehouse size and no greater than the configured DBU limit. Actual + capacity may be below this ceiling.""" max_num_clusters: Optional[int] = None """Maximum number of clusters that the autoscaler will create to handle concurrent queries. @@ -10027,8 +10029,9 @@ def create( :param instance_profile_arn: str (optional) Deprecated. Instance profile used to pass IAM role to the cluster :param max_dbu: int (optional) - Maximum compute capacity of a real-time SQL warehouse, in DBU per hour. Must be a multiple of 3 - between 3 and 960. The effective maximum cannot be lower than the warehouse's minimum capacity. + Maximum compute capacity of a real-time SQL warehouse, in DBU per hour. Must be at least the minimum + DBU for the selected warehouse size and no greater than the configured DBU limit. Actual capacity + may be below this ceiling. :param max_num_clusters: int (optional) Maximum number of clusters that the autoscaler will create to handle concurrent queries. @@ -10321,8 +10324,9 @@ def edit( :param instance_profile_arn: str (optional) Deprecated. Instance profile used to pass IAM role to the cluster :param max_dbu: int (optional) - Maximum compute capacity of a real-time SQL warehouse, in DBU per hour. Must be a multiple of 3 - between 3 and 960. The effective maximum cannot be lower than the warehouse's minimum capacity. + Maximum compute capacity of a real-time SQL warehouse, in DBU per hour. Must be at least the minimum + DBU for the selected warehouse size and no greater than the configured DBU limit. Actual capacity + may be below this ceiling. :param max_num_clusters: int (optional) Maximum number of clusters that the autoscaler will create to handle concurrent queries. diff --git a/databricks/sdk/version.py b/databricks/sdk/version.py index fca402585..812fce789 100644 --- a/databricks/sdk/version.py +++ b/databricks/sdk/version.py @@ -1 +1 @@ -__version__ = "0.158.0" +__version__ = "0.159.0" diff --git a/docs/account/iamv2/iam_v2.rst b/docs/account/iamv2/iam_v2.rst index cae389bfe..a90710fac 100644 --- a/docs/account/iamv2/iam_v2.rst +++ b/docs/account/iamv2/iam_v2.rst @@ -6,22 +6,23 @@ These APIs are used to manage identities and the workspace access of these identities in Databricks. - .. py:method:: create_direct_group_member(group_id: int, direct_group_member: DirectGroupMember) -> DirectGroupMember + .. py:method:: create_direct_group_member(parent: str, direct_group_member: DirectGroupMember) -> DirectGroupMember Creates a group membership (assigns a principal to a group). Authorization: the caller must be an account admin or a manager of the group (holds the ``roles/group.manager`` role on it). - :param group_id: int - Required. Internal ID of the group in Databricks. + :param parent: str + The parent resource name of the group to add the member to. Format: + accounts/{account_id}/groups/{group_id} :param direct_group_member: :class:`DirectGroupMember` Required. The direct group member to be added to the group. :returns: :class:`DirectGroupMember` - .. py:method:: create_group(group: Group) -> Group + .. py:method:: create_group(parent: str, group: Group) -> Group Creates a local group in the Databricks account and returns the created group. A local group is one that is not synced from the customer's identity provider, and can be created whether or not Account @@ -31,6 +32,8 @@ identity provider, resolve it by its external ID with ResolveGroup; to read an existing external identity, use the ExternalGroup resource. + :param parent: str + The parent resource name of the account in which to create the group. Format: accounts/{account_id} :param group: :class:`Group` Required. Group to be created in Databricks @@ -59,7 +62,7 @@ :returns: :class:`IdentityVisibilityFilter` - .. py:method:: create_service_principal(service_principal: ServicePrincipal) -> ServicePrincipal + .. py:method:: create_service_principal(parent: str, service_principal: ServicePrincipal) -> ServicePrincipal Creates a local service principal in the Databricks account and returns the created service principal. A local service principal is one that is not synced from the customer's identity provider, and can be @@ -69,13 +72,16 @@ identity provider, resolve it by its external ID with ResolveServicePrincipal; to read an existing external identity, use the ExternalServicePrincipal resource. + :param parent: str + The parent resource name of the account in which to create the service principal. Format: + accounts/{account_id} :param service_principal: :class:`ServicePrincipal` Required. Service principal to be created in Databricks :returns: :class:`ServicePrincipal` - .. py:method:: create_user(user: User) -> User + .. py:method:: create_user(parent: str, user: User) -> User Creates a local user in the Databricks account and returns the created user. A local user is one that is not synced from the customer's identity provider, and can be created whether or not Account @@ -85,21 +91,24 @@ identity provider, resolve it by its external ID with ResolveUser; to read an existing external identity, use the ExternalUser resource. + :param parent: str + The parent resource name of the account in which to create the user. Format: accounts/{account_id} :param user: :class:`User` Required. User to be created in Databricks :returns: :class:`User` - .. py:method:: create_workspace_assignment(workspace_id: int, workspace_assignment: WorkspaceAssignment) -> WorkspaceAssignment + .. py:method:: create_workspace_assignment(parent: str, workspace_assignment: WorkspaceAssignment) -> WorkspaceAssignment Creates a workspace assignment for a principal. Entitlements are granted one at a time rather than atomically. If the request fails partway through, the principal stays assigned to the workspace with only some of the requested entitlements. Get the assignment afterwards to confirm which entitlements were granted. - :param workspace_id: int - Required. The workspace ID for which the workspace assignment is being created. + :param parent: str + The parent resource name of the workspace to assign the principal to. Format: + accounts/{account_id}/workspaces/{workspace_id} :param workspace_assignment: :class:`WorkspaceAssignment` Required. Workspace assignment to be created in Databricks. @@ -118,30 +127,29 @@ :returns: :class:`WorkspaceAssignmentDetail` - .. py:method:: delete_direct_group_member(group_id: int, principal_id: int) + .. py:method:: delete_direct_group_member(name: str) Deletes a group membership (unassigns a principal from a group). Authorization: the caller must be an account admin or a manager of the group (holds the ``roles/group.manager`` role on it). - :param group_id: int - Required. Internal ID of the group in Databricks. - :param principal_id: int - Required. Internal ID of the principal to be unassigned from the group. + :param name: str + The resource name of the group membership. Format: + accounts/{account_id}/groups/{group_id}/direct-members/{principal_id} - .. py:method:: delete_group(group_id: str) + .. py:method:: delete_group(name: str) Deletes a group from the Databricks account by its internal ID. Authorization: the caller must be an account admin or a manager of the group (holds the ``roles/group.manager`` role on it). - :param group_id: str - Required. Internal ID of the group in Databricks. + :param name: str + The resource name of the group. Format: accounts/{account_id}/groups/{group_id} @@ -161,36 +169,36 @@ - .. py:method:: delete_service_principal(service_principal_id: str) + .. py:method:: delete_service_principal(name: str) Deletes a service principal from the Databricks account by its internal ID. - :param service_principal_id: str - Required. Internal ID of the service principal in Databricks. + :param name: str + The resource name of the service principal. Format: + accounts/{account_id}/service-principals/{service_principal_id} - .. py:method:: delete_user(user_id: str) + .. py:method:: delete_user(name: str) Deletes a user from the Databricks account by its internal ID. - :param user_id: str - Required. Internal ID of the user in Databricks. + :param name: str + The resource name of the user. Format: accounts/{account_id}/users/{user_id} - .. py:method:: delete_workspace_assignment(workspace_id: int, principal_id: int) + .. py:method:: delete_workspace_assignment(name: str) Deletes a workspace assignment for a principal, revoking all of its entitlements. Entitlements are revoked one at a time rather than atomically. If the request fails partway through, the principal stays assigned with some of its original entitlements. Retrying is safe. - :param workspace_id: int - The workspace ID where the principal has access. - :param principal_id: int - Required. ID of the principal in Databricks to delete workspace assignment for. + :param name: str + The resource name of the workspace assignment. Format: + accounts/{account_id}/workspaces/{workspace_id}/workspace-assignments/{principal_id} @@ -209,14 +217,13 @@ - .. py:method:: get_direct_group_member(group_id: int, principal_id: int) -> DirectGroupMember + .. py:method:: get_direct_group_member(name: str) -> DirectGroupMember Gets a provisioned direct member of a group. - :param group_id: int - Required. Internal ID of the group in Databricks. - :param principal_id: int - Required. Internal ID of the principal belonging to the group in Databricks. + :param name: str + The resource name of the group membership. Format: + accounts/{account_id}/groups/{group_id}/direct-members/{principal_id} :returns: :class:`DirectGroupMember` @@ -260,12 +267,12 @@ :returns: :class:`ExternalUser` - .. py:method:: get_group(group_id: str) -> Group + .. py:method:: get_group(name: str) -> Group Fetches a group from the Databricks account by its internal ID. - :param group_id: str - Required. Internal ID of the group in Databricks. + :param name: str + The resource name of the group. Format: accounts/{account_id}/groups/{group_id} :returns: :class:`Group` @@ -282,27 +289,28 @@ :returns: :class:`IdentityVisibilityFilter` - .. py:method:: get_service_principal(service_principal_id: str) -> ServicePrincipal + .. py:method:: get_service_principal(name: str) -> ServicePrincipal Fetches a service principal from the Databricks account by its internal ID. - :param service_principal_id: str - Required. Internal ID of the service principal in Databricks. + :param name: str + The resource name of the service principal. Format: + accounts/{account_id}/service-principals/{service_principal_id} :returns: :class:`ServicePrincipal` - .. py:method:: get_user(user_id: str) -> User + .. py:method:: get_user(name: str) -> User Fetches a user from the Databricks account by its internal ID. - :param user_id: str - Required. Internal ID of the user in Databricks. + :param name: str + The resource name of the user. Format: accounts/{account_id}/users/{user_id} :returns: :class:`User` - .. py:method:: get_workspace_access_detail(workspace_id: int, principal_id: int [, view: Optional[WorkspaceAccessDetailView]]) -> WorkspaceAccessDetail + .. py:method:: get_workspace_access_detail(name: str [, view: Optional[WorkspaceAccessDetailView]]) -> WorkspaceAccessDetail Returns the access details for a principal in a workspace. Allows for checking access details for any provisioned principal (user, service principal, or group) in a workspace. @@ -311,26 +319,22 @@ IdP or added explicitly to Databricks via SCIM/UI. Allows for passing in a "view" parameter to control what fields are returned (BASIC by default or FULL). - :param workspace_id: int - Required. The workspace ID for which the access details are being requested. - :param principal_id: int - Required. The internal ID of the principal (user/sp/group) for which the access details are being - requested. + :param name: str + The resource name of the workspace access detail. Format: + accounts/{account_id}/workspaces/{workspace_id}/workspace-access-details/{principal_id} :param view: :class:`WorkspaceAccessDetailView` (optional) Controls what fields are returned. :returns: :class:`WorkspaceAccessDetail` - .. py:method:: get_workspace_assignment(workspace_id: int, principal_id: int) -> WorkspaceAssignment + .. py:method:: get_workspace_assignment(name: str) -> WorkspaceAssignment Returns the assignment for a principal in a workspace. - :param workspace_id: int - Required. The workspace ID for which the assignment is being requested. - :param principal_id: int - Required. The internal ID of the principal (user/sp/group) for which the assignment is being - requested. + :param name: str + The resource name of the workspace assignment. Format: + accounts/{account_id}/workspaces/{workspace_id}/workspace-assignments/{principal_id} :returns: :class:`WorkspaceAssignment` @@ -348,13 +352,14 @@ :returns: :class:`WorkspaceAssignmentDetail` - .. py:method:: list_direct_group_members(group_id: int [, page_size: Optional[int], page_token: Optional[str]]) -> Iterator[DirectGroupMember] + .. py:method:: list_direct_group_members(parent: str [, page_size: Optional[int], page_token: Optional[str]]) -> Iterator[DirectGroupMember] Lists provisioned direct members of a group with their membership source (internal or from identity provider). - :param group_id: int - Required. Internal ID of the group in Databricks whose direct members are being listed. + :param parent: str + The parent resource name of the group whose direct members to list. Format: + accounts/{account_id}/groups/{group_id} :param page_size: int (optional) The maximum number of members to return. The service may return fewer than this value. If not provided, defaults to 1000, which is also the maximum allowed. Requests for more than the maximum @@ -366,11 +371,13 @@ :returns: Iterator over :class:`DirectGroupMember` - .. py:method:: list_groups( [, filter: Optional[str], page_size: Optional[int], page_token: Optional[str]]) -> Iterator[Group] + .. py:method:: list_groups(parent: str [, filter: Optional[str], page_size: Optional[int], page_token: Optional[str]]) -> Iterator[Group] Lists the groups in the Databricks account, returning one page per call. Supports filtering by group name or external ID. + :param parent: str + The parent resource name of the account whose groups to list. Format: accounts/{account_id} :param filter: str (optional) Optional. Allows filtering groups by group name or external id. :param page_size: int (optional) @@ -403,11 +410,14 @@ :returns: Iterator over :class:`IdentityVisibilityFilter` - .. py:method:: list_service_principals( [, filter: Optional[str], page_size: Optional[int], page_token: Optional[str]]) -> Iterator[ServicePrincipal] + .. py:method:: list_service_principals(parent: str [, filter: Optional[str], page_size: Optional[int], page_token: Optional[str]]) -> Iterator[ServicePrincipal] Lists the service principals in the Databricks account, returning one page per call. Supports filtering by application ID or external ID. + :param parent: str + The parent resource name of the account whose service principals to list. Format: + accounts/{account_id} :param filter: str (optional) Optional. Allows filtering service principals by application id or external id. :param page_size: int (optional) @@ -421,13 +431,13 @@ :returns: Iterator over :class:`ServicePrincipal` - .. py:method:: list_transitive_parent_groups(principal_id: int [, page_size: Optional[int], page_token: Optional[str]]) -> ListTransitiveParentGroupsResponse + .. py:method:: list_transitive_parent_groups(parent: str [, page_size: Optional[int], page_token: Optional[str]]) -> ListTransitiveParentGroupsResponse Lists all transitive parent groups of a principal. - :param principal_id: int - Required. Internal ID of the principal in Databricks whose transitive parent groups are being - listed. + :param parent: str + The parent resource name of the principal whose transitive parent groups to list. Format: + accounts/{account_id}/principals/{principal_id} :param page_size: int (optional) The maximum number of parent groups to return. The service may return fewer than this value. If not provided, defaults to 1000, which is also the maximum allowed. Requests for more than the maximum @@ -439,11 +449,13 @@ :returns: :class:`ListTransitiveParentGroupsResponse` - .. py:method:: list_users( [, filter: Optional[str], page_size: Optional[int], page_token: Optional[str]]) -> Iterator[User] + .. py:method:: list_users(parent: str [, filter: Optional[str], page_size: Optional[int], page_token: Optional[str]]) -> Iterator[User] Lists the users in the Databricks account, returning one page per call. Supports filtering by username or external ID. + :param parent: str + The parent resource name of the account whose users to list. Format: accounts/{account_id} :param filter: str (optional) Optional. Allows filtering users by username or external id. :param page_size: int (optional) @@ -475,14 +487,15 @@ :returns: Iterator over :class:`WorkspaceAssignmentDetail` - .. py:method:: list_workspace_assignments(workspace_id: int [, page_size: Optional[int], page_token: Optional[str]]) -> Iterator[WorkspaceAssignment] + .. py:method:: list_workspace_assignments(parent: str [, page_size: Optional[int], page_token: Optional[str]]) -> Iterator[WorkspaceAssignment] Lists workspace assignments for a workspace. The response omits the per-principal entitlement fields (``entitlements`` and ``effective_entitlements``). To read the entitlements for a single principal, get that principal's assignment. - :param workspace_id: int - Required. The workspace ID for which the workspace assignments are being fetched. + :param parent: str + The parent resource name of the workspace whose assignments to list. Format: + accounts/{account_id}/workspaces/{workspace_id} :param page_size: int (optional) The maximum number of workspace assignments to return. The service may return fewer than this value. If not provided, defaults to 1000, which is also the maximum allowed. Requests for more than the @@ -530,7 +543,7 @@ :returns: :class:`ResolveUserResponse` - .. py:method:: update_group(group_id: str, group: Group, update_mask: str) -> Group + .. py:method:: update_group(name: str, group: Group, update_mask: str) -> Group Updates an existing group in the Databricks account. Only the fields named in the update mask are modified. Returns the updated Group resource. @@ -543,8 +556,11 @@ Authorization: the caller must be an account admin or a manager of the group (holds the ``roles/group.manager`` role on it). - :param group_id: str - Required. Internal ID of the group in Databricks. + :param name: str + The resource name of the group. The format depends on the API that returned it: + + - Account-scoped: accounts/{account_id}/groups/{group_id} + - Workspace-scoped: groups/{group_id} :param group: :class:`Group` Required. Group to be updated in Databricks :param update_mask: str @@ -553,18 +569,22 @@ :returns: :class:`Group` - .. py:method:: update_service_principal(service_principal_id: str, service_principal: ServicePrincipal, update_mask: str) -> ServicePrincipal + .. py:method:: update_service_principal(name: str, service_principal: ServicePrincipal, update_mask: str) -> ServicePrincipal Updates an existing service principal in the Databricks account. Only the fields named in the update mask are modified. Returns the updated ServicePrincipal resource. - For locally managed service principals, is_agent must be set to true or false when updated. + To change is_agent on a locally managed service principal, set it to true or false. Naming is_agent in + the update mask without a value leaves it unchanged. When AIM is enabled and the service principal is an external identity (its external_id is set), only external_id can be updated; its other fields are sourced from your identity provider. - :param service_principal_id: str - Required. Internal ID of the service principal in Databricks. + :param name: str + The resource name of the service principal. The format depends on the API that returned it: + + - Account-scoped: accounts/{account_id}/service-principals/{service_principal_id} + - Workspace-scoped: service-principals/{service_principal_id} :param service_principal: :class:`ServicePrincipal` Required. Service Principal to be updated in Databricks :param update_mask: str @@ -573,7 +593,7 @@ :returns: :class:`ServicePrincipal` - .. py:method:: update_user(user_id: str, user: User, update_mask: str) -> User + .. py:method:: update_user(name: str, user: User, update_mask: str) -> User Updates an existing user in the Databricks account and returns the updated user. Only the fields named in the update mask are modified. The updatable fields are fullName.givenName, fullName.familyName, @@ -582,8 +602,11 @@ When AIM is enabled and the user is an external identity (its external_id is set), only external_id can be updated; its other fields are sourced from your identity provider. - :param user_id: str - Required. Internal ID of the user in Databricks. + :param name: str + The resource name of the user. The format depends on the API that returned it: + + - Account-scoped: accounts/{account_id}/users/{user_id} + - Workspace-scoped: users/{user_id} :param user: :class:`User` Required. User to be updated in Databricks :param update_mask: str @@ -592,16 +615,18 @@ :returns: :class:`User` - .. py:method:: update_workspace_assignment(workspace_id: int, principal_id: int, workspace_assignment: WorkspaceAssignment, update_mask: FieldMask) -> WorkspaceAssignment + .. py:method:: update_workspace_assignment(name: str, workspace_assignment: WorkspaceAssignment, update_mask: FieldMask) -> WorkspaceAssignment Updates the entitlements of a directly assigned principal in a workspace. Changes are applied one at a time rather than atomically. If the request fails partway through, only some of the requested changes take effect. Get the assignment afterwards to confirm the final state. - :param workspace_id: int - Required. The workspace ID for which the workspace assignment is being updated. - :param principal_id: int - Required. ID of the principal in Databricks. + :param name: str + The resource name of the workspace assignment. The format depends on the API that returned it: + + - Account-scoped: + accounts/{account_id}/workspaces/{workspace_id}/workspace-assignments/{principal_id} + - Workspace-scoped: workspace-assignments/{principal_id} :param workspace_assignment: :class:`WorkspaceAssignment` Required. Workspace assignment to be updated in Databricks. :param update_mask: FieldMask diff --git a/docs/account/settings/network_connectivity.rst b/docs/account/settings/network_connectivity.rst index 5816f9e4b..b6e2430a8 100644 --- a/docs/account/settings/network_connectivity.rst +++ b/docs/account/settings/network_connectivity.rst @@ -116,6 +116,20 @@ :returns: Iterator over :class:`NccPrivateEndpointRule` + .. py:method:: update_network_connectivity_config(network_connectivity_config_id: str, network_connectivity_config: NetworkConnectivityConfiguration, update_mask: FieldMask) -> NetworkConnectivityConfiguration + + Updates the name or billing workspace for a network connectivity configuration. + + :param network_connectivity_config_id: str + Your Network Connectivity Configuration ID. + :param network_connectivity_config: :class:`NetworkConnectivityConfiguration` + The network connectivity configuration fields to update. + :param update_mask: FieldMask + Must contain at least one of ``name`` or ``billing_workspace_id``. + + :returns: :class:`NetworkConnectivityConfiguration` + + .. py:method:: update_private_endpoint_rule(network_connectivity_config_id: str, private_endpoint_rule_id: str, private_endpoint_rule: UpdatePrivateEndpointRule, update_mask: str) -> NccPrivateEndpointRule Updates a private endpoint rule. Currently only a private endpoint rule to customer-managed resources diff --git a/docs/dbdataclasses/catalog.rst b/docs/dbdataclasses/catalog.rst index e31f602d4..0cfd318cb 100644 --- a/docs/dbdataclasses/catalog.rst +++ b/docs/dbdataclasses/catalog.rst @@ -1661,6 +1661,9 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: CREATE_EXTERNAL_VOLUME :value: "CREATE_EXTERNAL_VOLUME" + .. py:attribute:: CREATE_FEATURE + :value: "CREATE_FEATURE" + .. py:attribute:: CREATE_FOREIGN_CATALOG :value: "CREATE_FOREIGN_CATALOG" @@ -1730,6 +1733,9 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: MODIFY_CLEAN_ROOM :value: "MODIFY_CLEAN_ROOM" + .. py:attribute:: READ_FEATURE + :value: "READ_FEATURE" + .. py:attribute:: READ_FILES :value: "READ_FILES" @@ -2115,6 +2121,9 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: EXTERNAL_METADATA :value: "EXTERNAL_METADATA" + .. py:attribute:: FEATURE + :value: "FEATURE" + .. py:attribute:: FUNCTION :value: "FUNCTION" diff --git a/docs/dbdataclasses/computemigrations.rst b/docs/dbdataclasses/computemigrations.rst new file mode 100644 index 000000000..0a1de771d --- /dev/null +++ b/docs/dbdataclasses/computemigrations.rst @@ -0,0 +1,125 @@ +Compute Migrations +================== + +These dataclasses are used in the SDK to represent API requests and responses for services in the ``databricks.sdk.service.computemigrations`` module. + +.. py:currentmodule:: databricks.sdk.service.computemigrations +.. autoclass:: BakingConfig + :members: + :undoc-members: + +.. autoclass:: BakingReport + :members: + :undoc-members: + +.. autoclass:: ListMigrationsResponse + :members: + :undoc-members: + +.. autoclass:: Migration + :members: + :undoc-members: + +.. autoclass:: MigrationAsset + :members: + :undoc-members: + +.. py:class:: MigrationAssetType + + Type of a migration asset: the Classic asset being migrated, or an asset the migration agent created for it. Migrations can only be created for JOB assets. + + .. py:attribute:: FILE + :value: "FILE" + + .. py:attribute:: INTERACTIVE + :value: "INTERACTIVE" + + .. py:attribute:: JOB + :value: "JOB" + +.. autoclass:: MigrationInitiator + :members: + :undoc-members: + +.. py:class:: MigrationInitiatorRole + + Role of the migration initiator, determining approval authority. + + .. py:attribute:: ACCOUNT_ADMIN + :value: "ACCOUNT_ADMIN" + + .. py:attribute:: JOB_ADMIN + :value: "JOB_ADMIN" + + .. py:attribute:: SYSTEM + :value: "SYSTEM" + + .. py:attribute:: WORKSPACE_ADMIN + :value: "WORKSPACE_ADMIN" + +.. autoclass:: MigrationReport + :members: + :undoc-members: + +.. py:class:: MigrationState + + The state of a migration in its lifecycle. + + .. py:attribute:: CANCELLED + :value: "CANCELLED" + + .. py:attribute:: FAILED + :value: "FAILED" + + .. py:attribute:: MIGRATION_COMPLETE + :value: "MIGRATION_COMPLETE" + + .. py:attribute:: MIGRATION_OBSERVATION + :value: "MIGRATION_OBSERVATION" + + .. py:attribute:: MIGRATION_PENDING + :value: "MIGRATION_PENDING" + + .. py:attribute:: MIGRATION_REVIEW + :value: "MIGRATION_REVIEW" + + .. py:attribute:: MIGRATION_RUNNING + :value: "MIGRATION_RUNNING" + + .. py:attribute:: NOT_SUPPORTED + :value: "NOT_SUPPORTED" + + .. py:attribute:: PLANNING_PENDING + :value: "PLANNING_PENDING" + + .. py:attribute:: PLANNING_REVIEW + :value: "PLANNING_REVIEW" + + .. py:attribute:: PLANNING_RUNNING + :value: "PLANNING_RUNNING" + + .. py:attribute:: REVISION_PENDING + :value: "REVISION_PENDING" + + .. py:attribute:: TRANSLATION_CREATING + :value: "TRANSLATION_CREATING" + + .. py:attribute:: TRANSLATION_PENDING + :value: "TRANSLATION_PENDING" + + .. py:attribute:: TRANSLATION_READY + :value: "TRANSLATION_READY" + + .. py:attribute:: VALIDATION_RUNNING + :value: "VALIDATION_RUNNING" + + .. py:attribute:: VALIDATION_SCHEDULING + :value: "VALIDATION_SCHEDULING" + +.. autoclass:: MigrationStateTransition + :members: + :undoc-members: + +.. autoclass:: ValidationReport + :members: + :undoc-members: diff --git a/docs/dbdataclasses/genieagents.rst b/docs/dbdataclasses/genieagents.rst new file mode 100644 index 000000000..ba5283f03 --- /dev/null +++ b/docs/dbdataclasses/genieagents.rst @@ -0,0 +1,142 @@ +Genie Agents +============ + +These dataclasses are used in the SDK to represent API requests and responses for services in the ``databricks.sdk.service.genieagents`` module. + +.. py:currentmodule:: databricks.sdk.service.genieagents +.. autoclass:: GenieAgent + :members: + :undoc-members: + +.. autoclass:: GenieAgentBenchmarkAnswer + :members: + :undoc-members: + +.. py:class:: GenieAgentBenchmarkAnswerFormat + + The format of a benchmark answer. + + .. py:attribute:: GENIE_AGENT_BENCHMARK_ANSWER_FORMAT_SQL + :value: "GENIE_AGENT_BENCHMARK_ANSWER_FORMAT_SQL" + +.. autoclass:: GenieAgentBenchmarkQuestion + :members: + :undoc-members: + +.. autoclass:: GenieAgentBenchmarks + :members: + :undoc-members: + +.. autoclass:: GenieAgentColumnConfig + :members: + :undoc-members: + +.. autoclass:: GenieAgentConfig + :members: + :undoc-members: + +.. autoclass:: GenieAgentDataSources + :members: + :undoc-members: + +.. autoclass:: GenieAgentDefinition + :members: + :undoc-members: + +.. autoclass:: GenieAgentExampleQuestionSql + :members: + :undoc-members: + +.. autoclass:: GenieAgentInstructions + :members: + :undoc-members: + +.. py:class:: GenieAgentJoinRelationshipType + + How the rows of the two sides of a join relate. + + .. py:attribute:: GENIE_AGENT_JOIN_RELATIONSHIP_TYPE_MANY_TO_MANY + :value: "GENIE_AGENT_JOIN_RELATIONSHIP_TYPE_MANY_TO_MANY" + + .. py:attribute:: GENIE_AGENT_JOIN_RELATIONSHIP_TYPE_MANY_TO_ONE + :value: "GENIE_AGENT_JOIN_RELATIONSHIP_TYPE_MANY_TO_ONE" + + .. py:attribute:: GENIE_AGENT_JOIN_RELATIONSHIP_TYPE_ONE_TO_MANY + :value: "GENIE_AGENT_JOIN_RELATIONSHIP_TYPE_ONE_TO_MANY" + + .. py:attribute:: GENIE_AGENT_JOIN_RELATIONSHIP_TYPE_ONE_TO_ONE + :value: "GENIE_AGENT_JOIN_RELATIONSHIP_TYPE_ONE_TO_ONE" + +.. autoclass:: GenieAgentJoinSource + :members: + :undoc-members: + +.. autoclass:: GenieAgentJoinSpec + :members: + :undoc-members: + +.. autoclass:: GenieAgentMcpConnection + :members: + :undoc-members: + +.. autoclass:: GenieAgentSampleQuestion + :members: + :undoc-members: + +.. autoclass:: GenieAgentSqlExpression + :members: + :undoc-members: + +.. autoclass:: GenieAgentSqlFilter + :members: + :undoc-members: + +.. autoclass:: GenieAgentSqlFunction + :members: + :undoc-members: + +.. autoclass:: GenieAgentSqlMeasure + :members: + :undoc-members: + +.. autoclass:: GenieAgentSqlParameter + :members: + :undoc-members: + +.. autoclass:: GenieAgentSqlParameterDefaultRange + :members: + :undoc-members: + +.. autoclass:: GenieAgentSqlParameterDefaultValue + :members: + :undoc-members: + +.. autoclass:: GenieAgentSqlSnippets + :members: + :undoc-members: + +.. autoclass:: GenieAgentTable + :members: + :undoc-members: + +.. autoclass:: GenieAgentTextInstruction + :members: + :undoc-members: + +.. py:class:: GenieAgentView + + The fields returned for each Genie Agent in a list response. + + .. py:attribute:: GENIE_AGENT_VIEW_BASIC + :value: "GENIE_AGENT_VIEW_BASIC" + + .. py:attribute:: GENIE_AGENT_VIEW_FULL + :value: "GENIE_AGENT_VIEW_FULL" + +.. autoclass:: GenieAgentVolume + :members: + :undoc-members: + +.. autoclass:: ListGenieAgentsResponse + :members: + :undoc-members: diff --git a/docs/dbdataclasses/index.rst b/docs/dbdataclasses/index.rst index a8707c3aa..f77d8ee59 100644 --- a/docs/dbdataclasses/index.rst +++ b/docs/dbdataclasses/index.rst @@ -15,6 +15,7 @@ Dataclasses catalog cleanrooms compute + computemigrations dashboards database dataclassification @@ -23,6 +24,7 @@ Dataclasses domains environments files + genieagents iam iamv2 jobs diff --git a/docs/dbdataclasses/jobs.rst b/docs/dbdataclasses/jobs.rst index f03de4727..801f11f82 100644 --- a/docs/dbdataclasses/jobs.rst +++ b/docs/dbdataclasses/jobs.rst @@ -4,6 +4,51 @@ Jobs These dataclasses are used in the SDK to represent API requests and responses for services in the ``databricks.sdk.service.jobs`` module. .. py:currentmodule:: databricks.sdk.service.jobs +.. autoclass:: AiClassifyFunction + :members: + :undoc-members: + +.. autoclass:: AiClassifyFunctionInlineLabels + :members: + :undoc-members: + +.. autoclass:: AiClassifyFunctionLabel + :members: + :undoc-members: + +.. autoclass:: AiExtractFunction + :members: + :undoc-members: + +.. py:class:: AiExtractFunctionMode + + The extraction mode. + + .. py:attribute:: MODE_PRECISION + :value: "MODE_PRECISION" + +.. autoclass:: AiFunctionTask + :members: + :undoc-members: + +.. autoclass:: AiFunctionTaskOutput + :members: + :undoc-members: + +.. autoclass:: AiParseDocumentFunction + :members: + :undoc-members: + +.. py:class:: AiParseDocumentFunctionDescriptionElementType + + An element type that AI descriptions can be generated for. + + .. py:attribute:: DESCRIPTION_ELEMENT_TYPE_ALL + :value: "DESCRIPTION_ELEMENT_TYPE_ALL" + + .. py:attribute:: DESCRIPTION_ELEMENT_TYPE_FIGURE + :value: "DESCRIPTION_ELEMENT_TYPE_FIGURE" + .. autoclass:: AiRuntimeTask :members: :undoc-members: @@ -166,6 +211,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: CompletedJob + :members: + :undoc-members: + .. autoclass:: Compute :members: :undoc-members: @@ -344,6 +393,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: DeltaTableSource + :members: + :undoc-members: + .. autoclass:: DeploymentSpec :members: :undoc-members: @@ -454,6 +507,43 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: JobClusterServerlessCompute + :members: + :undoc-members: + +.. autoclass:: JobCompletionTriggerConfiguration + :members: + :undoc-members: + +.. py:class:: JobCompletionTriggerConfigurationIfStatus + + The status that all upstream jobs must have before a job completion trigger fires. + + .. py:attribute:: DONE + :value: "DONE" + + .. py:attribute:: FAILED + :value: "FAILED" + + .. py:attribute:: SUCCESS + :value: "SUCCESS" + +.. autoclass:: JobCompletionTriggerConfigurationMonitoredJob + :members: + :undoc-members: + +.. autoclass:: JobCompletionTriggerRunInfoDetails + :members: + :undoc-members: + +.. autoclass:: JobCompletionTriggerState + :members: + :undoc-members: + +.. autoclass:: JobCompletionTriggerStateMonitoredJobState + :members: + :undoc-members: + .. autoclass:: JobCompliance :members: :undoc-members: @@ -721,6 +811,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: QualifiedName + :members: + :undoc-members: + .. autoclass:: QueueDetails :members: :undoc-members: @@ -797,6 +891,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: ResolvedValuesAiFunctionTaskResolvedValues + :members: + :undoc-members: + .. autoclass:: ResolvedValuesAiRuntimeTaskResolvedValues :members: :undoc-members: @@ -1382,6 +1480,9 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: FILE_ARRIVAL :value: "FILE_ARRIVAL" + .. py:attribute:: JOB_COMPLETION + :value: "JOB_COMPLETION" + .. py:attribute:: ONE_TIME :value: "ONE_TIME" @@ -1424,6 +1525,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: DASHBOARDS :value: "DASHBOARDS" +.. autoclass:: VolumeSource + :members: + :undoc-members: + .. autoclass:: Webhook :members: :undoc-members: diff --git a/docs/dbdataclasses/sandbox.rst b/docs/dbdataclasses/sandbox.rst index ee65a6c7d..0aa768003 100644 --- a/docs/dbdataclasses/sandbox.rst +++ b/docs/dbdataclasses/sandbox.rst @@ -20,6 +20,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: DownloadFileResponse + :members: + :undoc-members: + .. autoclass:: EnvironmentSpec :members: :undoc-members: diff --git a/docs/index.rst b/docs/index.rst index 2172617d3..f338ec4ce 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -19,7 +19,7 @@ We are keen to hear feedback from you on these SDKs. Please `file GitHub issues logging dbutils dataplane - sandbox-command-execution + sandbox clients/workspace workspace/index clients/account diff --git a/docs/packages.py b/docs/packages.py index 14a7b645b..2f03c7f80 100644 --- a/docs/packages.py +++ b/docs/packages.py @@ -130,6 +130,7 @@ class Package: Package("catalog", "", ""), Package("cleanrooms", "", ""), Package("compute", "", ""), + Package("computemigrations", "Compute Migrations", "Manage migrations of Classic compute workloads to Serverless."), Package("dashboards", "", ""), Package( "database", @@ -158,6 +159,11 @@ class Package: "APIs to manage environment resources.\n\nThe Environments API provides management capabilities for different types of environments\nincluding workspace-level base environments that define the environment version and\ndependencies to be used in serverless notebooks and jobs.", ), Package("files", "", ""), + Package( + "genieagents", + "Genie Agents", + "Genie provides a no-code experience for business users, powered by AI/BI. Analysts set up Genie Agents that business users can use to ask questions using natural language. Genie uses data registered to Unity Catalog and requires at least CAN USE permission on a Pro or Serverless SQL warehouse. Also, Databricks Assistant must be enabled.", + ), Package("iam", "", ""), Package("iamv2", "", ""), Package("jobs", "", ""), diff --git a/docs/sandbox-command-execution.md b/docs/sandbox.md similarity index 61% rename from docs/sandbox-command-execution.md rename to docs/sandbox.md index ff985c5b7..e6418a01d 100644 --- a/docs/sandbox-command-execution.md +++ b/docs/sandbox.md @@ -1,13 +1,12 @@ -# Running commands in sandboxes +# Sandboxes Databricks Sandbox lets you run untrusted code in an isolated environment. For sandbox lifecycle operations (create, delete, etc.), see the [`WorkspaceClient.sandbox` reference](https://databricks-sdk-py.readthedocs.io/en/latest/workspace/sandbox/sandbox.html). -This guide shows how to run a command in a sandbox, read its output as it is -being produced, send it input, and get its exit code. +## Commands -## Prerequisites +### Prerequisites Install the SDK with the `sandbox` extra: @@ -20,7 +19,7 @@ value corresponds to the `name` field of the `Sandbox` object returned by `create_sandbox`, `get_sandbox`, or `list_sandboxes` (a bare `{sandbox_id}` is also accepted). If the sandbox is stopped, running a command starts it first. -## Run a command +### Run a command `execute_command` starts a command and returns a `CommandProcess`. You can read its standard output and standard error through `stdout` and `stderr`, and use @@ -46,7 +45,7 @@ print(process.wait()) # 3 `read(size)` returns up to `size` characters (or bytes). Output is `str` by default; pass `text=False` to `execute_command` to get `bytes`. -## Read output as it is produced +### Read output as it is produced Iterate `stdout` or `stderr` to read it line by line as the command writes it, or call `readline()` for a single line: @@ -62,7 +61,7 @@ For output that arrives without a trailing newline (a prompt, a progress bar), use `read1()` to receive output the instant it is written, without waiting for a full line. -## Send input +### Send input Write to `stdin`, then call `close()` to signal end of input: @@ -73,7 +72,7 @@ process.stdin.close() print(process.stdout.read()) # hello from stdin ``` -## Choose where output goes +### Choose where output goes The SDK keeps each stream's output in memory until you read it, so read long-running output as it is produced rather than buffering it all; if you don't @@ -99,7 +98,7 @@ process = w.sandbox.execute_command( print(process.wait()) # 0 ``` -## Stop waiting on a command +### Stop waiting on a command `wait(timeout=)` and `read(timeout=)` raise `TimeoutError` if the command does not finish in time. Call `close()` to stop receiving its output and release the @@ -117,7 +116,7 @@ with w.sandbox.execute_command(name, "/bin/sleep", args=["60"]) as process: # leaving the block closes the handle ``` -## Reconnect to a command +### Reconnect to a command If the connection is interrupted, the SDK reconnects and resumes the output where it left off. @@ -134,3 +133,45 @@ again = w.sandbox.connect_command(name, process.command_id) print(again.stdout.read()) # saved print(again.wait()) # 0 ``` + +## Files + +Pass the sandbox resource name (`sandboxes/{sandbox_id}`) to each method on `w.sandbox`: + +```python +from databricks.sdk import WorkspaceClient + +w = WorkspaceClient() +sandbox_name = "sandboxes/my-sandbox" + +w.sandbox.create_directory(sandbox_name, "/home/sandbox/data") +w.sandbox.rename_path(sandbox_name, "/home/sandbox/input.csv", "/home/sandbox/data/input.csv") +info = w.sandbox.get_metadata(sandbox_name, "/home/sandbox/data/input.csv").file +for child in w.sandbox.list_directory_contents(sandbox_name, "/home/sandbox/data"): + print(child.path, child.size) +w.sandbox.delete_path(sandbox_name, "/home/sandbox/data/input.csv") +``` + +- `create_directory` creates missing parents and succeeds if the directory exists. With `parents=False`, the parent must exist and the destination must be absent. +- `rename_path` moves to the exact destination path. It fails if the destination exists or the move crosses filesystems. +- `delete_path` removes files, symlinks, and empty directories. Use `recursive=True` for nonempty directories. Deleting a missing path succeeds. +- `get_metadata` describes the final symlink itself. The response's `file` field contains the metadata; `modify_time` preserves nanosecond precision. +- `list_directory_contents` includes dotfiles and fetches pages as you iterate. Concurrent directory changes can affect the results. + +Uploads and downloads stream bytes: + +```python +with open("report.bin", "rb") as source: + w.sandbox.upload_file(parent=sandbox_name, path="/home/sandbox/report.bin", contents=source) + +download = w.sandbox.download_file(parent=sandbox_name, path="/home/sandbox/report.bin") +with download.contents as source, open("report-copy.bin", "wb") as destination: + for chunk in iter(lambda: source.read(1024 * 1024), b""): + destination.write(chunk) +``` + +Uploads atomically replace the destination file. The parent directory must exist; `upload_file` has no `overwrite` option. + +`download_file` returns a `DownloadFileResponse` with a readable `contents` stream and HTTP metadata, including `content_length`, `content_type`, `etag`, and `last_modified`. Pass `byte_range="bytes=0-1023"` for a byte range and `if_range=etag` to require a matching file version. A mismatch returns the full file; an unsatisfiable range raises an SDK error. Close the contents stream after use. + +API and transport failures raise standard SDK exceptions. diff --git a/docs/workspace/catalog/credentials.rst b/docs/workspace/catalog/credentials.rst index fa2553ba9..79d533f6b 100644 --- a/docs/workspace/catalog/credentials.rst +++ b/docs/workspace/catalog/credentials.rst @@ -154,23 +154,22 @@ Validates a credential. - For service credentials (purpose is **SERVICE**), either the **credential_name** or the cloud-specific - credential must be provided. - For storage credentials (purpose is **STORAGE**), at least one of **external_location_name** and **url** need to be provided. If only one of them is provided, it will be used for validation. And if both are provided, the **url** will be used for validation, and **external_location_name** will be - ignored when checking overlapping urls. Either the **credential_name** or the cloud-specific - credential must be provided. + ignored when checking overlapping urls. The caller must be a metastore admin or the credential owner or have the required permission on the metastore and the credential (e.g., **CREATE_EXTERNAL_LOCATION** when purpose is **STORAGE**). :param aws_iam_role: :class:`AwsIamRole` (optional) + Deprecated: create a credential and pass its credential_name instead. :param azure_managed_identity: :class:`AzureManagedIdentity` (optional) + Deprecated: create a credential and pass its credential_name instead. :param credential_name: str (optional) Required. The name of an existing credential or long-lived cloud credential to validate. :param databricks_gcp_service_account: :class:`DatabricksGcpServiceAccount` (optional) + Deprecated: create a credential and pass its credential_name instead. :param external_location_name: str (optional) The name of an existing external location to validate. Only applicable for storage credentials (purpose is **STORAGE**.) diff --git a/docs/workspace/catalog/storage_credentials.rst b/docs/workspace/catalog/storage_credentials.rst index 2cd09f1e9..c1d7cbf92 100644 --- a/docs/workspace/catalog/storage_credentials.rst +++ b/docs/workspace/catalog/storage_credentials.rst @@ -229,21 +229,24 @@ the **url** will be used for validation, and **external_location_name** will be ignored when checking overlapping urls. - Either the **storage_credential_name** or the cloud-specific credential must be provided. - The caller must be a metastore admin or the storage credential owner or have the **CREATE_EXTERNAL_LOCATION** privilege on the metastore and the storage credential. :param aws_iam_role: :class:`AwsIamRoleRequest` (optional) - The AWS IAM role configuration. + Deprecated: create a storage credential and pass its storage_credential_name instead. The AWS IAM + role configuration. :param azure_managed_identity: :class:`AzureManagedIdentityRequest` (optional) - The Azure managed identity configuration. + Deprecated: create a storage credential and pass its storage_credential_name instead. The Azure + managed identity configuration. :param azure_service_principal: :class:`AzureServicePrincipal` (optional) - The Azure service principal configuration. + Deprecated: create a storage credential and pass its storage_credential_name instead. The Azure + service principal configuration. :param cloudflare_api_token: :class:`CloudflareApiToken` (optional) - The Cloudflare API token configuration. + Deprecated: create a storage credential and pass its storage_credential_name instead. The Cloudflare + API token configuration. :param databricks_gcp_service_account: :class:`DatabricksGcpServiceAccountRequest` (optional) - The Databricks created GCP service account configuration. + Deprecated: create a storage credential and pass its storage_credential_name instead. The Databricks + created GCP service account configuration. :param external_location_name: str (optional) The name of an existing external location to validate. :param read_only: bool (optional) diff --git a/docs/workspace/catalog/volumes.rst b/docs/workspace/catalog/volumes.rst index da9784bb0..1e56225bc 100644 --- a/docs/workspace/catalog/volumes.rst +++ b/docs/workspace/catalog/volumes.rst @@ -232,7 +232,7 @@ on the parent catalog and the **USE_SCHEMA** privilege on the parent schema. :param name: str - The three-level (fully qualified) name of the volume + The fully qualified name of the volume. :param include_browse: bool (optional) Whether to include volumes in the response for which the principal can only access selective metadata for diff --git a/docs/workspace/computemigrations/compute_migrations.rst b/docs/workspace/computemigrations/compute_migrations.rst new file mode 100644 index 000000000..e868749ec --- /dev/null +++ b/docs/workspace/computemigrations/compute_migrations.rst @@ -0,0 +1,201 @@ +``w.compute_migrations``: ComputeMigrations.v1 +============================================== +.. currentmodule:: databricks.sdk.service.computemigrations + +.. py:class:: ComputeMigrationsAPI + + Manage migrations of Classic compute workloads to Serverless. + + .. py:method:: approve_migration_execution(migration_id: str [, additional_feedback: Optional[str], etag: Optional[str]]) -> Migration + + Approves migration execution at the MIGRATION_REVIEW gate, advancing to execution. + + Valid only when the migration is in MIGRATION_REVIEW. If supplied, the ``etag`` must match the current + server value. + + Returns NOT_FOUND if the migration does not exist. Returns FAILED_PRECONDITION if the migration is not + in MIGRATION_REVIEW. Returns ABORTED if the etag does not match the current resource version. + + :param migration_id: str + The ID of the migration to approve for execution. + :param additional_feedback: str (optional) + Optional free-text feedback or instructions the user provides along with the approval, carried into + the next stage of the migration. + :param etag: str (optional) + Optional resource version for optimistic concurrency control. If supplied, it must match the current + server value. + + :returns: :class:`Migration` + + + .. py:method:: approve_migration_translation(migration_id: str [, additional_feedback: Optional[str], etag: Optional[str]]) -> Migration + + Approves the working copy at the TRANSLATION_READY gate. If the migration was created with + ``run_validation``, the working copy is then run to verify it executes successfully before the + migration is presented for execution review; otherwise it advances straight to MIGRATION_REVIEW. + + Valid only when the migration is in TRANSLATION_READY. If supplied, the ``etag`` must match the + current server value. + + Returns NOT_FOUND if the migration does not exist. Returns FAILED_PRECONDITION if the migration is not + in TRANSLATION_READY. Returns ABORTED if the etag does not match the current resource version. + + :param migration_id: str + The ID of the migration to approve. + :param additional_feedback: str (optional) + Optional free-text feedback or instructions the user provides along with the approval, carried into + the next stage of the migration. + :param etag: str (optional) + Optional resource version for optimistic concurrency control. If supplied, it must match the current + server value. + + :returns: :class:`Migration` + + + .. py:method:: cancel_migration(migration_id: str [, cancellation_reason: Optional[str], etag: Optional[str]]) -> Migration + + Cancels the migration, transitioning it to CANCELLED. + + Valid from any non-terminal state, unlike the review-gate transitions, which each require a specific + state. Any in-flight agent work is abandoned and the migration stops advancing. The resource is + retained for auditability. + + Returns NOT_FOUND if the migration does not exist. Returns FAILED_PRECONDITION if the migration is + already in a terminal state. Returns ABORTED if the etag does not match the current resource version. + + :param migration_id: str + The ID of the migration to cancel. + :param cancellation_reason: str (optional) + Optional free-text reason for the cancellation. + :param etag: str (optional) + Optional resource version for optimistic concurrency control. If supplied, it must match the current + server value. + + :returns: :class:`Migration` + + + .. py:method:: create(migration: Migration [, request_id: Optional[str]]) -> Migration + + Creates a new migration for a Classic compute asset. + + The migration is created in PLANNING_PENDING and advances through the state machine as planning work + is scheduled. The caller must have appropriate permissions: workspace admin for infrastructure-scoped + migrations, job owner/admin for individual workload migrations. + + At most one migration may hold an asset. Returns ALREADY_EXISTS if the asset has a migration in + flight, or one that already migrated it successfully; a migration that ended without succeeding, or + was deleted, releases the asset so it can be migrated again. Returns INVALID_PARAMETER_VALUE if the + asset is not eligible for migration. + + :param migration: :class:`Migration` + The migration to create. The server generates ``migration_id`` and sets ``state`` to + PLANNING_PENDING. Required fields: ``asset_type``, ``asset_uri``. + :param request_id: str (optional) + Client-generated request ID for idempotent request deduplication. The server de-duplicates requests + with the same request_id for at least 24 hours, returning the existing migration instead of creating + a duplicate. + + :returns: :class:`Migration` + + + .. py:method:: delete(migration_id: str [, etag: Optional[str]]) -> Migration + + Deletes a migration (soft-delete). + + Allowed from any state. Deleting a migration marks the record deleted only — it does not roll back a + migration that already ran, and resources created during the migration are not cleaned up. The + migration resource is retained for auditability. To stop an in-flight migration (agent teardown, + transition to CANCELLED), use CancelMigration instead. + + Returns NOT_FOUND if the migration does not exist, or was already deleted: a deleted migration is + hidden from every method (Get, Update, List, and Delete), so it cannot be deleted again. Returns + ABORTED if a supplied nonempty request etag does not match the resource version. + + :param migration_id: str + The ID of the migration to cancel. + :param etag: str (optional) + If nonempty, must match the current resource etag or the request fails with ABORTED. + + :returns: :class:`Migration` + + + .. py:method:: get(migration_id: str) -> Migration + + Retrieves a migration by its ID. + + Returns the current state, configuration, and output artifacts of the migration. Use this to poll + migration progress or inspect artifacts after state transitions. Returns NOT_FOUND if the migration + does not exist in the caller's workspace. + + :param migration_id: str + The ID of the migration to retrieve. + + :returns: :class:`Migration` + + + .. py:method:: list( [, asset_type: Optional[MigrationAssetType], page_size: Optional[int], page_token: Optional[str], states: Optional[List[MigrationState]]]) -> Iterator[Migration] + + Lists migrations in the caller's workspace, ordered by migration ID, with cursor-based pagination and + filtering by asset type and state. + + Eventually consistent: results are read from a cache, and pagination is not a point-in-time snapshot, + so a migration created during a List may not appear until a subsequent one. + + :param asset_type: :class:`MigrationAssetType` (optional) + Filter by asset type. + :param page_size: int (optional) + Maximum number of migrations to return. Default: 50, maximum: 1000. + :param page_token: str (optional) + Page token from a previous ListMigrations call. When paginating, all other parameters must match the + call that provided the page token. + :param states: List[:class:`MigrationState`] (optional) + Filter by migration state. Returns migrations in any of the specified states. + + :returns: Iterator over :class:`Migration` + + + .. py:method:: reject_migration_execution(migration_id: str, rejection_reason: str [, etag: Optional[str]]) -> Migration + + Rejects migration execution at the MIGRATION_REVIEW gate, sending the migration to REVISION_PENDING to + loop back for re-planning. + + Valid only when the migration is in MIGRATION_REVIEW. If supplied, the ``etag`` must match the current + server value. + + Returns NOT_FOUND if the migration does not exist. Returns FAILED_PRECONDITION if the migration is not + in MIGRATION_REVIEW. Returns ABORTED if the etag does not match the current resource version. + + :param migration_id: str + The ID of the migration to reject. + :param rejection_reason: str + Required free-text reason for the rejection, fed back into re-planning so the next plan can address + the user's concerns. + :param etag: str (optional) + Optional resource version for optimistic concurrency control. If supplied, it must match the current + server value. + + :returns: :class:`Migration` + + + .. py:method:: reject_migration_translation(migration_id: str, rejection_reason: str [, etag: Optional[str]]) -> Migration + + Rejects the working copy at the TRANSLATION_READY gate, sending the migration to REVISION_PENDING to + loop back for re-planning. + + Valid only when the migration is in TRANSLATION_READY. If supplied, the ``etag`` must match the + current server value. + + Returns NOT_FOUND if the migration does not exist. Returns FAILED_PRECONDITION if the migration is not + in TRANSLATION_READY. Returns ABORTED if the etag does not match the current resource version. + + :param migration_id: str + The ID of the migration to reject. + :param rejection_reason: str + Required free-text reason for the rejection, fed back into re-planning so the next plan can address + the user's concerns. + :param etag: str (optional) + Optional resource version for optimistic concurrency control. If supplied, it must match the current + server value. + + :returns: :class:`Migration` + \ No newline at end of file diff --git a/docs/workspace/computemigrations/index.rst b/docs/workspace/computemigrations/index.rst new file mode 100644 index 000000000..2ecc02098 --- /dev/null +++ b/docs/workspace/computemigrations/index.rst @@ -0,0 +1,10 @@ + +Compute Migrations +================== + +Manage migrations of Classic compute workloads to Serverless. + +.. toctree:: + :maxdepth: 1 + + compute_migrations \ No newline at end of file diff --git a/docs/workspace/genieagents/genie_agents.rst b/docs/workspace/genieagents/genie_agents.rst new file mode 100644 index 000000000..96437c800 --- /dev/null +++ b/docs/workspace/genieagents/genie_agents.rst @@ -0,0 +1,94 @@ +``w.genie_agents``: GenieAgents.v1 +================================== +.. currentmodule:: databricks.sdk.service.genieagents + +.. py:class:: GenieAgentsAPI + + Genie provides a no-code experience for business users, powered by AI/BI. Analysts set up Genie Agents + that business users can use to ask questions using natural language. Genie uses data registered to Unity + Catalog and requires at least CAN USE permission on a Pro or Serverless SQL warehouse. Also, Databricks + Assistant must be enabled. + + .. py:method:: create_genie_agent(genie_agent: GenieAgent [, genie_agent_id: Optional[str]]) -> GenieAgent + + Creates a Genie Agent. + + :param genie_agent: :class:`GenieAgent` + The Genie Agent to create. + :param genie_agent_id: str (optional) + The ID for the new Genie Agent: 32 lowercase hexadecimal characters (a UUID without dashes), not all + zeros. Assigned by the server if omitted. Fails with ALREADY_EXISTS if the ID is in use in the + workspace, including by a Genie Agent in the trash, or PERMISSION_DENIED if the caller cannot view + the Genie Agent using it. + + :returns: :class:`GenieAgent` + + + .. py:method:: delete_genie_agent(name: str [, etag: Optional[str]]) + + Deletes a Genie Agent, moving it to trash. + + :param name: str + The resource name of the Genie Agent to delete, in the form ``agents/{genie_agent_id}``. + :param etag: str (optional) + Etag from a previous read. On mismatch the delete fails with ABORTED and nothing is deleted. + + + + + .. py:method:: get_genie_agent(name: str) -> GenieAgent + + Gets a Genie Agent. + + :param name: str + The resource name of the Genie Agent to retrieve, in the form ``agents/{genie_agent_id}``. + + :returns: :class:`GenieAgent` + + + .. py:method:: list_genie_agents( [, page_size: Optional[int], page_token: Optional[str], view: Optional[GenieAgentView]]) -> Iterator[GenieAgent] + + Lists the Genie Agents in the workspace. Genie Agents created from an integration (for example, from a + dashboard) are not listed. + + :param page_size: int (optional) + Maximum number of Genie Agents to return per page. Defaults to 20 when unset or 0; values above 100 + are coerced to 100. Negative values are rejected with INVALID_PARAMETER_VALUE. With + ``GENIE_AGENT_VIEW_FULL``, a smaller per-page maximum applies. + :param page_token: str (optional) + Pagination token from a previous response's ``next_page_token``. A token that was not returned by + this method is rejected with INVALID_PARAMETER_VALUE. + :param view: :class:`GenieAgentView` (optional) + The fields to return for each Genie Agent. Defaults to ``GENIE_AGENT_VIEW_BASIC``. Use + ``GENIE_AGENT_VIEW_FULL`` only when the agent bodies are needed, since it returns smaller pages. + + :returns: Iterator over :class:`GenieAgent` + + + .. py:method:: update_genie_agent(name: str, genie_agent: GenieAgent, update_mask: FieldMask [, etag: Optional[str]]) -> GenieAgent + + Updates a Genie Agent. Only the fields named in update_mask are applied. + + :param name: str + The resource name of the Genie Agent, in the form ``agents/{genie_agent_id}``. Ignored on create. + :param genie_agent: :class:`GenieAgent` + The Genie Agent to update. Its ``name`` identifies the resource; fields named in update_mask are + applied. + :param update_mask: FieldMask + The set of fields to update. The mask is required and must name at least one field; an empty or + absent mask is rejected with INVALID_PARAMETER_VALUE, as is a path that is not an updatable field. + + A field named in the mask takes the value from ``genie_agent``, even when it is unset there: an + unset ``description`` is cleared; an unset ``display_name``, ``warehouse_id``, or ``definition`` is + rejected with INVALID_PARAMETER_VALUE; an unset ``parent_path`` is cleared without moving the Genie + Agent. ``definition`` is replaced as a whole; paths within it are not supported. + + The special value ``*`` replaces the whole Genie Agent, as if every updatable field were named. It + must be the only path in the mask. Prefer naming fields explicitly: with ``*``, any field added to + Genie Agent later is also replaced, so a client unaware of it may clear it. + :param etag: str (optional) + Etag from a previous read. On mismatch the update fails with ABORTED (metadata changes may already + be applied); get the agent and retry. May return TEMPORARILY_UNAVAILABLE. + + :returns: :class:`GenieAgent` + \ No newline at end of file diff --git a/docs/workspace/genieagents/index.rst b/docs/workspace/genieagents/index.rst new file mode 100644 index 000000000..f94f5e425 --- /dev/null +++ b/docs/workspace/genieagents/index.rst @@ -0,0 +1,10 @@ + +Genie Agents +============ + +Genie provides a no-code experience for business users, powered by AI/BI. Analysts set up Genie Agents that business users can use to ask questions using natural language. Genie uses data registered to Unity Catalog and requires at least CAN USE permission on a Pro or Serverless SQL warehouse. Also, Databricks Assistant must be enabled. + +.. toctree:: + :maxdepth: 1 + + genie_agents \ No newline at end of file diff --git a/docs/workspace/iamv2/workspace_iam_v2.rst b/docs/workspace/iamv2/workspace_iam_v2.rst index 16076261c..6d16aba82 100644 --- a/docs/workspace/iamv2/workspace_iam_v2.rst +++ b/docs/workspace/iamv2/workspace_iam_v2.rst @@ -6,19 +6,19 @@ These APIs are used to manage identities and the workspace access of these identities in Databricks. - .. py:method:: create_direct_group_member_proxy(group_id: int, direct_group_member: DirectGroupMember) -> DirectGroupMember + .. py:method:: create_direct_group_member_proxy(parent: str, direct_group_member: DirectGroupMember) -> DirectGroupMember Creates a group membership (assigns a principal to a group). - :param group_id: int - Required. Internal ID of the group in Databricks. + :param parent: str + The parent resource name of the group to add the member to. Format: groups/{group_id} :param direct_group_member: :class:`DirectGroupMember` Required. The group membership to create. :returns: :class:`DirectGroupMember` - .. py:method:: create_group_proxy(group: Group) -> Group + .. py:method:: create_group_proxy(group: Group [, parent: Optional[str]]) -> Group Creates a local group in the Databricks account that parents the calling workspace and returns the created group. A local group is one that is not synced from the customer's identity provider, and can @@ -30,11 +30,14 @@ :param group: :class:`Group` Required. Group to be created in Databricks + :param parent: str (optional) + The resource name of the account to create the group in, ``accounts/{account_id}``. Optional: + defaults to the account of the calling workspace, and any other account is rejected. :returns: :class:`Group` - .. py:method:: create_service_principal_proxy(service_principal: ServicePrincipal) -> ServicePrincipal + .. py:method:: create_service_principal_proxy(service_principal: ServicePrincipal [, parent: Optional[str]]) -> ServicePrincipal Creates a local service principal in the Databricks account that parents the calling workspace and returns the created service principal. A local service principal is one that is not synced from the @@ -47,11 +50,14 @@ :param service_principal: :class:`ServicePrincipal` Required. Service principal to be created in Databricks + :param parent: str (optional) + The resource name of the account to create the service principal in, ``accounts/{account_id}``. + Optional: defaults to the account of the calling workspace, and any other account is rejected. :returns: :class:`ServicePrincipal` - .. py:method:: create_user_proxy(user: User) -> User + .. py:method:: create_user_proxy(user: User [, parent: Optional[str]]) -> User Creates a local user in the Databricks account that parents the calling workspace and returns the created user. A local user is one that is not synced from the customer's identity provider, and can be @@ -63,6 +69,9 @@ :param user: :class:`User` Required. User to be created in Databricks + :param parent: str (optional) + The resource name of the account to create the user in, ``accounts/{account_id}``. Optional: + defaults to the account of the calling workspace, and any other account is rejected. :returns: :class:`User` @@ -77,7 +86,7 @@ :returns: :class:`WorkspaceAssignmentDetail` - .. py:method:: create_workspace_assignment_proxy(workspace_assignment: WorkspaceAssignment) -> WorkspaceAssignment + .. py:method:: create_workspace_assignment_proxy(workspace_assignment: WorkspaceAssignment [, parent: Optional[str]]) -> WorkspaceAssignment Creates a workspace assignment for a principal in the calling workspace. Entitlements are granted one at a time rather than atomically. If the request fails partway through, the principal stays assigned @@ -86,49 +95,51 @@ :param workspace_assignment: :class:`WorkspaceAssignment` Required. Workspace assignment to be created in Databricks. + :param parent: str (optional) + The resource name of the workspace to assign the principal to, + ``accounts/{account_id}/workspaces/{workspace_id}``. Optional: defaults to the calling workspace, + and any other workspace is rejected. :returns: :class:`WorkspaceAssignment` - .. py:method:: delete_direct_group_member_proxy(group_id: int, principal_id: int) + .. py:method:: delete_direct_group_member_proxy(name: str) Deletes a group membership (unassigns a principal from a group). - :param group_id: int - Required. Internal ID of the group in Databricks. - :param principal_id: int - Required. Internal ID of the principal to be unassigned from the group. + :param name: str + The resource name of the group membership. Format: groups/{group_id}/direct-members/{principal_id} - .. py:method:: delete_group_proxy(group_id: str) + .. py:method:: delete_group_proxy(name: str) Deletes a group by its internal ID from the Databricks account that parents the calling workspace. - :param group_id: str - Required. Internal ID of the group in Databricks. + :param name: str + The resource name of the group. Format: groups/{group_id} - .. py:method:: delete_service_principal_proxy(service_principal_id: str) + .. py:method:: delete_service_principal_proxy(name: str) Deletes a service principal by its internal ID from the Databricks account that parents the calling workspace. - :param service_principal_id: str - Required. Internal ID of the service principal in Databricks. + :param name: str + The resource name of the service principal. Format: service-principals/{service_principal_id} - .. py:method:: delete_user_proxy(user_id: str) + .. py:method:: delete_user_proxy(name: str) Deletes a user by its internal ID from the Databricks account that parents the calling workspace. - :param user_id: str - Required. Internal ID of the user in Databricks. + :param name: str + The resource name of the user. Format: users/{user_id} @@ -146,27 +157,25 @@ - .. py:method:: delete_workspace_assignment_proxy(principal_id: int) + .. py:method:: delete_workspace_assignment_proxy(name: str) Deletes a workspace assignment for a principal in the calling workspace, revoking all of its entitlements. Entitlements are revoked one at a time rather than atomically. If the request fails partway through, the principal stays assigned with some of its original entitlements. Retrying is safe. - :param principal_id: int - Required. ID of the principal in Databricks to delete workspace assignment for. + :param name: str + The resource name of the workspace assignment. Format: workspace-assignments/{principal_id} - .. py:method:: get_direct_group_member_proxy(group_id: int, principal_id: int) -> DirectGroupMember + .. py:method:: get_direct_group_member_proxy(name: str) -> DirectGroupMember Gets a provisioned direct member of a group. - :param group_id: int - Required. Internal ID of the group in Databricks. - :param principal_id: int - Required. Internal ID of the principal belonging to the group in Databricks. + :param name: str + The resource name of the group membership. Format: groups/{group_id}/direct-members/{principal_id} :returns: :class:`DirectGroupMember` @@ -210,38 +219,38 @@ :returns: :class:`ExternalUser` - .. py:method:: get_group_proxy(group_id: str) -> Group + .. py:method:: get_group_proxy(name: str) -> Group Fetches a group by its internal ID from the Databricks account that parents the calling workspace. - :param group_id: str - Required. Internal ID of the group in Databricks. + :param name: str + The resource name of the group. Format: groups/{group_id} :returns: :class:`Group` - .. py:method:: get_service_principal_proxy(service_principal_id: str) -> ServicePrincipal + .. py:method:: get_service_principal_proxy(name: str) -> ServicePrincipal Fetches a service principal by its internal ID from the Databricks account that parents the calling workspace. - :param service_principal_id: str - Required. Internal ID of the service principal in Databricks. + :param name: str + The resource name of the service principal. Format: service-principals/{service_principal_id} :returns: :class:`ServicePrincipal` - .. py:method:: get_user_proxy(user_id: str) -> User + .. py:method:: get_user_proxy(name: str) -> User Fetches a user by its internal ID from the Databricks account that parents the calling workspace. - :param user_id: str - Required. Internal ID of the user in Databricks. + :param name: str + The resource name of the user. Format: users/{user_id} :returns: :class:`User` - .. py:method:: get_workspace_access_detail_local(principal_id: int [, view: Optional[WorkspaceAccessDetailView]]) -> WorkspaceAccessDetail + .. py:method:: get_workspace_access_detail_local(name: str [, view: Optional[WorkspaceAccessDetailView]]) -> WorkspaceAccessDetail Returns the access details for a principal in the current workspace. Allows for checking access details for any provisioned principal (user, service principal, or group) in the current workspace. @@ -250,9 +259,8 @@ IdP or added explicitly to Databricks via SCIM/UI. Allows for passing in a "view" parameter to control what fields are returned (BASIC by default or FULL). - :param principal_id: int - Required. The internal ID of the principal (user/sp/group) for which the access details are being - requested. + :param name: str + The resource name of the workspace access detail. Format: workspace-access-details/{principal_id} :param view: :class:`WorkspaceAccessDetailView` (optional) Controls what fields are returned. @@ -270,35 +278,34 @@ :returns: :class:`WorkspaceAssignmentDetail` - .. py:method:: get_workspace_assignment_proxy(principal_id: int) -> WorkspaceAssignment + .. py:method:: get_workspace_assignment_proxy(name: str) -> WorkspaceAssignment Returns the assignment for a principal in the calling workspace. - :param principal_id: int - Required. The internal ID of the principal (user/sp/group) for which the assignment is being - requested. + :param name: str + The resource name of the workspace assignment. Format: workspace-assignments/{principal_id} :returns: :class:`WorkspaceAssignment` - .. py:method:: get_workspace_identity_detail(principal_id: int) -> WorkspaceIdentityDetail + .. py:method:: get_workspace_identity_detail(name: str) -> WorkspaceIdentityDetail Returns the identity details for a principal in a workspace. - :param principal_id: int - Required. The internal ID of the principal (user/sp/group) for which the identity details are being - requested. + :param name: str + The resource name of the workspace identity detail. Format: + principals/{principal_id}/workspace-identity-detail :returns: :class:`WorkspaceIdentityDetail` - .. py:method:: list_direct_group_members_proxy(group_id: int [, page_size: Optional[int], page_token: Optional[str]]) -> Iterator[DirectGroupMember] + .. py:method:: list_direct_group_members_proxy(parent: str [, page_size: Optional[int], page_token: Optional[str]]) -> Iterator[DirectGroupMember] Lists provisioned direct members of a group with their membership source (internal or from identity provider). - :param group_id: int - Required. Internal ID of the group in Databricks whose direct members are being listed. + :param parent: str + The parent resource name of the group whose direct members to list. Format: groups/{group_id} :param page_size: int (optional) The maximum number of members to return. The service may return fewer than this value. If not provided, defaults to 1000, which is also the maximum allowed. Requests for more than the maximum @@ -309,7 +316,7 @@ :returns: Iterator over :class:`DirectGroupMember` - .. py:method:: list_groups_proxy( [, filter: Optional[str], page_size: Optional[int], page_token: Optional[str]]) -> Iterator[Group] + .. py:method:: list_groups_proxy( [, filter: Optional[str], page_size: Optional[int], page_token: Optional[str], parent: Optional[str]]) -> Iterator[Group] Lists the groups in the Databricks account that parents the calling workspace, returning one page per call. Supports filtering by group name or external ID. @@ -323,11 +330,14 @@ :param page_token: str (optional) A page token, received from a previous ListGroups call. Provide this to retrieve the subsequent page. + :param parent: str (optional) + The resource name of the account to list groups in, ``accounts/{account_id}``. Optional: defaults to + the account of the calling workspace, and any other account is rejected. :returns: Iterator over :class:`Group` - .. py:method:: list_service_principals_proxy( [, filter: Optional[str], page_size: Optional[int], page_token: Optional[str]]) -> Iterator[ServicePrincipal] + .. py:method:: list_service_principals_proxy( [, filter: Optional[str], page_size: Optional[int], page_token: Optional[str], parent: Optional[str]]) -> Iterator[ServicePrincipal] Lists the service principals in the Databricks account that parents the calling workspace, returning one page per call. Supports filtering by application ID or external ID. @@ -341,17 +351,20 @@ :param page_token: str (optional) A page token, received from a previous ListServicePrincipals call. Provide this to retrieve the subsequent page. + :param parent: str (optional) + The resource name of the account to list service principals in, ``accounts/{account_id}``. Optional: + defaults to the account of the calling workspace, and any other account is rejected. :returns: Iterator over :class:`ServicePrincipal` - .. py:method:: list_transitive_parent_groups_proxy(principal_id: int [, page_size: Optional[int], page_token: Optional[str]]) -> ListTransitiveParentGroupsResponse + .. py:method:: list_transitive_parent_groups_proxy(parent: str [, page_size: Optional[int], page_token: Optional[str]]) -> ListTransitiveParentGroupsResponse Lists all transitive parent groups of a principal. - :param principal_id: int - Required. Internal ID of the principal in Databricks whose transitive parent groups are being - listed. + :param parent: str + The parent resource name of the principal whose transitive parent groups to list. Format: + principals/{principal_id} :param page_size: int (optional) The maximum number of parent groups to return. The service may return fewer than this value. If not provided, defaults to 1000, which is also the maximum allowed. Requests for more than the maximum @@ -363,7 +376,7 @@ :returns: :class:`ListTransitiveParentGroupsResponse` - .. py:method:: list_users_proxy( [, filter: Optional[str], page_size: Optional[int], page_token: Optional[str]]) -> Iterator[User] + .. py:method:: list_users_proxy( [, filter: Optional[str], page_size: Optional[int], page_token: Optional[str], parent: Optional[str]]) -> Iterator[User] Lists the users in the Databricks account that parents the calling workspace, returning one page per call. Supports filtering by username or external ID. @@ -376,6 +389,9 @@ are clamped to 1000. :param page_token: str (optional) A page token, received from a previous ListUsers call. Provide this to retrieve the subsequent page. + :param parent: str (optional) + The resource name of the account to list users in, ``accounts/{account_id}``. Optional: defaults to + the account of the calling workspace, and any other account is rejected. :returns: Iterator over :class:`User` @@ -396,7 +412,7 @@ :returns: Iterator over :class:`WorkspaceAssignmentDetail` - .. py:method:: list_workspace_assignments_proxy( [, page_size: Optional[int], page_token: Optional[str]]) -> Iterator[WorkspaceAssignment] + .. py:method:: list_workspace_assignments_proxy( [, page_size: Optional[int], page_token: Optional[str], parent: Optional[str]]) -> Iterator[WorkspaceAssignment] Lists workspace assignments for the calling workspace. The response omits the per-principal entitlement fields (``entitlements`` and ``effective_entitlements``). To read the entitlements for a @@ -408,6 +424,10 @@ maximum are clamped to 1000. :param page_token: str (optional) A page token from a previous list call. Provide this to retrieve the subsequent page. + :param parent: str (optional) + The resource name of the workspace to list assignments of, + ``accounts/{account_id}/workspaces/{workspace_id}``. Optional: defaults to the calling workspace, + and any other workspace is rejected. :returns: Iterator over :class:`WorkspaceAssignment` @@ -448,7 +468,7 @@ :returns: :class:`ResolveUserResponse` - .. py:method:: update_group_proxy(group_id: str, group: Group, update_mask: str) -> Group + .. py:method:: update_group_proxy(name: str, group: Group, update_mask: str) -> Group Updates an existing group in the Databricks account that parents the calling workspace. Only the fields named in the update mask are modified. Returns the updated Group resource. @@ -458,8 +478,11 @@ visibility filters, a new nonempty external_id must identify a visible group. Clearing external_id is allowed. - :param group_id: str - Required. Internal ID of the group in Databricks. + :param name: str + The resource name of the group. The format depends on the API that returned it: + + - Account-scoped: accounts/{account_id}/groups/{group_id} + - Workspace-scoped: groups/{group_id} :param group: :class:`Group` Required. Group to be updated in Databricks :param update_mask: str @@ -468,18 +491,22 @@ :returns: :class:`Group` - .. py:method:: update_service_principal_proxy(service_principal_id: str, service_principal: ServicePrincipal, update_mask: str) -> ServicePrincipal + .. py:method:: update_service_principal_proxy(name: str, service_principal: ServicePrincipal, update_mask: str) -> ServicePrincipal Updates an existing service principal in the Databricks account that parents the calling workspace. Only the fields named in the update mask are modified. Returns the updated ServicePrincipal resource. - For locally managed service principals, is_agent must be set to true or false when updated. + To change is_agent on a locally managed service principal, set it to true or false. Naming is_agent in + the update mask without a value leaves it unchanged. When AIM is enabled and the service principal is an external identity (its external_id is set), only external_id can be updated; its other fields are sourced from your identity provider. - :param service_principal_id: str - Required. Internal ID of the service principal in Databricks. + :param name: str + The resource name of the service principal. The format depends on the API that returned it: + + - Account-scoped: accounts/{account_id}/service-principals/{service_principal_id} + - Workspace-scoped: service-principals/{service_principal_id} :param service_principal: :class:`ServicePrincipal` Required. Service principal to be updated in Databricks :param update_mask: str @@ -488,7 +515,7 @@ :returns: :class:`ServicePrincipal` - .. py:method:: update_user_proxy(user_id: str, user: User, update_mask: str) -> User + .. py:method:: update_user_proxy(name: str, user: User, update_mask: str) -> User Updates an existing user in the Databricks account that parents the calling workspace and returns the updated user. Only the fields named in the update mask are modified. The updatable fields are @@ -497,8 +524,11 @@ When AIM is enabled and the user is an external identity (its external_id is set), only external_id can be updated; its other fields are sourced from your identity provider. - :param user_id: str - Required. Internal ID of the user in Databricks. + :param name: str + The resource name of the user. The format depends on the API that returned it: + + - Account-scoped: accounts/{account_id}/users/{user_id} + - Workspace-scoped: users/{user_id} :param user: :class:`User` Required. User to be updated in Databricks :param update_mask: str @@ -523,14 +553,18 @@ :returns: :class:`WorkspaceAssignmentDetail` - .. py:method:: update_workspace_assignment_proxy(principal_id: int, workspace_assignment: WorkspaceAssignment, update_mask: FieldMask) -> WorkspaceAssignment + .. py:method:: update_workspace_assignment_proxy(name: str, workspace_assignment: WorkspaceAssignment, update_mask: FieldMask) -> WorkspaceAssignment Updates the entitlements of a directly assigned principal in the calling workspace. Changes are applied one at a time rather than atomically. If the request fails partway through, only some of the requested changes take effect. Get the assignment afterwards to confirm the final state. - :param principal_id: int - Required. ID of the principal in Databricks. + :param name: str + The resource name of the workspace assignment. The format depends on the API that returned it: + + - Account-scoped: + accounts/{account_id}/workspaces/{workspace_id}/workspace-assignments/{principal_id} + - Workspace-scoped: workspace-assignments/{principal_id} :param workspace_assignment: :class:`WorkspaceAssignment` Required. Workspace assignment to be updated in Databricks. :param update_mask: FieldMask @@ -539,12 +573,13 @@ :returns: :class:`WorkspaceAssignment` - .. py:method:: update_workspace_identity_detail(principal_id: int, workspace_identity_detail: WorkspaceIdentityDetail, update_mask: FieldMask) -> WorkspaceIdentityDetail + .. py:method:: update_workspace_identity_detail(name: str, workspace_identity_detail: WorkspaceIdentityDetail, update_mask: FieldMask) -> WorkspaceIdentityDetail Updates a workspace identity detail for a principal. - :param principal_id: int - Required. ID of the principal in Databricks. + :param name: str + The resource name of the workspace identity detail. Format: + principals/{principal_id}/workspace-identity-detail :param workspace_identity_detail: :class:`WorkspaceIdentityDetail` Required. Workspace identity detail to be updated in Databricks. :param update_mask: FieldMask diff --git a/docs/workspace/index.rst b/docs/workspace/index.rst index 6d1f125c0..5fab702c1 100644 --- a/docs/workspace/index.rst +++ b/docs/workspace/index.rst @@ -16,6 +16,7 @@ These APIs are available from WorkspaceClient catalog/index cleanrooms/index compute/index + computemigrations/index dashboards/index database/index dataclassification/index @@ -23,6 +24,7 @@ These APIs are available from WorkspaceClient domains/index environments/index files/index + genieagents/index iam/index iamv2/index jobs/index diff --git a/docs/workspace/jobs/jobs.rst b/docs/workspace/jobs/jobs.rst index 686c1eb38..beab02956 100644 --- a/docs/workspace/jobs/jobs.rst +++ b/docs/workspace/jobs/jobs.rst @@ -203,10 +203,11 @@ serverless compute or a compute resource that uses Environments mode. For notebook tasks that use serverless compute or a compute resource that uses Environments mode, if - the environment_key is not specified, the notebook environment will be used if present. If a jobs - environment is specified, it will override the notebook environment. For other tasks that use - serverless compute or a compute resource that uses Environments mode, the task environment is - required to be specified using environment_key in the task settings. + the ``environment_key`` is not specified, the notebook environment will be used if present. If a + jobs environment is specified, it will override the notebook environment. + + For other tasks that use serverless compute or a compute resource that uses Environments mode, the + task environment is required to be specified using ``environment_key`` in the task settings. :param format: :class:`Format` (optional) Used to tell what is the format of the job. This field is ignored in Create/Update/Reset calls. When using the Jobs API 2.1 this value is always set to ``"MULTI_TASK"``. diff --git a/docs/workspace/ml/feature_engineering.rst b/docs/workspace/ml/feature_engineering.rst index b5a23818d..ef00b0610 100644 --- a/docs/workspace/ml/feature_engineering.rst +++ b/docs/workspace/ml/feature_engineering.rst @@ -11,9 +11,11 @@ Backfill features. :param feature_full_names: List[str] - Full names of the features to backfill. + Full names of the features to backfill. All features must share the same interval between values and + the same offset. The interval is the slide_duration of a sliding window or the window_duration of a + tumbling window. :param backfill_ranges: List[:class:`BackfillRange`] - Output ranges to backfill. + Time ranges of source data to backfill. Ranges must not overlap. :param budget_policy_id: str (optional) The budget policy ID, in UUID format, used to attribute the serverless compute cost of this backfill. If not specified, a default budget policy may be applied. @@ -374,15 +376,14 @@ .. py:method:: update_materialized_feature(materialized_feature_id: str, materialized_feature: MaterializedFeature, update_mask: str) -> MaterializedFeature - Update a materialized feature (pause/resume). + Updates a materialized feature. :param materialized_feature_id: str Server-assigned unique identifier for the materialized feature. :param materialized_feature: :class:`MaterializedFeature` The materialized feature to update. :param update_mask: str - Provide the materialization feature fields which should be updated. Currently, only the - pipeline_state field can be updated. + The fields to update. Only ``pipeline_schedule_state`` is supported. :returns: :class:`MaterializedFeature` diff --git a/docs/workspace/postgres/postgres.rst b/docs/workspace/postgres/postgres.rst index b75e23869..adda32f60 100644 --- a/docs/workspace/postgres/postgres.rst +++ b/docs/workspace/postgres/postgres.rst @@ -146,12 +146,16 @@ :param role: :class:`Role` The desired specification of a Role. :param replace_existing: bool (optional) - If true, update the role if it already exists instead of returning an error. - - When the role already exists, the provided ``role`` spec fully replaces the existing one: - ``membership_roles`` is overwritten, not merged. Leaving ``membership_roles`` empty clears all of - the role's existing memberships, including ``DATABRICKS_SUPERUSER``. Always send the complete - desired list of memberships when using this field. + If true, update the role if it already exists instead of returning an error. This requires + ``role_id`` to be set. Without ``role_id``, the request follows the normal create path and fails if + the role name already exists. + + When the role already exists, only mutable fields are updated. Attributes omitted from the request + retain their current values. Set an attribute explicitly to false to disable it. ``identity_type``, + ``auth_method``, and ``postgres_role`` retain their existing values; supplied values for these + fields are ignored. ``membership_roles`` is overwritten, not merged. Leaving ``membership_roles`` + empty clears all of the role's existing memberships, including ``DATABRICKS_SUPERUSER``. Always send + the complete desired list of memberships when using this field. :param role_id: str (optional) The ID to use for the Role, which will become the final component of the role's resource name. This ID becomes the role in Postgres. @@ -725,7 +729,12 @@ .. py:method:: update_role(name: str, role: Role, update_mask: FieldMask) -> UpdateRoleOperation - Update a role for a branch. + Updates the mutable fields of a Postgres role for a branch. + + The fields ``spec.identity_type``, ``spec.auth_method``, and ``spec.postgres_role`` are fixed when the + role is created. UpdateRole ignores values supplied for these fields in the request body and retains + their existing values. Explicitly naming one of these fields in ``update_mask`` returns an + INVALID_PARAMETER_VALUE error. :param name: str Output only. The full resource path of the role. Format: @@ -736,7 +745,17 @@ The role's ``name`` field is used to identify the role to update. Format: projects/{project_id}/branches/{branch_id}/roles/{role_id} :param update_mask: FieldMask - The list of fields to update. + The list of fields to update. Supported fields are ``spec.membership_roles``, + ``spec.attributes.createdb``, ``spec.attributes.createrole``, and ``spec.attributes.bypassrls``. + + ``spec.identity_type``, ``spec.auth_method``, and ``spec.postgres_role`` cannot be changed. Values + supplied for these fields in the request body are ignored. Including one of these fields as an + explicit path in update_mask returns an INVALID_PARAMETER_VALUE error. When ``spec`` is the only + path in the mask, only the supplied mutable attributes are updated. ``*`` by itself, or ``spec`` + combined with other supported paths, selects all supported mutable fields. For these masks, all + three attributes (``createdb``, ``createrole``, and ``bypassrls``) must be supplied; omitting any + returns an INVALID_PARAMETER_VALUE error. Each of these mask forms replaces ``membership_roles``. + Omitting it or supplying an empty list clears all memberships, including ``DATABRICKS_SUPERUSER``. :returns: :class:`Operation` diff --git a/docs/workspace/sandbox/sandbox.rst b/docs/workspace/sandbox/sandbox.rst index 715c78674..c1a936f66 100644 --- a/docs/workspace/sandbox/sandbox.rst +++ b/docs/workspace/sandbox/sandbox.rst @@ -68,6 +68,48 @@ + .. py:method:: download_file(parent: str, path: str [, byte_range: Optional[str], if_match: Optional[str], if_modified_since: Optional[str], if_none_match: Optional[str], if_range: Optional[str], if_unmodified_since: Optional[str]]) -> DownloadFileResponse + + Downloads a regular file from a sandbox as an application/octet-stream body. The sandbox starts + automatically if it is stopped. Filesystem access uses the sandbox user's permissions and follows + symbolic links. + + A full download returns 200 OK. A satisfiable single byte range returns 206 Partial Content. Multiple + byte ranges are unsupported and return 400. Invalid ranges are ignored; an unsatisfiable range returns + 416 with Content-Range describing the file length. Range is ignored for empty files. + + Conditional requests can return 304 Not Modified or 412 Precondition Failed. If-Range selects a + partial response only when its strong entity tag matches; otherwise, the entire file is returned. + HTTP-date validators in If-Range cause a full download. Reads are not snapshots: concurrent in-place + writes can affect the returned bytes. HEAD is unsupported. + + Range and conditional options may be supplied as query parameters or HTTP headers. A parameter and one + header may both be supplied when their values match exactly. Conflicting values or repeated parameters + return 400. + + :param parent: str + The sandbox resource name, in the form sandboxes/{sandbox_id}. + :param path: str + The absolute UTF-8 path of the file inside the sandbox. + :param byte_range: str (optional) + A single byte range, such as bytes=0-1023 or bytes=-1024. Equivalent to the Range header. + :param if_match: str (optional) + Entity tags or *. A mismatch returns 412 Precondition Failed. Equivalent to the If-Match header. + :param if_modified_since: str (optional) + An HTTP-date. Returns 304 if the file has not changed since this time. Ignored when If-None-Match is + present. Equivalent to the If-Modified-Since header. + :param if_none_match: str (optional) + Entity tags or *. A match returns 304 Not Modified. Equivalent to the If-None-Match header. + :param if_range: str (optional) + A strong entity tag. A mismatch, weak entity tag, or HTTP-date causes a full download. Equivalent to + the If-Range header. + :param if_unmodified_since: str (optional) + An HTTP-date. Returns 412 if the file has changed since this time. Ignored when If-Match is present. + Equivalent to the If-Unmodified-Since header. + + :returns: :class:`DownloadFileResponse` + + .. py:method:: execute_command(name: str, cmd: str [, args: Optional[list[str]], envs: Optional[Dict[str, str]], text: bool = True, stdout: StreamMode = StreamMode.PIPE, stderr: StreamMode = StreamMode.PIPE]) -> CommandProcess Start ``cmd`` in sandbox ``name`` and return a :class:`CommandProcess` @@ -218,4 +260,27 @@ - spec.compute.inactivity_timeout Any other path returns INVALID_PARAMETER_VALUE. :returns: :class:`Sandbox` + + + .. py:method:: upload_file(parent: str, contents: BinaryIO, path: str) + + Uploads an application/octet-stream body to a sandbox. The sandbox starts automatically if it is + stopped. Filesystem access uses the sandbox user's permissions and follows symbolic links. + + Creates the file or replaces its entire contents, creating parent directories if missing. The + completed upload becomes visible atomically. A failure before commit leaves an existing destination + unchanged; a lost response after commit does not undo the upload. Completion does not guarantee crash + durability. + + Returns 201 Created for a new file or 204 No Content for a replacement. Partial and conditional + uploads are unsupported: Content-Range, If-Match, If-None-Match, and If-Unmodified-Since return 400. + + :param parent: str + The sandbox resource name, in the form sandboxes/{sandbox_id}. + :param contents: BinaryIO + The complete file contents. An empty body creates an empty file. + :param path: str + The absolute UTF-8 destination path inside the sandbox. + + \ No newline at end of file diff --git a/docs/workspace/sql/warehouses.rst b/docs/workspace/sql/warehouses.rst index c5dfdac93..ce725c4a8 100644 --- a/docs/workspace/sql/warehouses.rst +++ b/docs/workspace/sql/warehouses.rst @@ -78,8 +78,9 @@ :param instance_profile_arn: str (optional) Deprecated. Instance profile used to pass IAM role to the cluster :param max_dbu: int (optional) - Maximum compute capacity of a real-time SQL warehouse, in DBU per hour. Must be a multiple of 3 - between 3 and 960. The effective maximum cannot be lower than the warehouse's minimum capacity. + Maximum compute capacity of a real-time SQL warehouse, in DBU per hour. Must be at least the minimum + DBU for the selected warehouse size and no greater than the configured DBU limit. Actual capacity + may be below this ceiling. :param max_num_clusters: int (optional) Maximum number of clusters that the autoscaler will create to handle concurrent queries. @@ -258,8 +259,9 @@ :param instance_profile_arn: str (optional) Deprecated. Instance profile used to pass IAM role to the cluster :param max_dbu: int (optional) - Maximum compute capacity of a real-time SQL warehouse, in DBU per hour. Must be a multiple of 3 - between 3 and 960. The effective maximum cannot be lower than the warehouse's minimum capacity. + Maximum compute capacity of a real-time SQL warehouse, in DBU per hour. Must be at least the minimum + DBU for the selected warehouse size and no greater than the configured DBU limit. Actual capacity + may be below this ceiling. :param max_num_clusters: int (optional) Maximum number of clusters that the autoscaler will create to handle concurrent queries. diff --git a/tests/test_files.py b/tests/test_files.py index a86b0df70..6b605e245 100644 --- a/tests/test_files.py +++ b/tests/test_files.py @@ -6,7 +6,6 @@ import os import platform import random -import re import threading import time from dataclasses import dataclass @@ -26,12 +25,10 @@ from databricks.sdk.core import Config from databricks.sdk.environments import Cloud, DatabricksEnvironment from databricks.sdk.errors.platform import ( - AlreadyExists, BadRequest, InternalError, NotImplemented, PermissionDenied, - TooManyRequests, ) from databricks.sdk.mixins.files import ( FallbackToDownloadUsingFilesApi, @@ -1534,7 +1531,7 @@ def __init__( custom_response_on_single_shot_upload: CustomResponse, # exception which is expected to be thrown (so upload is expected to have failed) expected_exception_type: Optional[Type[BaseException]], - # Whether abort is expected to be called for multipart/resumable upload, set to None if we don't care. + # Whether abort is expected to be called for multipart upload, set to None if we don't care. expected_multipart_upload_aborted: Optional[bool], expected_single_shot_upload: bool, use_storage_proxy: bool = False, @@ -1568,7 +1565,7 @@ def _expected_hostname(self) -> str: def customize_config(self, config: Config) -> None: pass - def create_multipart_upload_server_state(self) -> Union[MultipartUploadServerState, "ResumableUploadServerState"]: + def create_multipart_upload_server_state(self) -> MultipartUploadServerState: raise NotImplementedError def clear_state(self) -> None: @@ -1597,7 +1594,7 @@ def get_upload_file( raise ValueError(f"Unknown source type: {source_type}") def match_request_to_response( - self, request: requests.Request, server_state: Union[MultipartUploadServerState, "ResumableUploadServerState"] + self, request: requests.Request, server_state: MultipartUploadServerState ) -> Optional[requests.Response]: raise NotImplementedError @@ -2590,589 +2587,6 @@ def test_multipart_upload(config: Config, test_case: MultipartUploadTestCase) -> test_case.run(config) -class ResumableUploadServerState: - """This server state is updated on resumable upload (GCP)""" - - resumable_upload_url_prefix = "https://cloud_provider.com/resumable-upload/" - - def __init__(self, unconfirmed_delta: Union[int, list], expected_part_size: Optional[int]): - self.unconfirmed_delta = unconfirmed_delta - self.confirmed_last_byte: Optional[int] = None # inclusive - self.uploaded_parts = [] - self.session_token = "token-" + MultipartUploadServerState.randomstr() - self.file_content: Optional[FileContent] = None - self.aborted = False - self.expected_part_size = expected_part_size - - def save_part(self, start_offset: int, end_offset_incl: int, part_content: bytes, file_size_s: str) -> None: - assert not self.aborted - - assert len(part_content) > 0 - if self.expected_part_size is not None: - assert len(part_content) <= self.expected_part_size - - if self.confirmed_last_byte: - assert start_offset == self.confirmed_last_byte + 1 - else: - assert start_offset == 0 - - assert end_offset_incl == start_offset + len(part_content) - 1 - - is_last_part = file_size_s != "*" - if is_last_part: - assert int(file_size_s) == end_offset_incl + 1 - else: - assert not self.file_content # last part should not have been uploaded yet - - if isinstance(self.unconfirmed_delta, int): - unconfirmed_delta = self.unconfirmed_delta - elif len(self.uploaded_parts) < len(self.unconfirmed_delta): - unconfirmed_delta = self.unconfirmed_delta[len(self.uploaded_parts)] - else: - unconfirmed_delta = self.unconfirmed_delta[-1] # take the last delta - - if unconfirmed_delta >= len(part_content): - unconfirmed_delta = 0 # otherwise we never finish - - logger.info( - f"Saving part {len(self.uploaded_parts) + 1} of original size {len(part_content)} with unconfirmed delta {unconfirmed_delta}. is_last_part = {is_last_part}" - ) - - if unconfirmed_delta > 0: - part_content = part_content[:-unconfirmed_delta] - - with NamedTemporaryFile(mode="wb", delete=False) as f: - part_file = f.name - f.write(part_content) - - self.uploaded_parts.append(part_file) - - if is_last_part and unconfirmed_delta == 0: - size = 0 - sha256 = hashlib.sha256() - for part_path in self.uploaded_parts: - size += os.path.getsize(part_path) - with open(part_path, "rb") as f: - part_content = f.read() - sha256.update(part_content) - - assert size == end_offset_incl + 1 - self.file_content = FileContent(size, sha256.hexdigest()) - - self.confirmed_last_byte = end_offset_incl - unconfirmed_delta - - def cleanup(self) -> None: - for file in self.uploaded_parts: - os.remove(file) - - def get_file_content(self) -> Optional[FileContent]: - if self.aborted: - assert not self.file_content - - # content may be None even for a non-aborted upload, - # in case single-shot upload was used due to small stream size. - return self.file_content - - def abort_upload(self) -> None: - self.aborted = True - - -class ResumableUploadTestCase(UploadTestCase): - """Test case for resumable upload of a file. Resumable uploads are used on GCP. - - Resumable upload involves multiple HTTP requests: - - initiating upload (call to Databricks Files API) - - requesting resumable upload URL (call to Databricks Files API) - - uploading data in parts (calls to cloud storage provider or Databricks storage proxy) - - aborting the upload (call to cloud storage provider or Databricks storage proxy) - - Test case uses requests-mock library to mock all these requests. Within a test, mocks use - shared server state that tracks the upload. Mocks generate the "default" (successful) response. - - Response of each call can be modified by parameterising a respective `CustomResponse` object. - """ - - def __init__( - self, - name: str, - stream_size: int, - cloud: Cloud = None, - overwrite: bool = True, - source_type: Optional[List[UploadSourceType]] = None, - use_parallel: Optional[List[bool]] = None, - parallelism: Optional[int] = None, - multipart_upload_min_stream_size: int = 0, # disable single-shot uploads by default - multipart_upload_part_size: Optional[int] = None, - sdk_retry_timeout_seconds: Optional[int] = None, - multipart_upload_max_retries: Optional[int] = None, - # In resumable upload, when replying to part upload request, server returns - # (confirms) last accepted byte offset for the client to resume upload after. - # - # `unconfirmed_delta` defines offset from the end of the part that remains - # "unconfirmed", i.e. the last accepted offset would be (range_end - unconfirmed_delta). - # Can be int (same for all parts) or list (individual for each part). - unconfirmed_delta: Union[int, list] = 0, - custom_response_on_single_shot_upload: CustomResponse = CustomResponse(enabled=False), - custom_response_on_create_resumable_url: CustomResponse = CustomResponse(enabled=False), - custom_response_on_upload: CustomResponse = CustomResponse(enabled=False), - custom_response_on_status_check: CustomResponse = CustomResponse(enabled=False), - custom_response_on_abort: CustomResponse = CustomResponse(enabled=False), - # exception which is expected to be thrown (so upload is expected to have failed) - expected_exception_type: Optional[Type[BaseException]] = None, - # if abort is expected to be called - expected_multipart_upload_aborted: bool = False, - expected_single_shot_upload: bool = False, - expected_part_size: Optional[int] = None, - use_storage_proxy: bool = False, - ): - super().__init__( - name, - stream_size, - cloud, - overwrite, - source_type - or [UploadSourceType.FILE, UploadSourceType.SEEKABLE_STREAM, UploadSourceType.NONSEEKABLE_STREAM], - use_parallel - or [True, False], # Resumable Upload doesn't support parallel uploading of parts, but fallback should work - parallelism, - multipart_upload_min_stream_size, - multipart_upload_part_size, - sdk_retry_timeout_seconds, - multipart_upload_max_retries, - custom_response_on_single_shot_upload, - expected_exception_type, - expected_multipart_upload_aborted, - expected_single_shot_upload, - use_storage_proxy, - ) - - self.unconfirmed_delta = unconfirmed_delta - self.custom_response_on_create_resumable_url = copy.deepcopy(custom_response_on_create_resumable_url) - self.custom_response_on_upload = copy.deepcopy(custom_response_on_upload) - self.custom_response_on_status_check = copy.deepcopy(custom_response_on_status_check) - self.custom_response_on_abort = copy.deepcopy(custom_response_on_abort) - self.expected_exception_type = expected_exception_type - self.expected_part_size = expected_part_size - - def create_multipart_upload_server_state(self) -> ResumableUploadServerState: - return ResumableUploadServerState(self.unconfirmed_delta, self.expected_part_size) - - def clear_state(self) -> None: - super().clear_state() - self.custom_response_on_create_resumable_url.clear_state() - self.custom_response_on_upload.clear_state() - self.custom_response_on_status_check.clear_state() - self.custom_response_on_abort.clear_state() - - def match_request_to_response( - self, request: requests.Request, server_state: ResumableUploadServerState - ) -> Optional[requests.Response]: - request_url = urlparse(request.url) - request_query = parse_qs(request_url.query) - - # Storage proxy direct resumable upload. - if ( - self.use_storage_proxy - and request_url.hostname == self._expected_hostname - and request_query.get("upload_type") == ["resumable"] - and request.method == "PUT" - ): - assert UploadTestCase.is_auth_header_present(request) - - content_range_header = request.headers["Content-range"] - is_status_check_request = re.match("bytes \\*/\\*", content_range_header) - if is_status_check_request: - response_customizer = self.custom_response_on_status_check - else: - response_customizer = self.custom_response_on_upload - - def processor() -> list: - if not is_status_check_request: - body = request.body.read() - match = re.match("bytes (\\d+)-(\\d+)/(.+)", content_range_header) - [range_start_s, range_end_s, file_size_s] = match.groups() - server_state.save_part(int(range_start_s), int(range_end_s), body, file_size_s) - - if server_state.file_content: - return [200, "", {}] - else: - if server_state.confirmed_last_byte: - headers = {"Range": f"bytes=0-{server_state.confirmed_last_byte}"} - else: - headers = {} - return [308, "", headers] - - return response_customizer.generate_response(request, processor) - - # Initial request. - elif ( - request_url.hostname == self._expected_hostname - and request_url.path == f"/api/2.0/fs/files{self.path}" - and request_query.get("action") == ["initiate-upload"] - and request.method == "POST" - ): - assert UploadTestCase.is_auth_header_present(request) - assert request.text is None - - def processor() -> list: - response_json = {"resumable_upload": {"session_token": server_state.session_token}} - return [200, json.dumps(response_json), {}] - - # Different initiate error responses have been verified by test_multipart_upload(), - # so we're always generating a "success" response. - return CustomResponse(enabled=False).generate_response(request, processor) - - # Create resumable upload URL (presigned URL path only). - elif ( - request_url.hostname == self._expected_hostname - and request_url.path == "/api/2.0/fs/create-resumable-upload-url" - and request.method == "POST" - ): - # This should never be called when using the storage proxy. - assert not self.use_storage_proxy, "create-resumable-upload-url should not be called with storage proxy" - assert UploadTestCase.is_auth_header_present(request) - - request_json = request.json() - assert request_json.keys() == {"path", "session_token"} - assert request_json["path"] == self.path - assert request_json["session_token"] == server_state.session_token - - def processor() -> list: - resumable_upload_url = f"{ResumableUploadServerState.resumable_upload_url_prefix}{self.path}" - - response_json = { - "resumable_upload_url": { - "url": resumable_upload_url, - "headers": [self.cloud_provider_header], - } - } - return [200, json.dumps(response_json), {}] - - return self.custom_response_on_create_resumable_url.generate_response(request, processor) - - # Resumable upload, uploading part via presigned URL. - elif request.url.startswith(ResumableUploadServerState.resumable_upload_url_prefix) and request.method == "PUT": - assert not UploadTestCase.is_auth_header_present(request) - url_path = request.url[len(ResumableUploadServerState.resumable_upload_url_prefix) :] - assert url_path == self.path - - content_range_header = request.headers["Content-range"] - is_status_check_request = re.match("bytes \\*/\\*", content_range_header) - if not is_status_check_request: - h = self.cloud_provider_header - assert request.headers.get(h["name"]) == h["value"] - if is_status_check_request: - assert not request.body - response_customizer = self.custom_response_on_status_check - else: - response_customizer = self.custom_response_on_upload - - def processor() -> list: - if not is_status_check_request: - body = request.body.read() - - match = re.match("bytes (\\d+)-(\\d+)/(.+)", content_range_header) - [range_start_s, range_end_s, file_size_s] = match.groups() - - server_state.save_part(int(range_start_s), int(range_end_s), body, file_size_s) - - if server_state.file_content: - # Upload complete. - return [200, "", {}] - else: - # More data expected. - if server_state.confirmed_last_byte: - headers = {"Range": f"bytes=0-{server_state.confirmed_last_byte}"} - else: - headers = {} - return [308, "", headers] - - return response_customizer.generate_response(request, processor) - - # Abort upload via presigned URL. - elif ( - request.url.startswith(ResumableUploadServerState.resumable_upload_url_prefix) - and request.method == "DELETE" - ): - assert not UploadTestCase.is_auth_header_present(request) - h = self.cloud_provider_header - assert request.headers.get(h["name"]) == h["value"] - url_path = request.url[len(ResumableUploadServerState.resumable_upload_url_prefix) :] - assert url_path == self.path - - def processor() -> list: - server_state.abort_upload() - return [200, "", {}] - - return self.custom_response_on_abort.generate_response(request, processor) - - return None - - def __str__(self) -> str: - return self.name - - @staticmethod - def to_string(test_case: "ResumableUploadTestCase") -> str: - return str(test_case) - - -@pytest.mark.parametrize( - "test_case", - [ - # ------------------ failures on creating resumable upload URL ------------------ - ResumableUploadTestCase( - "Create resumable URL: 400 response is not retried and should fallback", - stream_size=1024 * 1024, - custom_response_on_create_resumable_url=CustomResponse( - code=400, - # 1 failure is enough - only_invocation=1, - ), - expected_multipart_upload_aborted=False, # upload didn't start - expected_single_shot_upload=True, - ), - ResumableUploadTestCase( - "Create resumable URL: 403 response is not retried and should fallback", - stream_size=1024 * 1024, - custom_response_on_create_resumable_url=CustomResponse(code=403, only_invocation=1), - expected_multipart_upload_aborted=False, # upload didn't start - expected_single_shot_upload=True, - ), - ResumableUploadTestCase( - "Create resumable URL: fallback to single-shot upload when presigned URLs are disabled", - stream_size=1024 * 1024, - custom_response_on_create_resumable_url=CustomResponse( - code=403, body=MultipartUploadTestCase.presigned_url_disabled_response, only_invocation=1 - ), - expected_multipart_upload_aborted=False, # upload didn't start - expected_single_shot_upload=True, - ), - ResumableUploadTestCase( - "Create resumable URL: fallback to single-shot upload when presigned URLs are not issued because of the NetworkZone is not populated to Filesystem service", - stream_size=1024 * 1024, - custom_response_on_create_resumable_url=CustomResponse( - code=500, - body=MultipartUploadTestCase.model_serving_presigned_url_internal_error_response, - only_invocation=1, - ), - expected_multipart_upload_aborted=False, # upload didn't start - expected_single_shot_upload=True, - ), - ResumableUploadTestCase( - "Create resumable URL: 500 response is not retried and should fallback", - stream_size=1024 * 1024, - custom_response_on_create_resumable_url=CustomResponse(code=500, only_invocation=1), - expected_multipart_upload_aborted=False, # upload didn't start - expected_single_shot_upload=True, - ), - ResumableUploadTestCase( - "Create resumable URL: non-JSON response is not retried", - stream_size=1024 * 1024, - custom_response_on_create_resumable_url=CustomResponse(body="Foo bar", only_invocation=1), - expected_multipart_upload_aborted=False, # upload didn't start - expected_single_shot_upload=True, - ), - ResumableUploadTestCase( - "Create resumable URL: meaningless JSON response falls back to single-shot", - stream_size=1024 * 1024, - custom_response_on_create_resumable_url=CustomResponse( - body='{"upload_part_urls":[{"url":""}]}', only_invocation=1 - ), - expected_multipart_upload_aborted=False, # upload didn't start - expected_single_shot_upload=True, - ), - ResumableUploadTestCase( - "Create resumable URL: node exists but URL is empty falls back to single-shot", - stream_size=1024 * 1024, - custom_response_on_create_resumable_url=CustomResponse( - body='{"resumable_upload_url":{"headers":[]}}', only_invocation=1 - ), - expected_multipart_upload_aborted=False, # upload didn't start - expected_single_shot_upload=True, - ), - ResumableUploadTestCase( - "Create resumable URL: permanent retryable status code", - stream_size=1024 * 1024, - custom_response_on_create_resumable_url=CustomResponse(code=429), - expected_multipart_upload_aborted=False, # upload didn't start - expected_single_shot_upload=True, - ), - ResumableUploadTestCase( - "Create resumable URL: intermittent retryable exception is retried", - stream_size=1024 * 1024, - custom_response_on_create_resumable_url=CustomResponse( - exception=requests.Timeout, - # 3 failures total - first_invocation=1, - last_invocation=3, - ), - expected_multipart_upload_aborted=False, # upload succeeds - ), - # ------------------ failures during upload ------------------ - ResumableUploadTestCase( - "Upload: retryable exception after file is uploaded", - stream_size=1024 * 1024, - custom_response_on_upload=CustomResponse( - exception=requests.ConnectionError, - # this makes server state change before exception is thrown - exception_happened_before_processing=False, - ), - # Despite the returned error, file has been uploaded. We'll discover that - # on the next status check and consider upload completed. - expected_multipart_upload_aborted=False, - ), - ResumableUploadTestCase( - "Upload: retryable exception before file is uploaded, not enough retries", - stream_size=1024 * 1024, - multipart_upload_max_retries=3, - custom_response_on_upload=CustomResponse( - exception=requests.ConnectionError, - # prevent server from saving this part - exception_happened_before_processing=True, - # fail 4 times, exceeding max_retries - first_invocation=1, - last_invocation=4, - ), - # File was never uploaded and we gave up retrying - expected_exception_type=requests.ConnectionError, - expected_multipart_upload_aborted=True, - ), - ResumableUploadTestCase( - "Upload: retryable exception before file is uploaded, enough retries", - stream_size=1024 * 1024, - multipart_upload_max_retries=4, - custom_response_on_upload=CustomResponse( - exception=requests.ConnectionError, - # prevent server from saving this part - exception_happened_before_processing=True, - # fail 4 times, not exceeding max_retries - first_invocation=1, - last_invocation=4, - ), - # File was uploaded after retries - expected_multipart_upload_aborted=False, - ), - ResumableUploadTestCase( - "Upload: intermittent 429 response: retried", - stream_size=100 * 1024 * 1024, - multipart_upload_part_size=7 * 1024 * 1024, - multipart_upload_max_retries=3, - custom_response_on_upload=CustomResponse( - code=429, - # 3 failures not exceeding max_retries - first_invocation=2, - last_invocation=4, - ), - expected_multipart_upload_aborted=False, # upload succeeded - ), - ResumableUploadTestCase( - "Upload: intermittent 429 response: retry exhausted", - stream_size=100 * 1024 * 1024, - multipart_upload_part_size=1 * 1024 * 1024, - multipart_upload_max_retries=3, - custom_response_on_upload=CustomResponse( - code=429, - # 4 failures exceeding max_retries - first_invocation=2, - last_invocation=5, - ), - expected_exception_type=TooManyRequests, - expected_multipart_upload_aborted=True, - ), - # -------------- abort failures -------------- - ResumableUploadTestCase( - "Abort: client error", - stream_size=1024 * 1024, - # prevent part from being uploaded - custom_response_on_upload=CustomResponse(code=403), - # internal server error does not prevent server state change - custom_response_on_abort=CustomResponse(code=501), - expected_exception_type=PermissionDenied, - # abort returned error but was actually processed - expected_multipart_upload_aborted=True, - ), - ResumableUploadTestCase( - "Abort: a failure is not retried", - stream_size=1024 * 1024, - custom_response_on_upload=CustomResponse(code=403), - custom_response_on_abort=CustomResponse(code=429, only_invocation=1), - expected_exception_type=PermissionDenied, - expected_multipart_upload_aborted=False, - ), - # -------------- file already exists -------------- - ResumableUploadTestCase( - "File already exists", - stream_size=1024 * 1024, - overwrite=False, - custom_response_on_upload=CustomResponse(code=412, only_invocation=1), - expected_exception_type=AlreadyExists, - expected_multipart_upload_aborted=True, - ), - # -------------- success cases -------------- - ResumableUploadTestCase( - "Multiple parts, zero unconfirmed delta", - stream_size=100 * 1024 * 1024, - multipart_upload_part_size=7 * 1024 * 1024 + 566, - # server accepts all the parts in full - unconfirmed_delta=0, - expected_multipart_upload_aborted=False, - expected_part_size=7 * 1024 * 1024 + 566, # chunk size is used - ), - ResumableUploadTestCase( - "Multiple small parts, zero unconfirmed delta", - stream_size=100 * 1024 * 1024, - multipart_upload_part_size=100 * 1024, - # server accepts all the parts in full - unconfirmed_delta=0, - expected_multipart_upload_aborted=False, - expected_part_size=100 * 1024, # chunk size is used - ), - ResumableUploadTestCase( - "Multiple parts, non-zero unconfirmed delta", - stream_size=100 * 1024 * 1024, - multipart_upload_part_size=7 * 1024 * 1024 + 566, - # for every part, server accepts all except last 239 bytes - unconfirmed_delta=239, - expected_multipart_upload_aborted=False, - expected_part_size=7 * 1024 * 1024 + 566, # chunk size is used - ), - ResumableUploadTestCase( - "Multiple parts, variable unconfirmed delta", - stream_size=100 * 1024 * 1024, - multipart_upload_part_size=7 * 1024 * 1024 + 566, - # for the first part, server accepts all except last 15Kib - # for the second part, server accepts it all - # for the 3rd part, server accepts all except last 25000 bytes - # for the 4th part, server accepts all except last 7 Mb - # for the 5th part onwards server accepts all except last 5 bytes - unconfirmed_delta=[15 * 1024, 0, 25000, 7 * 1024 * 1024, 5], - expected_multipart_upload_aborted=False, - expected_part_size=7 * 1024 * 1024 + 566, # chunk size is used - ), - ResumableUploadTestCase( - "Small stream, single-shot upload used", - stream_size=1024 * 1024, - multipart_upload_min_stream_size=1024 * 1024 + 1, - expected_multipart_upload_aborted=False, - expected_single_shot_upload=True, - ), - # -------------- storage proxy tests -------------- - ResumableUploadTestCase( - "Storage proxy: successful resumable upload", - stream_size=100 * 1024 * 1024, - multipart_upload_part_size=7 * 1024 * 1024 + 566, - unconfirmed_delta=0, - use_storage_proxy=True, - expected_multipart_upload_aborted=False, - expected_part_size=7 * 1024 * 1024 + 566, - ), - ], - ids=ResumableUploadTestCase.to_string, -) -def test_resumable_upload(config: Config, test_case: ResumableUploadTestCase) -> None: - test_case.run(config) - - @dataclass class CreateDownloadUrlResponseTestCase: data: Dict[str, Any] diff --git a/tests/test_files_walk.py b/tests/test_files_walk.py new file mode 100644 index 000000000..282e355f7 --- /dev/null +++ b/tests/test_files_walk.py @@ -0,0 +1,537 @@ +"""Which endpoints FilesExt transfers files through, and what a caller sees when they fail. + +FakeFilesServer stands in for the hosts a transfer can reach: the workspace (the Files API and the presigned URL +coordination API), the data-plane storage proxy, and cloud storage behind presigned URLs. It records every +operation per endpoint. + +The scenario matrix runs every combination of endpoint configuration, endpoint availability and transfer input, +and checks only what a caller sees: the file stored or downloaded, or the error raised. The outcome each scenario +must have comes from the rules in expected_upload and expected_download. +""" + +import collections +import io +import itertools +import json +import os +import platform +from dataclasses import dataclass +from tempfile import NamedTemporaryFile +from threading import Lock +from typing import Callable, Dict, List, Optional, Tuple, Union +from urllib.parse import parse_qs, urlparse + +import pytest +import requests +import requests_mock +from urllib3.exceptions import ProtocolError + +from databricks.sdk import WorkspaceClient +from databricks.sdk.core import Config + +from .test_files_utils import NonSeekableBuffer + +WORKSPACE = "localhost" +STORAGE_PROXY = "storage-proxy.databricks.com" +CLOUD = "cloud.example.com" +FILE_PATH = "/Volumes/catalog/schema/volume/file.bin" +LAST_MODIFIED = "Thu, 28 Nov 2024 16:39:14 GMT" +PART_SIZE = 1024 + +_ERROR_CODES = { + 400: "INVALID_PARAMETER_VALUE", + 403: "PERMISSION_DENIED", + 404: "NOT_FOUND", + 409: "ALREADY_EXISTS", + 500: "INTERNAL_ERROR", +} + +# Operations on presigned cloud URLs. They must never carry workspace credentials. +_PRESIGNED_OPS = {"offload.part", "offload.get"} + + +class BreakAfter: + """A download fault: the connection drops after `n` bytes of the body.""" + + def __init__(self, n: int): + self.n = n + + +class _BreakingBody(io.BytesIO): + """Serves the first `break_at` bytes, then fails every further read like a dropped connection. + + Covers every read method, since urllib3 may read through any of them. + """ + + def __init__(self, data: bytes, break_at: int): + super().__init__(data) + self._break_at = break_at + + def _limit(self, size: Optional[int]) -> int: + position = self.tell() + if position >= self._break_at: + raise ProtocolError("Connection broken: simulated") + remaining = self._break_at - position + return remaining if size is None or size < 0 else min(size, remaining) + + def read(self, size: Optional[int] = -1) -> bytes: + return super().read(self._limit(size)) + + def read1(self, size: Optional[int] = -1) -> bytes: + return super().read1(self._limit(size)) + + def readinto(self, buffer) -> int: + return super().readinto(memoryview(buffer)[: self._limit(len(buffer))]) + + def readinto1(self, buffer) -> int: + return super().readinto1(memoryview(buffer)[: self._limit(len(buffer))]) + + +Fault = Union[int, type, BreakAfter] + + +def always(fault: Fault) -> Callable[[int], Optional[Fault]]: + return lambda call: fault + + +class FakeFilesServer: + def __init__(self, faults: Optional[Dict[str, Callable[[int], Optional[Fault]]]] = None): + self.faults = faults or {} + self.download_content = b"" + self.calls: List[str] = [] + self.counts = collections.Counter() + self.parts: Dict[str, Dict[int, bytes]] = collections.defaultdict(dict) + self.stored: Dict[str, bytes] = {} + self.aborted = set() + self.ranges: Dict[str, List[Optional[str]]] = collections.defaultdict(list) + self.credentials: Dict[str, List[bool]] = collections.defaultdict(list) + self.request_headers: Dict[str, List[Dict[str, str]]] = collections.defaultdict(list) + self.urls: Dict[str, List[str]] = collections.defaultdict(list) + # Protocol checks that failed; a client that catches its request errors would otherwise hide them. + self.violations: List[str] = [] + self._lock = Lock() + + def stored_file(self) -> Optional[bytes]: + """The file as a caller would read it back, whichever endpoint stored it.""" + assert len(self.stored) <= 1, f"the file was stored by more than one endpoint: {sorted(self.stored)}" + return next(iter(self.stored.values()), None) + + def assert_credentials(self) -> None: + for op, sent in self.credentials.items(): + if op in _PRESIGNED_OPS: + assert not any(sent), f"workspace credentials were sent to a presigned URL ({op})" + else: + assert all(sent), f"workspace credentials were missing on {op}" + + def _operation(self, request: requests.Request) -> str: + url = urlparse(request.url) + query = parse_qs(url.query) + host, method, path = url.hostname, request.method, url.path + if host == WORKSPACE and path == "/api/2.0/preview/scim/v2/Me" and method == "GET": + return "scim" + if host == STORAGE_PROXY and path == "/api/2.0/fs/files/DatabricksInternal/Probe/fullstack/wis": + return "probe" + if path == f"/api/2.0/fs/files{FILE_PATH}": + assert host in (WORKSPACE, STORAGE_PROXY), f"Files API 2.0 request to {host}" + endpoint = "dp" if host == STORAGE_PROXY else None + action = query.get("action", [None])[0] + if method == "POST" and action == "initiate-upload": + return f"{endpoint or 'offload'}.initiate" + if method == "POST" and action == "complete-upload": + return f"{endpoint or 'offload'}.complete" + if method == "POST" and action == "abort-upload": + assert query.get("upload_type") == ["multipart"], f"abort-upload without upload_type: {request.url}" + return f"{endpoint or 'offload'}.abort" + if method == "PUT" and query.get("upload_type") == ["multipart"]: + return f"{endpoint}.part" + if method == "PUT": + return f"{endpoint or 'floor'}.single" + if method == "GET": + return f"{endpoint or 'floor'}.get" + if method == "HEAD" and host == WORKSPACE: + return "workspace.head" + if host == WORKSPACE and path == "/api/2.0/fs/create-upload-part-urls": + return "offload.part-urls" + if host == WORKSPACE and path == "/api/2.0/fs/create-download-url": + return "offload.download-url" + if host == CLOUD and path.startswith("/part/") and method == "PUT": + return "offload.part" + if host == CLOUD and path == "/download" and method == "GET": + return "offload.get" + raise AssertionError(f"Unexpected request {method} {request.url}") + + def handle(self, request: requests.Request) -> requests.Response: + op = self._operation(request) + with self._lock: + self.calls.append(op) + self.counts[op] += 1 + call = self.counts[op] + self.credentials[op].append(request.headers.get("Authorization") is not None) + self.request_headers[op].append(dict(request.headers)) + self.urls[op].append(request.url) + fault = self.faults.get(op, lambda _: None)(call) + if isinstance(fault, type): + raise fault("simulated network failure") + if isinstance(fault, int): + body = {"error_code": _ERROR_CODES.get(fault, "UNKNOWN"), "message": f"simulated {fault} on {op}"} + return requests_mock.create_response(request, status_code=fault, json=body) + return self._respond(op, request, fault) + + def _respond(self, op: str, request: requests.Request, fault: Optional[Fault]) -> requests.Response: + if op == "scim": + return requests_mock.create_response(request, json={}, headers={"X-Databricks-Org-Id": "12345"}) + if op == "probe": + return requests_mock.create_response(request) + endpoint, action = op.split(".", 1) + url = urlparse(request.url) + query = parse_qs(url.query) + if action == "initiate": + return requests_mock.create_response( + request, json={"multipart_upload": {"session_token": f"token-{endpoint}"}} + ) + if op == "offload.part-urls": + body = json.loads(request.body) + assert body["session_token"] == "token-offload" + start, count = body["start_part_number"], body["count"] + urls = [ + {"part_number": n, "url": f"https://{CLOUD}/part/{n}", "headers": []} + for n in range(start, start + count) + ] + return requests_mock.create_response(request, json={"upload_part_urls": urls}) + if action == "part": + part_number = int(url.path.split("/")[-1]) if endpoint == "offload" else int(query["part_number"][0]) + with self._lock: + self.parts[endpoint][part_number] = _read_body(request) + return requests_mock.create_response(request, headers={"ETag": f"etag-{endpoint}-{part_number}"}) + if action == "complete": + parts = json.loads(request.body)["parts"] + for part in parts: + assert part["etag"] == f"etag-{endpoint}-{part['part_number']}" + numbers = [part["part_number"] for part in parts] + assert numbers == sorted(self.parts[endpoint]), "complete-upload must list every uploaded part" + self.stored[endpoint] = b"".join(self.parts[endpoint][n] for n in numbers) + return requests_mock.create_response(request, json={}) + if action == "abort": + assert query.get("session_token") == [f"token-{endpoint}"] + self.aborted.add(endpoint) + return requests_mock.create_response(request, json={}) + if action == "single": + self.stored[endpoint] = _read_body(request) + return requests_mock.create_response(request, json={}) + if op == "offload.download-url": + return requests_mock.create_response(request, json={"url": f"https://{CLOUD}/download", "headers": []}) + if op == "workspace.head": + return requests_mock.create_response( + request, + headers={ + "Content-Length": str(len(self.download_content)), + "Content-Type": "application/octet-stream", + "Last-Modified": LAST_MODIFIED, + }, + ) + if action == "get": + return self._download(endpoint, request, fault) + raise AssertionError(f"No response for {op}") + + def _download(self, endpoint: str, request: requests.Request, fault: Optional[Fault]) -> requests.Response: + range_header = request.headers.get("Range") + self.ranges[endpoint].append(range_header) + data, status = self.download_content, 200 + if range_header: + start_s, end_s = range_header[len("bytes=") :].split("-") + start = int(start_s) + end = int(end_s) if end_s else len(data) - 1 + data, status = data[start : end + 1], 206 + headers = { + "Content-Length": str(len(data)), + "Content-Type": "application/octet-stream", + "Last-Modified": LAST_MODIFIED, + } + body = _BreakingBody(data, fault.n) if isinstance(fault, BreakAfter) else io.BytesIO(data) + return requests_mock.create_response(request, status_code=status, headers=headers, body=body) + + +def _read_body(request: requests.Request) -> bytes: + body = request.body + if body is None: + return b"" + if hasattr(body, "read"): + return body.read() + return body.encode() if isinstance(body, str) else bytes(body) + + +def _client(config: Config, storage_proxy: bool = True) -> WorkspaceClient: + config = config.copy() + config.experimental_files_ext_enable_storage_proxy = storage_proxy + config._header_factory = lambda: {"Authorization": "Bearer workspace-token"} + config.files_ext_multipart_upload_min_stream_size = PART_SIZE + config.files_ext_parallel_download_min_file_size = PART_SIZE + config.files_ext_parallel_download_default_part_size = PART_SIZE + config.files_ext_client_download_streaming_chunk_size = 512 + return WorkspaceClient(config=config) + + +SOURCES = ["file", "seekable", "non-seekable"] + + +def _upload( + w: WorkspaceClient, source: str, data: bytes, use_parallel: bool, tmp_path, overwrite: Optional[bool] = None +) -> None: + if source == "file": + local = tmp_path / "source.bin" + local.write_bytes(data) + w.files.upload_from(FILE_PATH, str(local), overwrite=overwrite, part_size=PART_SIZE, use_parallel=use_parallel) + return + stream = io.BytesIO(data) if source == "seekable" else NonSeekableBuffer(data) + w.files.upload(FILE_PATH, stream, overwrite=overwrite, part_size=PART_SIZE, use_parallel=use_parallel) + + +def _run(server: FakeFilesServer, fn: Callable[[], object]) -> object: + def handle(request: requests.Request) -> requests.Response: + try: + return server.handle(request) + except AssertionError as e: + server.violations.append(str(e)) + return requests_mock.create_response(request, status_code=400, json={"message": str(e)}) + + try: + with requests_mock.Mocker() as mock: + mock.add_matcher(handle) + return fn() + finally: + assert not server.violations, server.violations + + +LARGE = os.urandom(4 * PART_SIZE + 100) +SMALL = os.urandom(100) + + +# ======================== Scenario matrix (black box) ======================== +# +# Every combination of the endpoint configuration, which of its endpoints are up, and the transfer input, checked +# only on what the caller sees: the file stored or downloaded, or the error raised. The outcome a scenario must have +# comes from the rules in expected_upload and expected_download, not from the client. + + +@dataclass(frozen=True) +class EndpointList: + name: str + # The endpoints a transfer can use, best first. + endpoints: Tuple[str, ...] + storage_proxy: bool = False + + +ENDPOINT_LISTS = [ + EndpointList("built-in", ("presigned", "workspace")), + EndpointList("built-in with the storage proxy", ("proxy", "presigned", "workspace"), storage_proxy=True), +] + +# up: the endpoint works. down: every request to it fails, as when it can't be reached. rejects data: its probe and +# sessions answer, but every transfer fails. breaks mid-transfer: it takes the first part or sends the first bytes, +# then fails. +UP, DOWN, REJECTS, BREAKS = "up", "down", "rejects data", "breaks mid-transfer" + +_DATA = {"empty": b"", "small": SMALL, "one part": os.urandom(PART_SIZE), "several parts": LARGE} +SIZES = {size: len(data) for size, data in _DATA.items()} +UPLOAD_MODES = ["sequential", "parallel"] +DOWNLOAD_MODES = ["stream", "to-file", "parallel"] + +# What fails while an endpoint is down. For the workspace that is its Files API transfers: it still serves file +# metadata, opens presigned URL sessions and issues their URLs. +_DOWN_OPS = { + "proxy": ["probe", "dp.initiate", "dp.part", "dp.complete", "dp.abort", "dp.single", "dp.get"], + "presigned": ["offload.part", "offload.get"], + "workspace": ["floor.single", "floor.get"], +} +# The fake's name for each endpoint that can come first in a list. +_FAKE_NAMES = {"proxy": "dp", "presigned": "offload"} + +SUCCEEDS = "succeeds" + + +@dataclass(frozen=True) +class Scenario: + transfer: str + endpoint_list: EndpointList + availability: Tuple[Tuple[str, str], ...] + size: str + mode: str + source: Optional[str] = None + + def __str__(self) -> str: + transfer_input = ", ".join(x for x in (self.size, self.source, self.mode) if x) + states = ", ".join(f"{endpoint} {state}" for endpoint, state in self.availability) + return f"{transfer_input} | {self.endpoint_list.name} | {states}" + + def faults(self) -> Dict[str, Callable[[int], Optional[Fault]]]: + faults = {} + for endpoint, state in self.availability: + if state == DOWN: + faults.update({op: always(requests.ConnectionError) for op in _DOWN_OPS[endpoint]}) + elif state == REJECTS: + faults.update({f"{_FAKE_NAMES[endpoint]}.{op}": always(403) for op in ("part", "single", "get")}) + elif state == BREAKS and self.transfer == "upload": + # The first part lands, every later one fails. + faults[f"{_FAKE_NAMES[endpoint]}.part"] = lambda call: requests.ConnectionError if call > 1 else None + elif state == BREAKS and self.mode == "parallel": + # The first range lands, every later one fails. + faults[f"{_FAKE_NAMES[endpoint]}.get"] = lambda call: requests.ConnectionError if call > 1 else None + elif state == BREAKS: + # The response is cut off after some bytes reached the caller, and the endpoint is down after that. + faults[f"{_FAKE_NAMES[endpoint]}.get"] = lambda call: ( + BreakAfter(3 * 512) if call == 1 else requests.ConnectionError + ) + return faults + + +def _availabilities(endpoint_list: EndpointList) -> List[Tuple[Tuple[str, str], ...]]: + # An upload falls back to the workspace when the list has nothing else and a download doesn't, so the workspace is + # always on the axis. + axis = endpoint_list.endpoints + if "workspace" not in axis: + axis += ("workspace",) + combinations = [tuple(zip(axis, states)) for states in itertools.product([UP, DOWN], repeat=len(axis))] + first, rest = axis[0], axis[1:] + combinations.append(((first, BREAKS),) + tuple((endpoint, UP) for endpoint in rest)) + combinations += [ + ((first, REJECTS),) + tuple(zip(rest, states)) for states in itertools.product([UP, DOWN], repeat=len(rest)) + ] + return combinations + + +def _scenarios(transfer: str, modes: List[str], sources: List[Optional[str]]) -> List[Scenario]: + return [ + Scenario(transfer, endpoint_list, availability, size, mode, source) + for endpoint_list in ENDPOINT_LISTS + for availability in _availabilities(endpoint_list) + for size in SIZES + for source in sources + for mode in modes + # Only a transfer of several parts or chunks has data committed before a later failure. + if size == "several parts" or BREAKS not in dict(availability).values() + ] + + +UPLOAD_SCENARIOS = _scenarios("upload", UPLOAD_MODES, SOURCES) +DOWNLOAD_SCENARIOS = _scenarios("download", DOWNLOAD_MODES, [None]) + + +def _first_endpoint(scenario: Scenario) -> str: + """The endpoint a multipart upload or a download starts on: the storage proxy when it is enabled and its + probe reaches it, otherwise presigned URLs.""" + if scenario.endpoint_list.storage_proxy and dict(scenario.availability)["proxy"] != DOWN: + return "proxy" + return "presigned" + + +# Retries that run out raise TimeoutError from the API client's retry budget, which covers the workspace, and +# RuntimeError from the bounded retries of transfers to presigned URLs and the storage proxy. +def expected_upload(scenario: Scenario): + """The outcome an upload has. + + - A file under one part is sent in a single shot to the workspace. + - A non-seekable stream's size is known only once it is read. Read sequentially, a short one is sent in a + single shot as above. Read in parallel, an upload session is opened first; an empty stream is then sent in a + single shot to the workspace, and a short one as one part. + - Anything else is a multipart upload on the first endpoint. A failure there raises, except a rejected + transfer (403), which falls back to a single shot on the workspace. + - A single shot that follows a rejected part replays the part read plus the unread rest, which can't be + rewound, so the API client sends it once. An upload from a file in parallel reopens the file instead. + """ + states = dict(scenario.availability) + size = SIZES[scenario.size] + session_first = scenario.source == "non-seekable" and scenario.mode == "parallel" + workspace_up = states["workspace"] == UP + if size < PART_SIZE and not session_first: + return SUCCEEDS if workspace_up else TimeoutError + if session_first and size == 0: + return SUCCEEDS if workspace_up else requests.ConnectionError + state = states[_first_endpoint(scenario)] + if state == UP: + return SUCCEEDS + if state != REJECTS: + return RuntimeError + if workspace_up: + return SUCCEEDS + return TimeoutError if scenario.source == "file" and scenario.mode == "parallel" else requests.ConnectionError + + +def expected_download(scenario: Scenario): + """The outcome a download has. + + - It starts on the first endpoint. A failure there raises, except a rejected transfer (403), which falls back + to the Files API on the workspace. + - After bytes arrived, a stream resumes on that endpoint and, if that fails too, raises its broken connection; + a parallel download raises the failure of a later range. + """ + states = dict(scenario.availability) + state = states[_first_endpoint(scenario)] + if state == UP: + return SUCCEEDS + if state == BREAKS: + return RuntimeError if scenario.mode == "parallel" else requests.exceptions.ChunkedEncodingError + if state == REJECTS: + return SUCCEEDS if states["workspace"] == UP else TimeoutError + return RuntimeError + + +def _scenario_client(config: Config, scenario: Scenario) -> WorkspaceClient: + return _client(config, storage_proxy=scenario.endpoint_list.storage_proxy) + + +@pytest.mark.parametrize("scenario", UPLOAD_SCENARIOS, ids=str) +def test_upload_scenario(config, tmp_path, scenario): + server = FakeFilesServer(scenario.faults()) + w = _scenario_client(config, scenario) + data = _DATA[scenario.size] + expected = expected_upload(scenario) + + def upload(): + _upload(w, scenario.source, data, scenario.mode == "parallel", tmp_path) + + if expected == SUCCEEDS: + _run(server, upload) + assert server.stored_file() == data + else: + with pytest.raises(expected): + _run(server, upload) + assert server.stored_file() is None + server.assert_credentials() + + +def _download(w: WorkspaceClient) -> bytes: + return w.files.download(FILE_PATH).contents.read() + + +def _download_to(w: WorkspaceClient, use_parallel: bool, parallelism: Optional[int] = None) -> bytes: + if use_parallel and platform.system() == "Windows": + # FilesExt downloads sequentially on Windows, so the parallel rules don't apply there. + pytest.skip("Skipping parallel download tests on Windows") + with NamedTemporaryFile(delete=False) as f: + destination = f.name + try: + w.files.download_to(FILE_PATH, destination, use_parallel=use_parallel, parallelism=parallelism) + with open(destination, "rb") as f: + return f.read() + finally: + os.remove(destination) + + +@pytest.mark.parametrize("scenario", DOWNLOAD_SCENARIOS, ids=str) +def test_download_scenario(config, scenario): + server = FakeFilesServer(scenario.faults()) + server.download_content = _DATA[scenario.size] + w = _scenario_client(config, scenario) + download = { + "stream": lambda: _download(w), + "to-file": lambda: _download_to(w, use_parallel=False), + "parallel": lambda: _download_to(w, use_parallel=True), + }[scenario.mode] + expected = expected_download(scenario) + + if expected == SUCCEEDS: + assert _run(server, download) == server.download_content + else: + with pytest.raises(expected): + _run(server, download) + server.assert_credentials() diff --git a/tests/test_sandbox_files.py b/tests/test_sandbox_files.py new file mode 100644 index 000000000..27b105680 --- /dev/null +++ b/tests/test_sandbox_files.py @@ -0,0 +1,200 @@ +import io +from urllib.parse import parse_qs, urlsplit + +import pytest + +from databricks.sdk import WorkspaceClient +from databricks.sdk.errors import NotFound, PermissionDenied + +from .conftest import noop_credentials + + +_HOST = "https://example.cloud.databricks.com" +_FILES_URL = f"{_HOST}/api/2.0/sandbox-files/sandboxes/box" + + +def _sandbox(*, workspace_id=None): + workspace = WorkspaceClient( + host=_HOST, + workspace_id=workspace_id, + credentials_strategy=noop_credentials, + ) + return workspace.sandbox + + +def _query_parameters(request): + return parse_qs(urlsplit(request.url).query, keep_blank_values=True) + + +def _entry(name): + return { + "name": name, + "path": f"/tmp/{name}", + "file_type": "SANDBOX_FILE_TYPE_FILE", + "size": "1", + "mode": 0o644, + "permissions": "-rw-r--r--", + "owner": "sandbox", + "group": "sandbox", + "modify_time": "2023-11-14T22:13:20.123456789Z", + } + + +def test_unary_operations_use_generated_sandbox_methods(requests_mock): + requests_mock.post(f"{_FILES_URL}/paths:rename", json={}) + requests_mock.post(f"{_FILES_URL}/directories:make", json={}) + requests_mock.post(f"{_FILES_URL}/paths:remove", json={}) + requests_mock.get(f"{_FILES_URL}/paths:describe", json={"file": _entry("item")}) + requests_mock.get(f"{_FILES_URL}/directory-entries", json={"contents": [_entry("child")]}) + sandbox = _sandbox() + + sandbox.rename_path("sandboxes/box", "/tmp/source", "/tmp/destination") + sandbox.create_directory("sandboxes/box", "/tmp/directory", parents=False) + sandbox.delete_path("sandboxes/box", "/tmp/old", recursive=True) + metadata = sandbox.get_metadata("sandboxes/box", "/tmp/item") + contents = list(sandbox.list_directory_contents("sandboxes/box", "/tmp/directory", page_size=10)) + + assert metadata.file.path == "/tmp/item" + assert [entry.name for entry in contents] == ["child"] + rename, create, delete, metadata_request, listing = requests_mock.request_history + assert rename.json() == {"source": "/tmp/source", "destination": "/tmp/destination"} + assert create.json() == {"path": "/tmp/directory", "parents": False} + assert delete.json() == {"path": "/tmp/old", "recursive": True} + assert _query_parameters(metadata_request) == {"path": ["/tmp/item"]} + assert _query_parameters(listing) == {"path": ["/tmp/directory"], "page_size": ["10"]} + + +def test_list_directory_contents_follows_generated_pagination(requests_mock): + requests_mock.get( + f"{_FILES_URL}/directory-entries", + [ + {"json": {"contents": [_entry("first")], "next_page_token": "next"}}, + {"json": {"contents": [_entry("second")]}}, + ], + ) + + names = [entry.name for entry in _sandbox().list_directory_contents("sandboxes/box", "/tmp")] + + assert names == ["first", "second"] + assert _query_parameters(requests_mock.request_history[1])["page_token"] == ["next"] + + +def test_upload_file_streams_bytes_to_transfer_endpoint(requests_mock): + payload = b"raw file contents\x00\xff" + contents = io.BytesIO(payload) + uploaded = {} + + def capture_upload(request, context): + assert request.body is contents + uploaded["body"] = request.body.read() + context.status_code = 204 + return "" + + requests_mock.put(f"{_FILES_URL}/content", text=capture_upload) + + _sandbox(workspace_id="123").upload_file( + parent="sandboxes/box", + path="/tmp/report ?#%.bin", + contents=contents, + ) + + request = requests_mock.last_request + assert uploaded["body"] == payload + assert _query_parameters(request) == {"path": ["/tmp/report ?#%.bin"]} + assert request.headers["Content-Type"] == "application/octet-stream" + assert request.headers["X-Databricks-Workspace-Id"] == "123" + + +@pytest.mark.parametrize( + ("operation", "verb", "endpoint", "status_code", "error_code", "error_type"), + [ + ("get_metadata", "GET", "paths:describe", 404, "NOT_FOUND", NotFound), + ("download_file", "GET", "content", 404, "NOT_FOUND", NotFound), + ("upload_file", "PUT", "content", 403, "PERMISSION_DENIED", PermissionDenied), + ], +) +def test_file_operations_raise_sdk_errors( + requests_mock, operation, verb, endpoint, status_code, error_code, error_type +): + requests_mock.register_uri( + verb, + f"{_FILES_URL}/{endpoint}", + status_code=status_code, + json={"error_code": error_code, "message": "File access failed"}, + ) + + kwargs = {"contents": io.BytesIO(b"contents")} if operation == "upload_file" else {} + with pytest.raises(error_type, match="File access failed"): + getattr(_sandbox(), operation)(parent="sandboxes/box", path="/tmp/report.bin", **kwargs) + + +def test_download_file_returns_stream_and_transfer_metadata(requests_mock): + payload = b"downloaded contents" + headers = { + "Content-Length": str(len(payload)), + "Content-Type": "application/octet-stream", + "Last-Modified": "Wed, 10 Sep 2025 12:00:00 GMT", + "ETag": '"file-version"', + "Accept-Ranges": "bytes", + } + requests_mock.get(f"{_FILES_URL}/content", content=payload, headers=headers) + + response = _sandbox(workspace_id="123").download_file("sandboxes/box", "/tmp/report ?#%.bin") + + with response.contents as stream: + assert stream.read() == payload + request = requests_mock.last_request + assert _query_parameters(request) == {"path": ["/tmp/report ?#%.bin"]} + assert request.headers["Accept"] == "application/octet-stream" + assert request.headers["X-Databricks-Workspace-Id"] == "123" + for header in ("Range", "If-Range", "If-Match", "If-None-Match", "If-Modified-Since", "If-Unmodified-Since"): + assert header not in request.headers + assert response.content_length == len(payload) + assert response.content_type == "application/octet-stream" + assert response.last_modified == "Wed, 10 Sep 2025 12:00:00 GMT" + assert response.etag == '"file-version"' + assert response.accept_ranges == "bytes" + assert response.content_range is None + + +def test_download_file_returns_partial_stream_and_range_metadata(requests_mock): + requests_mock.get( + f"{_FILES_URL}/content", + status_code=206, + content=b"part", + headers={"Content-Length": "4", "Content-Range": "bytes 0-3/100", "ETag": '"file-version"'}, + ) + + response = _sandbox().download_file( + parent="sandboxes/box", path="/tmp/report.bin", byte_range="bytes=0-3", if_range='"file-version"' + ) + + with response.contents as stream: + assert stream.read() == b"part" + assert response.content_length == 4 + assert response.content_range == "bytes 0-3/100" + assert response.etag == '"file-version"' + assert _query_parameters(requests_mock.last_request) == { + "path": ["/tmp/report.bin"], + "byte_range": ["bytes=0-3"], + "if_range": ['"file-version"'], + } + + +@pytest.mark.parametrize( + ("argument", "value"), + [ + ("if_match", '"file-version"'), + ("if_none_match", '"file-version"'), + ("if_modified_since", "Wed, 10 Sep 2025 12:00:00 GMT"), + ("if_unmodified_since", "Wed, 10 Sep 2025 12:00:00 GMT"), + ], +) +def test_download_file_sends_conditional_query_parameters(requests_mock, argument, value): + requests_mock.get(f"{_FILES_URL}/content", content=b"contents") + + response = _sandbox().download_file("sandboxes/box", "/tmp/report.bin", **{argument: value}) + + with response.contents as stream: + assert stream.read() == b"contents" + assert _query_parameters(requests_mock.last_request) == {"path": ["/tmp/report.bin"], argument: [value]}