diff --git a/.codegen.json b/.codegen.json index 4a145805b..07d77d1ce 100644 --- a/.codegen.json +++ b/.codegen.json @@ -10,7 +10,7 @@ "tagging.py", "tagging.py.lock" ], - "allow_from_source": [ + "allow_from_packages": [ "**/*.{md,py,rst,png,svg,gif,tmpl,ps1}", "**/py.typed", ".codegen/{_last_sha,_openapi_sha}", diff --git a/.codegen/_last_sha b/.codegen/_last_sha index d720f4024..01458beb2 100644 --- a/.codegen/_last_sha +++ b/.codegen/_last_sha @@ -1 +1 @@ -8758d3b71cae7ab16f4bfb80629e5c156885ccbb \ No newline at end of file +33607074f4900fa4b72fdcd09c1c66429271844d \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md index d47b1e399..56dcde565 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,20 @@ # Version changelog +## Release v0.140.0 (2026-09-18) + +### API Changes +* Add `create_mcp_service_user_mapped_credential()`, `delete_mcp_service_user_mapped_credential()` and `get_mcp_service_user_mapped_credential()` methods for [w.ai_gateway](https://databricks-sdk-py.readthedocs.io/en/latest/workspace/catalog/ai_gateway.html) workspace-level service. +* Add `telemetry_export_destinations` field for `databricks.sdk.service.apps.AppUpdate`. +* Add `options` field for `databricks.sdk.service.catalog.McpServiceConfigSourceConnection`. +* Add `header_auth` field for `databricks.sdk.service.catalog.ModelProviderServiceConfigCustomProviderDirectConfig`. +* Add `timezone_id` field for `databricks.sdk.service.ml.CronSchedule`. +* Add `group_name` field for `databricks.sdk.service.pipelines.RunAs`. +* Add `avro_options` and `protobuf_options` fields for `databricks.sdk.service.pipelines.Transformer`. +* Add `tiktok_ads` enum value for `databricks.sdk.service.catalog.ConnectionType`. +* Add `agent_service` and `skill` enum values for `databricks.sdk.service.catalog.SecurableType`. +* Add `tiktok_ads` and `smartsheet` enum values for `databricks.sdk.service.pipelines.IngestionSourceType`. +* Add `avro` and `protobuf` enum values for `databricks.sdk.service.pipelines.TransformerFormat`. + ## Release v0.139.0 (2026-09-13) ### API Changes diff --git a/databricks/sdk/__init__.py b/databricks/sdk/__init__.py index 7e0ae3b6a..39883474f 100644 --- a/databricks/sdk/__init__.py +++ b/databricks/sdk/__init__.py @@ -1179,7 +1179,7 @@ def workspace_entity_tag_assignments(self) -> pkg_tags.WorkspaceEntityTagAssignm @property def workspace_iam_v2(self) -> pkg_iamv2.WorkspaceIamV2API: - """These APIs are used to manage identities and the workspace access of these identities in .""" + """These APIs are used to manage identities and the workspace access of these identities in Databricks.""" return self._workspace_iam_v2 @property @@ -1398,7 +1398,7 @@ def groups_v2(self) -> pkg_iam.AccountGroupsV2API: @property def iam_v2(self) -> pkg_iamv2.AccountIamV2API: - """These APIs are used to manage identities and the workspace access of these identities in .""" + """These APIs are used to manage identities and the workspace access of these identities in Databricks.""" return self._iam_v2 @property diff --git a/databricks/sdk/mixins/_grpc_transport.py b/databricks/sdk/mixins/_grpc_transport.py new file mode 100644 index 000000000..4abb573cc --- /dev/null +++ b/databricks/sdk/mixins/_grpc_transport.py @@ -0,0 +1,323 @@ +"""Generic gRPC streaming transport for hand-written SDK mixins. + +The databricks-sdk client is HTTP-only; its generated service clients speak +JSON over the shared ``ApiClient``. Some APIs are gRPC-streaming RPCs that the +SDK's codegen does not produce a client for, so a mixin for such an API needs a +gRPC channel instead of the HTTP transport. + +This module provides that channel, built from the SDK ``Config`` so it carries +the same credentials as the rest of the SDK: auth headers come from +``Config.authenticate()`` on every RPC, like the HTTP client's ``session.auth`` +hook, so whatever credential strategy the caller configured authenticates the +gRPC calls too. Any ``Config`` custom headers ride along as well, as they do on +the HTTP client's requests, and the channel advertises the SDK's User-Agent (the +same string the HTTP client sends) so gRPC traffic is attributable to the SDK. + +It is deliberately generic: it knows nothing about any particular service's +protos, headers, or responses. Those live in the mixin that uses the channel. + +``grpcio`` is imported lazily inside the functions that need it, so importing +this module never requires grpcio to be installed; grpcio ships only with the +optional ``sandbox`` extra. +""" + +from __future__ import annotations + +import logging +import time +import urllib.parse +from typing import TYPE_CHECKING, Callable, Dict, Iterator, List, Optional, Tuple, TypeVar + +if TYPE_CHECKING: + import grpc + + from databricks.sdk.config import Config + +_LOG = logging.getLogger("databricks.sdk") + +# 443 is the default HTTPS port, reused here as the default gRPC port (overridable via `port`). +_DEFAULT_GRPC_PORT = 443 + +# Retry defaults for transient gRPC failures - e.g. a backend still starting up +# refuses connections (UNAVAILABLE). 30 attempts x 3s bounds the wait at ~90s. +_DEFAULT_MAX_ATTEMPTS = 30 +_DEFAULT_RETRY_INTERVAL_SECONDS = 3.0 + +_T = TypeVar("_T") + +# Distinguishes "stream opened but yielded nothing" from a real first message. +_UNSET = object() + + +def _grpc_target(host: str, port: int = _DEFAULT_GRPC_PORT) -> str: + """Turn an SDK host (e.g. ``https://foo.databricks.com``) into a + ``host:port`` gRPC dial target. Keeps only host and port, dropping any path, + query, or fragment; appends ``port`` when the host has none.""" + if "://" not in host: + host = f"https://{host}" + netloc = urllib.parse.urlparse(host).netloc + if ":" not in netloc: + netloc = f"{netloc}:{port}" + return netloc + + +def _headers_to_metadata(headers: Dict[str, str]) -> List[Tuple[str, str]]: + """Convert HTTP-style auth headers into gRPC metadata. gRPC requires + lowercase metadata keys; values pass through unchanged.""" + return [(key.lower(), value) for key, value in headers.items()] + + +def _auth_metadata(cfg: "Config", callback: Callable[..., None]) -> None: + """Resolve the SDK's auth headers plus any configured custom headers and pass + them to gRPC's metadata ``callback``. + + Mirrors the HTTP client, which sends ``Config`` custom headers on every + request and layers auth on top; the two are merged here with auth winning on + a name collision. The merge is case-insensitive - gRPC metadata keys are + lowercase, so each source is lowercased before merging and every name ends up + at most once, with auth overriding a custom header of the same name whatever + its original casing. gRPC expects failures reported through ``callback``, not + raised, so a credential error is passed to the callback rather than raised. + Module-level (not a closure) so it stays testable without grpcio. + """ + try: + auth_headers = cfg.authenticate() + except Exception as e: # noqa: BLE001 - report the failure via callback, not by raising + callback((), e) + return + metadata: Dict[str, str] = {} + metadata.update(_headers_to_metadata(cfg._custom_headers)) + metadata.update(_headers_to_metadata(auth_headers)) + callback(tuple(metadata.items()), None) + + +def _channel_options(cfg: "Config") -> List[Tuple[str, str]]: + """gRPC channel options that carry the SDK's identity. Sets the channel + User-Agent to ``Config.user_agent`` - the same string the HTTP client sends + (``ApiClient`` is built with ``user_agent_base=cfg.user_agent``) - so gRPC + traffic is attributable to databricks-sdk exactly like the REST calls. gRPC + prepends it to its own agent token. Module-level so it stays testable + without grpcio.""" + return [("grpc.primary_user_agent", cfg.user_agent)] + + +def open_channel(cfg: "Config", *, port: int = _DEFAULT_GRPC_PORT) -> "grpc.Channel": + """Open a TLS gRPC channel to the workspace host, authenticated with the + SDK's configured credentials. + + Experimental: an internal helper for hand-written mixins, not public SDK + API. It may change or be removed without warning; do not call it directly. + + Auth headers are resolved from ``cfg.authenticate()`` on every RPC via a + gRPC call-credentials plugin (:func:`_auth_metadata`), so the channel + inherits whatever the SDK's credential strategy produces, refreshed per + call. The channel also advertises the SDK's User-Agent, the same string the + HTTP client sends (see :func:`_channel_options`). The caller owns the + returned channel and must ``close()`` it. + """ + import grpc # noqa: FlagLocalImports - lazy; grpcio ships only with the optional "sandbox" extra + + # Config.skip_verify is intentionally not honored: the SDK's HTTP client does not act + # on it either (it always verifies), so the gRPC transport keeps the same TLS behavior. + channel_credentials = grpc.composite_channel_credentials( + grpc.ssl_channel_credentials(), + grpc.metadata_call_credentials(lambda context, callback: _auth_metadata(cfg, callback)), + ) + return grpc.secure_channel(_grpc_target(cfg.host, port), channel_credentials, options=_channel_options(cfg)) + + +def _is_unavailable(error: Exception) -> bool: + """Default retry predicate: true only for gRPC ``UNAVAILABLE``, the code that + usually means the request never reached the server. Retrying re-issues the + call (at-least-once), so a caller needing other codes - or none - passes its + own ``is_retryable``.""" + import grpc # noqa: FlagLocalImports - lazy; grpcio ships only with the optional "sandbox" extra + + return isinstance(error, grpc.RpcError) and error.code() == grpc.StatusCode.UNAVAILABLE + + +def call_with_retry( + make_call: Callable[[], _T], + *, + is_retryable: Optional[Callable[[Exception], bool]] = None, + max_attempts: int = _DEFAULT_MAX_ATTEMPTS, + retry_interval_seconds: float = _DEFAULT_RETRY_INTERVAL_SECONDS, + sleep: Callable[[float], None] = time.sleep, +) -> _T: + """Invoke a gRPC call, retrying it on transient failures. + + Experimental: an internal helper for hand-written mixins, not public SDK + API. It may change or be removed without warning; do not call it directly. + + ``make_call`` performs the RPC and is re-invoked on each attempt. + ``is_retryable`` decides which errors retry (default: :func:`_is_unavailable`); + a non-retryable error propagates immediately, and once the attempt budget is + exhausted a ``TimeoutError`` is raised. For a server-streaming RPC use + :func:`open_stream_with_retry`. + """ + if is_retryable is None: + is_retryable = _is_unavailable + + last_error: Optional[Exception] = None + for attempt in range(max_attempts): + try: + return make_call() + except Exception as e: # noqa: BLE001 - broad exception catch by design. Retry on retryable errors and raise non-retryable ones. + if is_retryable(e): + last_error = e + _LOG.debug( + "gRPC call failed with a retryable error (attempt %d/%d); retrying in %ss", + attempt + 1, + max_attempts, + retry_interval_seconds, + ) + sleep(retry_interval_seconds) + continue + raise + waited = max_attempts * retry_interval_seconds + raise TimeoutError(f"gRPC call still failing after {waited:.0f}s ({max_attempts} attempts)") from last_error + + +class _ReplayFirstStream: + """Yields an already-read first item, then the rest of a server-streaming + gRPC call. The first item was read early to check the stream opened (see + :func:`open_stream_with_retry`); this hands it back before continuing. + + Attribute access delegates to the underlying stream, so a caller that + abandons iteration can still call ``cancel()`` (or any other stream method) + on the live RPC - which a plain chained iterator would hide. + """ + + def __init__(self, first: _T, stream: Iterator[_T]): + self._first = first + self._first_pending = True + self._stream = stream + + def __iter__(self) -> "Iterator[_T]": + return self + + def __next__(self) -> _T: + if self._first_pending: + self._first_pending = False + return self._first + return next(self._stream) + + def __getattr__(self, name: str): + # Reached only for attributes this proxy does not define (e.g. cancel); + # forward them to the underlying gRPC stream. + return getattr(self._stream, name) + + +def _open_stream( + make_stream: Callable[[], Iterator[_T]], + *, + is_retryable: Callable[[Exception], bool], + max_attempts: int, + retry_interval_seconds: float, + sleep: Callable[[float], None], +) -> Iterator[_T]: + """Open the stream and read its first item under :func:`call_with_retry`, + returning a stream that replays that item at the front (or the raw stream + when it is empty, so ``cancel()`` still works).""" + + def open_stream_and_read_first_item() -> Tuple[object, Iterator[_T]]: + stream = make_stream() + return next(stream, _UNSET), stream + + first, stream = call_with_retry( + open_stream_and_read_first_item, + is_retryable=is_retryable, + max_attempts=max_attempts, + retry_interval_seconds=retry_interval_seconds, + sleep=sleep, + ) + if first is _UNSET: + return stream + return _ReplayFirstStream(first, stream) + + +def _stream_reopening_on_failure( + first_stream: Iterator[_T], + make_stream: Callable[[], Iterator[_T]], + *, + is_retryable: Callable[[Exception], bool], + max_attempts: int, + retry_interval_seconds: float, + sleep: Callable[[float], None], +) -> Iterator[_T]: + """Yield from ``first_stream``, re-opening the whole call on a retryable + mid-stream failure (up to ``max_attempts`` re-opens). Each re-open restarts + from the first item, so the caller re-receives what it already got.""" + stream = first_stream + reopens = 0 + while True: + try: + item = next(stream) + except StopIteration: + return + except Exception as e: # noqa: BLE001 - re-open on retryable errors, raise otherwise + if not is_retryable(e) or reopens >= max_attempts: + raise + reopens += 1 + sleep(retry_interval_seconds) + stream = _open_stream( + make_stream, + is_retryable=is_retryable, + max_attempts=max_attempts, + retry_interval_seconds=retry_interval_seconds, + sleep=sleep, + ) + continue + yield item + + +def open_stream_with_retry( + make_stream: Callable[[], Iterator[_T]], + *, + is_retryable: Optional[Callable[[Exception], bool]] = None, + retry_mid_stream_failures: bool = False, + max_attempts: int = _DEFAULT_MAX_ATTEMPTS, + retry_interval_seconds: float = _DEFAULT_RETRY_INTERVAL_SECONDS, + sleep: Callable[[float], None] = time.sleep, +) -> Iterator[_T]: + """Open a server-streaming call and read its first item. + + Experimental: an internal helper for hand-written mixins, not public SDK + API. It may change or be removed without warning; do not call it directly. + + A gRPC stub sends the request only when the first item is read, so a + not-ready backend fails on that first read. This reads it inside + :func:`call_with_retry` (default predicate :func:`_is_unavailable`) and + replays it at the front, so the caller still receives the whole stream; the + returned object delegates attribute access to the live stream, so + ``cancel()`` still reaches the RPC. + + By default a failure once items are flowing is raised to the caller, since + gRPC cannot resume a broken stream. With ``retry_mid_stream_failures``, a + retryable mid-stream failure instead re-issues the whole call (up to + ``max_attempts`` re-opens); the new call restarts from the first item, so the + caller re-receives what it already got and the RPC runs again - use it only + for an idempotent, replay-tolerant consumer, and note the returned iterator + then does not expose ``cancel()``. ``is_retryable`` and ``sleep`` are + forwarded to :func:`call_with_retry`. + """ + if is_retryable is None: + is_retryable = _is_unavailable + + stream = _open_stream( + make_stream, + is_retryable=is_retryable, + max_attempts=max_attempts, + retry_interval_seconds=retry_interval_seconds, + sleep=sleep, + ) + if not retry_mid_stream_failures: + return stream + return _stream_reopening_on_failure( + stream, + make_stream, + is_retryable=is_retryable, + max_attempts=max_attempts, + retry_interval_seconds=retry_interval_seconds, + sleep=sleep, + ) diff --git a/databricks/sdk/service/aifunctions.py b/databricks/sdk/service/aifunctions.py index 04ab3dc9a..bc7b25568 100644 --- a/databricks/sdk/service/aifunctions.py +++ b/databricks/sdk/service/aifunctions.py @@ -710,8 +710,7 @@ def ai_extract(self, content: any, schema: any, *, options: Optional[AiExtractOp def ai_parse_document( self, content: str, *, options: Optional[AiParseDocumentOptions] = None ) -> AiParseDocumentResponse: - """Parse structured content from unstructured documents. For REST API requests, the default rate limit is - 120 pages per minute per workspace. Contact your Databricks account team to request a higher limit. + """Parse structured content from unstructured documents. :param content: str The document to parse, given as a Unity Catalog volume path to the source file (the REST API accepts diff --git a/databricks/sdk/service/apps.py b/databricks/sdk/service/apps.py index b5bedd687..1397914e0 100644 --- a/databricks/sdk/service/apps.py +++ b/databricks/sdk/service/apps.py @@ -1696,6 +1696,8 @@ class AppUpdate: status: Optional[AppUpdateUpdateStatus] = None + telemetry_export_destinations: Optional[List[TelemetryExportDestination]] = None + usage_policy_id: Optional[str] = None user_api_scopes: Optional[List[str]] = None @@ -1721,6 +1723,8 @@ def as_dict(self) -> dict: body["resources"] = [v.as_dict() for v in self.resources] if self.status: body["status"] = self.status.as_dict() + if self.telemetry_export_destinations: + body["telemetry_export_destinations"] = [v.as_dict() for v in self.telemetry_export_destinations] if self.usage_policy_id is not None: body["usage_policy_id"] = self.usage_policy_id if self.user_api_scopes: @@ -1748,6 +1752,8 @@ def as_shallow_dict(self) -> dict: body["resources"] = self.resources if self.status: body["status"] = self.status + if self.telemetry_export_destinations: + body["telemetry_export_destinations"] = self.telemetry_export_destinations if self.usage_policy_id is not None: body["usage_policy_id"] = self.usage_policy_id if self.user_api_scopes: @@ -1767,6 +1773,9 @@ def from_dict(cls, d: Dict[str, Any]) -> AppUpdate: git_repository=_from_dict(d, "git_repository", GitRepository), resources=_repeated_dict(d, "resources", AppResource), status=_from_dict(d, "status", AppUpdateUpdateStatus), + telemetry_export_destinations=_repeated_dict( + d, "telemetry_export_destinations", TelemetryExportDestination + ), usage_policy_id=d.get("usage_policy_id", None), user_api_scopes=d.get("user_api_scopes", None), ) @@ -1878,8 +1887,7 @@ class ComputeState(Enum): @dataclass class ComputeStatus: active_instances: Optional[int] = None - """The number of compute instances currently serving requests for this application. An instance is - considered active if it is reachable and ready to handle requests.""" + """The number of compute instances used and billed for this application.""" message: Optional[str] = None """Compute status message""" diff --git a/databricks/sdk/service/catalog.py b/databricks/sdk/service/catalog.py index 66a621934..c45e6ef0c 100644 --- a/databricks/sdk/service/catalog.py +++ b/databricks/sdk/service/catalog.py @@ -1918,6 +1918,7 @@ class ConnectionType(Enum): SQLDW = "SQLDW" SQLSERVER = "SQLSERVER" TERADATA = "TERADATA" + TIKTOK_ADS = "TIKTOK_ADS" UNKNOWN_CONNECTION_TYPE = "UNKNOWN_CONNECTION_TYPE" WORKDAY_RAAS = "WORKDAY_RAAS" ZENDESK = "ZENDESK" @@ -2910,6 +2911,27 @@ def from_dict(cls, d: Dict[str, Any]) -> DeleteCredentialResponse: return cls() +@dataclass +class DeleteMcpServiceUserMappedCredentialResponse: + """Delete returns no resource; a dedicated (empty) response keeps the revoke RPC's shape owned here + rather than google.protobuf.Empty.""" + + def as_dict(self) -> dict: + """Serializes the DeleteMcpServiceUserMappedCredentialResponse into a dictionary suitable for use as a JSON request body.""" + body = {} + return body + + def as_shallow_dict(self) -> dict: + """Serializes the DeleteMcpServiceUserMappedCredentialResponse into a shallow dictionary of its immediate attributes.""" + body = {} + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> DeleteMcpServiceUserMappedCredentialResponse: + """Deserializes the DeleteMcpServiceUserMappedCredentialResponse from a dictionary.""" + return cls() + + @dataclass class DeleteMonitorResponse: def as_dict(self) -> dict: @@ -6988,6 +7010,12 @@ class McpServiceConfigSourceConnection: callers can identify the broken dependency; tool invocation fails until the source connection is updated.""" + options: Optional[Dict[str, str]] = None + """Options needed to build the U2M authorize request, returned as a flat map. When set, it + includes: ``authorization_endpoint`` (OAuth authorize URL), ``token_endpoint`` (token-exchange + URL), ``oauth_scope`` (space-separated scopes to request), ``client_id`` (OAuth client id), and + ``oauth_provider`` (the OAuth provider).""" + def as_dict(self) -> dict: """Serializes the McpServiceConfigSourceConnection into a dictionary suitable for use as a JSON request body.""" body = {} @@ -6995,6 +7023,8 @@ def as_dict(self) -> dict: body["is_deleted"] = self.is_deleted if self.name is not None: body["name"] = self.name + if self.options: + body["options"] = self.options return body def as_shallow_dict(self) -> dict: @@ -7004,12 +7034,80 @@ def as_shallow_dict(self) -> dict: body["is_deleted"] = self.is_deleted if self.name is not None: body["name"] = self.name + if self.options: + body["options"] = self.options return body @classmethod def from_dict(cls, d: Dict[str, Any]) -> McpServiceConfigSourceConnection: """Deserializes the McpServiceConfigSourceConnection from a dictionary.""" - return cls(is_deleted=d.get("is_deleted", None), name=d.get("name", None)) + return cls(is_deleted=d.get("is_deleted", None), name=d.get("name", None), options=d.get("options", None)) + + +@dataclass +class McpServiceUserMappedCredential: + """A caller's per-user OAuth credential for an MCP service.""" + + options: Optional[Dict[str, str]] = None + """Token-expiry info for the credential, returned as a flat map: ``access_token_expiration`` + (always set) and ``refresh_token_expiration`` (set when the credential has a refresh token). + Both values are timestamps.""" + + provisioning_info: Optional[ProvisioningInfo] = None + """Provisioning state of the credential. ``ACTIVE`` means the caller is logged in and the + credential is usable; any other state means the login has not completed.""" + + def as_dict(self) -> dict: + """Serializes the McpServiceUserMappedCredential into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.options: + body["options"] = self.options + if self.provisioning_info: + body["provisioning_info"] = self.provisioning_info.as_dict() + return body + + def as_shallow_dict(self) -> dict: + """Serializes the McpServiceUserMappedCredential into a shallow dictionary of its immediate attributes.""" + body = {} + if self.options: + body["options"] = self.options + if self.provisioning_info: + body["provisioning_info"] = self.provisioning_info + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> McpServiceUserMappedCredential: + """Deserializes the McpServiceUserMappedCredential from a dictionary.""" + return cls( + options=d.get("options", None), provisioning_info=_from_dict(d, "provisioning_info", ProvisioningInfo) + ) + + +@dataclass +class McpServiceUserMappedCredentialLogin: + """Login input for an MCP service user credential. Carries the OAuth exchange fields as a flat map.""" + + options: Optional[Dict[str, str]] = None + """OAuth exchange fields: ``pkce_verifier``, ``authorization_code``, and ``oauth_redirect_uri``.""" + + def as_dict(self) -> dict: + """Serializes the McpServiceUserMappedCredentialLogin into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.options: + body["options"] = self.options + return body + + def as_shallow_dict(self) -> dict: + """Serializes the McpServiceUserMappedCredentialLogin into a shallow dictionary of its immediate attributes.""" + body = {} + if self.options: + body["options"] = self.options + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> McpServiceUserMappedCredentialLogin: + """Deserializes the McpServiceUserMappedCredentialLogin from a dictionary.""" + return cls(options=d.get("options", None)) @dataclass @@ -7815,6 +7913,47 @@ def from_dict(cls, d: Dict[str, Any]) -> ModelProviderServiceConfigAzureOpenAiPr ) +@dataclass +class ModelProviderServiceConfigCustomProviderApiKeyHeaderAuth: + """Header-based API-key authentication for a custom provider: the secret is forwarded on outbound + requests under a caller-chosen HTTP header, as ``: ``.""" + + api_key_name: Optional[str] = None + """HTTP header name that carries the API key on outbound requests (e.g., + ``Ocp-Apim-Subscription-Key``). The value forwarded under this header is supplied via + ``api_key_value``.""" + + api_key_value: Optional[ModelProviderServiceConfigProviderSecret] = None + """Secret value forwarded under the ``api_key_name`` header on outbound requests. Supplied as + inline plaintext via ``ProviderSecret.plaintext``.""" + + def as_dict(self) -> dict: + """Serializes the ModelProviderServiceConfigCustomProviderApiKeyHeaderAuth into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.api_key_name is not None: + body["api_key_name"] = self.api_key_name + if self.api_key_value: + body["api_key_value"] = self.api_key_value.as_dict() + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ModelProviderServiceConfigCustomProviderApiKeyHeaderAuth into a shallow dictionary of its immediate attributes.""" + body = {} + if self.api_key_name is not None: + body["api_key_name"] = self.api_key_name + if self.api_key_value: + body["api_key_value"] = self.api_key_value + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ModelProviderServiceConfigCustomProviderApiKeyHeaderAuth: + """Deserializes the ModelProviderServiceConfigCustomProviderApiKeyHeaderAuth from a dictionary.""" + return cls( + api_key_name=d.get("api_key_name", None), + api_key_value=_from_dict(d, "api_key_value", ModelProviderServiceConfigProviderSecret), + ) + + @dataclass class ModelProviderServiceConfigCustomProviderConfig: """Custom OpenAI-compatible provider configuration with bearer-token authentication.""" @@ -7844,8 +7983,9 @@ def from_dict(cls, d: Dict[str, Any]) -> ModelProviderServiceConfigCustomProvide @dataclass class ModelProviderServiceConfigCustomProviderDirectConfig: - """Direct form of a custom provider configuration. Set ``api_key`` to the bearer token sent in the - ``Authorization`` header.""" + """Direct form of a custom provider configuration. Set ``api_key`` to send the secret as an + ``Authorization`` bearer token, or ``header_auth`` to forward it under a caller-chosen HTTP + header.""" api_key: Optional[ModelProviderServiceConfigProviderSecret] = None """Bearer token forwarded in the ``Authorization`` header. Supply the value in @@ -7855,6 +7995,11 @@ class ModelProviderServiceConfigCustomProviderDirectConfig: """Endpoint URL of the OpenAI-compatible service (e.g., ``https://api.example.com/v1``). Required on Create.""" + header_auth: Optional[ModelProviderServiceConfigCustomProviderApiKeyHeaderAuth] = None + """Header-based API-key auth: the secret is forwarded on outbound requests under a caller-chosen + HTTP header rather than as an ``Authorization`` bearer token. Set this instead of ``api_key`` + for header auth.""" + def as_dict(self) -> dict: """Serializes the ModelProviderServiceConfigCustomProviderDirectConfig into a dictionary suitable for use as a JSON request body.""" body = {} @@ -7862,6 +8007,8 @@ def as_dict(self) -> dict: body["api_key"] = self.api_key.as_dict() if self.base_url is not None: body["base_url"] = self.base_url + if self.header_auth: + body["header_auth"] = self.header_auth.as_dict() return body def as_shallow_dict(self) -> dict: @@ -7871,13 +8018,17 @@ def as_shallow_dict(self) -> dict: body["api_key"] = self.api_key if self.base_url is not None: body["base_url"] = self.base_url + if self.header_auth: + body["header_auth"] = self.header_auth return body @classmethod def from_dict(cls, d: Dict[str, Any]) -> ModelProviderServiceConfigCustomProviderDirectConfig: """Deserializes the ModelProviderServiceConfigCustomProviderDirectConfig from a dictionary.""" return cls( - api_key=_from_dict(d, "api_key", ModelProviderServiceConfigProviderSecret), base_url=d.get("base_url", None) + api_key=_from_dict(d, "api_key", ModelProviderServiceConfigProviderSecret), + base_url=d.get("base_url", None), + header_auth=_from_dict(d, "header_auth", ModelProviderServiceConfigCustomProviderApiKeyHeaderAuth), ) @@ -11644,6 +11795,7 @@ def from_dict(cls, d: Dict[str, Any]) -> SecurablePermissions: class SecurableType(Enum): """The type of Unity Catalog securable.""" + AGENT_SERVICE = "AGENT_SERVICE" CATALOG = "CATALOG" CLEAN_ROOM = "CLEAN_ROOM" CONNECTION = "CONNECTION" @@ -11661,6 +11813,7 @@ class SecurableType(Enum): RECIPIENT = "RECIPIENT" SCHEMA = "SCHEMA" SHARE = "SHARE" + SKILL = "SKILL" STAGING_TABLE = "STAGING_TABLE" STORAGE_CREDENTIAL = "STORAGE_CREDENTIAL" TABLE = "TABLE" @@ -13856,6 +14009,36 @@ def create_mcp_service(self, mcp_service: McpService, parent: str, mcp_service_i res = self._api.do("POST", "/api/2.1/unity-catalog/mcp-services", query=query, body=body, headers=headers) return McpService.from_dict(res) + def create_mcp_service_user_mapped_credential( + self, name: str, login: McpServiceUserMappedCredentialLogin + ) -> McpServiceUserMappedCredential: + """Logs the caller in to an MCP service: creates their per-user OAuth credential, or re-authenticates it + if one already exists. The request body carries the OAuth exchange fields. + + You must be the owner of the MCP service or have ``EXECUTE`` on it, plus ``USE_CATALOG`` on the parent + catalog and ``USE_SCHEMA`` on the parent schema. + + :param name: str + Resource name of the MCP service. Format: ``mcp-services/{catalog}.{schema}.{mcp_service}``. + :param login: :class:`McpServiceUserMappedCredentialLogin` + + :returns: :class:`McpServiceUserMappedCredential` + """ + + body = login.as_dict() + query = {} + 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/2.1/unity-catalog/{name}/user-credentials", body=body, headers=headers) + return McpServiceUserMappedCredential.from_dict(res) + def create_model_provider_service( self, model_provider_service: ModelProviderService, parent: str, model_provider_service_id: str ) -> ModelProviderService: @@ -13974,6 +14157,29 @@ def delete_mcp_service(self, name: str, *, etag: Optional[str] = None): self._api.do("DELETE", f"/api/2.1/unity-catalog/{name}", query=query, headers=headers) + def delete_mcp_service_user_mapped_credential(self, name: str) -> DeleteMcpServiceUserMappedCredentialResponse: + """Revokes (deletes) the caller's per-user OAuth credential for an MCP service (logout). + + You must be the owner of the MCP service or have ``EXECUTE`` on it, plus ``USE_CATALOG`` on the parent + catalog and ``USE_SCHEMA`` on the parent schema. + + :param name: str + Resource name of the MCP service. Format: ``mcp-services/{catalog}.{schema}.{mcp_service}``. + + :returns: :class:`DeleteMcpServiceUserMappedCredentialResponse` + """ + + 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/2.1/unity-catalog/{name}/user-credentials", headers=headers) + return DeleteMcpServiceUserMappedCredentialResponse.from_dict(res) + def delete_model_provider_service(self, name: str, *, etag: Optional[str] = None): """Deletes the model provider service identified by its resource name. Optionally supply an ``etag`` to make the delete conditional on the model provider service not having changed since it was read. @@ -14061,6 +14267,32 @@ def get_mcp_service(self, name: str) -> McpService: res = self._api.do("GET", f"/api/2.1/unity-catalog/{name}", headers=headers) return McpService.from_dict(res) + def get_mcp_service_user_mapped_credential(self, name: str) -> McpServiceUserMappedCredential: + """Returns the caller's per-user OAuth login state for an MCP service. Read ``provisioning_info.state``: + ``ACTIVE`` means the caller is logged in and the credential is usable; any other state (for example a + failed or still-provisioning login) means the login has not completed and the caller should log in + again. If the caller has no credential yet, the RPC returns ``NOT_FOUND``. + + You must be the owner of the MCP service or have ``EXECUTE`` on it, plus ``USE_CATALOG`` on the parent + catalog and ``USE_SCHEMA`` on the parent schema. + + :param name: str + Resource name of the MCP service. Format: ``mcp-services/{catalog}.{schema}.{mcp_service}``. + + :returns: :class:`McpServiceUserMappedCredential` + """ + + 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.1/unity-catalog/{name}/user-credentials", headers=headers) + return McpServiceUserMappedCredential.from_dict(res) + def get_model_provider_service(self, name: str) -> ModelProviderService: """Returns the model provider service identified by its resource name. diff --git a/databricks/sdk/service/dashboards.py b/databricks/sdk/service/dashboards.py index 99805acd2..38cca7979 100644 --- a/databricks/sdk/service/dashboards.py +++ b/databricks/sdk/service/dashboards.py @@ -3806,7 +3806,8 @@ def create( ) -> Dashboard: """Create a draft dashboard. - Requires the Databricks SQL access entitlement. + Requires the `Databricks SQL access `__ + entitlement. Grant Databricks SQL access in addition to Workspace access. :param dashboard: :class:`Dashboard` :param dataset_catalog: str (optional) @@ -3842,6 +3843,9 @@ def create( def create_schedule(self, dashboard_id: str, schedule: Schedule) -> Schedule: """Create dashboard schedule. + Requires the `Databricks SQL access `__ + entitlement. Grant Databricks SQL access in addition to Workspace access. + :param dashboard_id: str UUID identifying the dashboard to which the schedule belongs. :param schedule: :class:`Schedule` @@ -3867,6 +3871,13 @@ def create_schedule(self, dashboard_id: str, schedule: Schedule) -> Schedule: def create_subscription(self, dashboard_id: str, schedule_id: str, subscription: Subscription) -> Subscription: """Create schedule subscription. + The caller must be a workspace user with one of the following `entitlements + `__: Workspace access, Databricks SQL access, + or Consumer access. + + Account-level users who are not members of the workspace cannot call this endpoint, even if the + dashboard has been shared with them. + :param dashboard_id: str UUID identifying the dashboard to which the subscription belongs. :param schedule_id: str @@ -3899,6 +3910,9 @@ def create_subscription(self, dashboard_id: str, schedule_id: str, subscription: def delete_schedule(self, dashboard_id: str, schedule_id: str, *, etag: Optional[str] = None): """Delete dashboard schedule. + Requires the `Databricks SQL access `__ + entitlement. Grant Databricks SQL access in addition to Workspace access. + :param dashboard_id: str UUID identifying the dashboard to which the schedule belongs. :param schedule_id: str @@ -3933,6 +3947,13 @@ def delete_subscription( ): """Delete schedule subscription. + The caller must be a workspace user with one of the following `entitlements + `__: Workspace access, Databricks SQL access, + or Consumer access. + + Account-level users who are not members of the workspace cannot call this endpoint, even if the + dashboard has been shared with them. + :param dashboard_id: str UUID identifying the dashboard which the subscription belongs. :param schedule_id: str @@ -3967,7 +3988,8 @@ def delete_subscription( def get(self, dashboard_id: str) -> Dashboard: """Get a draft dashboard. - Requires the Databricks SQL access entitlement. + Requires the `Databricks SQL access `__ + entitlement. Grant Databricks SQL access in addition to Workspace access. :param dashboard_id: str UUID identifying the dashboard. @@ -4015,6 +4037,13 @@ def get_published(self, dashboard_id: str) -> PublishedDashboard: def get_schedule(self, dashboard_id: str, schedule_id: str) -> Schedule: """Get dashboard schedule. + The caller must be a workspace user with one of the following `entitlements + `__: Workspace access, Databricks SQL access, + or Consumer access. + + Account-level users who are not members of the workspace cannot call this endpoint, even if the + dashboard has been shared with them. + :param dashboard_id: str UUID identifying the dashboard to which the schedule belongs. :param schedule_id: str @@ -4039,6 +4068,13 @@ def get_schedule(self, dashboard_id: str, schedule_id: str) -> Schedule: def get_subscription(self, dashboard_id: str, schedule_id: str, subscription_id: str) -> Subscription: """Get schedule subscription. + The caller must be a workspace user with one of the following `entitlements + `__: Workspace access, Databricks SQL access, + or Consumer access. + + Account-level users who are not members of the workspace cannot call this endpoint, even if the + dashboard has been shared with them. + :param dashboard_id: str UUID identifying the dashboard which the subscription belongs. :param schedule_id: str @@ -4074,7 +4110,8 @@ def list( ) -> Iterator[Dashboard]: """List dashboards. - Requires the Databricks SQL access entitlement. + Requires the `Databricks SQL access `__ + entitlement. Grant Databricks SQL access in addition to Workspace access. :param page_size: int (optional) The number of dashboards to return per page. @@ -4121,6 +4158,13 @@ def list_schedules( ) -> Iterator[Schedule]: """List dashboard schedules. + The caller must be a workspace user with one of the following `entitlements + `__: Workspace access, Databricks SQL access, + or Consumer access. + + Account-level users who are not members of the workspace cannot call this endpoint, even if the + dashboard has been shared with them. + :param dashboard_id: str UUID identifying the dashboard to which the schedules belongs. :param page_size: int (optional) @@ -4161,6 +4205,13 @@ def list_subscriptions( ) -> Iterator[Subscription]: """List schedule subscriptions. + The caller must be a workspace user with one of the following `entitlements + `__: Workspace access, Databricks SQL access, + or Consumer access. + + Account-level users who are not members of the workspace cannot call this endpoint, even if the + dashboard has been shared with them. + :param dashboard_id: str UUID identifying the dashboard which the subscriptions belongs. :param schedule_id: str @@ -4251,7 +4302,8 @@ def publish( ) -> PublishedDashboard: """Publish the current draft dashboard. - Requires the Databricks SQL access entitlement. + Requires the `Databricks SQL access `__ + entitlement. Grant Databricks SQL access in addition to Workspace access. :param dashboard_id: str UUID identifying the dashboard to be published. @@ -4284,7 +4336,8 @@ def publish( def revert(self, dashboard_id: str, *, etag: Optional[str] = None) -> RevertDashboardResponse: """Revert a dashboard's definition in draft mode to the last published version. - Requires the Databricks SQL access entitlement. + Requires the `Databricks SQL access `__ + entitlement. Grant Databricks SQL access in addition to Workspace access. :param dashboard_id: str UUID identifying the dashboard. @@ -4313,7 +4366,8 @@ def revert(self, dashboard_id: str, *, etag: Optional[str] = None) -> RevertDash def trash(self, dashboard_id: str): """Trash a dashboard. - Requires the Databricks SQL access entitlement. + Requires the `Databricks SQL access `__ + entitlement. Grant Databricks SQL access in addition to Workspace access. :param dashboard_id: str UUID identifying the dashboard. @@ -4334,7 +4388,8 @@ def trash(self, dashboard_id: str): def unpublish(self, dashboard_id: str): """Unpublish the dashboard. - Requires the Databricks SQL access entitlement. + Requires the `Databricks SQL access `__ + entitlement. Grant Databricks SQL access in addition to Workspace access. :param dashboard_id: str UUID identifying the published dashboard. @@ -4362,7 +4417,8 @@ def update( ) -> Dashboard: """Update a draft dashboard. - Requires the Databricks SQL access entitlement. + Requires the `Databricks SQL access `__ + entitlement. Grant Databricks SQL access in addition to Workspace access. :param dashboard_id: str UUID identifying the dashboard. @@ -4402,6 +4458,9 @@ def update( def update_schedule(self, dashboard_id: str, schedule_id: str, schedule: Schedule) -> Schedule: """Update dashboard schedule. + Requires the `Databricks SQL access `__ + entitlement. Grant Databricks SQL access in addition to Workspace access. + :param dashboard_id: str UUID identifying the dashboard to which the schedule belongs. :param schedule_id: str diff --git a/databricks/sdk/service/files.py b/databricks/sdk/service/files.py index f897417cd..5e606f2e3 100644 --- a/databricks/sdk/service/files.py +++ b/databricks/sdk/service/files.py @@ -787,7 +787,7 @@ class FilesAPI: The API supports `Unity Catalog volumes `__, where files and directories to operate on are specified using their volume URI path, which follows the format - /Volumes/<catalog_name>/<schema_name>/<volume_name>/<path_to_file>. + /Volumes////. The Files API has two distinct endpoints, one for working with files (``/fs/files``) and another one for working with directories (``/fs/directories``). Both endpoints use the standard HTTP methods GET, HEAD, diff --git a/databricks/sdk/service/iamv2.py b/databricks/sdk/service/iamv2.py index 8f18aa552..8b503394b 100644 --- a/databricks/sdk/service/iamv2.py +++ b/databricks/sdk/service/iamv2.py @@ -1336,7 +1336,7 @@ class WorkspacePermission(Enum): class AccountIamV2API: - """These APIs are used to manage identities and the workspace access of these identities in .""" + """These APIs are used to manage identities and the workspace access of these identities in Databricks.""" def __init__(self, api_client): self._api = api_client @@ -1377,7 +1377,7 @@ def create_group(self, group: Group) -> Group: identity, use the ExternalGroup resource. :param group: :class:`Group` - Required. Group to be created in + Required. Group to be created in Databricks :returns: :class:`Group` """ @@ -1404,7 +1404,7 @@ def create_service_principal(self, service_principal: ServicePrincipal) -> Servi external identity, use the ExternalServicePrincipal resource. :param service_principal: :class:`ServicePrincipal` - Required. Service principal to be created in + Required. Service principal to be created in Databricks :returns: :class:`ServicePrincipal` """ @@ -1431,7 +1431,7 @@ def create_user(self, user: User) -> User: identity, use the ExternalUser resource. :param user: :class:`User` - Required. User to be created in + Required. User to be created in Databricks :returns: :class:`User` """ @@ -1459,7 +1459,7 @@ def create_workspace_assignment( :param workspace_id: int Required. The workspace ID for which the workspace assignment is being created. :param workspace_assignment: :class:`WorkspaceAssignment` - Required. Workspace assignment to be created in . + Required. Workspace assignment to be created in Databricks. :returns: :class:`WorkspaceAssignment` """ @@ -1487,7 +1487,7 @@ def create_workspace_assignment_detail( :param workspace_id: int Required. The workspace ID for which the workspace assignment detail is being created. :param workspace_assignment_detail: :class:`WorkspaceAssignmentDetail` - Required. Workspace assignment detail to be created in . + Required. Workspace assignment detail to be created in Databricks. :returns: :class:`WorkspaceAssignmentDetail` """ @@ -2218,7 +2218,7 @@ def update_group(self, group_id: str, group: Group, update_mask: str) -> Group: :param group_id: str Required. Internal ID of the group in Databricks. :param group: :class:`Group` - Required. Group to be updated in + Required. Group to be updated in Databricks :param update_mask: str Optional. The list of fields to update. @@ -2255,7 +2255,7 @@ def update_service_principal( :param service_principal_id: str Required. Internal ID of the service principal in Databricks. :param service_principal: :class:`ServicePrincipal` - Required. Service Principal to be updated in + Required. Service Principal to be updated in Databricks :param update_mask: str Optional. The list of fields to update. @@ -2291,7 +2291,7 @@ def update_user(self, user_id: str, user: User, update_mask: str) -> User: :param user_id: str Required. Internal ID of the user in Databricks. :param user: :class:`User` - Required. User to be updated in + Required. User to be updated in Databricks :param update_mask: str Optional. The list of fields to update. @@ -2328,7 +2328,7 @@ def update_workspace_assignment( :param principal_id: int Required. ID of the principal in Databricks. :param workspace_assignment: :class:`WorkspaceAssignment` - Required. Workspace assignment to be updated in . + Required. Workspace assignment to be updated in Databricks. :param update_mask: FieldMask Required. The list of fields to update. @@ -2369,7 +2369,7 @@ def update_workspace_assignment_detail( :param principal_id: int Required. ID of the principal in Databricks. :param workspace_assignment_detail: :class:`WorkspaceAssignmentDetail` - Required. Workspace assignment detail to be updated in . + Required. Workspace assignment detail to be updated in Databricks. :param update_mask: FieldMask Required. The list of fields to update. @@ -2396,7 +2396,7 @@ def update_workspace_assignment_detail( class WorkspaceIamV2API: - """These APIs are used to manage identities and the workspace access of these identities in .""" + """These APIs are used to manage identities and the workspace access of these identities in Databricks.""" def __init__(self, api_client): self._api = api_client @@ -2438,7 +2438,7 @@ def create_group_proxy(self, group: Group) -> Group: identity, use the ExternalGroup resource. :param group: :class:`Group` - Required. Group to be created in + Required. Group to be created in Databricks :returns: :class:`Group` """ @@ -2468,7 +2468,7 @@ def create_service_principal_proxy(self, service_principal: ServicePrincipal) -> external identity, use the ExternalServicePrincipal resource. :param service_principal: :class:`ServicePrincipal` - Required. Service principal to be created in + Required. Service principal to be created in Databricks :returns: :class:`ServicePrincipal` """ @@ -2497,7 +2497,7 @@ def create_user_proxy(self, user: User) -> User: identity, use the ExternalUser resource. :param user: :class:`User` - Required. User to be created in + Required. User to be created in Databricks :returns: :class:`User` """ @@ -2522,7 +2522,7 @@ def create_workspace_assignment_detail_proxy( """Creates a workspace assignment detail for a principal in the calling workspace. :param workspace_assignment_detail: :class:`WorkspaceAssignmentDetail` - Required. Workspace assignment detail to be created in . + Required. Workspace assignment detail to be created in Databricks. :returns: :class:`WorkspaceAssignmentDetail` """ @@ -2548,7 +2548,7 @@ def create_workspace_assignment_proxy(self, workspace_assignment: WorkspaceAssig confirm which entitlements were granted. :param workspace_assignment: :class:`WorkspaceAssignment` - Required. Workspace assignment to be created in . + Required. Workspace assignment to be created in Databricks. :returns: :class:`WorkspaceAssignment` """ @@ -3315,7 +3315,7 @@ def update_group_proxy(self, group_id: str, group: Group, update_mask: str) -> G :param group_id: str Required. Internal ID of the group in Databricks. :param group: :class:`Group` - Required. Group to be updated in + Required. Group to be updated in Databricks :param update_mask: str Optional. The list of fields to update. @@ -3350,7 +3350,7 @@ def update_service_principal_proxy( :param service_principal_id: str Required. Internal ID of the service principal in Databricks. :param service_principal: :class:`ServicePrincipal` - Required. Service principal to be updated in + Required. Service principal to be updated in Databricks :param update_mask: str Optional. The list of fields to update. @@ -3390,7 +3390,7 @@ def update_user_proxy(self, user_id: str, user: User, update_mask: str) -> User: :param user_id: str Required. Internal ID of the user in Databricks. :param user: :class:`User` - Required. User to be updated in + Required. User to be updated in Databricks :param update_mask: str Optional. The list of fields to update. @@ -3423,7 +3423,7 @@ def update_workspace_assignment_detail_proxy( :param principal_id: int Required. ID of the principal in Databricks. :param workspace_assignment_detail: :class:`WorkspaceAssignmentDetail` - Required. Workspace assignment detail to be updated in . + Required. Workspace assignment detail to be updated in Databricks. :param update_mask: FieldMask Required. The list of fields to update. @@ -3462,7 +3462,7 @@ def update_workspace_assignment_proxy( :param principal_id: int Required. ID of the principal in Databricks. :param workspace_assignment: :class:`WorkspaceAssignment` - Required. Workspace assignment to be updated in . + Required. Workspace assignment to be updated in Databricks. :param update_mask: FieldMask Required. The list of fields to update. @@ -3495,7 +3495,7 @@ def update_workspace_identity_detail( :param principal_id: int Required. ID of the principal in Databricks. :param workspace_identity_detail: :class:`WorkspaceIdentityDetail` - Required. Workspace identity detail to be updated in . + Required. Workspace identity detail to be updated in Databricks. :param update_mask: FieldMask Required. The list of fields to update. diff --git a/databricks/sdk/service/jobs.py b/databricks/sdk/service/jobs.py index 13c158081..d2085b7ae 100644 --- a/databricks/sdk/service/jobs.py +++ b/databricks/sdk/service/jobs.py @@ -500,9 +500,11 @@ class BaseRun: the client-set performance target on the request depending on whether the performance mode is supported by the job type. - - ``STANDARD``: Enables cost-efficient execution of serverless workloads. - ``PERFORMANCE_OPTIMIZED``: Prioritizes fast startup and execution times through rapid scaling - and optimized cluster performance.""" + and optimized cluster performance. + - ``STANDARD``: Enables cost-efficient execution of serverless workloads. + - ``COST_OPTIMIZED``: Enables lower job costs by optimizing compute for your selected target + duration time. Must provide a duration target.""" effective_usage_policy_id: Optional[str] = None """The id of the usage policy used by this run for cost attribution purposes.""" @@ -3593,15 +3595,16 @@ def from_dict(cls, d: Dict[str, Any]) -> JobPermissionsDescription: @dataclass class JobRunAs: - """Write-only setting. Specifies the user or service principal that the job runs as. If not + """Write-only setting. Specifies the user, service principal, or group that the job runs as. If not specified, the job runs as the user who created the job. - Either ``user_name`` or ``service_principal_name`` should be specified. If not, an error is - thrown.""" + One of ``user_name``, ``service_principal_name``, or ``group_name`` should be specified. If not, + an error is thrown.""" group_name: Optional[str] = None - """Group name of an account group assigned to the workspace. Setting this field requires being a - member of the group.""" + """Group name of an account group assigned to the workspace. When set, all tasks run as the group + and the group's permissions are used for data access. Setting this field requires being a member + of the group, or having the ``Assume`` permission on the group.""" service_principal_name: Optional[str] = None """Application ID of an active service principal. Setting this field requires the @@ -3733,9 +3736,11 @@ class JobSettings: or cost-efficiency for the run. The performance target does not apply to tasks that run on Serverless GPU compute. - - ``STANDARD``: Enables cost-efficient execution of serverless workloads. - ``PERFORMANCE_OPTIMIZED``: Prioritizes fast startup and execution times through rapid scaling - and optimized cluster performance.""" + and optimized cluster performance. + - ``STANDARD``: Enables cost-efficient execution of serverless workloads. + - ``COST_OPTIMIZED``: Enables lower job costs by optimizing compute for your selected target + duration time. Must provide a duration target.""" queue: Optional[QueueSettings] = None """The queue settings of the job.""" @@ -5205,9 +5210,11 @@ class RepairHistoryItem: the client-set performance target on the request depending on whether the performance mode is supported by the job type. - - ``STANDARD``: Enables cost-efficient execution of serverless workloads. - ``PERFORMANCE_OPTIMIZED``: Prioritizes fast startup and execution times through rapid scaling - and optimized cluster performance.""" + and optimized cluster performance. + - ``STANDARD``: Enables cost-efficient execution of serverless workloads. + - ``COST_OPTIMIZED``: Enables lower job costs by optimizing compute for your selected target + duration time. Must provide a duration target.""" end_time: Optional[int] = None """The end time of the (repaired) run.""" @@ -5701,9 +5708,11 @@ class Run: the client-set performance target on the request depending on whether the performance mode is supported by the job type. - - ``STANDARD``: Enables cost-efficient execution of serverless workloads. - ``PERFORMANCE_OPTIMIZED``: Prioritizes fast startup and execution times through rapid scaling - and optimized cluster performance.""" + and optimized cluster performance. + - ``STANDARD``: Enables cost-efficient execution of serverless workloads. + - ``COST_OPTIMIZED``: Enables lower job costs by optimizing compute for your selected target + duration time. Must provide a duration target.""" effective_usage_policy_id: Optional[str] = None """The id of the usage policy used by this run for cost attribution purposes.""" @@ -6942,9 +6951,11 @@ class RunTask: the client-set performance target on the request depending on whether the performance mode is supported by the job type. - - ``STANDARD``: Enables cost-efficient execution of serverless workloads. - ``PERFORMANCE_OPTIMIZED``: Prioritizes fast startup and execution times through rapid scaling - and optimized cluster performance.""" + and optimized cluster performance. + - ``STANDARD``: Enables cost-efficient execution of serverless workloads. + - ``COST_OPTIMIZED``: Enables lower job costs by optimizing compute for your selected target + duration time. Must provide a duration target.""" effective_serverless_compute_id: Optional[str] = None """The id of the serverless compute this task ran on, either explicitly configured on the task or @@ -10437,9 +10448,11 @@ def create( cost-efficiency for the run. The performance target does not apply to tasks that run on Serverless GPU compute. - - ``STANDARD``: Enables cost-efficient execution of serverless workloads. - ``PERFORMANCE_OPTIMIZED``: Prioritizes fast startup and execution times through rapid scaling and optimized cluster performance. + - ``STANDARD``: Enables cost-efficient execution of serverless workloads. + - ``COST_OPTIMIZED``: Enables lower job costs by optimizing compute for your selected target + duration time. Must provide a duration target. :param queue: :class:`QueueSettings` (optional) The queue settings of the job. :param run_as: :class:`JobRunAs` (optional) @@ -10990,9 +11003,11 @@ def repair_run( performance or cost-efficiency for the run. This field overrides the performance target defined on the job level. - - ``STANDARD``: Enables cost-efficient execution of serverless workloads. - ``PERFORMANCE_OPTIMIZED``: Prioritizes fast startup and execution times through rapid scaling and optimized cluster performance. + - ``STANDARD``: Enables cost-efficient execution of serverless workloads. + - ``COST_OPTIMIZED``: Enables lower job costs by optimizing compute for your selected target + duration time. Must provide a duration target. :param pipeline_params: :class:`PipelineParams` (optional) Controls whether the pipeline should perform a full refresh :param python_named_params: Dict[str,str] (optional) @@ -11241,9 +11256,11 @@ def run_now( performance or cost-efficiency for the run. This field overrides the performance target defined on the job level. - - ``STANDARD``: Enables cost-efficient execution of serverless workloads. - ``PERFORMANCE_OPTIMIZED``: Prioritizes fast startup and execution times through rapid scaling and optimized cluster performance. + - ``STANDARD``: Enables cost-efficient execution of serverless workloads. + - ``COST_OPTIMIZED``: Enables lower job costs by optimizing compute for your selected target + duration time. Must provide a duration target. :param pipeline_params: :class:`PipelineParams` (optional) Controls whether the pipeline should perform a full refresh :param python_named_params: Dict[str,str] (optional) @@ -11468,9 +11485,11 @@ def submit( performance or cost-efficiency for the run. The performance target does not apply to tasks that run on Serverless GPU compute. - - ``STANDARD``: Enables cost-efficient execution of serverless workloads. - ``PERFORMANCE_OPTIMIZED``: Prioritizes fast startup and execution times through rapid scaling and optimized cluster performance. + - ``STANDARD``: Enables cost-efficient execution of serverless workloads. + - ``COST_OPTIMIZED``: Enables lower job costs by optimizing compute for your selected target + duration time. Must provide a duration target. :param queue: :class:`QueueSettings` (optional) The queue settings of the one-time run. :param run_as: :class:`JobRunAs` (optional) diff --git a/databricks/sdk/service/ml.py b/databricks/sdk/service/ml.py index 45373e198..b1c297bd3 100644 --- a/databricks/sdk/service/ml.py +++ b/databricks/sdk/service/ml.py @@ -1131,13 +1131,18 @@ class CronSchedule: cron_expression: Optional[str] = None """The cron expression defining the schedule (e.g., "0 0 * * *" for daily at midnight). The - schedule is interpreted in the UTC time zone. Required when mode is MANUAL (or unset). Left - empty when mode is DERIVED, where the service computes it (aligned to UTC) from the features' - window timing and fills it in on the response.""" + schedule is interpreted in timezone_id (defaults to UTC). Required when mode is MANUAL (or + unset). Left empty when mode is DERIVED, where the service computes it (aligned to UTC) from the + features' window timing and fills it in on the response.""" mode: Optional[CronScheduleMode] = None """How the schedule is determined. Defaults to MANUAL when unset.""" + timezone_id: Optional[str] = None + """A Java timezone ID. The schedule is resolved with respect to this timezone. Defaults to UTC when + omitted. Can only be configured for MANUAL schedules; DERIVED schedules are always aligned to + UTC.""" + def as_dict(self) -> dict: """Serializes the CronSchedule into a dictionary suitable for use as a JSON request body.""" body = {} @@ -1145,6 +1150,8 @@ def as_dict(self) -> dict: body["cron_expression"] = self.cron_expression if self.mode is not None: body["mode"] = self.mode.value + if self.timezone_id is not None: + body["timezone_id"] = self.timezone_id return body def as_shallow_dict(self) -> dict: @@ -1154,12 +1161,18 @@ def as_shallow_dict(self) -> dict: body["cron_expression"] = self.cron_expression if self.mode is not None: body["mode"] = self.mode + if self.timezone_id is not None: + body["timezone_id"] = self.timezone_id return body @classmethod def from_dict(cls, d: Dict[str, Any]) -> CronSchedule: """Deserializes the CronSchedule from a dictionary.""" - return cls(cron_expression=d.get("cron_expression", None), mode=_enum(d, "mode", CronScheduleMode)) + return cls( + cron_expression=d.get("cron_expression", None), + mode=_enum(d, "mode", CronScheduleMode), + timezone_id=d.get("timezone_id", None), + ) class CronScheduleMode(Enum): @@ -11803,6 +11816,10 @@ def publish_table(self, source_table_name: str, publish_spec: PublishSpec) -> Pu def update_online_store(self, name: str, online_store: OnlineStore, update_mask: str) -> OnlineStore: """Update an Online Feature Store. + This update is not guaranteed to be atomic: when a request changes multiple fields, some may be + applied while others fail. On a failed response, treat the update as partially applied and retry until + it succeeds. + :param name: str The name of the online store. This is the unique identifier for the online store. :param online_store: :class:`OnlineStore` diff --git a/databricks/sdk/service/pipelines.py b/databricks/sdk/service/pipelines.py index 591e7882d..0b93dbd5c 100644 --- a/databricks/sdk/service/pipelines.py +++ b/databricks/sdk/service/pipelines.py @@ -142,6 +142,58 @@ def from_dict(cls, d: Dict[str, Any]) -> AutoFullRefreshPolicy: return cls(enabled=d.get("enabled", None), min_interval_hours=d.get("min_interval_hours", None)) +@dataclass +class AvroTransformerOptions: + parse_mode: Optional[ParseMode] = None + """(Optional) Parse mode for Avro data. Valid values: FAILFAST, PERMISSIVE. Defaults to FAILFAST.""" + + schema: Optional[str] = None + """Inline Avro JSON schema string.""" + + schema_file_path: Optional[str] = None + """Path to a schema file (.avsc).""" + + schema_registry: Optional[SchemaRegistryConfig] = None + """(Optional) Schema registry to resolve the Avro schema at runtime instead of providing it inline + or via a file path.""" + + def as_dict(self) -> dict: + """Serializes the AvroTransformerOptions into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.parse_mode is not None: + body["parse_mode"] = self.parse_mode.value + if self.schema is not None: + body["schema"] = self.schema + if self.schema_file_path is not None: + body["schema_file_path"] = self.schema_file_path + if self.schema_registry: + body["schema_registry"] = self.schema_registry.as_dict() + return body + + def as_shallow_dict(self) -> dict: + """Serializes the AvroTransformerOptions into a shallow dictionary of its immediate attributes.""" + body = {} + if self.parse_mode is not None: + body["parse_mode"] = self.parse_mode + if self.schema is not None: + body["schema"] = self.schema + if self.schema_file_path is not None: + body["schema_file_path"] = self.schema_file_path + if self.schema_registry: + body["schema_registry"] = self.schema_registry + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> AvroTransformerOptions: + """Deserializes the AvroTransformerOptions from a dictionary.""" + return cls( + parse_mode=_enum(d, "parse_mode", ParseMode), + schema=d.get("schema", None), + schema_file_path=d.get("schema_file_path", None), + schema_registry=_from_dict(d, "schema_registry", SchemaRegistryConfig), + ) + + class CloneMode(Enum): """Enum to specify which mode of clone to execute""" @@ -200,6 +252,31 @@ def from_dict(cls, d: Dict[str, Any]) -> ConfluenceConnectorOptions: return cls(include_confluence_spaces=d.get("include_confluence_spaces", None)) +@dataclass +class ConfluentSchemaRegistryOptions: + subject: Optional[str] = None + """Required: subject name to resolve in the registry.""" + + def as_dict(self) -> dict: + """Serializes the ConfluentSchemaRegistryOptions into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.subject is not None: + body["subject"] = self.subject + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ConfluentSchemaRegistryOptions into a shallow dictionary of its immediate attributes.""" + body = {} + if self.subject is not None: + body["subject"] = self.subject + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ConfluentSchemaRegistryOptions: + """Deserializes the ConfluentSchemaRegistryOptions from a dictionary.""" + return cls(subject=d.get("subject", None)) + + @dataclass class ConnectionParameters: source_catalog: Optional[str] = None @@ -1733,8 +1810,10 @@ class IngestionSourceType(Enum): SALESFORCE = "SALESFORCE" SERVICENOW = "SERVICENOW" SHAREPOINT = "SHAREPOINT" + SMARTSHEET = "SMARTSHEET" SQLSERVER = "SQLSERVER" TERADATA = "TERADATA" + TIKTOK_ADS = "TIKTOK_ADS" WORKDAY_RAAS = "WORKDAY_RAAS" ZENDESK = "ZENDESK" @@ -2841,6 +2920,13 @@ def from_dict(cls, d: Dict[str, Any]) -> OutlookOptions: ) +class ParseMode(Enum): + """Determines how errors encountered while deserializing records are handled.""" + + FAILFAST = "FAILFAST" + PERMISSIVE = "PERMISSIVE" + + @dataclass class PathPattern: include: Optional[str] = None @@ -4100,6 +4186,69 @@ def from_dict(cls, d: Dict[str, Any]) -> PostgresSlotConfig: return cls(publication_name=d.get("publication_name", None), slot_name=d.get("slot_name", None)) +@dataclass +class ProtobufTransformerOptions: + desc_file_path: Optional[str] = None + """Required: path to the .desc file (dbfs:/... or /Volumes/...).""" + + message_name: Optional[str] = None + """Required: fully-qualified message type name.""" + + parse_mode: Optional[ParseMode] = None + """(Optional) Parse mode for Protobuf data. Valid values: FAILFAST, PERMISSIVE. Defaults to + FAILFAST.""" + + recursive_fields_max_depth: Optional[int] = None + """(Optional) Maximum expansion depth for recursive protobuf fields. Spark SQL does not natively + support recursive types, so recursive fields are expanded up to this depth and truncated beyond + it. Valid values: -1 (disallow recursive fields), 0 (drop), 1-10.""" + + schema_registry: Optional[SchemaRegistryConfig] = None + """(Optional) Schema registry to resolve the Protobuf schema at runtime instead of providing it via + desc_file_path.""" + + def as_dict(self) -> dict: + """Serializes the ProtobufTransformerOptions into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.desc_file_path is not None: + body["desc_file_path"] = self.desc_file_path + if self.message_name is not None: + body["message_name"] = self.message_name + if self.parse_mode is not None: + body["parse_mode"] = self.parse_mode.value + if self.recursive_fields_max_depth is not None: + body["recursive_fields_max_depth"] = self.recursive_fields_max_depth + if self.schema_registry: + body["schema_registry"] = self.schema_registry.as_dict() + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ProtobufTransformerOptions into a shallow dictionary of its immediate attributes.""" + body = {} + if self.desc_file_path is not None: + body["desc_file_path"] = self.desc_file_path + if self.message_name is not None: + body["message_name"] = self.message_name + if self.parse_mode is not None: + body["parse_mode"] = self.parse_mode + if self.recursive_fields_max_depth is not None: + body["recursive_fields_max_depth"] = self.recursive_fields_max_depth + if self.schema_registry: + body["schema_registry"] = self.schema_registry + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ProtobufTransformerOptions: + """Deserializes the ProtobufTransformerOptions from a dictionary.""" + return cls( + desc_file_path=d.get("desc_file_path", None), + message_name=d.get("message_name", None), + parse_mode=_enum(d, "parse_mode", ParseMode), + recursive_fields_max_depth=d.get("recursive_fields_max_depth", None), + schema_registry=_from_dict(d, "schema_registry", SchemaRegistryConfig), + ) + + class PublishingMode(Enum): """Enum representing the publishing mode of a pipeline.""" @@ -4465,6 +4614,11 @@ class RunAs: Only ``user_name`` or ``service_principal_name`` can be specified. If both are specified, an error is thrown.""" + group_name: Optional[str] = None + """Group name of an account group assigned to the workspace. When set, the pipeline runs as the + group and the group's permissions are used for data access. Setting this field requires being a + member of the group, or having the ``Assume`` permission on the group.""" + service_principal_name: Optional[str] = None """Application ID of an active service principal. Setting this field requires the ``servicePrincipal/user`` role.""" @@ -4475,6 +4629,8 @@ class RunAs: def as_dict(self) -> dict: """Serializes the RunAs into a dictionary suitable for use as a JSON request body.""" body = {} + if self.group_name is not None: + body["group_name"] = self.group_name if self.service_principal_name is not None: body["service_principal_name"] = self.service_principal_name if self.user_name is not None: @@ -4484,6 +4640,8 @@ def as_dict(self) -> dict: def as_shallow_dict(self) -> dict: """Serializes the RunAs into a shallow dictionary of its immediate attributes.""" body = {} + if self.group_name is not None: + body["group_name"] = self.group_name if self.service_principal_name is not None: body["service_principal_name"] = self.service_principal_name if self.user_name is not None: @@ -4493,7 +4651,57 @@ def as_shallow_dict(self) -> dict: @classmethod def from_dict(cls, d: Dict[str, Any]) -> RunAs: """Deserializes the RunAs from a dictionary.""" - return cls(service_principal_name=d.get("service_principal_name", None), user_name=d.get("user_name", None)) + return cls( + group_name=d.get("group_name", None), + service_principal_name=d.get("service_principal_name", None), + user_name=d.get("user_name", None), + ) + + +@dataclass +class SchemaRegistryConfig: + confluent_options: Optional[ConfluentSchemaRegistryOptions] = None + """Required: Confluent-compatible schema registry options.""" + + connection_name: Optional[str] = None + """(Optional) UC connection for registry authentication. Specify if different from the top-level + source connection.""" + + protobuf_message_name: Optional[str] = None + """(Optional, Protobuf only) Selects a specific message from a schema that defines multiple + Protobuf messages. Simple ("Location") or fully-qualified ("com.example.protos.Location"). + Defaults to the first message.""" + + def as_dict(self) -> dict: + """Serializes the SchemaRegistryConfig into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.confluent_options: + body["confluent_options"] = self.confluent_options.as_dict() + if self.connection_name is not None: + body["connection_name"] = self.connection_name + if self.protobuf_message_name is not None: + body["protobuf_message_name"] = self.protobuf_message_name + return body + + def as_shallow_dict(self) -> dict: + """Serializes the SchemaRegistryConfig into a shallow dictionary of its immediate attributes.""" + body = {} + if self.confluent_options: + body["confluent_options"] = self.confluent_options + if self.connection_name is not None: + body["connection_name"] = self.connection_name + if self.protobuf_message_name is not None: + body["protobuf_message_name"] = self.protobuf_message_name + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> SchemaRegistryConfig: + """Deserializes the SchemaRegistryConfig from a dictionary.""" + return cls( + confluent_options=_from_dict(d, "confluent_options", ConfluentSchemaRegistryOptions), + connection_name=d.get("connection_name", None), + protobuf_message_name=d.get("protobuf_message_name", None), + ) @dataclass @@ -5357,6 +5565,8 @@ class TikTokAdsOptionsTikTokReportType(Enum): class Transformer: """Specifies how to transform binary data into structured data.""" + avro_options: Optional[AvroTransformerOptions] = None + format: Optional[TransformerFormat] = None """Required: the wire format of the data.""" @@ -5370,9 +5580,13 @@ class Transformer: """Optional output column name. When set, the transformed result is written to this column instead of replacing the input column.""" + protobuf_options: Optional[ProtobufTransformerOptions] = None + def as_dict(self) -> dict: """Serializes the Transformer into a dictionary suitable for use as a JSON request body.""" body = {} + if self.avro_options: + body["avro_options"] = self.avro_options.as_dict() if self.format is not None: body["format"] = self.format.value if self.input_column is not None: @@ -5381,11 +5595,15 @@ def as_dict(self) -> dict: body["json_options"] = self.json_options.as_dict() if self.output_column is not None: body["output_column"] = self.output_column + if self.protobuf_options: + body["protobuf_options"] = self.protobuf_options.as_dict() return body def as_shallow_dict(self) -> dict: """Serializes the Transformer into a shallow dictionary of its immediate attributes.""" body = {} + if self.avro_options: + body["avro_options"] = self.avro_options if self.format is not None: body["format"] = self.format if self.input_column is not None: @@ -5394,21 +5612,27 @@ def as_shallow_dict(self) -> dict: body["json_options"] = self.json_options if self.output_column is not None: body["output_column"] = self.output_column + if self.protobuf_options: + body["protobuf_options"] = self.protobuf_options return body @classmethod def from_dict(cls, d: Dict[str, Any]) -> Transformer: """Deserializes the Transformer from a dictionary.""" return cls( + avro_options=_from_dict(d, "avro_options", AvroTransformerOptions), format=_enum(d, "format", TransformerFormat), input_column=d.get("input_column", None), json_options=_from_dict(d, "json_options", JsonTransformerOptions), output_column=d.get("output_column", None), + protobuf_options=_from_dict(d, "protobuf_options", ProtobufTransformerOptions), ) class TransformerFormat(Enum): + AVRO = "AVRO" JSON = "JSON" + PROTOBUF = "PROTOBUF" STRING = "STRING" diff --git a/databricks/sdk/version.py b/databricks/sdk/version.py index 99ac4c403..120dcc32b 100644 --- a/databricks/sdk/version.py +++ b/databricks/sdk/version.py @@ -1 +1 @@ -__version__ = "0.139.0" +__version__ = "0.140.0" diff --git a/docs/account/iamv2/iam_v2.rst b/docs/account/iamv2/iam_v2.rst index 8abfe2e2e..83a145c1c 100644 --- a/docs/account/iamv2/iam_v2.rst +++ b/docs/account/iamv2/iam_v2.rst @@ -4,7 +4,7 @@ .. py:class:: AccountIamV2API - These APIs are used to manage identities and the workspace access of these identities in . + 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 @@ -29,7 +29,7 @@ identity, use the ExternalGroup resource. :param group: :class:`Group` - Required. Group to be created in + Required. Group to be created in Databricks :returns: :class:`Group` @@ -45,7 +45,7 @@ external identity, use the ExternalServicePrincipal resource. :param service_principal: :class:`ServicePrincipal` - Required. Service principal to be created in + Required. Service principal to be created in Databricks :returns: :class:`ServicePrincipal` @@ -61,7 +61,7 @@ identity, use the ExternalUser resource. :param user: :class:`User` - Required. User to be created in + Required. User to be created in Databricks :returns: :class:`User` @@ -76,7 +76,7 @@ :param workspace_id: int Required. The workspace ID for which the workspace assignment is being created. :param workspace_assignment: :class:`WorkspaceAssignment` - Required. Workspace assignment to be created in . + Required. Workspace assignment to be created in Databricks. :returns: :class:`WorkspaceAssignment` @@ -88,7 +88,7 @@ :param workspace_id: int Required. The workspace ID for which the workspace assignment detail is being created. :param workspace_assignment_detail: :class:`WorkspaceAssignmentDetail` - Required. Workspace assignment detail to be created in . + Required. Workspace assignment detail to be created in Databricks. :returns: :class:`WorkspaceAssignmentDetail` @@ -464,7 +464,7 @@ :param group_id: str Required. Internal ID of the group in Databricks. :param group: :class:`Group` - Required. Group to be updated in + Required. Group to be updated in Databricks :param update_mask: str Optional. The list of fields to update. @@ -482,7 +482,7 @@ :param service_principal_id: str Required. Internal ID of the service principal in Databricks. :param service_principal: :class:`ServicePrincipal` - Required. Service Principal to be updated in + Required. Service Principal to be updated in Databricks :param update_mask: str Optional. The list of fields to update. @@ -501,7 +501,7 @@ :param user_id: str Required. Internal ID of the user in Databricks. :param user: :class:`User` - Required. User to be updated in + Required. User to be updated in Databricks :param update_mask: str Optional. The list of fields to update. @@ -519,7 +519,7 @@ :param principal_id: int Required. ID of the principal in Databricks. :param workspace_assignment: :class:`WorkspaceAssignment` - Required. Workspace assignment to be updated in . + Required. Workspace assignment to be updated in Databricks. :param update_mask: FieldMask Required. The list of fields to update. @@ -537,7 +537,7 @@ :param principal_id: int Required. ID of the principal in Databricks. :param workspace_assignment_detail: :class:`WorkspaceAssignmentDetail` - Required. Workspace assignment detail to be updated in . + Required. Workspace assignment detail to be updated in Databricks. :param update_mask: FieldMask Required. The list of fields to update. diff --git a/docs/dbdataclasses/catalog.rst b/docs/dbdataclasses/catalog.rst index dca2c2481..53486b808 100644 --- a/docs/dbdataclasses/catalog.rst +++ b/docs/dbdataclasses/catalog.rst @@ -374,6 +374,9 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: TERADATA :value: "TERADATA" + .. py:attribute:: TIKTOK_ADS + :value: "TIKTOK_ADS" + .. py:attribute:: UNKNOWN_CONNECTION_TYPE :value: "UNKNOWN_CONNECTION_TYPE" @@ -624,6 +627,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: DeleteMcpServiceUserMappedCredentialResponse + :members: + :undoc-members: + .. autoclass:: DeleteMonitorResponse :members: :undoc-members: @@ -1111,6 +1118,14 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: McpServiceUserMappedCredential + :members: + :undoc-members: + +.. autoclass:: McpServiceUserMappedCredentialLogin + :members: + :undoc-members: + .. autoclass:: MetastoreAssignment :members: :undoc-members: @@ -1159,6 +1174,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: ModelProviderServiceConfigCustomProviderApiKeyHeaderAuth + :members: + :undoc-members: + .. autoclass:: ModelProviderServiceConfigCustomProviderConfig :members: :undoc-members: @@ -2053,6 +2072,9 @@ These dataclasses are used in the SDK to represent API requests and responses fo The type of Unity Catalog securable. + .. py:attribute:: AGENT_SERVICE + :value: "AGENT_SERVICE" + .. py:attribute:: CATALOG :value: "CATALOG" @@ -2104,6 +2126,9 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: SHARE :value: "SHARE" + .. py:attribute:: SKILL + :value: "SKILL" + .. py:attribute:: STAGING_TABLE :value: "STAGING_TABLE" diff --git a/docs/dbdataclasses/pipelines.rst b/docs/dbdataclasses/pipelines.rst index 44e7d8f1c..645aaeaa8 100644 --- a/docs/dbdataclasses/pipelines.rst +++ b/docs/dbdataclasses/pipelines.rst @@ -20,6 +20,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: AvroTransformerOptions + :members: + :undoc-members: + .. py:class:: CloneMode Enum to specify which mode of clone to execute @@ -35,6 +39,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: ConfluentSchemaRegistryOptions + :members: + :undoc-members: + .. autoclass:: ConnectionParameters :members: :undoc-members: @@ -325,12 +333,18 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: SHAREPOINT :value: "SHAREPOINT" + .. py:attribute:: SMARTSHEET + :value: "SMARTSHEET" + .. py:attribute:: SQLSERVER :value: "SQLSERVER" .. py:attribute:: TERADATA :value: "TERADATA" + .. py:attribute:: TIKTOK_ADS + :value: "TIKTOK_ADS" + .. py:attribute:: WORKDAY_RAAS :value: "WORKDAY_RAAS" @@ -486,6 +500,16 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. py:class:: ParseMode + + Determines how errors encountered while deserializing records are handled. + + .. py:attribute:: FAILFAST + :value: "FAILFAST" + + .. py:attribute:: PERMISSIVE + :value: "PERMISSIVE" + .. autoclass:: PathPattern :members: :undoc-members: @@ -621,6 +645,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: ProtobufTransformerOptions + :members: + :undoc-members: + .. py:class:: PublishingMode Enum representing the publishing mode of a pipeline. @@ -667,6 +695,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: SchemaRegistryConfig + :members: + :undoc-members: + .. autoclass:: SchemaSpec :members: :undoc-members: @@ -819,9 +851,15 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:class:: TransformerFormat + .. py:attribute:: AVRO + :value: "AVRO" + .. py:attribute:: JSON :value: "JSON" + .. py:attribute:: PROTOBUF + :value: "PROTOBUF" + .. py:attribute:: STRING :value: "STRING" diff --git a/docs/workspace/aifunctions/ai_functions.rst b/docs/workspace/aifunctions/ai_functions.rst index d2384cd07..423f82ab4 100644 --- a/docs/workspace/aifunctions/ai_functions.rst +++ b/docs/workspace/aifunctions/ai_functions.rst @@ -48,8 +48,7 @@ .. py:method:: ai_parse_document(content: str [, options: Optional[AiParseDocumentOptions]]) -> AiParseDocumentResponse - Parse structured content from unstructured documents. For REST API requests, the default rate limit is - 120 pages per minute per workspace. Contact your Databricks account team to request a higher limit. + Parse structured content from unstructured documents. :param content: str The document to parse, given as a Unity Catalog volume path to the source file (the REST API accepts diff --git a/docs/workspace/catalog/ai_gateway.rst b/docs/workspace/catalog/ai_gateway.rst index 23e44bdf9..79a3863d6 100644 --- a/docs/workspace/catalog/ai_gateway.rst +++ b/docs/workspace/catalog/ai_gateway.rst @@ -31,6 +31,21 @@ :returns: :class:`McpService` + .. py:method:: create_mcp_service_user_mapped_credential(name: str, login: McpServiceUserMappedCredentialLogin) -> McpServiceUserMappedCredential + + Logs the caller in to an MCP service: creates their per-user OAuth credential, or re-authenticates it + if one already exists. The request body carries the OAuth exchange fields. + + You must be the owner of the MCP service or have ``EXECUTE`` on it, plus ``USE_CATALOG`` on the parent + catalog and ``USE_SCHEMA`` on the parent schema. + + :param name: str + Resource name of the MCP service. Format: ``mcp-services/{catalog}.{schema}.{mcp_service}``. + :param login: :class:`McpServiceUserMappedCredentialLogin` + + :returns: :class:`McpServiceUserMappedCredential` + + .. py:method:: create_model_provider_service(model_provider_service: ModelProviderService, parent: str, model_provider_service_id: str) -> ModelProviderService Creates a model provider service in a Unity Catalog schema. A model provider service stores @@ -99,6 +114,19 @@ + .. py:method:: delete_mcp_service_user_mapped_credential(name: str) -> DeleteMcpServiceUserMappedCredentialResponse + + Revokes (deletes) the caller's per-user OAuth credential for an MCP service (logout). + + You must be the owner of the MCP service or have ``EXECUTE`` on it, plus ``USE_CATALOG`` on the parent + catalog and ``USE_SCHEMA`` on the parent schema. + + :param name: str + Resource name of the MCP service. Format: ``mcp-services/{catalog}.{schema}.{mcp_service}``. + + :returns: :class:`DeleteMcpServiceUserMappedCredentialResponse` + + .. py:method:: delete_model_provider_service(name: str [, etag: Optional[str]]) Deletes the model provider service identified by its resource name. Optionally supply an ``etag`` to @@ -152,6 +180,22 @@ :returns: :class:`McpService` + .. py:method:: get_mcp_service_user_mapped_credential(name: str) -> McpServiceUserMappedCredential + + Returns the caller's per-user OAuth login state for an MCP service. Read ``provisioning_info.state``: + ``ACTIVE`` means the caller is logged in and the credential is usable; any other state (for example a + failed or still-provisioning login) means the login has not completed and the caller should log in + again. If the caller has no credential yet, the RPC returns ``NOT_FOUND``. + + You must be the owner of the MCP service or have ``EXECUTE`` on it, plus ``USE_CATALOG`` on the parent + catalog and ``USE_SCHEMA`` on the parent schema. + + :param name: str + Resource name of the MCP service. Format: ``mcp-services/{catalog}.{schema}.{mcp_service}``. + + :returns: :class:`McpServiceUserMappedCredential` + + .. py:method:: get_model_provider_service(name: str) -> ModelProviderService Returns the model provider service identified by its resource name. diff --git a/docs/workspace/dashboards/lakeview.rst b/docs/workspace/dashboards/lakeview.rst index d3399b210..e6bc33eb1 100644 --- a/docs/workspace/dashboards/lakeview.rst +++ b/docs/workspace/dashboards/lakeview.rst @@ -11,7 +11,8 @@ Create a draft dashboard. - Requires the Databricks SQL access entitlement. + Requires the `Databricks SQL access `__ + entitlement. Grant Databricks SQL access in addition to Workspace access. :param dashboard: :class:`Dashboard` :param dataset_catalog: str (optional) @@ -30,6 +31,9 @@ Create dashboard schedule. + Requires the `Databricks SQL access `__ + entitlement. Grant Databricks SQL access in addition to Workspace access. + :param dashboard_id: str UUID identifying the dashboard to which the schedule belongs. :param schedule: :class:`Schedule` @@ -42,6 +46,13 @@ Create schedule subscription. + The caller must be a workspace user with one of the following `entitlements + `__: Workspace access, Databricks SQL access, + or Consumer access. + + Account-level users who are not members of the workspace cannot call this endpoint, even if the + dashboard has been shared with them. + :param dashboard_id: str UUID identifying the dashboard to which the subscription belongs. :param schedule_id: str @@ -56,6 +67,9 @@ Delete dashboard schedule. + Requires the `Databricks SQL access `__ + entitlement. Grant Databricks SQL access in addition to Workspace access. + :param dashboard_id: str UUID identifying the dashboard to which the schedule belongs. :param schedule_id: str @@ -71,6 +85,13 @@ Delete schedule subscription. + The caller must be a workspace user with one of the following `entitlements + `__: Workspace access, Databricks SQL access, + or Consumer access. + + Account-level users who are not members of the workspace cannot call this endpoint, even if the + dashboard has been shared with them. + :param dashboard_id: str UUID identifying the dashboard which the subscription belongs. :param schedule_id: str @@ -88,7 +109,8 @@ Get a draft dashboard. - Requires the Databricks SQL access entitlement. + Requires the `Databricks SQL access `__ + entitlement. Grant Databricks SQL access in addition to Workspace access. :param dashboard_id: str UUID identifying the dashboard. @@ -116,6 +138,13 @@ Get dashboard schedule. + The caller must be a workspace user with one of the following `entitlements + `__: Workspace access, Databricks SQL access, + or Consumer access. + + Account-level users who are not members of the workspace cannot call this endpoint, even if the + dashboard has been shared with them. + :param dashboard_id: str UUID identifying the dashboard to which the schedule belongs. :param schedule_id: str @@ -128,6 +157,13 @@ Get schedule subscription. + The caller must be a workspace user with one of the following `entitlements + `__: Workspace access, Databricks SQL access, + or Consumer access. + + Account-level users who are not members of the workspace cannot call this endpoint, even if the + dashboard has been shared with them. + :param dashboard_id: str UUID identifying the dashboard which the subscription belongs. :param schedule_id: str @@ -142,7 +178,8 @@ List dashboards. - Requires the Databricks SQL access entitlement. + Requires the `Databricks SQL access `__ + entitlement. Grant Databricks SQL access in addition to Workspace access. :param page_size: int (optional) The number of dashboards to return per page. @@ -162,6 +199,13 @@ List dashboard schedules. + The caller must be a workspace user with one of the following `entitlements + `__: Workspace access, Databricks SQL access, + or Consumer access. + + Account-level users who are not members of the workspace cannot call this endpoint, even if the + dashboard has been shared with them. + :param dashboard_id: str UUID identifying the dashboard to which the schedules belongs. :param page_size: int (optional) @@ -177,6 +221,13 @@ List schedule subscriptions. + The caller must be a workspace user with one of the following `entitlements + `__: Workspace access, Databricks SQL access, + or Consumer access. + + Account-level users who are not members of the workspace cannot call this endpoint, even if the + dashboard has been shared with them. + :param dashboard_id: str UUID identifying the dashboard which the subscriptions belongs. :param schedule_id: str @@ -212,7 +263,8 @@ Publish the current draft dashboard. - Requires the Databricks SQL access entitlement. + Requires the `Databricks SQL access `__ + entitlement. Grant Databricks SQL access in addition to Workspace access. :param dashboard_id: str UUID identifying the dashboard to be published. @@ -229,7 +281,8 @@ Revert a dashboard's definition in draft mode to the last published version. - Requires the Databricks SQL access entitlement. + Requires the `Databricks SQL access `__ + entitlement. Grant Databricks SQL access in addition to Workspace access. :param dashboard_id: str UUID identifying the dashboard. @@ -244,7 +297,8 @@ Trash a dashboard. - Requires the Databricks SQL access entitlement. + Requires the `Databricks SQL access `__ + entitlement. Grant Databricks SQL access in addition to Workspace access. :param dashboard_id: str UUID identifying the dashboard. @@ -256,7 +310,8 @@ Unpublish the dashboard. - Requires the Databricks SQL access entitlement. + Requires the `Databricks SQL access `__ + entitlement. Grant Databricks SQL access in addition to Workspace access. :param dashboard_id: str UUID identifying the published dashboard. @@ -268,7 +323,8 @@ Update a draft dashboard. - Requires the Databricks SQL access entitlement. + Requires the `Databricks SQL access `__ + entitlement. Grant Databricks SQL access in addition to Workspace access. :param dashboard_id: str UUID identifying the dashboard. @@ -289,6 +345,9 @@ Update dashboard schedule. + Requires the `Databricks SQL access `__ + entitlement. Grant Databricks SQL access in addition to Workspace access. + :param dashboard_id: str UUID identifying the dashboard to which the schedule belongs. :param schedule_id: str diff --git a/docs/workspace/files/files.rst b/docs/workspace/files/files.rst old mode 100755 new mode 100644 index f3e2f3aa6..8b6f2e2a2 --- a/docs/workspace/files/files.rst +++ b/docs/workspace/files/files.rst @@ -11,7 +11,7 @@ The API supports `Unity Catalog volumes `__, where files and directories to operate on are specified using their volume URI path, which follows the format - /Volumes/<catalog_name>/<schema_name>/<volume_name>/<path_to_file>. + /Volumes////. The Files API has two distinct endpoints, one for working with files (``/fs/files``) and another one for working with directories (``/fs/directories``). Both endpoints use the standard HTTP methods GET, HEAD, diff --git a/docs/workspace/iamv2/workspace_iam_v2.rst b/docs/workspace/iamv2/workspace_iam_v2.rst index 363ed5b51..ff36b5bbd 100644 --- a/docs/workspace/iamv2/workspace_iam_v2.rst +++ b/docs/workspace/iamv2/workspace_iam_v2.rst @@ -4,7 +4,7 @@ .. py:class:: WorkspaceIamV2API - These APIs are used to manage identities and the workspace access of these identities in . + 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 @@ -29,7 +29,7 @@ identity, use the ExternalGroup resource. :param group: :class:`Group` - Required. Group to be created in + Required. Group to be created in Databricks :returns: :class:`Group` @@ -46,7 +46,7 @@ external identity, use the ExternalServicePrincipal resource. :param service_principal: :class:`ServicePrincipal` - Required. Service principal to be created in + Required. Service principal to be created in Databricks :returns: :class:`ServicePrincipal` @@ -62,7 +62,7 @@ identity, use the ExternalUser resource. :param user: :class:`User` - Required. User to be created in + Required. User to be created in Databricks :returns: :class:`User` @@ -72,7 +72,7 @@ Creates a workspace assignment detail for a principal in the calling workspace. :param workspace_assignment_detail: :class:`WorkspaceAssignmentDetail` - Required. Workspace assignment detail to be created in . + Required. Workspace assignment detail to be created in Databricks. :returns: :class:`WorkspaceAssignmentDetail` @@ -85,7 +85,7 @@ confirm which entitlements were granted. :param workspace_assignment: :class:`WorkspaceAssignment` - Required. Workspace assignment to be created in . + Required. Workspace assignment to be created in Databricks. :returns: :class:`WorkspaceAssignment` @@ -459,7 +459,7 @@ :param group_id: str Required. Internal ID of the group in Databricks. :param group: :class:`Group` - Required. Group to be updated in + Required. Group to be updated in Databricks :param update_mask: str Optional. The list of fields to update. @@ -477,7 +477,7 @@ :param service_principal_id: str Required. Internal ID of the service principal in Databricks. :param service_principal: :class:`ServicePrincipal` - Required. Service principal to be updated in + Required. Service principal to be updated in Databricks :param update_mask: str Optional. The list of fields to update. @@ -496,7 +496,7 @@ :param user_id: str Required. Internal ID of the user in Databricks. :param user: :class:`User` - Required. User to be updated in + Required. User to be updated in Databricks :param update_mask: str Optional. The list of fields to update. @@ -512,7 +512,7 @@ :param principal_id: int Required. ID of the principal in Databricks. :param workspace_assignment_detail: :class:`WorkspaceAssignmentDetail` - Required. Workspace assignment detail to be updated in . + Required. Workspace assignment detail to be updated in Databricks. :param update_mask: FieldMask Required. The list of fields to update. @@ -528,7 +528,7 @@ :param principal_id: int Required. ID of the principal in Databricks. :param workspace_assignment: :class:`WorkspaceAssignment` - Required. Workspace assignment to be updated in . + Required. Workspace assignment to be updated in Databricks. :param update_mask: FieldMask Required. The list of fields to update. @@ -542,7 +542,7 @@ :param principal_id: int Required. ID of the principal in Databricks. :param workspace_identity_detail: :class:`WorkspaceIdentityDetail` - Required. Workspace identity detail to be updated in . + Required. Workspace identity detail to be updated in Databricks. :param update_mask: FieldMask Required. The list of fields to update. diff --git a/docs/workspace/jobs/jobs.rst b/docs/workspace/jobs/jobs.rst index c21b8be02..c4a51fbbe 100644 --- a/docs/workspace/jobs/jobs.rst +++ b/docs/workspace/jobs/jobs.rst @@ -237,9 +237,11 @@ cost-efficiency for the run. The performance target does not apply to tasks that run on Serverless GPU compute. - - ``STANDARD``: Enables cost-efficient execution of serverless workloads. - ``PERFORMANCE_OPTIMIZED``: Prioritizes fast startup and execution times through rapid scaling and optimized cluster performance. + - ``STANDARD``: Enables cost-efficient execution of serverless workloads. + - ``COST_OPTIMIZED``: Enables lower job costs by optimizing compute for your selected target + duration time. Must provide a duration target. :param queue: :class:`QueueSettings` (optional) The queue settings of the job. :param run_as: :class:`JobRunAs` (optional) @@ -737,9 +739,11 @@ performance or cost-efficiency for the run. This field overrides the performance target defined on the job level. - - ``STANDARD``: Enables cost-efficient execution of serverless workloads. - ``PERFORMANCE_OPTIMIZED``: Prioritizes fast startup and execution times through rapid scaling and optimized cluster performance. + - ``STANDARD``: Enables cost-efficient execution of serverless workloads. + - ``COST_OPTIMIZED``: Enables lower job costs by optimizing compute for your selected target + duration time. Must provide a duration target. :param pipeline_params: :class:`PipelineParams` (optional) Controls whether the pipeline should perform a full refresh :param python_named_params: Dict[str,str] (optional) @@ -957,9 +961,11 @@ performance or cost-efficiency for the run. This field overrides the performance target defined on the job level. - - ``STANDARD``: Enables cost-efficient execution of serverless workloads. - ``PERFORMANCE_OPTIMIZED``: Prioritizes fast startup and execution times through rapid scaling and optimized cluster performance. + - ``STANDARD``: Enables cost-efficient execution of serverless workloads. + - ``COST_OPTIMIZED``: Enables lower job costs by optimizing compute for your selected target + duration time. Must provide a duration target. :param pipeline_params: :class:`PipelineParams` (optional) Controls whether the pipeline should perform a full refresh :param python_named_params: Dict[str,str] (optional) @@ -1105,9 +1111,11 @@ performance or cost-efficiency for the run. The performance target does not apply to tasks that run on Serverless GPU compute. - - ``STANDARD``: Enables cost-efficient execution of serverless workloads. - ``PERFORMANCE_OPTIMIZED``: Prioritizes fast startup and execution times through rapid scaling and optimized cluster performance. + - ``STANDARD``: Enables cost-efficient execution of serverless workloads. + - ``COST_OPTIMIZED``: Enables lower job costs by optimizing compute for your selected target + duration time. Must provide a duration target. :param queue: :class:`QueueSettings` (optional) The queue settings of the one-time run. :param run_as: :class:`JobRunAs` (optional) diff --git a/docs/workspace/ml/feature_store.rst b/docs/workspace/ml/feature_store.rst old mode 100755 new mode 100644 index 40ba8dc3b..54aef3b0e --- a/docs/workspace/ml/feature_store.rst +++ b/docs/workspace/ml/feature_store.rst @@ -79,6 +79,10 @@ Update an Online Feature Store. + This update is not guaranteed to be atomic: when a request changes multiple fields, some may be + applied while others fail. On a failed response, treat the update as partially applied and retry until + it succeeds. + :param name: str The name of the online store. This is the unique identifier for the online store. :param online_store: :class:`OnlineStore` diff --git a/pyproject.toml b/pyproject.toml index 1d168e512..93c2024a5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,6 +53,7 @@ dev = [ "openai", "langchain-openai", "httpx", + "grpcio<2", # the sandbox extra's dep; also needed to run the gRPC transport unit tests "build", # some integration tests depend on the databricks-sdk-py wheel "check-manifest", ] @@ -65,6 +66,17 @@ openai = [ "langchain-openai", "httpx", ] +sandbox = [ + # gRPC transport for the sandbox command-execution mixin. Optional so + # users who don't use the sandbox APIs don't pull the large native grpcio + # wheel; imported lazily in databricks/sdk/mixins/_grpc_transport.py. + # + # `<2` caps the major version to guard against a breaking grpcio 2.x. No + # lower bound is needed: the transport uses only long-stable grpc APIs, and + # pip installs the newest grpcio with wheels for the user's Python (the SDK + # requires >=3.10), so a too-old release is never selected. + "grpcio<2", +] [tool.setuptools.dynamic] version = { attr = "databricks.sdk.version.__version__" } diff --git a/tests/test_compute_mixins.py b/tests/test_compute_mixins.py index ec895b022..d0e4858b2 100644 --- a/tests/test_compute_mixins.py +++ b/tests/test_compute_mixins.py @@ -1,6 +1,8 @@ import pytest +from databricks.sdk import WorkspaceClient from databricks.sdk.mixins.compute import SemVer +from databricks.sdk.service import compute @pytest.mark.parametrize( @@ -40,3 +42,171 @@ def test_sorting_semver(): SemVer(1, 0, 0), SemVer(12, 0, 0), ] + + +def test_select_spark_version_filters_and_selects_latest_lts_from_http(config, requests_mock): + requests_mock.get( + "http://localhost/api/2.1/clusters/spark-versions", + json={ + "versions": [ + { + "key": "14.3.x-scala2.12", + "name": "14.3 LTS (includes Apache Spark 3.5.0, Scala 2.12)", + }, + { + "key": "15.4.x-scala2.12", + "name": "15.4 LTS (includes Apache Spark 3.5.0, Scala 2.12)", + }, + { + "key": "16.0.x-scala2.12", + "name": "16.0 Beta (includes Apache Spark 4.0.0, Scala 2.12)", + }, + { + "key": "15.4.x-photon-scala2.12", + "name": "15.4 LTS Photon (includes Apache Spark 3.5.0, Scala 2.12)", + }, + ] + }, + ) + workspace = WorkspaceClient(config=config) + + selected = workspace.clusters.select_spark_version(long_term_support=True) + + assert selected == "15.4.x-scala2.12" + assert requests_mock.last_request.method == "GET" + + +def test_select_node_type_filters_diskless_and_unavailable_nodes_from_http(config, requests_mock): + def node_type( + node_type_id, + *, + memory_mb, + num_cores, + local_disks, + local_disk_size_gb, + node_info=None, + ): + node = { + "category": "General Purpose", + "description": node_type_id, + "instance_type_id": node_type_id, + "is_deprecated": False, + "memory_mb": memory_mb, + "node_instance_type": { + "instance_type_id": node_type_id, + "local_disk_size_gb": local_disk_size_gb, + "local_disks": local_disks, + "local_nvme_disk_size_gb": 0, + "local_nvme_disks": 0, + }, + "node_type_id": node_type_id, + "num_cores": num_cores, + "num_gpus": 0, + } + if node_info is not None: + node["node_info"] = node_info + return node + + requests_mock.get( + "http://localhost/api/2.1/clusters/list-node-types", + json={ + "node_types": [ + node_type( + "unavailable-local", + memory_mb=4096, + num_cores=2, + local_disks=1, + local_disk_size_gb=100, + node_info={"status": ["NotAvailableInRegion"]}, + ), + node_type( + "diskless-small", + memory_mb=4096, + num_cores=2, + local_disks=0, + local_disk_size_gb=0, + ), + node_type( + "local-medium", + memory_mb=8192, + num_cores=4, + local_disks=1, + local_disk_size_gb=100, + ), + node_type( + "local-large", + memory_mb=16384, + num_cores=8, + local_disks=1, + local_disk_size_gb=200, + ), + ] + }, + ) + workspace = WorkspaceClient(config=config) + + selected = workspace.clusters.select_node_type(local_disk=True) + + assert selected == "local-medium" + assert requests_mock.last_request.method == "GET" + + +def test_ensure_cluster_is_running_waits_for_termination_then_starts(config, monkeypatch, requests_mock): + cluster_url = "http://localhost/api/2.1/clusters/get?cluster_id=cluster-1" + requests_mock.register_uri( + "GET", + cluster_url, + [ + {"json": {"cluster_id": "cluster-1", "cluster_name": "fixture", "state": "TERMINATING"}}, + {"json": {"cluster_id": "cluster-1", "cluster_name": "fixture", "state": "TERMINATED"}}, + {"json": {"cluster_id": "cluster-1", "cluster_name": "fixture", "state": "PENDING"}}, + {"json": {"cluster_id": "cluster-1", "cluster_name": "fixture", "state": "RUNNING"}}, + ], + ) + requests_mock.post("http://localhost/api/2.1/clusters/start", json={}) + monkeypatch.setattr(compute.time, "sleep", lambda _: None) + workspace = WorkspaceClient(config=config) + + workspace.clusters.ensure_cluster_is_running("cluster-1") + + assert [request.method for request in requests_mock.request_history] == ["GET", "GET", "POST", "GET", "GET"] + assert requests_mock.request_history[2].json() == {"cluster_id": "cluster-1"} + + +def test_cluster_events_reposts_next_page_request_at_http_boundary(config, requests_mock): + next_page = { + "cluster_id": "cluster-1", + "event_types": ["STARTING", "TERMINATING"], + "limit": 1, + "offset": 1, + } + requests_mock.register_uri( + "POST", + "http://localhost/api/2.1/clusters/events", + [ + { + "json": { + "events": [{"cluster_id": "cluster-1", "timestamp": 1000, "type": "STARTING"}], + "next_page": next_page, + } + }, + {"json": {"events": [{"cluster_id": "cluster-1", "timestamp": 2000, "type": "TERMINATING"}]}}, + ], + ) + workspace = WorkspaceClient(config=config) + + events = list( + workspace.clusters.events( + "cluster-1", + event_types=[compute.EventType.STARTING, compute.EventType.TERMINATING], + limit=1, + ) + ) + + assert [event.type for event in events] == [compute.EventType.STARTING, compute.EventType.TERMINATING] + assert requests_mock.request_history[0].json() == { + "cluster_id": "cluster-1", + "event_types": ["STARTING", "TERMINATING"], + "limit": 1, + } + assert requests_mock.request_history[1].json() == next_page diff --git a/tests/test_data_plane.py b/tests/test_data_plane.py index 1cade49d3..499bba649 100644 --- a/tests/test_data_plane.py +++ b/tests/test_data_plane.py @@ -2,7 +2,9 @@ from unittest.mock import patch from urllib import parse -from databricks.sdk import data_plane, oauth +from databricks.sdk import WorkspaceClient, data_plane, oauth +from databricks.sdk.config import Config +from databricks.sdk.credentials_provider import OAuthCredentialsProvider, OauthCredentialsStrategy from databricks.sdk.oauth import Token cp_token = Token(access_token="control plane token", token_type="type", expiry=datetime.now() + timedelta(hours=1)) @@ -67,3 +69,65 @@ def test_token_source_get_token_existing(config): retrieve_token.assert_not_called() assert result_token.access_token == another_token.access_token + + +def test_serving_endpoint_data_plane_discovers_exchanges_and_queries_over_http(requests_mock): + control_plane_token = Token(access_token="control-plane-token", token_type="Bearer") + credentials_strategy = OauthCredentialsStrategy( + "fixture-oauth", + lambda _: OAuthCredentialsProvider( + lambda: {"Authorization": "Bearer control-plane-token"}, + lambda: control_plane_token, + ), + ) + config = Config( + host="http://localhost", + workspace_id="12345", + credentials_strategy=credentials_strategy, + disable_async_token_refresh=True, + ) + endpoint_name = "route-optimized" + data_plane_url = "https://data-plane.example/invocations" + authorization_details = '{"type":"query","resource":"route-optimized"}' + requests_mock.get( + f"http://localhost/api/2.0/serving-endpoints/{endpoint_name}", + json={ + "name": endpoint_name, + "data_plane_info": { + "query_info": { + "endpoint_url": data_plane_url, + "authorization_details": authorization_details, + } + }, + }, + ) + requests_mock.post( + "http://localhost/oidc/v1/token", + json={"access_token": "data-plane-token", "token_type": "Bearer", "expires_in": 3600}, + ) + requests_mock.post( + data_plane_url, + json={"predictions": [0.75]}, + headers={"served-model-name": "model-v1"}, + ) + + response = WorkspaceClient(config=config).serving_endpoints_data_plane.query( + name=endpoint_name, + dataframe_records=[{"col": 1.0}], + ) + + discovery_request, exchange_request, query_request = requests_mock.request_history + assert response.predictions == [0.75] + assert response.served_model_name == "model-v1" + assert discovery_request.method == "GET" + assert discovery_request.headers["Authorization"] == "Bearer control-plane-token" + assert parse.parse_qs(exchange_request.text) == { + "assertion": ["control-plane-token"], + "authorization_details": [authorization_details], + "grant_type": ["urn:ietf:params:oauth:grant-type:jwt-bearer"], + } + assert exchange_request.headers["Content-Type"] == "application/x-www-form-urlencoded" + assert query_request.url == data_plane_url + assert query_request.headers["Authorization"] == "Bearer data-plane-token" + assert query_request.headers["X-Databricks-Workspace-Id"] == "12345" + assert query_request.json() == {"dataframe_records": [{"col": 1.0}]} diff --git a/tests/test_dbfs_mixins.py b/tests/test_dbfs_mixins.py index 5d785585d..4b0d3a240 100644 --- a/tests/test_dbfs_mixins.py +++ b/tests/test_dbfs_mixins.py @@ -1,64 +1,121 @@ +import base64 +import io + import pytest -from databricks.sdk.errors import NotFound +from databricks.sdk import WorkspaceClient from databricks.sdk.mixins.files import DbfsExt, _DbfsPath, _LocalPath, _VolumesPath -def test_moving_dbfs_file_to_local_dir(config, tmp_path, mocker): - from databricks.sdk import WorkspaceClient - from databricks.sdk.service.files import FileInfo, ReadResponse - - get_status = mocker.patch( - "databricks.sdk.service.files.DbfsAPI.get_status", - return_value=FileInfo(path="a", is_dir=False, file_size=4), +def test_dbfs_upload_chunks_payload_through_http_api(config, requests_mock): + chunk_size = 1024 * 1024 + payload = b"a" * chunk_size + b"tail" + requests_mock.post("http://localhost/api/2.0/dbfs/create", json={"handle": 123}) + requests_mock.post("http://localhost/api/2.0/dbfs/add-block", json={}) + requests_mock.post("http://localhost/api/2.0/dbfs/close", json={}) + + WorkspaceClient(config=config).dbfs.upload("/tmp/payload.bin", io.BytesIO(payload), overwrite=True) + + create_request, first_block, second_block, close_request = requests_mock.request_history + assert create_request.json() == {"overwrite": True, "path": "/tmp/payload.bin"} + assert first_block.url == "http://localhost/api/2.0/dbfs/add-block" + assert second_block.url == "http://localhost/api/2.0/dbfs/add-block" + assert base64.b64decode(first_block.json()["data"]) == payload[:chunk_size] + assert base64.b64decode(second_block.json()["data"]) == payload[chunk_size:] + assert first_block.json()["handle"] == 123 + assert second_block.json()["handle"] == 123 + assert close_request.json() == {"handle": 123} + + +def test_copy_routes_dbfs_source_to_volumes_http_api(config, requests_mock): + source_path = "/tmp/source.txt" + target_path = "/Volumes/main/default/volume/copied.txt" + target_url = f"http://localhost/api/2.0/fs/files{target_path}" + target_directory_url = f"http://localhost/api/2.0/fs/directories{target_path}" + payload = b"hello" + not_found = {"error_code": "RESOURCE_DOES_NOT_EXIST", "message": "not found"} + requests_mock.head(target_url, status_code=404, json=not_found) + requests_mock.head(target_directory_url, status_code=404, json=not_found) + requests_mock.get( + "http://localhost/api/2.0/dbfs/get-status", + json={"path": source_path, "is_dir": False, "file_size": len(payload)}, ) - def fake_read(path: str, *, length: int = None, offset: int = None): - assert path == "a" - assert length == 1048576 - if not offset: - return ReadResponse(bytes_read=4, data="aGVsbG8=") - return ReadResponse(bytes_read=0) - - mocker.patch("databricks.sdk.service.files.DbfsAPI.read", wraps=fake_read) - delete = mocker.patch("databricks.sdk.service.files.DbfsAPI.delete") - - w = WorkspaceClient(config=config) - w.dbfs.move_("a", f"file:{tmp_path}", recursive=True) - - get_status.assert_called_with("a") - delete.assert_called_with("a", recursive=True) - - with (tmp_path / "a").open("rb") as f: - assert f.read() == b"hello" - + def read_chunk(request, _context): + offset = int(request.qs["offset"][0]) + length = int(request.qs["length"][0]) + chunk = payload[offset : offset + length] + return {"bytes_read": len(chunk), "data": base64.b64encode(chunk).decode()} + + requests_mock.get("http://localhost/api/2.0/dbfs/read", json=read_chunk) + uploaded = {} + + def capture_upload(request, _context): + body = request.body + uploaded["body"] = body.read() if hasattr(body, "read") else body + return "" + + requests_mock.put(target_url, text=capture_upload) + + WorkspaceClient(config=config).dbfs.copy(source_path, target_path) + + status_requests = [request for request in requests_mock.request_history if "/dbfs/get-status" in request.url] + read_requests = [request for request in requests_mock.request_history if "/dbfs/read" in request.url] + upload_request = requests_mock.request_history[-1] + assert len(status_requests) == 2 + assert all(request.qs == {"path": [source_path]} for request in status_requests) + assert [request.qs["offset"] for request in read_requests] == [["0"], [str(len(payload))]] + assert all(request.qs["length"] == [str(1024 * 1024)] for request in read_requests) + assert upload_request.method == "PUT" + assert upload_request.url == f"{target_url}?overwrite=false" + assert upload_request.headers["Content-Type"] == "application/octet-stream" + assert uploaded["body"] == payload + + +def test_move_routes_volume_to_local_and_deletes_remote_source(config, requests_mock, tmp_path): + source_path = "/Volumes/main/default/volume/source.txt" + source_url = f"http://localhost/api/2.0/fs/files{source_path}" + source_directory_url = f"http://localhost/api/2.0/fs/directories{source_path}" + payload = b"moved from a volume" + requests_mock.head( + source_directory_url, + status_code=404, + json={"error_code": "RESOURCE_DOES_NOT_EXIST", "message": "not a directory"}, + ) + requests_mock.get( + source_url, + content=payload, + headers={ + "Content-Length": str(len(payload)), + "Content-Type": "application/octet-stream", + "Last-Modified": "Wed, 10 Sep 2025 12:00:00 GMT", + }, + ) + requests_mock.delete(source_url, status_code=204) -def test_moving_local_dir_to_dbfs(config, tmp_path, mocker): - from databricks.sdk import WorkspaceClient - from databricks.sdk.service.files import CreateResponse + WorkspaceClient(config=config).dbfs.move_(source_path, f"file:{tmp_path}", recursive=True) - with (tmp_path / "a").open("wb") as f: - f.write(b"hello") + assert (tmp_path / "source.txt").read_bytes() == payload + assert [request.method for request in requests_mock.request_history] == [ + "HEAD", + "HEAD", + "GET", + "HEAD", + "DELETE", + ] + assert requests_mock.request_history[-1].url == source_url - mocker.patch( - "databricks.sdk.service.files.DbfsAPI.create", - return_value=CreateResponse(123), - ) - get_status = mocker.patch( - "databricks.sdk.service.files.DbfsAPI.get_status", - side_effect=NotFound(), - ) - add_block = mocker.patch("databricks.sdk.service.files.DbfsAPI.add_block") - close = mocker.patch("databricks.sdk.service.files.DbfsAPI.close") +def test_move_within_dbfs_uses_dbfs_move_http_api(config, requests_mock): + requests_mock.post("http://localhost/api/2.0/dbfs/move", json={}) - w = WorkspaceClient(config=config) - w.dbfs.move_(f"file:{tmp_path}", "a", recursive=True) + WorkspaceClient(config=config).dbfs.move_("/tmp/source", "/tmp/destination") - get_status.assert_called_with("a") - close.assert_called_with(123) - add_block.assert_called_with(123, "aGVsbG8=") - assert not (tmp_path / "a").exists() + request = requests_mock.last_request + assert request.json() == { + "destination_path": "/tmp/destination", + "source_path": "/tmp/source", + } @pytest.mark.parametrize( @@ -85,33 +142,36 @@ def test_fs_path_invalid(config): def test_dbfs_local_path_mkdir(config, tmp_path): - from databricks.sdk import WorkspaceClient - - w = WorkspaceClient(config=config) - w.dbfs._path(f"file:{tmp_path}/test_dir").mkdir() - assert w.dbfs.exists(f"file:{tmp_path}/test_dir") - + client = WorkspaceClient(config=config) + client.dbfs._path(f"file:{tmp_path}/test_dir").mkdir() + assert client.dbfs.exists(f"file:{tmp_path}/test_dir") -def test_dbfs_exists(config, mocker): - from databricks.sdk import WorkspaceClient - get_status = mocker.patch( - "databricks.sdk.service.files.DbfsAPI.get_status", - side_effect=NotFound(), +def test_dbfs_exists_maps_not_found_http_response_to_false(config, requests_mock): + requests_mock.get( + "http://localhost/api/2.0/dbfs/get-status", + status_code=404, + json={"error_code": "RESOURCE_DOES_NOT_EXIST", "message": "not found"}, ) - client = WorkspaceClient(config=config) - client.dbfs.exists("/abc/def/ghi") - - get_status.assert_called_with("/abc/def/ghi") + exists = WorkspaceClient(config=config).dbfs.exists("/abc/def/ghi") + assert exists is False + assert requests_mock.last_request.qs == {"path": ["/abc/def/ghi"]} -def test_volume_exists(config, mocker): - from databricks.sdk import WorkspaceClient - get_metadata = mocker.patch("databricks.sdk.service.files.FilesAPI.get_metadata") +def test_volume_exists_reads_file_metadata_over_http(config, requests_mock): + path = "/Volumes/abc/def/ghi" + requests_mock.head( + f"http://localhost/api/2.0/fs/files{path}", + headers={ + "Content-Length": "4", + "Content-Type": "application/octet-stream", + "Last-Modified": "Wed, 10 Sep 2025 12:00:00 GMT", + }, + ) - client = WorkspaceClient(config=config) - client.dbfs.exists("/Volumes/abc/def/ghi") + exists = WorkspaceClient(config=config).dbfs.exists(path) - get_metadata.assert_called_with("/Volumes/abc/def/ghi") + assert exists is True + assert requests_mock.last_request.method == "HEAD" diff --git a/tests/test_files.py b/tests/test_files.py index 25c7c9689..63ab60715 100644 --- a/tests/test_files.py +++ b/tests/test_files.py @@ -41,6 +41,53 @@ logger = logging.getLogger(__name__) +def test_files_http_wire_format_preserves_bytes_metadata_and_stream_lifecycle(config: Config, requests_mock): + config = config.copy() + config.disable_experimental_files_api_client = True + file_path = "/Volumes/main/default/volume/report ?#%.bin" + file_url = "http://localhost/api/2.0/fs/files/Volumes/main/default/volume/report%20%3F%23%25.bin" + payload = b"raw file contents\x00\xff" + uploaded = {} + + def capture_upload(request, _context): + body = request.body + uploaded["body"] = body.read() if hasattr(body, "read") else body + return "" + + requests_mock.put(file_url, text=capture_upload) + response_headers = { + "Content-Length": str(len(payload)), + "Content-Type": "application/octet-stream", + "Last-Modified": "Wed, 10 Sep 2025 12:00:00 GMT", + } + requests_mock.head(file_url, headers=response_headers) + requests_mock.get(file_url, content=payload, headers=response_headers) + files = WorkspaceClient(config=config).files + + files.upload(file_path, io.BytesIO(payload), overwrite=False) + metadata = files.get_metadata(file_path) + download = files.download(file_path) + with download.contents as stream: + assert stream.read() == payload + with pytest.raises(ValueError, match="I/O operation on closed file"): + download.contents.read() + + upload_request, metadata_request, download_request = requests_mock.request_history + assert uploaded["body"] == payload + assert upload_request.url == f"{file_url}?overwrite=false" + assert upload_request.headers["Content-Type"] == "application/octet-stream" + assert metadata.content_length == len(payload) + assert metadata.content_type == "application/octet-stream" + assert metadata.last_modified == "Wed, 10 Sep 2025 12:00:00 GMT" + assert metadata_request.method == "HEAD" + assert metadata_request.url == file_url + assert download.content_length == len(payload) + assert download.content_type == "application/octet-stream" + assert download.last_modified == "Wed, 10 Sep 2025 12:00:00 GMT" + assert download_request.method == "GET" + assert download_request.url == file_url + + class CustomResponse: """Custom response allows to override the "default" response generated by the server with the "custom" response to simulate failure error code, unexpected response body or diff --git a/tests/test_grpc_transport.py b/tests/test_grpc_transport.py new file mode 100644 index 000000000..89ff2e45b --- /dev/null +++ b/tests/test_grpc_transport.py @@ -0,0 +1,477 @@ +import pytest + +from databricks.sdk.mixins._grpc_transport import ( + _ReplayFirstStream, + _auth_metadata, + _channel_options, + _grpc_target, + _headers_to_metadata, + _is_unavailable, + call_with_retry, + open_channel, + open_stream_with_retry, +) + + +class _FakeConfig: + """Minimal stand-in for databricks.sdk.config.Config: the transport reads + ``host``, ``user_agent``, ``_custom_headers``, and calls ``authenticate()``.""" + + def __init__( + self, + *, + headers=None, + host="https://foo.cloud.databricks.com", + error=None, + custom_headers=None, + user_agent="databricks-sdk-py/1.0.0 python/3.12 os/linux auth/pat", + ): + self._headers = headers or {} + self._error = error + self.host = host + self._custom_headers = custom_headers or {} + self.user_agent = user_agent + + def authenticate(self): + if self._error is not None: + raise self._error + return self._headers + + +@pytest.mark.parametrize( + "host,expected", + [ + ("https://foo.cloud.databricks.com", "foo.cloud.databricks.com:443"), + ("http://foo.cloud.databricks.com", "foo.cloud.databricks.com:443"), + ("foo.cloud.databricks.com", "foo.cloud.databricks.com:443"), + ("https://foo.cloud.databricks.com/", "foo.cloud.databricks.com:443"), + ("https://foo.cloud.databricks.com/?x=1", "foo.cloud.databricks.com:443"), + ("https://foo.cloud.databricks.com/some/path", "foo.cloud.databricks.com:443"), + ("https://foo.cloud.databricks.com#frag", "foo.cloud.databricks.com:443"), + # An explicit port is preserved rather than having :443 appended. + ("https://localhost:8443", "localhost:8443"), + ("localhost:8443", "localhost:8443"), + ], +) +def test_grpc_target(host, expected): + assert _grpc_target(host) == expected + + +def test_grpc_target_custom_port(): + assert _grpc_target("foo.cloud.databricks.com", port=9090) == "foo.cloud.databricks.com:9090" + + +def test_headers_to_metadata_lowercases_keys(): + headers = {"Authorization": "Bearer tok", "X-Databricks-Org-Id": "123"} + assert _headers_to_metadata(headers) == [ + ("authorization", "Bearer tok"), + ("x-databricks-org-id", "123"), + ] + + +def test_headers_to_metadata_empty(): + assert _headers_to_metadata({}) == [] + + +class _Retryable(Exception): + pass + + +class _Fatal(Exception): + pass + + +def _is_retryable(err: BaseException) -> bool: + return isinstance(err, _Retryable) + + +def test_retry_succeeds_after_transient_failures(): + calls = {"n": 0} + slept: list[float] = [] + + def make_call(): + calls["n"] += 1 + if calls["n"] < 3: + raise _Retryable() + return "ok" + + result = call_with_retry( + make_call, + is_retryable=_is_retryable, + max_attempts=5, + retry_interval_seconds=1.5, + sleep=slept.append, + ) + + assert result == "ok" + assert calls["n"] == 3 + # Slept once before each of the two retries, never after success. + assert slept == [1.5, 1.5] + + +def test_retry_reraises_non_retryable(): + calls = {"n": 0} + + def make_call(): + calls["n"] += 1 + raise _Fatal() + + with pytest.raises(_Fatal): + call_with_retry(make_call, is_retryable=_is_retryable, max_attempts=5, sleep=lambda _: None) + # A non-retryable error propagates on the first attempt, no retries. + assert calls["n"] == 1 + + +def test_retry_exhausts_attempts(): + calls = {"n": 0} + + def make_call(): + calls["n"] += 1 + raise _Retryable() + + with pytest.raises(TimeoutError): + call_with_retry( + make_call, + is_retryable=_is_retryable, + max_attempts=4, + retry_interval_seconds=2.0, + sleep=lambda _: None, + ) + assert calls["n"] == 4 + + +def test_stream_retry_opens_and_yields_full_stream(): + slept: list[float] = [] + opens = {"n": 0} + + def make_stream(): + opens["n"] += 1 + if opens["n"] < 3: + # A not-ready backend fails on the first pull, so raise from the + # iterator rather than from make_stream itself. + def failing(): + raise _Retryable() + yield # pragma: no cover - unreachable, makes this a generator + + return failing() + return iter(["start", "data", "end"]) + + stream = open_stream_with_retry( + make_stream, + is_retryable=_is_retryable, + max_attempts=5, + retry_interval_seconds=1.0, + sleep=slept.append, + ) + + assert list(stream) == ["start", "data", "end"] + assert opens["n"] == 3 + assert slept == [1.0, 1.0] + + +def test_stream_retry_empty_stream(): + stream = open_stream_with_retry(lambda: iter(()), is_retryable=_is_retryable, sleep=lambda _: None) + assert list(stream) == [] + + +def test_stream_retry_does_not_retry_mid_stream_failure(): + opens = {"n": 0} + + def make_stream(): + opens["n"] += 1 + + def gen(): + yield "start" + raise _Retryable() + + return gen() + + stream = open_stream_with_retry(make_stream, is_retryable=_is_retryable, sleep=lambda _: None) + + assert next(stream) == "start" + with pytest.raises(_Retryable): + next(stream) + assert opens["n"] == 1 + + +def test_stream_mid_stream_retry_reopens_and_replays_from_start(): + opens = {"n": 0} + slept: list[float] = [] + + def make_stream(): + opens["n"] += 1 + if opens["n"] == 1: + + def failing(): + yield "start" + raise _Retryable() + + return failing() + return iter(["start", "data", "end"]) + + stream = open_stream_with_retry( + make_stream, + is_retryable=_is_retryable, + retry_mid_stream_failures=True, + max_attempts=5, + retry_interval_seconds=1.0, + sleep=slept.append, + ) + + # The mid-stream failure re-issues the whole call, which restarts from the + # first item, so "start" is delivered twice. + assert list(stream) == ["start", "start", "data", "end"] + assert opens["n"] == 2 + assert slept == [1.0] + + +def test_stream_mid_stream_retry_raises_non_retryable(): + opens = {"n": 0} + + def make_stream(): + opens["n"] += 1 + + def gen(): + yield "start" + raise _Fatal() + + return gen() + + stream = open_stream_with_retry( + make_stream, + is_retryable=_is_retryable, + retry_mid_stream_failures=True, + sleep=lambda _: None, + ) + + assert next(stream) == "start" + with pytest.raises(_Fatal): + next(stream) + assert opens["n"] == 1 + + +def test_stream_mid_stream_retry_bounded_by_max_attempts(): + opens = {"n": 0} + slept: list[float] = [] + + def make_stream(): + opens["n"] += 1 + + def gen(): + yield "start" + raise _Retryable() + + return gen() + + stream = open_stream_with_retry( + make_stream, + is_retryable=_is_retryable, + retry_mid_stream_failures=True, + max_attempts=3, + retry_interval_seconds=2.0, + sleep=slept.append, + ) + + collected = [] + with pytest.raises(_Retryable): + collected.extend(stream) + # One initial open plus max_attempts re-opens, each yielding "start". + assert collected == ["start", "start", "start", "start"] + assert opens["n"] == 4 + assert slept == [2.0, 2.0, 2.0] + + +def test_auth_metadata_forwards_headers_as_lowercased_metadata(): + captured = {} + + def callback(metadata, error): + captured["metadata"] = metadata + captured["error"] = error + + cfg = _FakeConfig(headers={"Authorization": "Bearer tok", "X-Databricks-Org-Id": "42"}) + _auth_metadata(cfg, callback) + + assert captured["error"] is None + assert captured["metadata"] == (("authorization", "Bearer tok"), ("x-databricks-org-id", "42")) + + +def test_auth_metadata_forwards_credential_error(): + captured = {} + + def callback(metadata, error): + captured["metadata"] = metadata + captured["error"] = error + + boom = RuntimeError("token refresh failed") + _auth_metadata(_FakeConfig(error=boom), callback) + + # On failure gRPC's contract is empty metadata + the error, not a raise. + assert captured["metadata"] == () + assert captured["error"] is boom + + +def test_auth_metadata_merges_custom_headers_with_auth_winning(): + captured = {} + + def callback(metadata, error): + captured["metadata"] = dict(metadata) + captured["error"] = error + + cfg = _FakeConfig( + headers={"Authorization": "Bearer tok"}, + custom_headers={"X-Custom": "v", "Authorization": "should-be-overridden"}, + ) + _auth_metadata(cfg, callback) + + assert captured["error"] is None + # Custom headers ride along; auth headers win on a key collision. + assert captured["metadata"] == {"x-custom": "v", "authorization": "Bearer tok"} + + +def test_auth_metadata_merge_is_case_insensitive(): + captured = {} + + def callback(metadata, error): + captured["metadata"] = list(metadata) + + # Custom header and auth header name the same logical header with different + # casing: the result must carry a single lowercased key, with auth winning. + cfg = _FakeConfig( + headers={"Authorization": "Bearer tok"}, + custom_headers={"authorization": "should-be-overridden"}, + ) + _auth_metadata(cfg, callback) + + assert captured["metadata"] == [("authorization", "Bearer tok")] + + +def test_is_unavailable_matches_only_unavailable(): + grpc = pytest.importorskip("grpc") + + class _RpcError(grpc.RpcError): + def __init__(self, code): + self._code = code + + def code(self): + return self._code + + assert _is_unavailable(_RpcError(grpc.StatusCode.UNAVAILABLE)) + assert not _is_unavailable(_RpcError(grpc.StatusCode.DEADLINE_EXCEEDED)) + assert not _is_unavailable(RuntimeError("boom")) + + +def test_call_with_retry_default_retries_unavailable_not_deadline(): + grpc = pytest.importorskip("grpc") + + class _RpcError(grpc.RpcError): + def __init__(self, code): + self._code = code + + def code(self): + return self._code + + unavailable_calls = {"n": 0} + + def raise_unavailable(): + unavailable_calls["n"] += 1 + raise _RpcError(grpc.StatusCode.UNAVAILABLE) + + with pytest.raises(TimeoutError): + call_with_retry(raise_unavailable, max_attempts=3, sleep=lambda _: None) + assert unavailable_calls["n"] == 3 + + deadline_calls = {"n": 0} + + def raise_deadline(): + deadline_calls["n"] += 1 + raise _RpcError(grpc.StatusCode.DEADLINE_EXCEEDED) + + with pytest.raises(grpc.RpcError): + call_with_retry(raise_deadline, max_attempts=3, sleep=lambda _: None) + assert deadline_calls["n"] == 1 + + +def test_call_with_retry_caller_can_opt_into_deadline(): + grpc = pytest.importorskip("grpc") + + class _RpcError(grpc.RpcError): + def __init__(self, code): + self._code = code + + def code(self): + return self._code + + calls = {"n": 0} + + def raise_deadline(): + calls["n"] += 1 + raise _RpcError(grpc.StatusCode.DEADLINE_EXCEEDED) + + def retry_deadline(err): + return isinstance(err, grpc.RpcError) and err.code() == grpc.StatusCode.DEADLINE_EXCEEDED + + with pytest.raises(TimeoutError): + call_with_retry(raise_deadline, is_retryable=retry_deadline, max_attempts=2, sleep=lambda _: None) + assert calls["n"] == 2 + + +def test_replay_first_stream_yields_first_item_then_delegates_attributes(): + class _FakeStream: + def __init__(self): + self._rest = iter(["data", "end"]) + self.cancelled = False + + def __next__(self): + return next(self._rest) + + def cancel(self): + self.cancelled = True + + underlying = _FakeStream() + wrapped = _ReplayFirstStream("start", underlying) + + assert list(wrapped) == ["start", "data", "end"] + # cancel() (and any other attribute) reaches the underlying gRPC stream. + wrapped.cancel() + assert underlying.cancelled + + +def test_stream_retry_default_predicate_excludes_deadline_exceeded(): + grpc = pytest.importorskip("grpc") + + class _RpcError(grpc.RpcError): + def __init__(self, code): + self._code = code + + def code(self): + return self._code + + def make_stream_raising(code): + def failing(): + raise _RpcError(code) + yield # pragma: no cover - unreachable, makes this a generator + + return failing() + + with pytest.raises(grpc.RpcError): + open_stream_with_retry( + lambda: make_stream_raising(grpc.StatusCode.DEADLINE_EXCEEDED), + max_attempts=3, + sleep=lambda _: None, + ) + + +def test_channel_options_carry_the_sdk_user_agent(): + cfg = _FakeConfig(user_agent="databricks-sdk-py/9.9.9 python/3.12 os/linux auth/pat") + # The gRPC channel advertises the same User-Agent the HTTP client sends. + assert _channel_options(cfg) == [("grpc.primary_user_agent", cfg.user_agent)] + + +def test_open_channel_returns_secure_channel(): + grpc = pytest.importorskip("grpc") + + channel = open_channel(_FakeConfig(headers={"Authorization": "Bearer tok"})) + try: + assert isinstance(channel, grpc.Channel) + finally: + channel.close() diff --git a/tests/test_jobs_mixin.py b/tests/test_jobs_mixin.py index c14c7dd83..96339be3d 100644 --- a/tests/test_jobs_mixin.py +++ b/tests/test_jobs_mixin.py @@ -3,6 +3,7 @@ from typing import Optional, Pattern from databricks.sdk import WorkspaceClient +from databricks.sdk.service import compute, jobs def make_getrun_path_pattern(run_id: int, page_token: Optional[str] = None) -> Pattern[str]: @@ -53,8 +54,6 @@ def test_get_run_with_no_pagination(config, requests_mock): def test_get_run_pagination_with_tasks(config, requests_mock): - from databricks.sdk.service import compute, jobs - cluster_spec = compute.ClusterSpec( spark_version="11.3.x-scala2.12", custom_tags={"ResourceClass": "SingleNode"}, @@ -199,8 +198,6 @@ def test_get_job_forwards_include_trigger_state(config, requests_mock): def test_get_job_pagination_with_tasks(config, requests_mock): - from databricks.sdk.service import compute, jobs - cluster_spec = compute.ClusterSpec( spark_version="11.3.x-scala2.12", custom_tags={"ResourceClass": "SingleNode"}, @@ -374,8 +371,6 @@ def test_list_jobs_without_task_expansion(config, requests_mock): def test_list_jobs_with_many_tasks(config, requests_mock): - from databricks.sdk.service import compute, jobs - cluster_spec = compute.ClusterSpec( spark_version="11.3.x-scala2.12", custom_tags={"ResourceClass": "SingleNode"}, @@ -816,3 +811,55 @@ def test_list_runs(config, requests_mock): history = requests_mock.request_history assert all("300" not in request.qs.get("run_id", [""]) for request in history) assert all("200" not in request.qs.get("run_id", [""]) for request in history) + + +def test_submit_waiter_polls_real_jobs_service_until_terminal_state(config, monkeypatch, requests_mock): + requests_mock.post( + "http://localhost/api/2.2/jobs/runs/submit", + json={"run_id": 123}, + ) + requests_mock.register_uri( + "GET", + make_getrun_path_pattern(123), + [ + { + "json": { + "run_id": 123, + "run_name": "fixture-run", + "state": {"life_cycle_state": "PENDING", "state_message": "Waiting for compute"}, + } + }, + { + "json": { + "run_id": 123, + "run_name": "fixture-run", + "state": {"life_cycle_state": "RUNNING", "state_message": "Running"}, + } + }, + { + "json": { + "run_id": 123, + "run_name": "fixture-run", + "state": { + "life_cycle_state": "TERMINATED", + "result_state": "SUCCESS", + "state_message": "Completed", + }, + } + }, + ], + ) + monkeypatch.setattr(jobs.time, "sleep", lambda _: None) + workspace = WorkspaceClient(config=config) + observed_states = [] + + run = workspace.jobs.submit(run_name="fixture-run").result( + callback=lambda polled: observed_states.append(polled.state.life_cycle_state) + ) + + assert run.run_id == 123 + assert run.state.life_cycle_state == jobs.RunLifeCycleState.TERMINATED + assert run.state.result_state == jobs.RunResultState.SUCCESS + assert observed_states == [jobs.RunLifeCycleState.PENDING, jobs.RunLifeCycleState.RUNNING] + assert requests_mock.request_history[0].json() == {"run_name": "fixture-run"} + assert [request.method for request in requests_mock.request_history] == ["POST", "GET", "GET", "GET"] diff --git a/tests/test_workspace.py b/tests/test_workspace.py new file mode 100644 index 000000000..c7c71e479 --- /dev/null +++ b/tests/test_workspace.py @@ -0,0 +1,92 @@ +import io +from urllib.parse import parse_qs, urlsplit + +import pytest + +from databricks.sdk import WorkspaceClient +from databricks.sdk.service.workspace import ExportFormat, ObjectType + + +def test_recursive_list_traverses_directories_through_workspace_http_api(config, requests_mock): + requests_mock.get( + "http://localhost/api/2.0/workspace/list", + response_list=[ + { + "json": { + "objects": [ + {"object_type": "DIRECTORY", "path": "/root/nested"}, + {"object_type": "NOTEBOOK", "path": "/root/top.py"}, + ] + } + }, + { + "json": { + "objects": [ + {"object_type": "DIRECTORY", "path": "/root/nested/empty"}, + {"object_type": "FILE", "path": "/root/nested/data.bin"}, + ] + } + }, + {"json": {"objects": []}}, + ], + ) + workspace = WorkspaceClient(config=config).workspace + + objects = list(workspace.list("/root", notebooks_modified_after=123, recursive=True)) + + assert [(obj.path, obj.object_type) for obj in objects] == [ + ("/root/top.py", ObjectType.NOTEBOOK), + ("/root/nested/data.bin", ObjectType.FILE), + ] + assert [parse_qs(urlsplit(request.url).query) for request in requests_mock.request_history] == [ + {"notebooks_modified_after": ["123"], "path": ["/root"]}, + {"notebooks_modified_after": ["123"], "path": ["/root/nested"]}, + {"notebooks_modified_after": ["123"], "path": ["/root/nested/empty"]}, + ] + + +def test_upload_and_download_use_workspace_http_wire_format(config, requests_mock): + uploaded = {} + + def capture_upload(request, _context): + uploaded["body"] = request.body + return "" + + requests_mock.post("http://localhost/api/2.0/workspace/import", text=capture_upload) + requests_mock.get( + "http://localhost/api/2.0/workspace/export", + content=b"print('downloaded')\n", + ) + workspace = WorkspaceClient(config=config).workspace + + workspace.upload( + "/Users/sdk/example.py", + io.BytesIO(b"print('uploaded')\n"), + overwrite=True, + ) + download = workspace.download("/Users/sdk/example.py", format=ExportFormat.AUTO) + with download as stream: + assert stream.read() == b"print('downloaded')\n" + with pytest.raises(ValueError, match="I/O operation on closed file"): + download.read() + + upload_request, download_request = requests_mock.request_history + assert upload_request.method == "POST" + assert upload_request.url == "http://localhost/api/2.0/workspace/import" + assert upload_request.headers["Content-Type"].startswith("multipart/form-data; boundary=") + assert b'name="path"' in uploaded["body"] + assert b"/Users/sdk/example.py" in uploaded["body"] + assert b'name="language"' in uploaded["body"] + assert b"PYTHON" in uploaded["body"] + assert b'name="overwrite"' in uploaded["body"] + assert b"true" in uploaded["body"] + assert b'name="content"' in uploaded["body"] + assert b"print('uploaded')\n" in uploaded["body"] + + assert download_request.method == "GET" + assert download_request.url.startswith("http://localhost/api/2.0/workspace/export?") + assert parse_qs(urlsplit(download_request.url).query) == { + "direct_download": ["true"], + "format": ["AUTO"], + "path": ["/Users/sdk/example.py"], + }