diff --git a/.codegen/_last_sha b/.codegen/_last_sha index fe07bcb12..7be7d3677 100644 --- a/.codegen/_last_sha +++ b/.codegen/_last_sha @@ -1 +1 @@ -b936af3d07da4d0b9f2fcea823a97dd17e65fbcf \ No newline at end of file +52efa6fe60dbb9e77eaee68e53c4199e2235b95e \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md index 2a709af8c..2fd912a95 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,20 @@ # Version changelog +## Release v0.161.0 (2026-11-04) + +### New Features and Improvements +* `WorkspaceClient.files` uploads and downloads move to the next transfer endpoint whenever an attempt fails before any data was committed: presigned URLs, then the Files API on the workspace. Once data was committed, the transfer stays on its endpoint and raises. Before commit, a failure is retried only on the last endpoint. +* With the experimental `experimental_files_ext_enable_endpoints_api` config option, `WorkspaceClient.files` follows the transfer endpoints the workspace lists (`GET /api/2.0/fs/endpoints`), including the storage proxy and Files API 3.0 endpoints. +* `experimental_files_ext_enable_storage_proxy` is deprecated and has no effect: the storage proxy is used when the workspace lists it, with `experimental_files_ext_enable_endpoints_api`. +* `WorkspaceClient.files.download_to` writes a new file through a temp file next to it, so a failed download leaves no partial file. An existing destination (a file, a symlink, a pipe or a device) is written in place as before, and a parallel download to anything other than a regular file is sequential. + +### API Changes +* Add `databricks.sdk.service.lakewatch` package. +* Add [w.lakewatch](https://databricks-sdk-py.readthedocs.io/en/latest/workspace/lakewatch/lakewatch.html) workspace-level service. +* Add `cancel_batch()`, `create_batch()`, `get_batch()` and `list_batches()` methods for [w.ai_functions](https://databricks-sdk-py.readthedocs.io/en/latest/workspace/aifunctions/ai_functions.html) workspace-level service. +* Add `create_feature_spec()`, `delete_feature_spec()` and `get_feature_spec()` methods for [w.feature_engineering](https://databricks-sdk-py.readthedocs.io/en/latest/workspace/ml/feature_engineering.html) workspace-level service. +* [Breaking] Remove `critical` enum value for `databricks.sdk.service.jobs.AiRuntimeTaskPriorityClass`. + ## Release v0.160.0 (2026-11-03) ### API Changes diff --git a/databricks/sdk/__init__.py b/databricks/sdk/__init__.py index 6bf05feb6..ed37e83d6 100644 --- a/databricks/sdk/__init__.py +++ b/databricks/sdk/__init__.py @@ -33,6 +33,7 @@ from databricks.sdk.service import iamv2 as pkg_iamv2 from databricks.sdk.service import jobs as pkg_jobs from databricks.sdk.service import knowledgeassistants as pkg_knowledgeassistants +from databricks.sdk.service import lakewatch as pkg_lakewatch from databricks.sdk.service import marketplace as pkg_marketplace from databricks.sdk.service import ml as pkg_ml from databricks.sdk.service import networking as pkg_networking @@ -155,6 +156,7 @@ from databricks.sdk.service.knowledgeassistants import KnowledgeAssistantsAPI from databricks.sdk.service.dashboards import LakeviewAPI from databricks.sdk.service.dashboards import LakeviewEmbeddedAPI +from databricks.sdk.service.lakewatch import LakewatchAPI from databricks.sdk.service.compute import LibrariesAPI from databricks.sdk.service.settings import LlmProxyPartnerPoweredAccountAPI from databricks.sdk.service.settings import LlmProxyPartnerPoweredEnforceAPI @@ -436,6 +438,7 @@ def __init__( self._knowledge_assistants = pkg_knowledgeassistants.KnowledgeAssistantsAPI(self._api_client) self._lakeview = pkg_dashboards.LakeviewAPI(self._api_client) self._lakeview_embedded = pkg_dashboards.LakeviewEmbeddedAPI(self._api_client) + self._lakewatch = pkg_lakewatch.LakewatchAPI(self._api_client) self._libraries = pkg_compute.LibrariesAPI(self._api_client) self._materialized_features = pkg_ml.MaterializedFeaturesAPI(self._api_client) self._metastores = pkg_catalog.MetastoresAPI(self._api_client) @@ -858,6 +861,11 @@ def lakeview_embedded(self) -> pkg_dashboards.LakeviewEmbeddedAPI: """Token-based Lakeview APIs for embedding dashboards in external applications.""" return self._lakeview_embedded + @property + def lakewatch(self) -> pkg_lakewatch.LakewatchAPI: + """Manage Lakewatch: security-analytics data sources, detection rules, and per-catalog configuration.""" + return self._lakewatch + @property def libraries(self) -> pkg_compute.LibrariesAPI: """The Libraries API allows you to install and uninstall libraries and get the status of libraries on a cluster.""" diff --git a/databricks/sdk/config.py b/databricks/sdk/config.py index bf7f268bf..a92a8925f 100644 --- a/databricks/sdk/config.py +++ b/databricks/sdk/config.py @@ -260,12 +260,12 @@ class Config: # is hit first will stop the retry loop. experimental_files_ext_cloud_api_max_retries: int = 3 - # Whether to enable the storage proxy for file operations. - # When enabled, the SDK will probe the storage proxy and use it if available. + # Deprecated and unused: the storage proxy is used when the workspace lists it, see + # experimental_files_ext_enable_endpoints_api. experimental_files_ext_enable_storage_proxy: bool = False # Whether to ask the workspace which endpoints to use for file transfers (GET /api/2.0/fs/endpoints). - # When disabled, or when the server cannot answer, the SDK uses a built-in order: Presigned URL, then + # When disabled, or when the server cannot answer, the SDK uses a built-in order: presigned URLs, then # Files API 2.0 on the workspace host. experimental_files_ext_enable_endpoints_api: bool = False diff --git a/databricks/sdk/mixins/files.py b/databricks/sdk/mixins/files.py index 6f4bdb8a4..271b79261 100644 --- a/databricks/sdk/mixins/files.py +++ b/databricks/sdk/mixins/files.py @@ -778,6 +778,7 @@ class _ResolvedEndpoint(Enum): IO_OFFLOAD = "IO_OFFLOAD" FILES_API_V2 = "FILES_API_V2" + FILES_API_V3 = "FILES_API_V3" # Maps the server's (api, api_version) to a protocol compiled into this SDK. Pairs not listed here are @@ -786,6 +787,7 @@ class _ResolvedEndpoint(Enum): _RECOGNIZED_ENDPOINTS: dict[tuple[str, Optional[str]], _ResolvedEndpoint] = { ("IO_OFFLOAD", None): _ResolvedEndpoint.IO_OFFLOAD, ("FILES_API", "2.0"): _ResolvedEndpoint.FILES_API_V2, + ("FILES_API", "3.0"): _ResolvedEndpoint.FILES_API_V3, } # The Files API path each protocol initiates, completes, uploads and downloads through. IO_OFFLOAD initiates @@ -793,8 +795,12 @@ class _ResolvedEndpoint(Enum): _FILES_API_PATHS: dict[_ResolvedEndpoint, str] = { _ResolvedEndpoint.IO_OFFLOAD: "/api/2.0/fs/files", _ResolvedEndpoint.FILES_API_V2: "/api/2.0/fs/files", + _ResolvedEndpoint.FILES_API_V3: "/api/3.0/fs/files", } +# Files API 3.0 rejects multipart parts above 4000 MiB, below the largest part size the SDK may pick. +_FILES_API_V3_MAX_PART_SIZE = 4000 * 1024 * 1024 + _ENDPOINT_SCHEMES = ("http", "https") @@ -913,14 +919,43 @@ def build_download_url(self, path: str, expire_time: str) -> CreateDownloadUrlRe ) -class FilesExt(files.FilesAPI): - __doc__ = files.FilesAPI.__doc__ +class _FilesApiV3RequestBuilder: + """Builds Files API 3.0 multipart requests, sent directly to the endpoint host. + + 3.0 identifies a multipart upload by the `upload_id` returned from initiate-upload, which callers pass + as `session_token`. Parts go straight to the host, so there is no URL to presign or expire. + """ - # Storage-proxy hostname for data plane file operations. - _STORAGE_PROXY_HOSTNAME: str = "http://storage-proxy.databricks.com" + def __init__(self, hostname: str): + self._hostname = hostname - # Timeout in seconds for the storage-proxy health check probe. - _STORAGE_PROXY_PROBE_TIMEOUT: float = 3.0 + def _file_url(self, path: str) -> str: + return f"{self._hostname}/api/3.0/fs/files{_files_api_file_path(path)}" + + def build_upload_part_urls( + self, path: str, session_token: str, start_part_number: int, count: int, expire_time: str + ) -> list[_PresignedUrl]: + """Builds URLs for uploading multipart parts.""" + base = self._file_url(path) + results = [] + for part_number in range(start_part_number, start_part_number + count): + query = parse.urlencode({"upload_id": session_token, "part_number": part_number}) + results.append( + _PresignedUrl( + url=f"{base}?{query}", + headers={"Content-Type": "application/octet-stream"}, + part_number=part_number, + ) + ) + return results + + def build_download_url(self, path: str, expire_time: str) -> CreateDownloadUrlResponse: + """Builds a URL for downloading a file directly from the host.""" + return CreateDownloadUrlResponse(url=self._file_url(path), headers={}) + + +class FilesExt(files.FilesAPI): + __doc__ = files.FilesAPI.__doc__ _ENDPOINTS_API_PATH: str = "/api/2.0/fs/endpoints" @@ -934,8 +969,6 @@ def __init__(self, api_client, config: Config): self._multipart_upload_read_ahead_bytes = 1 self._cached_cloud_provider_session: Optional[requests.Session] = None self._cloud_provider_session_lock = Lock() - self._dp_hostname_available: Optional[bool] = None - self._storage_proxy_probe_lock = Lock() self._cached_authenticated_session: Optional[requests.Session] = None self._authenticated_session_lock = Lock() self._cached_endpoints: Optional[list[_Endpoint]] = None @@ -956,41 +989,18 @@ def _resolve_endpoints(self) -> list[_Endpoint]: return list(self._cached_endpoints) def _default_endpoints(self) -> list[_Endpoint]: - """The built-in order: presigned URLs, or the storage proxy when experimental_files_ext_enable_storage_proxy - is set and the probe finds it can serve this workspace; then the Files API on the workspace host, which takes - uploads only in a single shot. + """The built-in order: presigned URLs, then the Files API on the workspace host, which takes uploads only in a + single shot. """ host = self._config.host - if self._storage_proxy_available(): - first = _Endpoint( - _ResolvedEndpoint.FILES_API_V2, self._STORAGE_PROXY_HOSTNAME, multipart_upload_supported=True - ) - else: - first = _Endpoint(_ResolvedEndpoint.IO_OFFLOAD, host, multipart_upload_supported=True) - return [first, _Endpoint(_ResolvedEndpoint.FILES_API_V2, host, multipart_upload_supported=False)] + return [ + _Endpoint(_ResolvedEndpoint.IO_OFFLOAD, host, multipart_upload_supported=True), + _Endpoint(_ResolvedEndpoint.FILES_API_V2, host, multipart_upload_supported=False), + ] def _transfer_endpoints(self) -> list[_Endpoint]: """The endpoints a transfer walks, best first.""" - return self._default_endpoints() - - def _storage_proxy_available(self) -> bool: - """Whether the storage-proxy flag is set and the proxy can serve this workspace. - - The probe runs once per client and its result, including a failure, is cached. - """ - if not self._config.experimental_files_ext_enable_storage_proxy: - return False - available = self._dp_hostname_available - if available is not None: - return available - with self._storage_proxy_probe_lock: - if self._dp_hostname_available is None: - self._dp_hostname_available = self._probe_storage_proxy() - if self._dp_hostname_available: - _LOG.info("Storage proxy is available, will use it for file operations.") - else: - _LOG.info("Storage proxy is not available, will use presigned URLs.") - return self._dp_hostname_available + return self._resolve_endpoints() def _fetch_endpoints(self) -> list[_Endpoint]: """Fetches the endpoint list from the server, keeping only entries this SDK can speak. @@ -1072,15 +1082,41 @@ def _endpoint_kind(entry) -> Optional[_ResolvedEndpoint]: return None return _RECOGNIZED_ENDPOINTS.get((api, api_version)) - @staticmethod - def _parse_endpoint(entry: dict, kind: _ResolvedEndpoint) -> Optional[_Endpoint]: + def _parse_endpoint(self, entry: dict, kind: _ResolvedEndpoint) -> Optional[_Endpoint]: scheme = entry.get("scheme") hostname = entry.get("hostname") if scheme not in _ENDPOINT_SCHEMES or not isinstance(hostname, str) or not hostname: return None + # A Files API 3.0 host is reached over the internet with workspace credentials, so it must use TLS. + if kind == _ResolvedEndpoint.FILES_API_V3 and scheme != "https": + return None + base_url = f"{scheme}://{hostname}" + # Only a bare host and port that requests sends to as written. Transfers route by the host parsed here, so it + # must be the host the request reaches, and a parse error mid-transfer is not something every walk survives. + try: + parsed = parse.urlsplit(base_url) + origin = _origin(base_url) + prepared = requests.PreparedRequest() + prepared.prepare_url(base_url, None) + sent_as_written = _origin(prepared.url) == origin + except ValueError: + return None + if not sent_as_written or parsed.netloc != hostname or parsed.username is not None: + return None + host = origin[1] + # Plain HTTP carries workspace credentials in the clear. The workspace host may use it only if the SDK is + # configured to reach the workspace that way; another host only for Files API 2.0, i.e. the in-cluster + # storage proxy. + if scheme == "http": + workspace_scheme, workspace_host, _ = _origin(self._config.host) + if host == workspace_host: + if workspace_scheme != "http": + return None + elif kind != _ResolvedEndpoint.FILES_API_V2: + return None # Mirrors the proto getter: an unset optional bool reads as false. multipart_upload_supported = entry.get("multipart_upload_supported") is True - return _Endpoint(kind, f"{scheme}://{hostname}", multipart_upload_supported) + return _Endpoint(kind, base_url, multipart_upload_supported) def _uses_api_client(self, endpoint: _Endpoint) -> bool: """Whether transfers to the endpoint go through the workspace API client. @@ -1105,46 +1141,10 @@ def _request_builder_for(self, endpoint: _Endpoint): """ if endpoint.kind == _ResolvedEndpoint.IO_OFFLOAD: return _PresignedUrlRequestBuilder(self._api, endpoint.base_url) + if endpoint.kind == _ResolvedEndpoint.FILES_API_V3: + return _FilesApiV3RequestBuilder(endpoint.base_url) return _StorageProxyRequestBuilder(endpoint.base_url) - def _resolve_workspace_id(self) -> Optional[str]: - """Resolves the workspace ID by querying the SCIM Me endpoint.""" - try: - response = self._api.do("GET", "/api/2.0/preview/scim/v2/Me", response_headers=["X-Databricks-Org-Id"]) - workspace_id = response.get("X-Databricks-Org-Id") - if workspace_id: - return str(workspace_id) - except Exception: - _LOG.debug("Failed to resolve workspace ID for storage proxy probe.") - return None - - def _probe_storage_proxy(self) -> bool: - """Probes the storage proxy to check if it can serve this workspace. - - Uses a HEAD request to the probe endpoint with the workspace ID as a - query parameter. The proxy returns 200 if it can serve this workspace, - 403 for a wrong workspace, and connection errors if unreachable. - Returns False if the workspace ID cannot be resolved. - """ - workspace_id = self._resolve_workspace_id() - if not workspace_id: - _LOG.debug("Workspace ID not available, skipping storage proxy probe.") - return False - proxy_host = self._STORAGE_PROXY_HOSTNAME - probe_url = f"{proxy_host}/api/2.0/fs/files/DatabricksInternal/Probe/fullstack/wis?ew={workspace_id}" - try: - headers = self._config.authenticate() - session = self._cloud_provider_session() - response = session.request( - "HEAD", - probe_url, - headers=headers, - timeout=self._STORAGE_PROXY_PROBE_TIMEOUT, - ) - return response.status_code == 200 - except Exception: - return False - def _authenticated_session(self) -> requests.Session: """Returns an HTTP session with SDK auth for direct Files API data requests. @@ -1501,6 +1501,8 @@ def _clients_for(self, endpoints: list[_Endpoint]) -> list: for endpoint in endpoints: if endpoint.kind == _ResolvedEndpoint.IO_OFFLOAD: clients.append(_PresignedEndpointClient(self, endpoint)) + elif endpoint.kind == _ResolvedEndpoint.FILES_API_V3: + clients.append(_FilesApiV3EndpointClient(self, endpoint)) else: clients.append(_FilesApiV2EndpointClient(self, endpoint)) return clients @@ -2180,6 +2182,27 @@ def _abort_query(self, session_token: str) -> dict: return {"action": "abort-upload", "upload_type": "multipart", "session_token": session_token} +class _FilesApiV3EndpointClient(_FilesApiV2EndpointClient): + """FILES_API 3.0 (RIG): the 2.0 protocol with `upload_id` in place of `session_token`, and parts of at most + _FILES_API_V3_MAX_PART_SIZE.""" + + @property + def max_part_size_opt(self) -> Optional[int]: + return _FILES_API_V3_MAX_PART_SIZE + + def _session_from(self, response: dict, path: str, overwrite_opt: Optional[bool]): + upload_id = response.get("upload_id") + if not upload_id: + raise ValueError(f"Unexpected server response: {response}") + return _FilesApiMultipartSession(self, path, upload_id) + + def _complete_query(self, session_token: str) -> dict: + return {"action": "complete-upload", "upload_id": session_token} + + def _abort_query(self, session_token: str) -> dict: + return {"action": "abort-upload", "upload_id": session_token} + + class _FilesApiMultipartSession(_MultipartUploadSession): """A multipart upload whose parts go straight to the Files API host, so there is no URL to presign or expire.""" diff --git a/databricks/sdk/service/aifunctions.py b/databricks/sdk/service/aifunctions.py index 5b559c42f..91064dfc9 100644 --- a/databricks/sdk/service/aifunctions.py +++ b/databricks/sdk/service/aifunctions.py @@ -5,16 +5,20 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Dict, List, Any, Optional +from enum import Enum +from typing import Dict, List, Any, Iterator, Optional +from google.protobuf.timestamp_pb2 import Timestamp import logging from databricks.sdk.service._internal import ( + _enum, _from_dict, _int64, _repeated_dict, _repeated_int64, + _timestamp, ) @@ -773,6 +777,206 @@ def from_dict(cls, d: Dict[str, Any]) -> AiParseDocumentResponseMetadata: ) +@dataclass +class Batch: + """An asynchronous run of one AI function over a file of requests.""" + + endpoint: str + """The AI function to run for every request, for example ``ai_extract``.""" + + input_ref: str + """Path of the JSONL input file on a Unity Catalog volume. Each line is one request.""" + + create_time: Optional[Timestamp] = None + """Output only. Time the batch was created.""" + + name: Optional[str] = None + """ID of the batch, in the form ``batch_``. The service assigns it.""" + + output_ref: Optional[str] = None + """Path on a Unity Catalog volume to write the JSONL output to. When it is omitted, the output + stays on Databricks-managed storage and a finished batch returns a short-lived download URL in + this field instead.""" + + params: Optional[Dict[str, any]] = None + """Function options applied to every request in the batch, for example the extraction schema.""" + + percent_complete: Optional[float] = None + """Output only. Percentage of the requests that finished, from 0 to 100. Returned when one batch is + read, not when batches are listed.""" + + request_counts: Optional[BatchRequestCounts] = None + """Output only. Request counts. Returned when one batch is read, not when batches are listed.""" + + status: Optional[BatchStatus] = None + """Output only. Status of the batch.""" + + def as_dict(self) -> dict: + """Serializes the Batch into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.create_time is not None: + body["create_time"] = self.create_time.ToJsonString() + if self.endpoint is not None: + body["endpoint"] = self.endpoint + if self.input_ref is not None: + body["input_ref"] = self.input_ref + if self.name is not None: + body["name"] = self.name + if self.output_ref is not None: + body["output_ref"] = self.output_ref + if self.params: + body["params"] = self.params + if self.percent_complete is not None: + body["percent_complete"] = self.percent_complete + if self.request_counts: + body["request_counts"] = self.request_counts.as_dict() + if self.status is not None: + body["status"] = self.status.value + return body + + def as_shallow_dict(self) -> dict: + """Serializes the Batch into a shallow dictionary of its immediate attributes.""" + body = {} + if self.create_time is not None: + body["create_time"] = self.create_time + if self.endpoint is not None: + body["endpoint"] = self.endpoint + if self.input_ref is not None: + body["input_ref"] = self.input_ref + if self.name is not None: + body["name"] = self.name + if self.output_ref is not None: + body["output_ref"] = self.output_ref + if self.params: + body["params"] = self.params + if self.percent_complete is not None: + body["percent_complete"] = self.percent_complete + if self.request_counts: + body["request_counts"] = self.request_counts + if self.status is not None: + body["status"] = self.status + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> Batch: + """Deserializes the Batch from a dictionary.""" + return cls( + create_time=_timestamp(d, "create_time"), + endpoint=d.get("endpoint", None), + input_ref=d.get("input_ref", None), + name=d.get("name", None), + output_ref=d.get("output_ref", None), + params=d.get("params", None), + percent_complete=d.get("percent_complete", None), + request_counts=_from_dict(d, "request_counts", BatchRequestCounts), + status=_enum(d, "status", BatchStatus), + ) + + +@dataclass +class BatchRequestCounts: + """Counts of the requests in a batch.""" + + failed: Optional[int] = None + """Output only. Number of requests that failed and will not be retried.""" + + succeeded: Optional[int] = None + """Output only. Number of requests that finished successfully.""" + + total: Optional[int] = None + """Output only. Number of requests in the input file.""" + + def as_dict(self) -> dict: + """Serializes the BatchRequestCounts into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.failed is not None: + body["failed"] = self.failed + if self.succeeded is not None: + body["succeeded"] = self.succeeded + if self.total is not None: + body["total"] = self.total + return body + + def as_shallow_dict(self) -> dict: + """Serializes the BatchRequestCounts into a shallow dictionary of its immediate attributes.""" + body = {} + if self.failed is not None: + body["failed"] = self.failed + if self.succeeded is not None: + body["succeeded"] = self.succeeded + if self.total is not None: + body["total"] = self.total + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> BatchRequestCounts: + """Deserializes the BatchRequestCounts from a dictionary.""" + return cls(failed=_int64(d, "failed"), succeeded=_int64(d, "succeeded"), total=_int64(d, "total")) + + +class BatchStatus(Enum): + """The status of a batch.""" + + CANCELLED = "CANCELLED" + CANCELLING = "CANCELLING" + COMPLETED = "COMPLETED" + EXPIRED = "EXPIRED" + FAILED = "FAILED" + FINALIZING = "FINALIZING" + IN_PROGRESS = "IN_PROGRESS" + PENDING = "PENDING" + + +@dataclass +class CancelBatchResponse: + def as_dict(self) -> dict: + """Serializes the CancelBatchResponse into a dictionary suitable for use as a JSON request body.""" + body = {} + return body + + def as_shallow_dict(self) -> dict: + """Serializes the CancelBatchResponse into a shallow dictionary of its immediate attributes.""" + body = {} + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> CancelBatchResponse: + """Deserializes the CancelBatchResponse from a dictionary.""" + return cls() + + +@dataclass +class ListBatchesResponse: + batches: Optional[List[Batch]] = None + """The batches in this page.""" + + next_page_token: Optional[str] = None + """Token that reads the next page. It is empty on the last page.""" + + def as_dict(self) -> dict: + """Serializes the ListBatchesResponse into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.batches: + body["batches"] = [v.as_dict() for v in self.batches] + if self.next_page_token is not None: + body["next_page_token"] = self.next_page_token + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ListBatchesResponse into a shallow dictionary of its immediate attributes.""" + body = {} + if self.batches: + body["batches"] = self.batches + if self.next_page_token is not None: + body["next_page_token"] = self.next_page_token + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ListBatchesResponse: + """Deserializes the ListBatchesResponse from a dictionary.""" + return cls(batches=_repeated_dict(d, "batches", Batch), next_page_token=d.get("next_page_token", None)) + + class AiFunctionsAPI: """Transform and enrich data with AI on Databricks.""" @@ -1000,3 +1204,103 @@ def ai_parse_document( res = self._api.do("POST", "/api/2.0/ai-functions/ai-parse-document", body=body, headers=headers) return AiParseDocumentResponse.from_dict(res) + + def cancel_batch(self, name: str) -> CancelBatchResponse: + """Cancel a batch. Cancellation is asynchronous: requests in flight drain first, so read the batch to + observe the ``CANCELLING`` and ``CANCELLED`` statuses. + + :param name: str + ID of the batch, in the form ``batch_``. + + :returns: :class:`CancelBatchResponse` + """ + + body = {} + 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.0/ai-functions/batches/{name}/cancel", body=body, headers=headers) + return CancelBatchResponse.from_dict(res) + + def create_batch(self, batch: Batch) -> Batch: + """Run an AI function over every request in a file on a Unity Catalog volume. The batch runs + asynchronously; read it to follow its progress. + + :param batch: :class:`Batch` + The batch to create. + + :returns: :class:`Batch` + """ + + body = batch.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", "/api/2.0/ai-functions/batches", body=body, headers=headers) + return Batch.from_dict(res) + + def get_batch(self, name: str) -> Batch: + """Read one batch, with its status and progress. + + :param name: str + ID of the batch, in the form ``batch_``. + + :returns: :class:`Batch` + """ + + headers = { + "Accept": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do("GET", f"/api/2.0/ai-functions/batches/{name}", headers=headers) + return Batch.from_dict(res) + + def list_batches(self, *, page_size: Optional[int] = None, page_token: Optional[str] = None) -> Iterator[Batch]: + """List the batches the caller created in this workspace. + + :param page_size: int (optional) + Maximum number of batches to return in one page. + :param page_token: str (optional) + Page token from an earlier response. Omit it to read the first page. + + :returns: Iterator over :class:`Batch` + """ + + query = {} + if page_size is not None: + query["page_size"] = page_size + if page_token is not None: + query["page_token"] = page_token + headers = { + "Accept": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + while True: + json = self._api.do("GET", "/api/2.0/ai-functions/batches", query=query, headers=headers) + if "batches" in json: + for v in json["batches"]: + yield Batch.from_dict(v) + if "next_page_token" not in json or not json["next_page_token"]: + return + query["page_token"] = json["next_page_token"] diff --git a/databricks/sdk/service/compute.py b/databricks/sdk/service/compute.py index 117c06862..a909f71e4 100644 --- a/databricks/sdk/service/compute.py +++ b/databricks/sdk/service/compute.py @@ -4453,7 +4453,7 @@ class GetInstancePool: in this cluster. For example, the Spark nodes can be provisioned and optimized for memory or compute intensive workloads. A list of available node types can be retrieved by using the `clusters/listNodeTypes `__ - API call.""" + API call. The node type cannot be changed after the pool is created.""" preloaded_docker_images: Optional[List[DockerImage]] = None """Custom Docker Image BYOC""" @@ -5298,7 +5298,7 @@ class InstancePoolAndStats: in this cluster. For example, the Spark nodes can be provisioned and optimized for memory or compute intensive workloads. A list of available node types can be retrieved by using the `clusters/listNodeTypes `__ - API call.""" + API call. The node type cannot be changed after the pool is created.""" preloaded_docker_images: Optional[List[DockerImage]] = None """Custom Docker Image BYOC""" @@ -10922,7 +10922,7 @@ def create( this cluster. For example, the Spark nodes can be provisioned and optimized for memory or compute intensive workloads. A list of available node types can be retrieved by using the `clusters/listNodeTypes `__ API - call. + call. The node type cannot be changed after the pool is created. :param aws_attributes: :class:`InstancePoolAwsAttributes` (optional) Attributes related to instance pools running on Amazon Web Services. If not specified at pool creation, a set of default values will be used. @@ -11067,7 +11067,7 @@ def edit( this cluster. For example, the Spark nodes can be provisioned and optimized for memory or compute intensive workloads. A list of available node types can be retrieved by using the `clusters/listNodeTypes `__ API - call. + call. The node type cannot be changed after the pool is created. :param custom_tags: Dict[str,str] (optional) Additional tags for pool resources. Databricks will tag all pool resources (e.g., AWS instances and EBS volumes) with these tags in addition to ``default_tags``. Notes: diff --git a/databricks/sdk/service/dataquality.py b/databricks/sdk/service/dataquality.py index cee0ea541..71a9f4f14 100644 --- a/databricks/sdk/service/dataquality.py +++ b/databricks/sdk/service/dataquality.py @@ -659,10 +659,10 @@ class NotificationDestination: """Destination of the data quality monitoring notification.""" destination_ids: Optional[List[str]] = None - """The notification-destination IDs (Slack / Teams / webhook) to send the notification to. The - destinations must already exist in the workspace; a workspace admin creates them. A maximum of - 25 destination IDs is supported. Only supported on schema monitors, via - ``anomaly_detection_config``; rejected on table monitors.""" + """The notification-destination IDs (Slack / Teams / webhook / Asana / ServiceNow / AWS + EventBridge) to send the notification to. The destinations must already exist in the workspace; + a workspace admin creates them. A maximum of 25 destination IDs is supported. Only supported on + schema monitors, via ``anomaly_detection_config``; rejected on table monitors.""" destination_user_ids: Optional[List[int]] = None """The IDs of the workspace users to notify by email. A maximum of 25 user IDs is supported. Only diff --git a/databricks/sdk/service/jobs.py b/databricks/sdk/service/jobs.py index c0ee023cc..41345c9d2 100644 --- a/databricks/sdk/service/jobs.py +++ b/databricks/sdk/service/jobs.py @@ -633,7 +633,6 @@ class AiRuntimeTaskPriorityClass(Enum): ranks pending work.""" BEST_EFFORT = "BEST_EFFORT" - CRITICAL = "CRITICAL" NORMAL = "NORMAL" diff --git a/databricks/sdk/service/lakewatch.py b/databricks/sdk/service/lakewatch.py new file mode 100644 index 000000000..6b3ce9f3d --- /dev/null +++ b/databricks/sdk/service/lakewatch.py @@ -0,0 +1,3282 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. +# ruff: noqa: F811, F841 +# F401 is intentionally NOT covered: `make fmt` uses `ruff check --fix-only` +# to strip the fat-import header below; ignoring F401 would defeat that. + +from __future__ import annotations +from dataclasses import dataclass +from enum import Enum +from typing import Dict, List, Any, Iterator, Optional + +from google.protobuf.timestamp_pb2 import Timestamp + +import logging + +from databricks.sdk.service._internal import ( + _enum, + _from_dict, + _int64, + _repeated_dict, + _timestamp, +) +from databricks.sdk.common.types.fieldmask import FieldMask + + +_LOG = logging.getLogger("databricks.sdk") + + +# all definitions in this file are in alphabetical order + + +@dataclass +class AlertConfig: + """Per-catalog alert-related config. Singleton. Name format: ``catalogs/{catalog}/alertConfig``.""" + + default_security_alert_collation_window: Optional[str] = None + """Window for collating Signals into security alerts. Duration string — compact + day/hour/minute/second components, e.g. ``1d``, ``12h`` (see ``Schedule.at_least_every`` in + common.proto for the grammar). Reads return the value as supplied. When unset or empty, the + window is ``1h``; the window in effect is reported in + ``effective_default_security_alert_collation_window``.""" + + default_system_alert_collation_window: Optional[str] = None + """Window for collating SystemEvents into system alerts. Duration string — compact + day/hour/minute/second components, e.g. ``1d``, ``12h`` (see ``Schedule.at_least_every`` in + common.proto for the grammar). Reads return the value as supplied. When unset or empty, the + window is ``1d``; the window in effect is reported in + ``effective_default_system_alert_collation_window``.""" + + effective_default_security_alert_collation_window: Optional[str] = None + """The security alert collation window in effect: ``default_security_alert_collation_window`` when + set, otherwise ``1h``.""" + + effective_default_system_alert_collation_window: Optional[str] = None + """The system alert collation window in effect: ``default_system_alert_collation_window`` when set, + otherwise ``1d``.""" + + etag: Optional[str] = None + """Opaque entity tag for optimistic concurrency. Returned on every read; pass it back as the Update + request's ``etag`` to make the update conditional on the resource being unchanged.""" + + name: Optional[str] = None + + security_alert_export_config: Optional[ExportConfig] = None + """Export destination for security alerts. When unset, export is off and reads return ``{"enabled": + false}``.""" + + system_alert_export_config: Optional[ExportConfig] = None + """Export destination for system alerts. When unset, export is off and reads return ``{"enabled": + false}``.""" + + update_time: Optional[Timestamp] = None + + updated_by: Optional[Principal] = None + + version: Optional[int] = None + """Monotonic version, bumped on each update. Informational; use ``etag`` for optimistic + concurrency.""" + + def as_dict(self) -> dict: + """Serializes the AlertConfig into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.default_security_alert_collation_window is not None: + body["default_security_alert_collation_window"] = self.default_security_alert_collation_window + if self.default_system_alert_collation_window is not None: + body["default_system_alert_collation_window"] = self.default_system_alert_collation_window + if self.effective_default_security_alert_collation_window is not None: + body["effective_default_security_alert_collation_window"] = ( + self.effective_default_security_alert_collation_window + ) + if self.effective_default_system_alert_collation_window is not None: + body["effective_default_system_alert_collation_window"] = ( + self.effective_default_system_alert_collation_window + ) + if self.etag is not None: + body["etag"] = self.etag + if self.name is not None: + body["name"] = self.name + if self.security_alert_export_config: + body["security_alert_export_config"] = self.security_alert_export_config.as_dict() + if self.system_alert_export_config: + body["system_alert_export_config"] = self.system_alert_export_config.as_dict() + if self.update_time is not None: + body["update_time"] = self.update_time.ToJsonString() + if self.updated_by: + body["updated_by"] = self.updated_by.as_dict() + if self.version is not None: + body["version"] = self.version + return body + + def as_shallow_dict(self) -> dict: + """Serializes the AlertConfig into a shallow dictionary of its immediate attributes.""" + body = {} + if self.default_security_alert_collation_window is not None: + body["default_security_alert_collation_window"] = self.default_security_alert_collation_window + if self.default_system_alert_collation_window is not None: + body["default_system_alert_collation_window"] = self.default_system_alert_collation_window + if self.effective_default_security_alert_collation_window is not None: + body["effective_default_security_alert_collation_window"] = ( + self.effective_default_security_alert_collation_window + ) + if self.effective_default_system_alert_collation_window is not None: + body["effective_default_system_alert_collation_window"] = ( + self.effective_default_system_alert_collation_window + ) + if self.etag is not None: + body["etag"] = self.etag + if self.name is not None: + body["name"] = self.name + if self.security_alert_export_config: + body["security_alert_export_config"] = self.security_alert_export_config + if self.system_alert_export_config: + body["system_alert_export_config"] = self.system_alert_export_config + if self.update_time is not None: + body["update_time"] = self.update_time + if self.updated_by: + body["updated_by"] = self.updated_by + if self.version is not None: + body["version"] = self.version + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> AlertConfig: + """Deserializes the AlertConfig from a dictionary.""" + return cls( + default_security_alert_collation_window=d.get("default_security_alert_collation_window", None), + default_system_alert_collation_window=d.get("default_system_alert_collation_window", None), + effective_default_security_alert_collation_window=d.get( + "effective_default_security_alert_collation_window", None + ), + effective_default_system_alert_collation_window=d.get( + "effective_default_system_alert_collation_window", None + ), + etag=d.get("etag", None), + name=d.get("name", None), + security_alert_export_config=_from_dict(d, "security_alert_export_config", ExportConfig), + system_alert_export_config=_from_dict(d, "system_alert_export_config", ExportConfig), + update_time=_timestamp(d, "update_time"), + updated_by=_from_dict(d, "updated_by", Principal), + version=_int64(d, "version"), + ) + + +@dataclass +class AutoLoader: + """Ingests files from cloud storage with Auto Loader. Mutually exclusive with + ``unity_catalog_connection`` and ``existing_table``.""" + + format: AutoLoaderFormat + """File format of the source files. Required.""" + + location: str + """URI of the source files: a path in a Unity Catalog volume or external location, or a cloud + storage URL such as an Amazon S3 path. Required, including when new files are discovered through + ``queue_url``. Cannot be changed after creation.""" + + raw_table: TableName + """Table that stores the raw data before parsers schematize it. Every Auto Loader datasource writes + its raw rows to this table before any parser reads them. The table must be in the datasource's + catalog and schema; an omitted ``catalog_name`` or ``schema_name`` defaults to them. Cannot be + changed after creation. Reads return the value as supplied; the table the data is written to is + reported in ``effective_raw_table``.""" + + aws_secret_scope: Optional[str] = None + """Name of the workspace secret scope that holds the AWS access key ID and secret access key, under + the keys ``awsAccessKeyId`` and ``awsSecretAccessKey``.""" + + effective_include_existing_files: Optional[bool] = None + """Whether ingestion started with the files already present at ``location``. Populated on get, + create, and update responses.""" + + effective_raw_table: Optional[TableName] = None + """The table that the raw data is written to: ``raw_table`` in the datasource's catalog and schema, + under their current Unity Catalog names. Populated on get, create, and update responses.""" + + envelope: Optional[Envelope] = None + """How to unwrap a transport envelope around the source records before parsing. Supported only for + the JSON and JSONL formats. If unset, each source record is treated as one event.""" + + format_options: Optional[Dict[str, str]] = None + """Format-specific parsing options, passed to Auto Loader. Supported for the CSV, JSON, JSONL, and + XML formats, each of which accepts a fixed set of options. The XML format requires the + ``rowTag`` option.""" + + include_existing_files: Optional[bool] = None + """Whether to ingest the files already present at ``location`` when ingestion first starts. Set + true to ingest the existing files and then every new file. Set false, or leave unset, to ingest + only files that arrive after the datasource starts. Defaults to false. Cannot be changed after + creation. Reads return the value as supplied; the setting in effect is reported in + ``effective_include_existing_files``.""" + + override_partition_columns: Optional[List[str]] = None + """Hive-style partition columns to derive from the source file paths. If unset, no partition + columns are derived.""" + + queue_url: Optional[str] = None + """URL of a customer-managed Amazon SQS queue that receives file-arrival notifications. Set it to + discover new files through the queue. If unset, new files are discovered through the file events + of the source's Unity Catalog external location when it has them configured, and by directory + listing otherwise. Cannot be changed after creation.""" + + service_credentials: Optional[str] = None + """Name of the Unity Catalog service credential to use.""" + + def as_dict(self) -> dict: + """Serializes the AutoLoader into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.aws_secret_scope is not None: + body["aws_secret_scope"] = self.aws_secret_scope + if self.effective_include_existing_files is not None: + body["effective_include_existing_files"] = self.effective_include_existing_files + if self.effective_raw_table: + body["effective_raw_table"] = self.effective_raw_table.as_dict() + if self.envelope: + body["envelope"] = self.envelope.as_dict() + if self.format is not None: + body["format"] = self.format.value + if self.format_options: + body["format_options"] = self.format_options + if self.include_existing_files is not None: + body["include_existing_files"] = self.include_existing_files + if self.location is not None: + body["location"] = self.location + if self.override_partition_columns: + body["override_partition_columns"] = [v for v in self.override_partition_columns] + if self.queue_url is not None: + body["queue_url"] = self.queue_url + if self.raw_table: + body["raw_table"] = self.raw_table.as_dict() + if self.service_credentials is not None: + body["service_credentials"] = self.service_credentials + return body + + def as_shallow_dict(self) -> dict: + """Serializes the AutoLoader into a shallow dictionary of its immediate attributes.""" + body = {} + if self.aws_secret_scope is not None: + body["aws_secret_scope"] = self.aws_secret_scope + if self.effective_include_existing_files is not None: + body["effective_include_existing_files"] = self.effective_include_existing_files + if self.effective_raw_table: + body["effective_raw_table"] = self.effective_raw_table + if self.envelope: + body["envelope"] = self.envelope + if self.format is not None: + body["format"] = self.format + if self.format_options: + body["format_options"] = self.format_options + if self.include_existing_files is not None: + body["include_existing_files"] = self.include_existing_files + if self.location is not None: + body["location"] = self.location + if self.override_partition_columns: + body["override_partition_columns"] = self.override_partition_columns + if self.queue_url is not None: + body["queue_url"] = self.queue_url + if self.raw_table: + body["raw_table"] = self.raw_table + if self.service_credentials is not None: + body["service_credentials"] = self.service_credentials + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> AutoLoader: + """Deserializes the AutoLoader from a dictionary.""" + return cls( + aws_secret_scope=d.get("aws_secret_scope", None), + effective_include_existing_files=d.get("effective_include_existing_files", None), + effective_raw_table=_from_dict(d, "effective_raw_table", TableName), + envelope=_from_dict(d, "envelope", Envelope), + format=_enum(d, "format", AutoLoaderFormat), + format_options=d.get("format_options", None), + include_existing_files=d.get("include_existing_files", None), + location=d.get("location", None), + override_partition_columns=d.get("override_partition_columns", None), + queue_url=d.get("queue_url", None), + raw_table=_from_dict(d, "raw_table", TableName), + service_credentials=d.get("service_credentials", None), + ) + + +class AutoLoaderFormat(Enum): + """File format of the source files that Auto Loader reads.""" + + AUTOLOADER_FORMAT_CSV = "AUTOLOADER_FORMAT_CSV" + AUTOLOADER_FORMAT_JSON = "AUTOLOADER_FORMAT_JSON" + AUTOLOADER_FORMAT_JSONL = "AUTOLOADER_FORMAT_JSONL" + AUTOLOADER_FORMAT_PARQUET = "AUTOLOADER_FORMAT_PARQUET" + AUTOLOADER_FORMAT_XML = "AUTOLOADER_FORMAT_XML" + + +@dataclass +class CloudWatchLogsEnvelope: + """Unwraps an AWS CloudWatch Logs delivery envelope: each element of ``logEvents`` becomes its own + row, with the element's ``message`` parsed as the event body.""" + + header_retention: Optional[bool] = None + """Whether to keep the envelope's ``owner``, ``logGroup``, and ``logStream`` values in the + ``data.lw_header`` field of each unwrapped row. Must be set explicitly to ``true`` or ``false``.""" + + def as_dict(self) -> dict: + """Serializes the CloudWatchLogsEnvelope into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.header_retention is not None: + body["header_retention"] = self.header_retention + return body + + def as_shallow_dict(self) -> dict: + """Serializes the CloudWatchLogsEnvelope into a shallow dictionary of its immediate attributes.""" + body = {} + if self.header_retention is not None: + body["header_retention"] = self.header_retention + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> CloudWatchLogsEnvelope: + """Deserializes the CloudWatchLogsEnvelope from a dictionary.""" + return cls(header_retention=d.get("header_retention", None)) + + +@dataclass +class CommonConfig: + """Catch-all per-catalog config — settings that don't belong to any one resource family. + Singleton. Name format: ``catalogs/{catalog}/commonConfig``.""" + + custom_content_pack_path: Optional[str] = None + """UC volume path holding user-authored custom content packs, e.g. + ``/Volumes/{catalog}/{schema}/{volume}``. User-provided with no default: when unset, custom + content packs are unavailable (referencing one is rejected) and only built-in content packs + apply.""" + + custom_presets_path: Optional[str] = None + """UC volume path holding user-authored custom presets, e.g. + ``/Volumes/{catalog}/{schema}/{volume}``. User-provided with no default: when unset, custom + presets are unavailable and only built-in presets apply.""" + + etag: Optional[str] = None + """Opaque entity tag for optimistic concurrency. Returned on every read; pass it back as the Update + request's ``etag`` to make the update conditional on the resource being unchanged.""" + + name: Optional[str] = None + + per_resource_permissions_enabled: Optional[bool] = None + """Whether users may grant permissions on individual resources rather than only at the Folder + level.""" + + update_time: Optional[Timestamp] = None + + updated_by: Optional[Principal] = None + + version: Optional[int] = None + """Monotonic version, bumped on each update. Informational; use ``etag`` for optimistic + concurrency.""" + + def as_dict(self) -> dict: + """Serializes the CommonConfig into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.custom_content_pack_path is not None: + body["custom_content_pack_path"] = self.custom_content_pack_path + if self.custom_presets_path is not None: + body["custom_presets_path"] = self.custom_presets_path + if self.etag is not None: + body["etag"] = self.etag + if self.name is not None: + body["name"] = self.name + if self.per_resource_permissions_enabled is not None: + body["per_resource_permissions_enabled"] = self.per_resource_permissions_enabled + if self.update_time is not None: + body["update_time"] = self.update_time.ToJsonString() + if self.updated_by: + body["updated_by"] = self.updated_by.as_dict() + if self.version is not None: + body["version"] = self.version + return body + + def as_shallow_dict(self) -> dict: + """Serializes the CommonConfig into a shallow dictionary of its immediate attributes.""" + body = {} + if self.custom_content_pack_path is not None: + body["custom_content_pack_path"] = self.custom_content_pack_path + if self.custom_presets_path is not None: + body["custom_presets_path"] = self.custom_presets_path + if self.etag is not None: + body["etag"] = self.etag + if self.name is not None: + body["name"] = self.name + if self.per_resource_permissions_enabled is not None: + body["per_resource_permissions_enabled"] = self.per_resource_permissions_enabled + if self.update_time is not None: + body["update_time"] = self.update_time + if self.updated_by: + body["updated_by"] = self.updated_by + if self.version is not None: + body["version"] = self.version + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> CommonConfig: + """Deserializes the CommonConfig from a dictionary.""" + return cls( + custom_content_pack_path=d.get("custom_content_pack_path", None), + custom_presets_path=d.get("custom_presets_path", None), + etag=d.get("etag", None), + name=d.get("name", None), + per_resource_permissions_enabled=d.get("per_resource_permissions_enabled", None), + update_time=_timestamp(d, "update_time"), + updated_by=_from_dict(d, "updated_by", Principal), + version=_int64(d, "version"), + ) + + +@dataclass +class ContentPackReference: + """Reference to the ContentPack (and pack version) a resource was originally installed from.""" + + uuid: Optional[str] = None + """UUID of the content pack.""" + + version: Optional[str] = None + """Pack version the resource was installed from.""" + + def as_dict(self) -> dict: + """Serializes the ContentPackReference into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.uuid is not None: + body["uuid"] = self.uuid + if self.version is not None: + body["version"] = self.version + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ContentPackReference into a shallow dictionary of its immediate attributes.""" + body = {} + if self.uuid is not None: + body["uuid"] = self.uuid + if self.version is not None: + body["version"] = self.version + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ContentPackReference: + """Deserializes the ContentPackReference from a dictionary.""" + return cls(uuid=d.get("uuid", None), version=d.get("version", None)) + + +@dataclass +class Datasource: + """A Lakewatch datasource: a pipeline that ingests security data and writes it to Unity Catalog + tables. + + What a datasource ingests is defined by its ingestion method (``autoloader``, + ``unity_catalog_connection``, or ``existing_table``), which sets how the data arrives, together + with its ``parsers``, which define the schemas that the arriving data is matched against and + written as.""" + + catalog_name: str + """Name of the Unity Catalog catalog that contains the datasource. Required on create and cannot be + changed afterward. Reads return the value as supplied; the catalog's current name is reported in + ``effective_catalog_name``.""" + + schema_name: str + """Name of the Unity Catalog schema that contains the datasource. Required on create and cannot be + changed afterward. Reads return the value as supplied; the schema's current name is reported in + ``effective_schema_name``.""" + + datasource_name: str + """Name of the datasource in Unity Catalog, unique within its schema. Required on create and cannot + be changed afterward. Reads return the value as supplied; the datasource's current Unity Catalog + name, which reflects renames, is reported in ``effective_datasource_name``.""" + + schedule: Schedule + """Schedule on which the datasource's ingestion job runs. A datasource cannot be paused, so a + request that sets ``schedule.enabled`` to ``false`` is rejected. Continuous schedules are not + supported for Unity Catalog connection datasources.""" + + warehouse_id: str + """ID of the SQL warehouse that runs the checks Lakewatch performs to validate a create or update + request, such as source and destination access checks. Required and non-blank on every create + and update request; there is no default warehouse. The requester must be able to use the + warehouse. Used only for those checks: it is not stored, is not returned on reads, and does not + affect the compute that runs the ingestion job.""" + + annotations: Optional[Dict[str, str]] = None + """Free-form key-value annotations on the datasource, managed by the client.""" + + autoloader: Optional[AutoLoader] = None + + compute_mode: Optional[DatasourceComputeMode] = None + """Compute mode for the datasource's ingestion job. If unset, the ``default_compute_mode`` of the + catalog's datasource configuration applies, or ``DS_COMPUTE_PERFORMANCE_OPTIMIZED`` if the + catalog has none.""" + + compute_workspace_id: Optional[str] = None + """ID of the workspace whose compute runs this datasource's ingestion job and owns its associated + resources, such as the storage the ingested data lands in. Server-assigned.""" + + create_time: Optional[Timestamp] = None + + created_by: Optional[Principal] = None + + datasource_id: Optional[str] = None + """Unity Catalog securable ID (a UUID) of the datasource. It never changes, including when the + datasource is renamed in Unity Catalog, and is the final segment of ``name``.""" + + description: Optional[str] = None + """Free-form description of the datasource.""" + + effective_catalog_name: Optional[str] = None + """Current name of the Unity Catalog catalog that contains the datasource, which reflects catalog + renames without a change to ``catalog_name``. Populated on get, create, and update responses; + list responses omit it.""" + + effective_datasource_name: Optional[str] = None + """Current UC leaf name of the datasource (server-owned). Tracks UC renames so callers can observe + the live name without having to update the client-owned ``datasource_name``; unset until + resolved.""" + + effective_run_as_principal: Optional[str] = None + """The identity the datasource runs as: the resolved ``run_as_principal``, or the requester's + identity if ``run_as_principal`` was omitted on create or cleared on update. Given as a user + email or service principal application ID, or as the numeric principal ID if the identity can no + longer be resolved. Populated on get, create, and update responses; list responses omit it.""" + + effective_schema_name: Optional[str] = None + """Current name of the Unity Catalog schema that contains the datasource, which reflects schema + renames without a change to ``schema_name``. Populated on get, create, and update responses; + list responses omit it.""" + + etag: Optional[str] = None + """Opaque entity tag for optimistic concurrency control. Returned by get, create, and update + responses and by list responses in the ``DATASOURCE_VIEW_FULL`` view. To make an update + conditional on the datasource not having changed since it was read, pass this value as the + update request's ``etag``. Changes to ``status`` do not change the etag.""" + + existing_table: Optional[ExistingTable] = None + + name: Optional[str] = None + """Resource name of the datasource, in the format ``datasources/{datasource_id}``. Unlike the + datasource's Unity Catalog location (``catalog_name``, ``schema_name``, and + ``datasource_name``), the name does not change when the datasource or its catalog or schema is + renamed in Unity Catalog.""" + + owner: Optional[Principal] = None + """Unity Catalog owner of the datasource, assigned by the server. Create responses omit it.""" + + parsers: Optional[List[ParserRef]] = None + """Parsers that schematize the ingested data. Each record is matched against the parsers in order + and written to the parsed table of the first parser that it matches. Existing-table and Unity + Catalog connection datasources require at least one parser; an Auto Loader datasource with no + parsers writes only its raw table.""" + + run_as_principal: Optional[str] = None + """The user or service principal that the datasource runs as, given as a user email or service + principal application ID. When it is supplied on create or update, the requester must be + authorized to run the datasource as that principal, or the request is rejected. When it is + omitted on create, the datasource runs as the requester. On update, a mask that includes + ``run_as_principal`` without a value clears it, so the datasource runs as the requester; a mask + that omits it keeps the current identity. Reads return the value as supplied; the identity the + datasource runs as is reported in ``effective_run_as_principal``.""" + + status: Optional[DatasourceStatus] = None + """Runtime status of the jobs and pipelines that run this datasource.""" + + unity_catalog_connection: Optional[UnityCatalogConnection] = None + + update_time: Optional[Timestamp] = None + + updated_by: Optional[Principal] = None + + version: Optional[int] = None + """Version number of the datasource, incremented on each update. Informational only; use ``etag`` + for optimistic concurrency control.""" + + def as_dict(self) -> dict: + """Serializes the Datasource into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.annotations: + body["annotations"] = self.annotations + if self.autoloader: + body["autoloader"] = self.autoloader.as_dict() + if self.catalog_name is not None: + body["catalog_name"] = self.catalog_name + if self.compute_mode is not None: + body["compute_mode"] = self.compute_mode.value + if self.compute_workspace_id is not None: + body["compute_workspace_id"] = self.compute_workspace_id + if self.create_time is not None: + body["create_time"] = self.create_time.ToJsonString() + if self.created_by: + body["created_by"] = self.created_by.as_dict() + if self.datasource_id is not None: + body["datasource_id"] = self.datasource_id + if self.datasource_name is not None: + body["datasource_name"] = self.datasource_name + if self.description is not None: + body["description"] = self.description + if self.effective_catalog_name is not None: + body["effective_catalog_name"] = self.effective_catalog_name + if self.effective_datasource_name is not None: + body["effective_datasource_name"] = self.effective_datasource_name + if self.effective_run_as_principal is not None: + body["effective_run_as_principal"] = self.effective_run_as_principal + if self.effective_schema_name is not None: + body["effective_schema_name"] = self.effective_schema_name + if self.etag is not None: + body["etag"] = self.etag + if self.existing_table: + body["existing_table"] = self.existing_table.as_dict() + if self.name is not None: + body["name"] = self.name + if self.owner: + body["owner"] = self.owner.as_dict() + if self.parsers: + body["parsers"] = [v.as_dict() for v in self.parsers] + if self.run_as_principal is not None: + body["run_as_principal"] = self.run_as_principal + if self.schedule: + body["schedule"] = self.schedule.as_dict() + if self.schema_name is not None: + body["schema_name"] = self.schema_name + if self.status: + body["status"] = self.status.as_dict() + if self.unity_catalog_connection: + body["unity_catalog_connection"] = self.unity_catalog_connection.as_dict() + if self.update_time is not None: + body["update_time"] = self.update_time.ToJsonString() + if self.updated_by: + body["updated_by"] = self.updated_by.as_dict() + if self.version is not None: + body["version"] = self.version + if self.warehouse_id is not None: + body["warehouse_id"] = self.warehouse_id + return body + + def as_shallow_dict(self) -> dict: + """Serializes the Datasource into a shallow dictionary of its immediate attributes.""" + body = {} + if self.annotations: + body["annotations"] = self.annotations + if self.autoloader: + body["autoloader"] = self.autoloader + if self.catalog_name is not None: + body["catalog_name"] = self.catalog_name + if self.compute_mode is not None: + body["compute_mode"] = self.compute_mode + if self.compute_workspace_id is not None: + body["compute_workspace_id"] = self.compute_workspace_id + if self.create_time is not None: + body["create_time"] = self.create_time + if self.created_by: + body["created_by"] = self.created_by + if self.datasource_id is not None: + body["datasource_id"] = self.datasource_id + if self.datasource_name is not None: + body["datasource_name"] = self.datasource_name + if self.description is not None: + body["description"] = self.description + if self.effective_catalog_name is not None: + body["effective_catalog_name"] = self.effective_catalog_name + if self.effective_datasource_name is not None: + body["effective_datasource_name"] = self.effective_datasource_name + if self.effective_run_as_principal is not None: + body["effective_run_as_principal"] = self.effective_run_as_principal + if self.effective_schema_name is not None: + body["effective_schema_name"] = self.effective_schema_name + if self.etag is not None: + body["etag"] = self.etag + if self.existing_table: + body["existing_table"] = self.existing_table + if self.name is not None: + body["name"] = self.name + if self.owner: + body["owner"] = self.owner + if self.parsers: + body["parsers"] = self.parsers + if self.run_as_principal is not None: + body["run_as_principal"] = self.run_as_principal + if self.schedule: + body["schedule"] = self.schedule + if self.schema_name is not None: + body["schema_name"] = self.schema_name + if self.status: + body["status"] = self.status + if self.unity_catalog_connection: + body["unity_catalog_connection"] = self.unity_catalog_connection + if self.update_time is not None: + body["update_time"] = self.update_time + if self.updated_by: + body["updated_by"] = self.updated_by + if self.version is not None: + body["version"] = self.version + if self.warehouse_id is not None: + body["warehouse_id"] = self.warehouse_id + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> Datasource: + """Deserializes the Datasource from a dictionary.""" + return cls( + annotations=d.get("annotations", None), + autoloader=_from_dict(d, "autoloader", AutoLoader), + catalog_name=d.get("catalog_name", None), + compute_mode=_enum(d, "compute_mode", DatasourceComputeMode), + compute_workspace_id=d.get("compute_workspace_id", None), + create_time=_timestamp(d, "create_time"), + created_by=_from_dict(d, "created_by", Principal), + datasource_id=d.get("datasource_id", None), + datasource_name=d.get("datasource_name", None), + description=d.get("description", None), + effective_catalog_name=d.get("effective_catalog_name", None), + effective_datasource_name=d.get("effective_datasource_name", None), + effective_run_as_principal=d.get("effective_run_as_principal", None), + effective_schema_name=d.get("effective_schema_name", None), + etag=d.get("etag", None), + existing_table=_from_dict(d, "existing_table", ExistingTable), + name=d.get("name", None), + owner=_from_dict(d, "owner", Principal), + parsers=_repeated_dict(d, "parsers", ParserRef), + run_as_principal=d.get("run_as_principal", None), + schedule=_from_dict(d, "schedule", Schedule), + schema_name=d.get("schema_name", None), + status=_from_dict(d, "status", DatasourceStatus), + unity_catalog_connection=_from_dict(d, "unity_catalog_connection", UnityCatalogConnection), + update_time=_timestamp(d, "update_time"), + updated_by=_from_dict(d, "updated_by", Principal), + version=_int64(d, "version"), + warehouse_id=d.get("warehouse_id", None), + ) + + +class DatasourceComputeMode(Enum): + """Compute mode for the datasource's ingestion job.""" + + DS_COMPUTE_PERFORMANCE_OPTIMIZED = "DS_COMPUTE_PERFORMANCE_OPTIMIZED" + DS_COMPUTE_STANDARD = "DS_COMPUTE_STANDARD" + + +@dataclass +class DatasourceConfig: + """Per-catalog datasource-related config. Singleton. Name format: + ``catalogs/{catalog}/datasourceConfig``.""" + + default_compute_mode: Optional[DatasourceComputeMode] = None + """Default compute mode applied to datasource ingestion jobs when the individual datasource doesn't + override it. v1 ``datasources.defaultComputeMode``.""" + + default_folder: Optional[str] = None + """Schema name of the registered folder (in this catalog) that a new datasource is placed in when + the request does not specify a schema. Must name a schema registered as a folder. Omit to clear. + Reads return the value as supplied; the folder's current schema name is reported in + ``effective_default_folder``.""" + + effective_default_folder: Optional[str] = None + """The current schema name of the ``default_folder`` folder. Tracks UC schema renames. Unset when + no default is configured, or when the folder's schema no longer exists or is not visible to the + caller.""" + + etag: Optional[str] = None + """Opaque entity tag for optimistic concurrency. Returned on every read; pass it back as the Update + request's ``etag`` to make the update conditional on the resource being unchanged.""" + + name: Optional[str] = None + + update_time: Optional[Timestamp] = None + + updated_by: Optional[Principal] = None + + version: Optional[int] = None + """Monotonic version, bumped on each update. Informational; use ``etag`` for optimistic + concurrency.""" + + def as_dict(self) -> dict: + """Serializes the DatasourceConfig into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.default_compute_mode is not None: + body["default_compute_mode"] = self.default_compute_mode.value + if self.default_folder is not None: + body["default_folder"] = self.default_folder + if self.effective_default_folder is not None: + body["effective_default_folder"] = self.effective_default_folder + if self.etag is not None: + body["etag"] = self.etag + if self.name is not None: + body["name"] = self.name + if self.update_time is not None: + body["update_time"] = self.update_time.ToJsonString() + if self.updated_by: + body["updated_by"] = self.updated_by.as_dict() + if self.version is not None: + body["version"] = self.version + return body + + def as_shallow_dict(self) -> dict: + """Serializes the DatasourceConfig into a shallow dictionary of its immediate attributes.""" + body = {} + if self.default_compute_mode is not None: + body["default_compute_mode"] = self.default_compute_mode + if self.default_folder is not None: + body["default_folder"] = self.default_folder + if self.effective_default_folder is not None: + body["effective_default_folder"] = self.effective_default_folder + if self.etag is not None: + body["etag"] = self.etag + if self.name is not None: + body["name"] = self.name + if self.update_time is not None: + body["update_time"] = self.update_time + if self.updated_by: + body["updated_by"] = self.updated_by + if self.version is not None: + body["version"] = self.version + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> DatasourceConfig: + """Deserializes the DatasourceConfig from a dictionary.""" + return cls( + default_compute_mode=_enum(d, "default_compute_mode", DatasourceComputeMode), + default_folder=d.get("default_folder", None), + effective_default_folder=d.get("effective_default_folder", None), + etag=d.get("etag", None), + name=d.get("name", None), + update_time=_timestamp(d, "update_time"), + updated_by=_from_dict(d, "updated_by", Principal), + version=_int64(d, "version"), + ) + + +@dataclass +class DatasourceStatus: + """Read-only runtime status of a datasource, with the status of each job and pipeline that runs it.""" + + ingestion_job: Optional[ResourceStatus] = None + """Status of the datasource's ingestion job.""" + + ingestion_pipeline: Optional[ResourceStatus] = None + """Status of the managed ingestion pipeline. Populated only when ``unity_catalog_connection`` is + set, with ``pipeline_id`` identifying the pipeline.""" + + def as_dict(self) -> dict: + """Serializes the DatasourceStatus into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.ingestion_job: + body["ingestion_job"] = self.ingestion_job.as_dict() + if self.ingestion_pipeline: + body["ingestion_pipeline"] = self.ingestion_pipeline.as_dict() + return body + + def as_shallow_dict(self) -> dict: + """Serializes the DatasourceStatus into a shallow dictionary of its immediate attributes.""" + body = {} + if self.ingestion_job: + body["ingestion_job"] = self.ingestion_job + if self.ingestion_pipeline: + body["ingestion_pipeline"] = self.ingestion_pipeline + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> DatasourceStatus: + """Deserializes the DatasourceStatus from a dictionary.""" + return cls( + ingestion_job=_from_dict(d, "ingestion_job", ResourceStatus), + ingestion_pipeline=_from_dict(d, "ingestion_pipeline", ResourceStatus), + ) + + +class DatasourceView(Enum): + """Which fields list responses return for each datasource.""" + + DATASOURCE_VIEW_BASIC = "DATASOURCE_VIEW_BASIC" + DATASOURCE_VIEW_FULL = "DATASOURCE_VIEW_FULL" + + +@dataclass +class Envelope: + """Unwraps a transport envelope from the source data so that each contained event becomes its own + row. Applied when the source data is read, before parsing.""" + + cloud_watch_logs: Optional[CloudWatchLogsEnvelope] = None + + json_array: Optional[JsonArrayEnvelope] = None + + def as_dict(self) -> dict: + """Serializes the Envelope into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.cloud_watch_logs: + body["cloud_watch_logs"] = self.cloud_watch_logs.as_dict() + if self.json_array: + body["json_array"] = self.json_array.as_dict() + return body + + def as_shallow_dict(self) -> dict: + """Serializes the Envelope into a shallow dictionary of its immediate attributes.""" + body = {} + if self.cloud_watch_logs: + body["cloud_watch_logs"] = self.cloud_watch_logs + if self.json_array: + body["json_array"] = self.json_array + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> Envelope: + """Deserializes the Envelope from a dictionary.""" + return cls( + cloud_watch_logs=_from_dict(d, "cloud_watch_logs", CloudWatchLogsEnvelope), + json_array=_from_dict(d, "json_array", JsonArrayEnvelope), + ) + + +@dataclass +class ExistingTable: + """Uses an existing Unity Catalog Delta table or view as the datasource's source. Mutually + exclusive with ``autoloader`` and ``unity_catalog_connection``.""" + + name: str + """Name of the source table or view. Required on create and cannot be changed. Reads return the + value as supplied; the table's current name is reported in ``effective_name``.""" + + catalog: Optional[str] = None + """Name of the catalog the existing table lives in. Optional on create and defaults to the + datasource's own catalog. Cannot be changed. Reads return the value as supplied; the catalog in + use is reported in ``effective_catalog``.""" + + effective_catalog: Optional[str] = None + """Current Unity Catalog catalog name of the source table, resolved from its stable table ID, so a + rename of the source table is reported here without changing ``catalog``.""" + + effective_name: Optional[str] = None + """Current Unity Catalog name of the source table, resolved from its stable table ID.""" + + effective_schema: Optional[str] = None + """Current Unity Catalog schema name of the source table, resolved from its stable table ID.""" + + include_existing_rows: Optional[bool] = None + """Whether to ingest the rows already present in the source table when ingestion first starts. Set + true to ingest the existing rows and then every new commit. Set false, or leave unset, to ingest + only rows written after the datasource starts. Defaults to false, and must be true for a Delta + Sharing source table. Takes effect only at first start and cannot be changed. Reads return the + value as supplied.""" + + schema: Optional[str] = None + """Name of the schema the existing table lives in. Optional, but required when ``catalog`` is set; + supplying ``catalog`` without ``schema`` is rejected. Defaults to the datasource's own schema. + Cannot be changed. Reads return the value as supplied; the schema in use is reported in + ``effective_schema``.""" + + def as_dict(self) -> dict: + """Serializes the ExistingTable into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.catalog is not None: + body["catalog"] = self.catalog + if self.effective_catalog is not None: + body["effective_catalog"] = self.effective_catalog + if self.effective_name is not None: + body["effective_name"] = self.effective_name + if self.effective_schema is not None: + body["effective_schema"] = self.effective_schema + if self.include_existing_rows is not None: + body["include_existing_rows"] = self.include_existing_rows + if self.name is not None: + body["name"] = self.name + if self.schema is not None: + body["schema"] = self.schema + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ExistingTable into a shallow dictionary of its immediate attributes.""" + body = {} + if self.catalog is not None: + body["catalog"] = self.catalog + if self.effective_catalog is not None: + body["effective_catalog"] = self.effective_catalog + if self.effective_name is not None: + body["effective_name"] = self.effective_name + if self.effective_schema is not None: + body["effective_schema"] = self.effective_schema + if self.include_existing_rows is not None: + body["include_existing_rows"] = self.include_existing_rows + if self.name is not None: + body["name"] = self.name + if self.schema is not None: + body["schema"] = self.schema + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ExistingTable: + """Deserializes the ExistingTable from a dictionary.""" + return cls( + catalog=d.get("catalog", None), + effective_catalog=d.get("effective_catalog", None), + effective_name=d.get("effective_name", None), + effective_schema=d.get("effective_schema", None), + include_existing_rows=d.get("include_existing_rows", None), + name=d.get("name", None), + schema=d.get("schema", None), + ) + + +@dataclass +class ExportConfig: + """Configuration for delivering events / alerts to an external destination. + + Exactly one of the destination variants is set; the server picks it via the ``oneof``.""" + + enabled: Optional[bool] = None + """Whether export is currently active.""" + + jira: Optional[JiraExportConfig] = None + + service_now: Optional[ServiceNowExportConfig] = None + + slack: Optional[SlackExportConfig] = None + + tines: Optional[TinesExportConfig] = None + + webhook: Optional[WebhookExportConfig] = None + + workspace_id: Optional[str] = None + """Deprecated and ignored: exports always egress through the request's own workspace. Kept so + configs stored by older builds still round-trip.""" + + def as_dict(self) -> dict: + """Serializes the ExportConfig into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.enabled is not None: + body["enabled"] = self.enabled + if self.jira: + body["jira"] = self.jira.as_dict() + if self.service_now: + body["service_now"] = self.service_now.as_dict() + if self.slack: + body["slack"] = self.slack.as_dict() + if self.tines: + body["tines"] = self.tines.as_dict() + if self.webhook: + body["webhook"] = self.webhook.as_dict() + if self.workspace_id is not None: + body["workspace_id"] = self.workspace_id + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ExportConfig into a shallow dictionary of its immediate attributes.""" + body = {} + if self.enabled is not None: + body["enabled"] = self.enabled + if self.jira: + body["jira"] = self.jira + if self.service_now: + body["service_now"] = self.service_now + if self.slack: + body["slack"] = self.slack + if self.tines: + body["tines"] = self.tines + if self.webhook: + body["webhook"] = self.webhook + if self.workspace_id is not None: + body["workspace_id"] = self.workspace_id + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ExportConfig: + """Deserializes the ExportConfig from a dictionary.""" + return cls( + enabled=d.get("enabled", None), + jira=_from_dict(d, "jira", JiraExportConfig), + service_now=_from_dict(d, "service_now", ServiceNowExportConfig), + slack=_from_dict(d, "slack", SlackExportConfig), + tines=_from_dict(d, "tines", TinesExportConfig), + webhook=_from_dict(d, "webhook", WebhookExportConfig), + workspace_id=d.get("workspace_id", None), + ) + + +@dataclass +class JiraExportConfig: + instance_url: str + """Jira instance URL (e.g. ``https://acme.atlassian.net``).""" + + project_key: str + """Jira project key the exported issues are created under.""" + + user_email: str + """User email paired with the API token for Jira Cloud basic auth.""" + + default_issue_type: str + """Default issue type for created tickets (e.g. ``Task``, ``Incident``). Jira issue creation + requires an issue type, so it is validated at config-save time rather than failing the export at + Jira later.""" + + api_token_uc_secret: str + """Three-level Unity Catalog name (``catalog.schema.secret``) of the secret holding the Jira API + token, paired with ``user_email`` as the Jira Cloud basic-auth credential (``email:token``). + Only the name is stored; the token is read when an alert is exported, so rotating the secret + takes effect without editing this config.""" + + def as_dict(self) -> dict: + """Serializes the JiraExportConfig into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.api_token_uc_secret is not None: + body["api_token_uc_secret"] = self.api_token_uc_secret + if self.default_issue_type is not None: + body["default_issue_type"] = self.default_issue_type + if self.instance_url is not None: + body["instance_url"] = self.instance_url + if self.project_key is not None: + body["project_key"] = self.project_key + if self.user_email is not None: + body["user_email"] = self.user_email + return body + + def as_shallow_dict(self) -> dict: + """Serializes the JiraExportConfig into a shallow dictionary of its immediate attributes.""" + body = {} + if self.api_token_uc_secret is not None: + body["api_token_uc_secret"] = self.api_token_uc_secret + if self.default_issue_type is not None: + body["default_issue_type"] = self.default_issue_type + if self.instance_url is not None: + body["instance_url"] = self.instance_url + if self.project_key is not None: + body["project_key"] = self.project_key + if self.user_email is not None: + body["user_email"] = self.user_email + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> JiraExportConfig: + """Deserializes the JiraExportConfig from a dictionary.""" + return cls( + api_token_uc_secret=d.get("api_token_uc_secret", None), + default_issue_type=d.get("default_issue_type", None), + instance_url=d.get("instance_url", None), + project_key=d.get("project_key", None), + user_email=d.get("user_email", None), + ) + + +@dataclass +class JsonArrayEnvelope: + """Unwraps a JSON array of events into one row per array element.""" + + path: str + """Name of the root-level key that holds the array of events, for example ``Records``. Nested paths + are not supported.""" + + def as_dict(self) -> dict: + """Serializes the JsonArrayEnvelope into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.path is not None: + body["path"] = self.path + return body + + def as_shallow_dict(self) -> dict: + """Serializes the JsonArrayEnvelope into a shallow dictionary of its immediate attributes.""" + body = {} + if self.path is not None: + body["path"] = self.path + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> JsonArrayEnvelope: + """Deserializes the JsonArrayEnvelope from a dictionary.""" + return cls(path=d.get("path", None)) + + +@dataclass +class ListDatasourcesResponse: + datasources: Optional[List[Datasource]] = None + + next_page_token: Optional[str] = None + """Opaque token for fetching the next page; pass it as ``page_token`` on the next call. It is empty + only when there are no more results. Page size is best-effort, so a short or even empty page + does not mean the end; keep paging until this token is empty, which is the only reliable + end-of-results signal.""" + + def as_dict(self) -> dict: + """Serializes the ListDatasourcesResponse into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.datasources: + body["datasources"] = [v.as_dict() for v in self.datasources] + if self.next_page_token is not None: + body["next_page_token"] = self.next_page_token + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ListDatasourcesResponse into a shallow dictionary of its immediate attributes.""" + body = {} + if self.datasources: + body["datasources"] = self.datasources + if self.next_page_token is not None: + body["next_page_token"] = self.next_page_token + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ListDatasourcesResponse: + """Deserializes the ListDatasourcesResponse from a dictionary.""" + return cls( + datasources=_repeated_dict(d, "datasources", Datasource), next_page_token=d.get("next_page_token", None) + ) + + +@dataclass +class ListRulesResponse: + next_page_token: Optional[str] = None + """Token to pass as ``page_token`` in a subsequent ``ListRules`` call to fetch the next page, or + empty when there are no more results.""" + + rules: Optional[List[Rule]] = None + """The rules in this page of results.""" + + def as_dict(self) -> dict: + """Serializes the ListRulesResponse into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.next_page_token is not None: + body["next_page_token"] = self.next_page_token + if self.rules: + body["rules"] = [v.as_dict() for v in self.rules] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ListRulesResponse into a shallow dictionary of its immediate attributes.""" + body = {} + if self.next_page_token is not None: + body["next_page_token"] = self.next_page_token + if self.rules: + body["rules"] = self.rules + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ListRulesResponse: + """Deserializes the ListRulesResponse from a dictionary.""" + return cls(next_page_token=d.get("next_page_token", None), rules=_repeated_dict(d, "rules", Rule)) + + +@dataclass +class ParserRef: + """Reference to a parser that the datasource uses. Each parser defines the schema of the parsed + table that it writes.""" + + id: str + """Identifier of the parser. Stable across renames, unlike a parser name.""" + + version: str + """Version of the parser to use, given as an explicit version number.""" + + effective_target_table: Optional[TableName] = None + """The table this parser's data is written to: ``target_table`` when supplied, otherwise the table + derived from the datasource's catalog and schema and the parser name. Fixed once resolved, so a + later parser rename does not move it. Populated on get, create, and update responses with the + resolved table.""" + + target_table: Optional[TableName] = None + """Table that this parser's data is written to. The table must be in the datasource's catalog and + schema. If unset, the table name is derived from the parser name. Reads return the value as + supplied; the table the data is written to is reported in ``effective_target_table``.""" + + def as_dict(self) -> dict: + """Serializes the ParserRef into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.effective_target_table: + body["effective_target_table"] = self.effective_target_table.as_dict() + if self.id is not None: + body["id"] = self.id + if self.target_table: + body["target_table"] = self.target_table.as_dict() + if self.version is not None: + body["version"] = self.version + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ParserRef into a shallow dictionary of its immediate attributes.""" + body = {} + if self.effective_target_table: + body["effective_target_table"] = self.effective_target_table + if self.id is not None: + body["id"] = self.id + if self.target_table: + body["target_table"] = self.target_table + if self.version is not None: + body["version"] = self.version + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ParserRef: + """Deserializes the ParserRef from a dictionary.""" + return cls( + effective_target_table=_from_dict(d, "effective_target_table", TableName), + id=d.get("id", None), + target_table=_from_dict(d, "target_table", TableName), + version=d.get("version", None), + ) + + +@dataclass +class Principal: + """A principal — user or service-principal — used as the value type for ``created_by`` / + ``updated_by`` / ``assignee`` etc. across v2 resources.""" + + display_name: Optional[str] = None + """Display name (user or SP name).""" + + email: Optional[str] = None + """Email (regular users only).""" + + id: Optional[str] = None + """User / service-principal ID.""" + + principal_type: Optional[PrincipalType] = None + """Whether this principal is a regular user, service principal, or system principal.""" + + def as_dict(self) -> dict: + """Serializes the Principal into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.display_name is not None: + body["display_name"] = self.display_name + if self.email is not None: + body["email"] = self.email + if self.id is not None: + body["id"] = self.id + if self.principal_type is not None: + body["principal_type"] = self.principal_type.value + return body + + def as_shallow_dict(self) -> dict: + """Serializes the Principal into a shallow dictionary of its immediate attributes.""" + body = {} + if self.display_name is not None: + body["display_name"] = self.display_name + if self.email is not None: + body["email"] = self.email + if self.id is not None: + body["id"] = self.id + if self.principal_type is not None: + body["principal_type"] = self.principal_type + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> Principal: + """Deserializes the Principal from a dictionary.""" + return cls( + display_name=d.get("display_name", None), + email=d.get("email", None), + id=d.get("id", None), + principal_type=_enum(d, "principal_type", PrincipalType), + ) + + +class PrincipalType(Enum): + """Type of a principal (user / SP / system).""" + + REGULAR_USER = "REGULAR_USER" + SERVICE_PRINCIPAL = "SERVICE_PRINCIPAL" + SYSTEM_USER = "SYSTEM_USER" + + +@dataclass +class ResourceStatus: + """Read-only runtime status of one job or pipeline backing a resource. + + Shared across resources: a Rule's single job, a Datasource's main / health / + historical-ingestion jobs, and a Datasource's ingestion pipeline (the ``pipeline_id`` field is + set only for the pipeline). For a datasource's historical ingestion, the shard count is simply + the number of ``historical_job_statuses`` entries — there is no separate count field. + + Two orthogonal axes are modeled separately (v1 conflated them into a single ``jobStatus`` enum): + ``lifecycle`` (is it configured / does it still exist) and ``run_state`` (what the most recent + or in-flight run did). + + Every field that holds a ``ResourceStatus`` is OUTPUT_ONLY at its use site (Datasource + ingest/health/historical/pipeline statuses, Rule status), so the inner fields don't repeat the + annotation.""" + + create_time: Optional[Timestamp] = None + + job_id: Optional[str] = None + """Databricks Jobs API job ID. Unset for a pure ingestion pipeline.""" + + last_complete_time: Optional[Timestamp] = None + """When the most recent run completed. Pairs with ``run_state = SUCCEEDED``.""" + + last_error_message: Optional[str] = None + """Most recent error message, if any. Pairs with ``run_state = FAILED``.""" + + last_start_time: Optional[Timestamp] = None + """When the most recent run started. Pairs with ``run_state = RUNNING``.""" + + last_status_change_time: Optional[Timestamp] = None + + lifecycle_state: Optional[ResourceStatusLifecycle] = None + """Scheduling lifecycle of the resource's underlying job/pipeline: whether it is unscheduled, + scheduled, stopped, or deleted. Distinct from ``run_state``, which reports the outcome of the + most recent run.""" + + pipeline_id: Optional[str] = None + """DLT ingestion-pipeline ID. Set only when this status represents the datasource's ingestion + pipeline (unset for job statuses).""" + + run_state: Optional[ResourceStatusRunState] = None + + def as_dict(self) -> dict: + """Serializes the ResourceStatus into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.create_time is not None: + body["create_time"] = self.create_time.ToJsonString() + if self.job_id is not None: + body["job_id"] = self.job_id + if self.last_complete_time is not None: + body["last_complete_time"] = self.last_complete_time.ToJsonString() + if self.last_error_message is not None: + body["last_error_message"] = self.last_error_message + if self.last_start_time is not None: + body["last_start_time"] = self.last_start_time.ToJsonString() + if self.last_status_change_time is not None: + body["last_status_change_time"] = self.last_status_change_time.ToJsonString() + if self.lifecycle_state is not None: + body["lifecycle_state"] = self.lifecycle_state.value + if self.pipeline_id is not None: + body["pipeline_id"] = self.pipeline_id + if self.run_state is not None: + body["run_state"] = self.run_state.value + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ResourceStatus into a shallow dictionary of its immediate attributes.""" + body = {} + if self.create_time is not None: + body["create_time"] = self.create_time + if self.job_id is not None: + body["job_id"] = self.job_id + if self.last_complete_time is not None: + body["last_complete_time"] = self.last_complete_time + if self.last_error_message is not None: + body["last_error_message"] = self.last_error_message + if self.last_start_time is not None: + body["last_start_time"] = self.last_start_time + if self.last_status_change_time is not None: + body["last_status_change_time"] = self.last_status_change_time + if self.lifecycle_state is not None: + body["lifecycle_state"] = self.lifecycle_state + if self.pipeline_id is not None: + body["pipeline_id"] = self.pipeline_id + if self.run_state is not None: + body["run_state"] = self.run_state + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ResourceStatus: + """Deserializes the ResourceStatus from a dictionary.""" + return cls( + create_time=_timestamp(d, "create_time"), + job_id=d.get("job_id", None), + last_complete_time=_timestamp(d, "last_complete_time"), + last_error_message=d.get("last_error_message", None), + last_start_time=_timestamp(d, "last_start_time"), + last_status_change_time=_timestamp(d, "last_status_change_time"), + lifecycle_state=_enum(d, "lifecycle_state", ResourceStatusLifecycle), + pipeline_id=d.get("pipeline_id", None), + run_state=_enum(d, "run_state", ResourceStatusRunState), + ) + + +class ResourceStatusLifecycle(Enum): + """Configured scheduling / existence of the job or pipeline.""" + + DELETED = "DELETED" + SCHEDULED = "SCHEDULED" + STOPPED = "STOPPED" + UNSCHEDULED = "UNSCHEDULED" + + +class ResourceStatusRunState(Enum): + """Outcome / progress of the most recent (or in-flight) run.""" + + CANCELED = "CANCELED" + FAILED = "FAILED" + RUNNING = "RUNNING" + SUCCEEDED = "SUCCEEDED" + + +@dataclass +class Rule: + """A Lakewatch detection rule: a query that evaluates ingested data and generates a signal for each + match. Each rule is a securable in a Unity Catalog schema.""" + + catalog_name: str + """Name of the Unity Catalog catalog that contains the rule. Required. Cannot be changed after + creation. Reads return the value as supplied; the catalog's current name is reported in + ``effective_catalog_name``.""" + + schema_name: str + """Name of the Unity Catalog schema that contains the rule. Required. Cannot be changed after + creation. Reads return the value as supplied; the schema's current name is reported in + ``effective_schema_name``.""" + + rule_name: str + """Name of the rule within its Unity Catalog schema. Required. Cannot be changed after creation. + Reads return the value as supplied; the current UC name is reported in ``effective_rule_name``.""" + + severity: str + """Free-form severity label, such as ``Informational``, ``Medium``, ``High``, or ``Critical``.""" + + fidelity: RuleFidelity + """Signal quality for this rule.""" + + schedule: RuleSchedule + """When the rule executes.""" + + output: RuleOutput + """What to record in each signal.""" + + annotations: Optional[Dict[str, str]] = None + """Free-form client-managed key/value annotations.""" + + batch: Optional[RuleInputBatch] = None + """Batch SQL rule input.""" + + category: Optional[str] = None + """Free-form rule category.""" + + compute_workspace_id: Optional[str] = None + """ID of the workspace whose compute runs this rule's jobs and owns its associated resources, such + as the SQL warehouse a scheduled rule runs on. Server-assigned; a rule's warehouse can only be + changed from this workspace.""" + + content_resource_id: Optional[str] = None + """Immutable UUIDv4 identifying the authored content resource installed as this rule.""" + + create_time: Optional[Timestamp] = None + """Time the rule was created.""" + + created_by: Optional[Principal] = None + """Principal that created the rule.""" + + deduplication: Optional[RuleDeduplication] = None + """Signal deduplication mode.""" + + deduplication_time_window: Optional[str] = None + """Time window for deduplication, such as ``1h``, ``24h``, ``7d``, or ``automatic``.""" + + description: Optional[str] = None + """Free-form description.""" + + effective_catalog_name: Optional[str] = None + """Current name of the Unity Catalog catalog that contains the rule. Reflects catalog renames, + which do not change ``catalog_name``. Unset until resolved.""" + + effective_rule_name: Optional[str] = None + """Current UC leaf name of the rule (server-owned). Tracks UC renames so callers can observe the + live name without having to update the client-owned ``rule_name``; unset until resolved.""" + + effective_run_as_principal: Optional[str] = None + """The identity this rule's job runs as: the resolved ``run_as_principal``, or the requester's + identity when it was omitted on create or cleared on update. The value is a user email or + service-principal application ID, falling back to the numeric principal ID if that identity can + no longer be resolved. Populated on single-rule reads (get/create/update); list responses omit + it.""" + + effective_schema_name: Optional[str] = None + """Current name of the Unity Catalog schema that contains the rule. Reflects schema renames, which + do not change ``schema_name``. Unset until resolved.""" + + etag: Optional[str] = None + """Opaque entity tag for optimistic concurrency. Content-version ETags change on user edits and + remain stable across runtime status updates.""" + + mitre: Optional[List[RuleMitreMapping]] = None + """Up to 64 MITRE ATT&CK mappings, producing at most 192 flattened alert tags.""" + + name: Optional[str] = None + """Resource name, in the form ``rules/{rule_id}``. Assigned by the server on create; required on + update.""" + + objective: Optional[str] = None + """Long-form description of what this rule detects.""" + + original_content_pack: Optional[ContentPackReference] = None + """Reference to the content pack this rule was originally installed from.""" + + owner: Optional[Principal] = None + """Unity Catalog owner of the rule. Assigned by the server. Populated by ``GetRule``.""" + + references: Optional[List[str]] = None + """Ordered reference URLs carried from the authored rule metadata. A maximum of three references of + up to 512 characters each is supported. Requests exceeding these limits are rejected.""" + + rule_id: Optional[str] = None + """Unity Catalog securable UUID of the rule, and the last segment of ``name``. Never changes, + including when the rule's catalog or schema is renamed.""" + + run_as_principal: Optional[str] = None + """The principal whose identity this rule's job runs as, given as a user email or service-principal + application ID. When supplied on create/update, Lakewatch resolves the principal, records it on + the rule, and creates the delegation to run as that principal — the requester must be + authorized to delegate to it, or the request is rejected. When omitted on create, the + requester's own identity is used. An update whose mask does not include ``run_as_principal`` + keeps the current identity, and the requester must still be authorized to delegate to it; an + editor who is not must supply a run_as_principal they can delegate to (for example, their own + identity), or clear it, either of which reassigns the rule. An update whose mask includes + ``run_as_principal`` but supplies no value clears it: the rule then runs as the requester, + exactly as when it is omitted on create. Reads return the value as supplied; the identity the + rule runs as is reported in ``effective_run_as_principal``.""" + + runbook: Optional[str] = None + """Long-form operational guidance carried from the authored rule metadata. Supports up to 4096 + characters. Requests exceeding this limit are rejected.""" + + status: Optional[ResourceStatus] = None + """Read-only lifecycle status.""" + + suppress_until_time: Optional[Timestamp] = None + """If set, this rule's signals do not open new alerts until the given time.""" + + timerange: Optional[RuleTimeRange] = None + """Time-range configuration for batch rules. If unset, the execution time range is extended by a + fixed ``15m`` lookback.""" + + update_time: Optional[Timestamp] = None + """Time the rule was last updated.""" + + updated_by: Optional[Principal] = None + """Principal that last updated the rule.""" + + version: Optional[int] = None + """Monotonic version, bumped on each update. Informational; use ``etag`` for optimistic + concurrency.""" + + def as_dict(self) -> dict: + """Serializes the Rule into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.annotations: + body["annotations"] = self.annotations + if self.batch: + body["batch"] = self.batch.as_dict() + if self.catalog_name is not None: + body["catalog_name"] = self.catalog_name + if self.category is not None: + body["category"] = self.category + if self.compute_workspace_id is not None: + body["compute_workspace_id"] = self.compute_workspace_id + if self.content_resource_id is not None: + body["content_resource_id"] = self.content_resource_id + if self.create_time is not None: + body["create_time"] = self.create_time.ToJsonString() + if self.created_by: + body["created_by"] = self.created_by.as_dict() + if self.deduplication is not None: + body["deduplication"] = self.deduplication.value + if self.deduplication_time_window is not None: + body["deduplication_time_window"] = self.deduplication_time_window + if self.description is not None: + body["description"] = self.description + if self.effective_catalog_name is not None: + body["effective_catalog_name"] = self.effective_catalog_name + if self.effective_rule_name is not None: + body["effective_rule_name"] = self.effective_rule_name + if self.effective_run_as_principal is not None: + body["effective_run_as_principal"] = self.effective_run_as_principal + if self.effective_schema_name is not None: + body["effective_schema_name"] = self.effective_schema_name + if self.etag is not None: + body["etag"] = self.etag + if self.fidelity is not None: + body["fidelity"] = self.fidelity.value + if self.mitre: + body["mitre"] = [v.as_dict() for v in self.mitre] + if self.name is not None: + body["name"] = self.name + if self.objective is not None: + body["objective"] = self.objective + if self.original_content_pack: + body["original_content_pack"] = self.original_content_pack.as_dict() + if self.output: + body["output"] = self.output.as_dict() + if self.owner: + body["owner"] = self.owner.as_dict() + if self.references: + body["references"] = [v for v in self.references] + if self.rule_id is not None: + body["rule_id"] = self.rule_id + if self.rule_name is not None: + body["rule_name"] = self.rule_name + if self.run_as_principal is not None: + body["run_as_principal"] = self.run_as_principal + if self.runbook is not None: + body["runbook"] = self.runbook + if self.schedule: + body["schedule"] = self.schedule.as_dict() + if self.schema_name is not None: + body["schema_name"] = self.schema_name + if self.severity is not None: + body["severity"] = self.severity + if self.status: + body["status"] = self.status.as_dict() + if self.suppress_until_time is not None: + body["suppress_until_time"] = self.suppress_until_time.ToJsonString() + if self.timerange: + body["timerange"] = self.timerange.as_dict() + if self.update_time is not None: + body["update_time"] = self.update_time.ToJsonString() + if self.updated_by: + body["updated_by"] = self.updated_by.as_dict() + if self.version is not None: + body["version"] = self.version + return body + + def as_shallow_dict(self) -> dict: + """Serializes the Rule into a shallow dictionary of its immediate attributes.""" + body = {} + if self.annotations: + body["annotations"] = self.annotations + if self.batch: + body["batch"] = self.batch + if self.catalog_name is not None: + body["catalog_name"] = self.catalog_name + if self.category is not None: + body["category"] = self.category + if self.compute_workspace_id is not None: + body["compute_workspace_id"] = self.compute_workspace_id + if self.content_resource_id is not None: + body["content_resource_id"] = self.content_resource_id + if self.create_time is not None: + body["create_time"] = self.create_time + if self.created_by: + body["created_by"] = self.created_by + if self.deduplication is not None: + body["deduplication"] = self.deduplication + if self.deduplication_time_window is not None: + body["deduplication_time_window"] = self.deduplication_time_window + if self.description is not None: + body["description"] = self.description + if self.effective_catalog_name is not None: + body["effective_catalog_name"] = self.effective_catalog_name + if self.effective_rule_name is not None: + body["effective_rule_name"] = self.effective_rule_name + if self.effective_run_as_principal is not None: + body["effective_run_as_principal"] = self.effective_run_as_principal + if self.effective_schema_name is not None: + body["effective_schema_name"] = self.effective_schema_name + if self.etag is not None: + body["etag"] = self.etag + if self.fidelity is not None: + body["fidelity"] = self.fidelity + if self.mitre: + body["mitre"] = self.mitre + if self.name is not None: + body["name"] = self.name + if self.objective is not None: + body["objective"] = self.objective + if self.original_content_pack: + body["original_content_pack"] = self.original_content_pack + if self.output: + body["output"] = self.output + if self.owner: + body["owner"] = self.owner + if self.references: + body["references"] = self.references + if self.rule_id is not None: + body["rule_id"] = self.rule_id + if self.rule_name is not None: + body["rule_name"] = self.rule_name + if self.run_as_principal is not None: + body["run_as_principal"] = self.run_as_principal + if self.runbook is not None: + body["runbook"] = self.runbook + if self.schedule: + body["schedule"] = self.schedule + if self.schema_name is not None: + body["schema_name"] = self.schema_name + if self.severity is not None: + body["severity"] = self.severity + if self.status: + body["status"] = self.status + if self.suppress_until_time is not None: + body["suppress_until_time"] = self.suppress_until_time + if self.timerange: + body["timerange"] = self.timerange + if self.update_time is not None: + body["update_time"] = self.update_time + if self.updated_by: + body["updated_by"] = self.updated_by + if self.version is not None: + body["version"] = self.version + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> Rule: + """Deserializes the Rule from a dictionary.""" + return cls( + annotations=d.get("annotations", None), + batch=_from_dict(d, "batch", RuleInputBatch), + catalog_name=d.get("catalog_name", None), + category=d.get("category", None), + compute_workspace_id=d.get("compute_workspace_id", None), + content_resource_id=d.get("content_resource_id", None), + create_time=_timestamp(d, "create_time"), + created_by=_from_dict(d, "created_by", Principal), + deduplication=_enum(d, "deduplication", RuleDeduplication), + deduplication_time_window=d.get("deduplication_time_window", None), + description=d.get("description", None), + effective_catalog_name=d.get("effective_catalog_name", None), + effective_rule_name=d.get("effective_rule_name", None), + effective_run_as_principal=d.get("effective_run_as_principal", None), + effective_schema_name=d.get("effective_schema_name", None), + etag=d.get("etag", None), + fidelity=_enum(d, "fidelity", RuleFidelity), + mitre=_repeated_dict(d, "mitre", RuleMitreMapping), + name=d.get("name", None), + objective=d.get("objective", None), + original_content_pack=_from_dict(d, "original_content_pack", ContentPackReference), + output=_from_dict(d, "output", RuleOutput), + owner=_from_dict(d, "owner", Principal), + references=d.get("references", None), + rule_id=d.get("rule_id", None), + rule_name=d.get("rule_name", None), + run_as_principal=d.get("run_as_principal", None), + runbook=d.get("runbook", None), + schedule=_from_dict(d, "schedule", RuleSchedule), + schema_name=d.get("schema_name", None), + severity=d.get("severity", None), + status=_from_dict(d, "status", ResourceStatus), + suppress_until_time=_timestamp(d, "suppress_until_time"), + timerange=_from_dict(d, "timerange", RuleTimeRange), + update_time=_timestamp(d, "update_time"), + updated_by=_from_dict(d, "updated_by", Principal), + version=_int64(d, "version"), + ) + + +@dataclass +class RuleConfig: + """Per-catalog rule-related config. Singleton. Name format: ``catalogs/{catalog}/ruleConfig``.""" + + default_folder: Optional[str] = None + """Schema name of the registered folder (in this catalog) that a new rule is placed in when the + request does not specify a schema. Must name a schema registered as a folder. Omit to clear. + Reads return the value as supplied; the folder's current schema name is reported in + ``effective_default_folder``.""" + + default_schedule: Optional[RuleSchedule] = None + """Default schedule for rules that don't specify one.""" + + default_signal_collation_window: Optional[str] = None + """Default duration string for signal deduplication when a rule doesn't override it — compact + day/hour/minute/second components, e.g. ``5d``, ``1h``. Reads return the value as supplied; + there is no server default, so it is unset until set. An empty value clears it.""" + + detection_categories: Optional[List[str]] = None + """Available detection categories a rule can be tagged with. Reads return the value as supplied. + When empty, rules can be tagged with the default categories: ``Static Signature``, ``Fuzzy + Signature``, ``Behavioral``, ``Statistical``, and ``Learning (AI/ML)``. The categories in effect + are reported in ``effective_detection_categories``.""" + + effective_default_folder: Optional[str] = None + """The current schema name of the ``default_folder`` folder. Tracks UC schema renames. Unset when + no default is configured, or when the folder's schema no longer exists or is not visible to the + caller.""" + + effective_detection_categories: Optional[List[str]] = None + """The detection categories in effect: ``detection_categories`` when set, otherwise the default + categories.""" + + etag: Optional[str] = None + """Opaque entity tag for optimistic concurrency. Returned on every read; pass it back as the Update + request's ``etag`` to make the update conditional on the resource being unchanged.""" + + name: Optional[str] = None + + update_time: Optional[Timestamp] = None + + updated_by: Optional[Principal] = None + + version: Optional[int] = None + """Monotonic version, bumped on each update. Informational; use ``etag`` for optimistic + concurrency.""" + + def as_dict(self) -> dict: + """Serializes the RuleConfig into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.default_folder is not None: + body["default_folder"] = self.default_folder + if self.default_schedule: + body["default_schedule"] = self.default_schedule.as_dict() + if self.default_signal_collation_window is not None: + body["default_signal_collation_window"] = self.default_signal_collation_window + if self.detection_categories: + body["detection_categories"] = [v for v in self.detection_categories] + if self.effective_default_folder is not None: + body["effective_default_folder"] = self.effective_default_folder + if self.effective_detection_categories: + body["effective_detection_categories"] = [v for v in self.effective_detection_categories] + if self.etag is not None: + body["etag"] = self.etag + if self.name is not None: + body["name"] = self.name + if self.update_time is not None: + body["update_time"] = self.update_time.ToJsonString() + if self.updated_by: + body["updated_by"] = self.updated_by.as_dict() + if self.version is not None: + body["version"] = self.version + return body + + def as_shallow_dict(self) -> dict: + """Serializes the RuleConfig into a shallow dictionary of its immediate attributes.""" + body = {} + if self.default_folder is not None: + body["default_folder"] = self.default_folder + if self.default_schedule: + body["default_schedule"] = self.default_schedule + if self.default_signal_collation_window is not None: + body["default_signal_collation_window"] = self.default_signal_collation_window + if self.detection_categories: + body["detection_categories"] = self.detection_categories + if self.effective_default_folder is not None: + body["effective_default_folder"] = self.effective_default_folder + if self.effective_detection_categories: + body["effective_detection_categories"] = self.effective_detection_categories + if self.etag is not None: + body["etag"] = self.etag + if self.name is not None: + body["name"] = self.name + if self.update_time is not None: + body["update_time"] = self.update_time + if self.updated_by: + body["updated_by"] = self.updated_by + if self.version is not None: + body["version"] = self.version + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> RuleConfig: + """Deserializes the RuleConfig from a dictionary.""" + return cls( + default_folder=d.get("default_folder", None), + default_schedule=_from_dict(d, "default_schedule", RuleSchedule), + default_signal_collation_window=d.get("default_signal_collation_window", None), + detection_categories=d.get("detection_categories", None), + effective_default_folder=d.get("effective_default_folder", None), + effective_detection_categories=d.get("effective_detection_categories", None), + etag=d.get("etag", None), + name=d.get("name", None), + update_time=_timestamp(d, "update_time"), + updated_by=_from_dict(d, "updated_by", Principal), + version=_int64(d, "version"), + ) + + +class RuleDeduplication(Enum): + """Signal deduplication mode for a detection rule.""" + + DEDUPLICATION_DEFAULT = "DEDUPLICATION_DEFAULT" + DEDUPLICATION_DISABLED = "DEDUPLICATION_DISABLED" + DEDUPLICATION_ENABLED = "DEDUPLICATION_ENABLED" + + +class RuleFidelity(Enum): + """Maturity of a detection rule.""" + + FIDELITY_DEVELOPMENT = "FIDELITY_DEVELOPMENT" + FIDELITY_HIGH = "FIDELITY_HIGH" + FIDELITY_INVESTIGATIVE = "FIDELITY_INVESTIGATIVE" + FIDELITY_MEDIUM = "FIDELITY_MEDIUM" + + +@dataclass +class RuleInputBatch: + """Batch input for a detection rule.""" + + custom: Optional[RuleInputCustomNotebook] = None + """Custom-notebook user function. The output is used as input data for the rule.""" + + sql: Optional[RuleInputBatchSql] = None + """SQL query evaluated as the rule input.""" + + warehouse_id: Optional[str] = None + """SQL warehouse used to run this batch rule.""" + + def as_dict(self) -> dict: + """Serializes the RuleInputBatch into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.custom: + body["custom"] = self.custom.as_dict() + if self.sql: + body["sql"] = self.sql.as_dict() + if self.warehouse_id is not None: + body["warehouse_id"] = self.warehouse_id + return body + + def as_shallow_dict(self) -> dict: + """Serializes the RuleInputBatch into a shallow dictionary of its immediate attributes.""" + body = {} + if self.custom: + body["custom"] = self.custom + if self.sql: + body["sql"] = self.sql + if self.warehouse_id is not None: + body["warehouse_id"] = self.warehouse_id + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> RuleInputBatch: + """Deserializes the RuleInputBatch from a dictionary.""" + return cls( + custom=_from_dict(d, "custom", RuleInputCustomNotebook), + sql=_from_dict(d, "sql", RuleInputBatchSql), + warehouse_id=d.get("warehouse_id", None), + ) + + +@dataclass +class RuleInputBatchSql: + """Batch SQL input for a detection rule.""" + + query: str + """SQL query evaluated as the rule input.""" + + default_catalog_name: Optional[str] = None + """Default catalog used to resolve unqualified identifiers in ``query``. If unset, the rule's + ``catalog_name`` is used.""" + + default_schema_name: Optional[str] = None + """Default schema used to resolve unqualified identifiers in ``query``. If unset, the rule's + ``schema_name`` is used.""" + + def as_dict(self) -> dict: + """Serializes the RuleInputBatchSql into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.default_catalog_name is not None: + body["default_catalog_name"] = self.default_catalog_name + if self.default_schema_name is not None: + body["default_schema_name"] = self.default_schema_name + if self.query is not None: + body["query"] = self.query + return body + + def as_shallow_dict(self) -> dict: + """Serializes the RuleInputBatchSql into a shallow dictionary of its immediate attributes.""" + body = {} + if self.default_catalog_name is not None: + body["default_catalog_name"] = self.default_catalog_name + if self.default_schema_name is not None: + body["default_schema_name"] = self.default_schema_name + if self.query is not None: + body["query"] = self.query + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> RuleInputBatchSql: + """Deserializes the RuleInputBatchSql from a dictionary.""" + return cls( + default_catalog_name=d.get("default_catalog_name", None), + default_schema_name=d.get("default_schema_name", None), + query=d.get("query", None), + ) + + +@dataclass +class RuleInputCustomNotebook: + """Reference to a custom-user-function notebook for stream or batch input.""" + + notebook: str + """Fully-qualified notebook path, or a path relative to the configured base notebook directory.""" + + options: Optional[Dict[str, str]] = None + """Optional arguments supplied when invoking the custom function.""" + + def as_dict(self) -> dict: + """Serializes the RuleInputCustomNotebook into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.notebook is not None: + body["notebook"] = self.notebook + if self.options: + body["options"] = self.options + return body + + def as_shallow_dict(self) -> dict: + """Serializes the RuleInputCustomNotebook into a shallow dictionary of its immediate attributes.""" + body = {} + if self.notebook is not None: + body["notebook"] = self.notebook + if self.options: + body["options"] = self.options + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> RuleInputCustomNotebook: + """Deserializes the RuleInputCustomNotebook from a dictionary.""" + return cls(notebook=d.get("notebook", None), options=d.get("options", None)) + + +@dataclass +class RuleMitreMapping: + """MITRE ATT&CK mapping for a detection rule.""" + + sub_technique: Optional[str] = None + """Human-readable MITRE sub-technique name.""" + + sub_technique_id: Optional[str] = None + """MITRE sub-technique ID, such as ``T1078.001``.""" + + tactic: Optional[str] = None + """Adversary objective during this attack, such as ``Command and Control``.""" + + taxonomy: Optional[str] = None + """Structured framework that organizes adversary behaviors, such as ``mitre-attack``.""" + + technique: Optional[str] = None + """Human-readable MITRE technique name.""" + + technique_id: Optional[str] = None + """MITRE technique ID, such as ``T1078``.""" + + def as_dict(self) -> dict: + """Serializes the RuleMitreMapping into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.sub_technique is not None: + body["sub_technique"] = self.sub_technique + if self.sub_technique_id is not None: + body["sub_technique_id"] = self.sub_technique_id + if self.tactic is not None: + body["tactic"] = self.tactic + if self.taxonomy is not None: + body["taxonomy"] = self.taxonomy + if self.technique is not None: + body["technique"] = self.technique + if self.technique_id is not None: + body["technique_id"] = self.technique_id + return body + + def as_shallow_dict(self) -> dict: + """Serializes the RuleMitreMapping into a shallow dictionary of its immediate attributes.""" + body = {} + if self.sub_technique is not None: + body["sub_technique"] = self.sub_technique + if self.sub_technique_id is not None: + body["sub_technique_id"] = self.sub_technique_id + if self.tactic is not None: + body["tactic"] = self.tactic + if self.taxonomy is not None: + body["taxonomy"] = self.taxonomy + if self.technique is not None: + body["technique"] = self.technique + if self.technique_id is not None: + body["technique_id"] = self.technique_id + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> RuleMitreMapping: + """Deserializes the RuleMitreMapping from a dictionary.""" + return cls( + sub_technique=d.get("sub_technique", None), + sub_technique_id=d.get("sub_technique_id", None), + tactic=d.get("tactic", None), + taxonomy=d.get("taxonomy", None), + technique=d.get("technique", None), + technique_id=d.get("technique_id", None), + ) + + +@dataclass +class RuleOutput: + """What to record in each signal that this rule generates.""" + + context: Optional[Dict[str, str]] = None + """Key/value context to record from the input row. Values may contain SQL expressions in braces.""" + + default_context: Optional[bool] = None + """If true, also record every column of the triggering input row. Custom context wins on key + collision.""" + + event_time_column: Optional[str] = None + """Column in the rule output that holds the event timestamp.""" + + summary: Optional[str] = None + """Signal summary. May contain SQL expressions in braces that the server substitutes from the input + row.""" + + def as_dict(self) -> dict: + """Serializes the RuleOutput into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.context: + body["context"] = self.context + if self.default_context is not None: + body["default_context"] = self.default_context + if self.event_time_column is not None: + body["event_time_column"] = self.event_time_column + if self.summary is not None: + body["summary"] = self.summary + return body + + def as_shallow_dict(self) -> dict: + """Serializes the RuleOutput into a shallow dictionary of its immediate attributes.""" + body = {} + if self.context: + body["context"] = self.context + if self.default_context is not None: + body["default_context"] = self.default_context + if self.event_time_column is not None: + body["event_time_column"] = self.event_time_column + if self.summary is not None: + body["summary"] = self.summary + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> RuleOutput: + """Deserializes the RuleOutput from a dictionary.""" + return cls( + context=d.get("context", None), + default_context=d.get("default_context", None), + event_time_column=d.get("event_time_column", None), + summary=d.get("summary", None), + ) + + +@dataclass +class RuleSchedule: + """Defines when a detection rule executes. + + A batch rule sets exactly one of ``at_least_every`` or ``cron_expression``. A streaming rule + runs continuously and sets no timing; it accepts only ``enabled``, and an omitted schedule means + enabled. The server rejects any other combination. ``enabled`` is orthogonal to the timing.""" + + at_least_every: Optional[str] = None + """Defines the minimum cadence as a duration string. It may contain components for days, hours, + minutes, and seconds (e.g. ``5d``, ``5d3h``, ``3d2h5m12s``, ``12m16s``).""" + + continuous: Optional[bool] = None + """Runs the rule continuously. Detection rules do not accept it: batch rules reject it, and + streaming rules run continuously without it.""" + + cron_expression: Optional[str] = None + """Defines the schedule using a Quartz cron expression.""" + + enabled: Optional[bool] = None + """Sets whether the job is scheduled to run. Defaults to true.""" + + def as_dict(self) -> dict: + """Serializes the RuleSchedule into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.at_least_every is not None: + body["at_least_every"] = self.at_least_every + if self.continuous is not None: + body["continuous"] = self.continuous + if self.cron_expression is not None: + body["cron_expression"] = self.cron_expression + if self.enabled is not None: + body["enabled"] = self.enabled + return body + + def as_shallow_dict(self) -> dict: + """Serializes the RuleSchedule into a shallow dictionary of its immediate attributes.""" + body = {} + if self.at_least_every is not None: + body["at_least_every"] = self.at_least_every + if self.continuous is not None: + body["continuous"] = self.continuous + if self.cron_expression is not None: + body["cron_expression"] = self.cron_expression + if self.enabled is not None: + body["enabled"] = self.enabled + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> RuleSchedule: + """Deserializes the RuleSchedule from a dictionary.""" + return cls( + at_least_every=d.get("at_least_every", None), + continuous=d.get("continuous", None), + cron_expression=d.get("cron_expression", None), + enabled=d.get("enabled", None), + ) + + +@dataclass +class RuleTimeRange: + """Time-range configuration for a batch detection rule.""" + + extend_lookback: Optional[str] = None + """Extends the start of the execution time range by a compact duration with integer days, hours, + minutes, and seconds in that order, such as ``1h30m`` or ``15m``. Zero duration disables + extension. Unset or empty uses a fixed ``15m`` at execution time. Only the start moves; the end + is unchanged. Has no effect on batch SQL that does not use ``IN_TIMERANGE``. Values must fit + signed 64-bit microseconds.""" + + def as_dict(self) -> dict: + """Serializes the RuleTimeRange into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.extend_lookback is not None: + body["extend_lookback"] = self.extend_lookback + return body + + def as_shallow_dict(self) -> dict: + """Serializes the RuleTimeRange into a shallow dictionary of its immediate attributes.""" + body = {} + if self.extend_lookback is not None: + body["extend_lookback"] = self.extend_lookback + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> RuleTimeRange: + """Deserializes the RuleTimeRange from a dictionary.""" + return cls(extend_lookback=d.get("extend_lookback", None)) + + +class RuleView(Enum): + """Selects how much of each rule a list response returns.""" + + RULE_VIEW_BASIC = "RULE_VIEW_BASIC" + RULE_VIEW_FULL = "RULE_VIEW_FULL" + + +@dataclass +class Schedule: + """Defines when a recurring job executes. + + Exactly one of the timing modes (``at_least_every``, ``cron_expression``, ``continuous``) must + be set; combining them or setting none is rejected by the server. ``compute_group`` and + ``enabled`` are orthogonal to the timing.""" + + at_least_every: Optional[str] = None + """Duration string defining the minimum cadence. May contain components for days, hours, minutes, + seconds (e.g. ``5d``, ``5d3h``, ``3d2h5m12s``, ``12m16s``).""" + + compute_group: Optional[str] = None + """``dedicated``, ``inherited``, ``automatic``, or a custom group name. Reads return the value as + supplied. A rule whose schedule omits it runs on ``dedicated``.""" + + continuous: Optional[bool] = None + """Run continuously (sentinel — must be ``true`` when set in this branch).""" + + cron_expression: Optional[str] = None + """Quartz cron expression.""" + + enabled: Optional[bool] = None + """Whether the job is scheduled to run. Defaults to true.""" + + def as_dict(self) -> dict: + """Serializes the Schedule into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.at_least_every is not None: + body["at_least_every"] = self.at_least_every + if self.compute_group is not None: + body["compute_group"] = self.compute_group + if self.continuous is not None: + body["continuous"] = self.continuous + if self.cron_expression is not None: + body["cron_expression"] = self.cron_expression + if self.enabled is not None: + body["enabled"] = self.enabled + return body + + def as_shallow_dict(self) -> dict: + """Serializes the Schedule into a shallow dictionary of its immediate attributes.""" + body = {} + if self.at_least_every is not None: + body["at_least_every"] = self.at_least_every + if self.compute_group is not None: + body["compute_group"] = self.compute_group + if self.continuous is not None: + body["continuous"] = self.continuous + if self.cron_expression is not None: + body["cron_expression"] = self.cron_expression + if self.enabled is not None: + body["enabled"] = self.enabled + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> Schedule: + """Deserializes the Schedule from a dictionary.""" + return cls( + at_least_every=d.get("at_least_every", None), + compute_group=d.get("compute_group", None), + continuous=d.get("continuous", None), + cron_expression=d.get("cron_expression", None), + enabled=d.get("enabled", None), + ) + + +@dataclass +class ServiceNowExportConfig: + """ServiceNow destination. Alerts are delivered to the instance's Table API, authorized by an OAuth + 2.0 client-credentials token minted from the registered application's client ID and secret.""" + + instance_id: str + """ServiceNow instance ID — the host label only, not the full URL. For an instance reachable at + ``https://acme.service-now.com``, this is ``acme``.""" + + client_id: str + """OAuth client ID of the registered ServiceNow application. An identifier rather than secret + material — ServiceNow shows it on the application record — so it is stored and read back + verbatim.""" + + client_secret_uc_secret: str + """Three-level Unity Catalog name (``catalog.schema.secret``) of the secret holding the OAuth + client secret. Only the name is stored; the secret is read when an alert is exported, so + rotating it takes effect without editing this config.""" + + def as_dict(self) -> dict: + """Serializes the ServiceNowExportConfig into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.client_id is not None: + body["client_id"] = self.client_id + if self.client_secret_uc_secret is not None: + body["client_secret_uc_secret"] = self.client_secret_uc_secret + if self.instance_id is not None: + body["instance_id"] = self.instance_id + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ServiceNowExportConfig into a shallow dictionary of its immediate attributes.""" + body = {} + if self.client_id is not None: + body["client_id"] = self.client_id + if self.client_secret_uc_secret is not None: + body["client_secret_uc_secret"] = self.client_secret_uc_secret + if self.instance_id is not None: + body["instance_id"] = self.instance_id + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ServiceNowExportConfig: + """Deserializes the ServiceNowExportConfig from a dictionary.""" + return cls( + client_id=d.get("client_id", None), + client_secret_uc_secret=d.get("client_secret_uc_secret", None), + instance_id=d.get("instance_id", None), + ) + + +@dataclass +class SlackExportConfig: + channel: str + """Channel the alert is posted to: a channel ID (e.g. ``C0123456789``) or a channel name without + the leading ``#`` (e.g. ``soc-alerts``). The value is stored as sent, so a leading ``#``, + surrounding whitespace, or a value naming more than one channel is rejected. For private + channels, prefer the channel ID — Slack does not guarantee name resolution for those. + + The Slack app must be a member of the channel and hold the ``chat:write`` scope.""" + + bot_token_uc_secret: str + """Three-level Unity Catalog name (``catalog.schema.secret``) of the secret holding the Slack bot + token. Only the name is stored; the token is read when an alert is exported, so rotating the + secret takes effect without editing this config.""" + + def as_dict(self) -> dict: + """Serializes the SlackExportConfig into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.bot_token_uc_secret is not None: + body["bot_token_uc_secret"] = self.bot_token_uc_secret + if self.channel is not None: + body["channel"] = self.channel + return body + + def as_shallow_dict(self) -> dict: + """Serializes the SlackExportConfig into a shallow dictionary of its immediate attributes.""" + body = {} + if self.bot_token_uc_secret is not None: + body["bot_token_uc_secret"] = self.bot_token_uc_secret + if self.channel is not None: + body["channel"] = self.channel + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> SlackExportConfig: + """Deserializes the SlackExportConfig from a dictionary.""" + return cls(bot_token_uc_secret=d.get("bot_token_uc_secret", None), channel=d.get("channel", None)) + + +@dataclass +class TableName: + """A fully-qualified Unity Catalog table reference by catalog / schema / table name.""" + + table_name: str + """Unity Catalog table name.""" + + catalog_name: Optional[str] = None + """Name of the parent Unity Catalog catalog.""" + + schema_name: Optional[str] = None + """Name of the parent Unity Catalog schema.""" + + def as_dict(self) -> dict: + """Serializes the TableName into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.catalog_name is not None: + body["catalog_name"] = self.catalog_name + if self.schema_name is not None: + body["schema_name"] = self.schema_name + if self.table_name is not None: + body["table_name"] = self.table_name + return body + + def as_shallow_dict(self) -> dict: + """Serializes the TableName into a shallow dictionary of its immediate attributes.""" + body = {} + if self.catalog_name is not None: + body["catalog_name"] = self.catalog_name + if self.schema_name is not None: + body["schema_name"] = self.schema_name + if self.table_name is not None: + body["table_name"] = self.table_name + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> TableName: + """Deserializes the TableName from a dictionary.""" + return cls( + catalog_name=d.get("catalog_name", None), + schema_name=d.get("schema_name", None), + table_name=d.get("table_name", None), + ) + + +@dataclass +class TinesExportConfig: + """Tines destination. Alerts are delivered to a Tines webhook action, authorized by a secret sent + in the ``Authorization`` header.""" + + webhook_uri: str + """Tines webhook action URL, typically of the form ``https://.tines.com/webhook/``. + Any HTTPS host is accepted, so that a Tines tenant on a custom domain works. + + The secret is sent as a header, so it must not be embedded in the URL: a trailing secret path + segment (``/webhook//``) and a ``secret`` query parameter are both rejected. + API-path-prefix URLs (``/api/public/...``) are not supported.""" + + webhook_secret_uc_secret: str + """Three-level Unity Catalog name (``catalog.schema.secret``) of the secret holding the Tines + webhook secret. Only the name is stored; the secret is read when an alert is exported, so + rotating it takes effect without editing this config. Required: without it the destination would + be an unauthenticated endpoint.""" + + export_updates: Optional[bool] = None + """Whether a later change to an already-exported alert is delivered too, as a second + ``lakewatch.alert.updated`` event to the same webhook action (a close is such an update, + carrying a terminal ``alert_state``). Creation export is unaffected either way. + + Opt-in, and unset means off: Tines is an event transport rather than a system whose record is + edited in place, so an update is an extra event the receiving story has to correlate to the + alert it already handled. A destination configured before this event existed must not start + receiving more without its owner asking.""" + + def as_dict(self) -> dict: + """Serializes the TinesExportConfig into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.export_updates is not None: + body["export_updates"] = self.export_updates + if self.webhook_secret_uc_secret is not None: + body["webhook_secret_uc_secret"] = self.webhook_secret_uc_secret + if self.webhook_uri is not None: + body["webhook_uri"] = self.webhook_uri + return body + + def as_shallow_dict(self) -> dict: + """Serializes the TinesExportConfig into a shallow dictionary of its immediate attributes.""" + body = {} + if self.export_updates is not None: + body["export_updates"] = self.export_updates + if self.webhook_secret_uc_secret is not None: + body["webhook_secret_uc_secret"] = self.webhook_secret_uc_secret + if self.webhook_uri is not None: + body["webhook_uri"] = self.webhook_uri + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> TinesExportConfig: + """Deserializes the TinesExportConfig from a dictionary.""" + return cls( + export_updates=d.get("export_updates", None), + webhook_secret_uc_secret=d.get("webhook_secret_uc_secret", None), + webhook_uri=d.get("webhook_uri", None), + ) + + +@dataclass +class UnityCatalogConnection: + """Ingests data through a managed connector that uses a Unity Catalog connection. Mutually + exclusive with ``autoloader`` and ``existing_table``.""" + + connection_name: str + """Name of the Unity Catalog connection to ingest through.""" + + configuration: Optional[Dict[str, str]] = None + """Free-form key-value configuration for the managed ingestion pipeline.""" + + connector: Optional[str] = None + """Connector type of the connection, such as ``SLACK_AUDIT_LOGS``. Resolved by the server from + ``connection_name`` through Unity Catalog; not accepted as input.""" + + effective_start_time: Optional[Timestamp] = None + """The start time ingestion uses: the supplied ``start_time``, or the creation time when it was + omitted. Always populated.""" + + start_time: Optional[Timestamp] = None + """Earliest source event time to ingest. Leave unset to start at creation time, with no history + backfilled. Rejected if in the future, beyond the connector's history window, or for a connector + without start-time support. Takes effect only at first start and cannot be changed. Reads return + the value as supplied; the start time ingestion uses is reported in ``effective_start_time``.""" + + def as_dict(self) -> dict: + """Serializes the UnityCatalogConnection into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.configuration: + body["configuration"] = self.configuration + if self.connection_name is not None: + body["connection_name"] = self.connection_name + if self.connector is not None: + body["connector"] = self.connector + if self.effective_start_time is not None: + body["effective_start_time"] = self.effective_start_time.ToJsonString() + if self.start_time is not None: + body["start_time"] = self.start_time.ToJsonString() + return body + + def as_shallow_dict(self) -> dict: + """Serializes the UnityCatalogConnection into a shallow dictionary of its immediate attributes.""" + body = {} + if self.configuration: + body["configuration"] = self.configuration + if self.connection_name is not None: + body["connection_name"] = self.connection_name + if self.connector is not None: + body["connector"] = self.connector + if self.effective_start_time is not None: + body["effective_start_time"] = self.effective_start_time + if self.start_time is not None: + body["start_time"] = self.start_time + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> UnityCatalogConnection: + """Deserializes the UnityCatalogConnection from a dictionary.""" + return cls( + configuration=d.get("configuration", None), + connection_name=d.get("connection_name", None), + connector=d.get("connector", None), + effective_start_time=_timestamp(d, "effective_start_time"), + start_time=_timestamp(d, "start_time"), + ) + + +@dataclass +class WebhookExportConfig: + export_updates: Optional[bool] = None + """Whether a later change to an already-exported alert is delivered too, as a second + ``lakewatch.alert.updated`` event to the same webhook URL (a close is such an update, carrying a + terminal ``alert_state``). Creation export is unaffected either way. + + Opt-in, and unset means off: a webhook is an event transport rather than a system whose record + is edited in place, so an update is an extra event the receiving consumer has to correlate to + the alert it already handled. A destination configured before this event existed must not start + receiving more without its owner asking.""" + + url: Optional[str] = None + """Webhook destination the server POSTs payloads to, in plaintext. Userinfo (``user:password@``) is + sent as HTTP Basic auth.""" + + def as_dict(self) -> dict: + """Serializes the WebhookExportConfig into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.export_updates is not None: + body["export_updates"] = self.export_updates + if self.url is not None: + body["url"] = self.url + return body + + def as_shallow_dict(self) -> dict: + """Serializes the WebhookExportConfig into a shallow dictionary of its immediate attributes.""" + body = {} + if self.export_updates is not None: + body["export_updates"] = self.export_updates + if self.url is not None: + body["url"] = self.url + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> WebhookExportConfig: + """Deserializes the WebhookExportConfig from a dictionary.""" + return cls(export_updates=d.get("export_updates", None), url=d.get("url", None)) + + +class LakewatchAPI: + """Manage Lakewatch: security-analytics data sources, detection rules, and per-catalog configuration.""" + + def __init__(self, api_client): + self._api = api_client + + def create_datasource(self, datasource: Datasource) -> Datasource: + """Creates a new datasource. + + :param datasource: :class:`Datasource` + + :returns: :class:`Datasource` + """ + + body = datasource.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", "/api/2.0/lakewatch/datasources", body=body, headers=headers) + return Datasource.from_dict(res) + + def create_rule(self, rule: Rule) -> Rule: + """Creates a new rule. + + :param rule: :class:`Rule` + The rule to create. Server-assigned fields (``name``, ``rule_id``, timestamps, ``etag``, …) are + ignored on the request. + + :returns: :class:`Rule` + """ + + body = rule.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", "/api/2.0/lakewatch/rules", body=body, headers=headers) + return Rule.from_dict(res) + + def delete_datasource(self, name: str): + """Deletes a datasource by its resource name. + + :param name: str + + + """ + + headers = { + "Accept": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + self._api.do("DELETE", f"/api/2.0/lakewatch/{name}", headers=headers) + + def delete_rule(self, name: str): + """Deletes a rule by its resource name. + + :param name: str + Resource name of the rule to delete. Format: ``rules/{rule_id}``. + + + """ + + headers = { + "Accept": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + self._api.do("DELETE", f"/api/2.0/lakewatch/{name}", headers=headers) + + def get_alert_config(self, name: str) -> AlertConfig: + """Returns the alert delivery configuration for a catalog. + + :param name: str + + :returns: :class:`AlertConfig` + """ + + headers = { + "Accept": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do("GET", f"/api/2.0/lakewatch/{name}", headers=headers) + return AlertConfig.from_dict(res) + + def get_common_config(self, name: str) -> CommonConfig: + """Returns the common Lakewatch configuration for a catalog. + + :param name: str + + :returns: :class:`CommonConfig` + """ + + headers = { + "Accept": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do("GET", f"/api/2.0/lakewatch/{name}", headers=headers) + return CommonConfig.from_dict(res) + + def get_datasource(self, name: str) -> Datasource: + """Returns a single datasource by its resource name. + + :param name: str + + :returns: :class:`Datasource` + """ + + headers = { + "Accept": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do("GET", f"/api/2.0/lakewatch/{name}", headers=headers) + return Datasource.from_dict(res) + + def get_datasource_config(self, name: str) -> DatasourceConfig: + """Returns the datasource ingestion configuration for a catalog. + + :param name: str + + :returns: :class:`DatasourceConfig` + """ + + headers = { + "Accept": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do("GET", f"/api/2.0/lakewatch/{name}", headers=headers) + return DatasourceConfig.from_dict(res) + + def get_rule(self, name: str) -> Rule: + """Returns a single rule by its resource name. + + :param name: str + Full resource name of the detection rule, in the form ``rules/{rule_id}``. + + :returns: :class:`Rule` + """ + + headers = { + "Accept": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do("GET", f"/api/2.0/lakewatch/{name}", headers=headers) + return Rule.from_dict(res) + + def get_rule_config(self, name: str) -> RuleConfig: + """Returns the detection-rule configuration for a catalog. + + :param name: str + + :returns: :class:`RuleConfig` + """ + + headers = { + "Accept": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do("GET", f"/api/2.0/lakewatch/{name}", headers=headers) + return RuleConfig.from_dict(res) + + def list_datasources( + self, + catalog_name: str, + *, + filter: Optional[str] = None, + page_size: Optional[int] = None, + page_token: Optional[str] = None, + schema_name: Optional[str] = None, + view: Optional[DatasourceView] = None, + ) -> Iterator[Datasource]: + """Lists the datasources in a Unity Catalog catalog, optionally narrowed to one schema or filtered by + name. Results are paginated. Datasources that are still being created or deleted are not listed. + + :param catalog_name: str + Name of the Unity Catalog catalog to list datasources from. Required; each request lists the + datasources of a single catalog. + :param filter: str (optional) + Filter expression that restricts the results to datasources with a given name, for example + ``datasource_name = "my_source"``. Only a single exact-match comparison on ``datasource_name`` is + supported, and it matches the datasource's current Unity Catalog name, which reads report in + ``effective_datasource_name``. Filtered results are returned in a single page. + :param page_size: int (optional) + Maximum number of results to return in one page. This is a best-effort upper bound, not a guarantee: + the server may return fewer items than requested (even zero) while more results still exist. If + unset or 0, the server picks a default. + :param page_token: str (optional) + :param schema_name: str (optional) + Name of a Unity Catalog schema in ``catalog_name``. Optional; if set, only the datasources in this + schema are returned, and if omitted, the datasources in all schemas of the catalog are returned. + :param view: :class:`DatasourceView` (optional) + Which fields to return for each datasource. Defaults to ``DATASOURCE_VIEW_FULL``. + ``DATASOURCE_VIEW_BASIC`` returns only the identity and Unity Catalog location fields. + + :returns: Iterator over :class:`Datasource` + """ + + query = {} + if catalog_name is not None: + query["catalog_name"] = catalog_name + if filter is not None: + query["filter"] = filter + if page_size is not None: + query["page_size"] = page_size + if page_token is not None: + query["page_token"] = page_token + if schema_name is not None: + query["schema_name"] = schema_name + if view is not None: + query["view"] = view.value + headers = { + "Accept": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + while True: + json = self._api.do("GET", "/api/2.0/lakewatch/datasources", query=query, headers=headers) + if "datasources" in json: + for v in json["datasources"]: + yield Datasource.from_dict(v) + if "next_page_token" not in json or not json["next_page_token"]: + return + query["page_token"] = json["next_page_token"] + + def list_rules( + self, + catalog_name: str, + *, + filter: Optional[str] = None, + page_size: Optional[int] = None, + page_token: Optional[str] = None, + schema_name: Optional[str] = None, + view: Optional[RuleView] = None, + ) -> Iterator[Rule]: + """Lists rules, with optional filtering and pagination. + + :param catalog_name: str + Name of the Unity Catalog catalog to list rules from. Required. + :param filter: str (optional) + Filter expression that restricts the results, such as ``rule_name = "my_rule"`` to return only the + rule with that name. The name is matched against the current UC name, which reads report in + ``effective_rule_name``. + :param page_size: int (optional) + Maximum number of rules to return. The server may return fewer. If unset or 0, the server picks a + default page size. + :param page_token: str (optional) + Page token from a previous ``ListRules`` response, used to fetch the next page of results. + :param schema_name: str (optional) + Name of a Unity Catalog schema in ``catalog_name``. If set, only rules in this schema are returned. + If unset, rules in all schemas of the catalog are returned. + :param view: :class:`RuleView` (optional) + How much of each rule to return. Defaults to ``RULE_VIEW_FULL``. ``RULE_VIEW_BASIC`` returns only + each rule's identity and Unity Catalog location. + + :returns: Iterator over :class:`Rule` + """ + + query = {} + if catalog_name is not None: + query["catalog_name"] = catalog_name + if filter is not None: + query["filter"] = filter + if page_size is not None: + query["page_size"] = page_size + if page_token is not None: + query["page_token"] = page_token + if schema_name is not None: + query["schema_name"] = schema_name + if view is not None: + query["view"] = view.value + headers = { + "Accept": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + while True: + json = self._api.do("GET", "/api/2.0/lakewatch/rules", query=query, headers=headers) + if "rules" in json: + for v in json["rules"]: + yield Rule.from_dict(v) + if "next_page_token" not in json or not json["next_page_token"]: + return + query["page_token"] = json["next_page_token"] + + def update_alert_config( + self, name: str, alert_config: AlertConfig, update_mask: FieldMask, *, etag: Optional[str] = None + ) -> AlertConfig: + """Updates the alert delivery configuration for a catalog. + + :param name: str + :param alert_config: :class:`AlertConfig` + :param update_mask: FieldMask + The field mask must be a single string, with multiple fields separated by commas (no spaces). The + field path is relative to the resource object, using a dot (``.``) to navigate sub-fields (e.g., + ``author.given_name``). Specification of elements in sequence or map fields is not allowed, as only + the entire collection field can be specified. Field names must exactly match the resource field + names. + + A field mask of ``*`` indicates full replacement. It’s recommended to always explicitly list the + fields being updated and avoid using ``*`` wildcards, as it can lead to unintended results if the + API changes in the future. + :param etag: str (optional) + Optimistic-concurrency precondition: if set and it does not match the resource's current ``etag``, + the update is rejected as a conflict (HTTP 409). + + :returns: :class:`AlertConfig` + """ + + body = alert_config.as_dict() + query = {} + if etag is not None: + query["etag"] = etag + if update_mask is not None: + query["update_mask"] = update_mask.ToJsonString() + headers = { + "Accept": "application/json", + "Content-Type": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do("PATCH", f"/api/2.0/lakewatch/{name}", query=query, body=body, headers=headers) + return AlertConfig.from_dict(res) + + def update_common_config( + self, name: str, common_config: CommonConfig, update_mask: FieldMask, *, etag: Optional[str] = None + ) -> CommonConfig: + """Updates the common Lakewatch configuration for a catalog. + + :param name: str + :param common_config: :class:`CommonConfig` + :param update_mask: FieldMask + The field mask must be a single string, with multiple fields separated by commas (no spaces). The + field path is relative to the resource object, using a dot (``.``) to navigate sub-fields (e.g., + ``author.given_name``). Specification of elements in sequence or map fields is not allowed, as only + the entire collection field can be specified. Field names must exactly match the resource field + names. + + A field mask of ``*`` indicates full replacement. It’s recommended to always explicitly list the + fields being updated and avoid using ``*`` wildcards, as it can lead to unintended results if the + API changes in the future. + :param etag: str (optional) + Optimistic-concurrency precondition: if set and it does not match the resource's current ``etag``, + the update is rejected as a conflict (HTTP 409). + + :returns: :class:`CommonConfig` + """ + + body = common_config.as_dict() + query = {} + if etag is not None: + query["etag"] = etag + if update_mask is not None: + query["update_mask"] = update_mask.ToJsonString() + headers = { + "Accept": "application/json", + "Content-Type": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do("PATCH", f"/api/2.0/lakewatch/{name}", query=query, body=body, headers=headers) + return CommonConfig.from_dict(res) + + def update_datasource( + self, name: str, datasource: Datasource, update_mask: FieldMask, *, etag: Optional[str] = None + ) -> Datasource: + """Updates an existing datasource. + + :param name: str + Resource name of the datasource, in the format ``datasources/{datasource_id}``. Unlike the + datasource's Unity Catalog location (``catalog_name``, ``schema_name``, and ``datasource_name``), + the name does not change when the datasource or its catalog or schema is renamed in Unity Catalog. + :param datasource: :class:`Datasource` + :param update_mask: FieldMask + The field mask must be a single string, with multiple fields separated by commas (no spaces). The + field path is relative to the resource object, using a dot (``.``) to navigate sub-fields (e.g., + ``author.given_name``). Specification of elements in sequence or map fields is not allowed, as only + the entire collection field can be specified. Field names must exactly match the resource field + names. + + A field mask of ``*`` indicates full replacement. It’s recommended to always explicitly list the + fields being updated and avoid using ``*`` wildcards, as it can lead to unintended results if the + API changes in the future. + :param etag: str (optional) + Optimistic-concurrency precondition: if set to a non-empty value that does not match the resource's + current ``etag``, the update is rejected as a conflict (HTTP 409). Omit this field or send an empty + string to update unconditionally. + + :returns: :class:`Datasource` + """ + + body = datasource.as_dict() + query = {} + if etag is not None: + query["etag"] = etag + if update_mask is not None: + query["update_mask"] = update_mask.ToJsonString() + headers = { + "Accept": "application/json", + "Content-Type": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do("PATCH", f"/api/2.0/lakewatch/{name}", query=query, body=body, headers=headers) + return Datasource.from_dict(res) + + def update_datasource_config( + self, name: str, datasource_config: DatasourceConfig, update_mask: FieldMask, *, etag: Optional[str] = None + ) -> DatasourceConfig: + """Updates the datasource ingestion configuration for a catalog. + + :param name: str + :param datasource_config: :class:`DatasourceConfig` + :param update_mask: FieldMask + The field mask must be a single string, with multiple fields separated by commas (no spaces). The + field path is relative to the resource object, using a dot (``.``) to navigate sub-fields (e.g., + ``author.given_name``). Specification of elements in sequence or map fields is not allowed, as only + the entire collection field can be specified. Field names must exactly match the resource field + names. + + A field mask of ``*`` indicates full replacement. It’s recommended to always explicitly list the + fields being updated and avoid using ``*`` wildcards, as it can lead to unintended results if the + API changes in the future. + :param etag: str (optional) + Optimistic-concurrency precondition: if set and it does not match the resource's current ``etag``, + the update is rejected as a conflict (HTTP 409). + + :returns: :class:`DatasourceConfig` + """ + + body = datasource_config.as_dict() + query = {} + if etag is not None: + query["etag"] = etag + if update_mask is not None: + query["update_mask"] = update_mask.ToJsonString() + headers = { + "Accept": "application/json", + "Content-Type": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do("PATCH", f"/api/2.0/lakewatch/{name}", query=query, body=body, headers=headers) + return DatasourceConfig.from_dict(res) + + def update_rule(self, name: str, rule: Rule, update_mask: FieldMask, *, etag: Optional[str] = None) -> Rule: + """Updates an existing rule. + + :param name: str + Resource name, in the form ``rules/{rule_id}``. Assigned by the server on create; required on + update. + :param rule: :class:`Rule` + The rule with its new field values. ``rule.name`` identifies the rule to update. + :param update_mask: FieldMask + The set of fields to update. Only fields named in the mask are modified. + :param etag: str (optional) + Optimistic-concurrency token from the latest read. A non-empty token that does not match the current + ``etag`` rejects the update as a conflict (HTTP 409). Omit this field or provide an empty string to + update unconditionally. + + :returns: :class:`Rule` + """ + + body = rule.as_dict() + query = {} + if etag is not None: + query["etag"] = etag + if update_mask is not None: + query["update_mask"] = update_mask.ToJsonString() + headers = { + "Accept": "application/json", + "Content-Type": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do("PATCH", f"/api/2.0/lakewatch/{name}", query=query, body=body, headers=headers) + return Rule.from_dict(res) + + def update_rule_config( + self, name: str, rule_config: RuleConfig, update_mask: FieldMask, *, etag: Optional[str] = None + ) -> RuleConfig: + """Updates the detection-rule configuration for a catalog. + + :param name: str + :param rule_config: :class:`RuleConfig` + :param update_mask: FieldMask + The field mask must be a single string, with multiple fields separated by commas (no spaces). The + field path is relative to the resource object, using a dot (``.``) to navigate sub-fields (e.g., + ``author.given_name``). Specification of elements in sequence or map fields is not allowed, as only + the entire collection field can be specified. Field names must exactly match the resource field + names. + + A field mask of ``*`` indicates full replacement. It’s recommended to always explicitly list the + fields being updated and avoid using ``*`` wildcards, as it can lead to unintended results if the + API changes in the future. + :param etag: str (optional) + Optimistic-concurrency precondition: if set and it does not match the resource's current ``etag``, + the update is rejected as a conflict (HTTP 409). + + :returns: :class:`RuleConfig` + """ + + body = rule_config.as_dict() + query = {} + if etag is not None: + query["etag"] = etag + if update_mask is not None: + query["update_mask"] = update_mask.ToJsonString() + headers = { + "Accept": "application/json", + "Content-Type": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do("PATCH", f"/api/2.0/lakewatch/{name}", query=query, body=body, headers=headers) + return RuleConfig.from_dict(res) diff --git a/databricks/sdk/service/ml.py b/databricks/sdk/service/ml.py index d7bb273c7..db399071a 100644 --- a/databricks/sdk/service/ml.py +++ b/databricks/sdk/service/ml.py @@ -1496,6 +1496,67 @@ def from_dict(cls, d: Dict[str, Any]) -> DatasetInput: return cls(dataset=_from_dict(d, "dataset", Dataset), tags=_repeated_dict(d, "tags", InputTag)) +@dataclass +class DeclarativeFeature: + """A reference to a Feature.""" + + feature: str + """Three-part Unity Catalog name (catalog.schema.feature) of the Feature.""" + + def as_dict(self) -> dict: + """Serializes the DeclarativeFeature into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.feature is not None: + body["feature"] = self.feature + return body + + def as_shallow_dict(self) -> dict: + """Serializes the DeclarativeFeature into a shallow dictionary of its immediate attributes.""" + body = {} + if self.feature is not None: + body["feature"] = self.feature + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> DeclarativeFeature: + """Deserializes the DeclarativeFeature from a dictionary.""" + return cls(feature=d.get("feature", None)) + + +@dataclass +class DefaultValue: + """A default value for a looked-up feature.""" + + feature: str + """Name of the feature: its output name if it is renamed, otherwise its name in the feature table.""" + + default_value: Optional[str] = None + """Default value for the feature.""" + + def as_dict(self) -> dict: + """Serializes the DefaultValue into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.default_value is not None: + body["default_value"] = self.default_value + if self.feature is not None: + body["feature"] = self.feature + return body + + def as_shallow_dict(self) -> dict: + """Serializes the DefaultValue into a shallow dictionary of its immediate attributes.""" + body = {} + if self.default_value is not None: + body["default_value"] = self.default_value + if self.feature is not None: + body["feature"] = self.feature + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> DefaultValue: + """Deserializes the DefaultValue from a dictionary.""" + return cls(default_value=d.get("default_value", None), feature=d.get("feature", None)) + + @dataclass class DeleteCommentResponse: def as_dict(self) -> dict: @@ -2527,6 +2588,85 @@ def from_dict(cls, d: Dict[str, Any]) -> Feature: ) +@dataclass +class FeatureFunction: + """A Python UDF that computes a feature.""" + + udf: str + """Three-part name of the Python UDF.""" + + input_bindings: Optional[List[FeatureFunctionInputBinding]] = None + """Bindings of the UDF's parameters to features.""" + + output: Optional[str] = None + """Output name of the feature that the UDF computes.""" + + def as_dict(self) -> dict: + """Serializes the FeatureFunction into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.input_bindings: + body["input_bindings"] = [v.as_dict() for v in self.input_bindings] + if self.output is not None: + body["output"] = self.output + if self.udf is not None: + body["udf"] = self.udf + return body + + def as_shallow_dict(self) -> dict: + """Serializes the FeatureFunction into a shallow dictionary of its immediate attributes.""" + body = {} + if self.input_bindings: + body["input_bindings"] = self.input_bindings + if self.output is not None: + body["output"] = self.output + if self.udf is not None: + body["udf"] = self.udf + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> FeatureFunction: + """Deserializes the FeatureFunction from a dictionary.""" + return cls( + input_bindings=_repeated_dict(d, "input_bindings", FeatureFunctionInputBinding), + output=d.get("output", None), + udf=d.get("udf", None), + ) + + +@dataclass +class FeatureFunctionInputBinding: + """Binds one parameter of the UDF to a feature.""" + + parameter: str + """Name of the UDF parameter.""" + + bound_feature: str + """Name of the feature bound to the parameter.""" + + def as_dict(self) -> dict: + """Serializes the FeatureFunctionInputBinding into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.bound_feature is not None: + body["bound_feature"] = self.bound_feature + if self.parameter is not None: + body["parameter"] = self.parameter + return body + + def as_shallow_dict(self) -> dict: + """Serializes the FeatureFunctionInputBinding into a shallow dictionary of its immediate attributes.""" + body = {} + if self.bound_feature is not None: + body["bound_feature"] = self.bound_feature + if self.parameter is not None: + body["parameter"] = self.parameter + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> FeatureFunctionInputBinding: + """Deserializes the FeatureFunctionInputBinding from a dictionary.""" + return cls(bound_feature=d.get("bound_feature", None), parameter=d.get("parameter", None)) + + @dataclass class FeatureLineage: feature_specs: Optional[List[FeatureLineageFeatureSpec]] = None @@ -2685,6 +2825,83 @@ def from_dict(cls, d: Dict[str, Any]) -> FeatureList: return cls(features=_repeated_dict(d, "features", LinkedFeature)) +@dataclass +class FeatureLookup: + """Features to look up from a feature table.""" + + table: str + """Three-part name of the feature table.""" + + lookup_keys: List[str] + """Keys to use when joining with the feature table.""" + + default_values: Optional[List[DefaultValue]] = None + """Default values for features.""" + + feature_names: Optional[List[str]] = None + """Names of the features to look up, or empty to look up every feature in the table.""" + + lookback_window: Optional[Duration] = None + """Lookback window to use when performing point-in-time lookup.""" + + rename_outputs: Optional[List[RenameOutput]] = None + """Output names for features, in place of their names in the feature table.""" + + timestamp_lookup_keys: Optional[List[str]] = None + """Keys to use when performing point-in-time lookup.""" + + def as_dict(self) -> dict: + """Serializes the FeatureLookup into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.default_values: + body["default_values"] = [v.as_dict() for v in self.default_values] + if self.feature_names: + body["feature_names"] = [v for v in self.feature_names] + if self.lookback_window is not None: + body["lookback_window"] = self.lookback_window.ToJsonString() + if self.lookup_keys: + body["lookup_keys"] = [v for v in self.lookup_keys] + if self.rename_outputs: + body["rename_outputs"] = [v.as_dict() for v in self.rename_outputs] + if self.table is not None: + body["table"] = self.table + if self.timestamp_lookup_keys: + body["timestamp_lookup_keys"] = [v for v in self.timestamp_lookup_keys] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the FeatureLookup into a shallow dictionary of its immediate attributes.""" + body = {} + if self.default_values: + body["default_values"] = self.default_values + if self.feature_names: + body["feature_names"] = self.feature_names + if self.lookback_window is not None: + body["lookback_window"] = self.lookback_window + if self.lookup_keys: + body["lookup_keys"] = self.lookup_keys + if self.rename_outputs: + body["rename_outputs"] = self.rename_outputs + if self.table is not None: + body["table"] = self.table + if self.timestamp_lookup_keys: + body["timestamp_lookup_keys"] = self.timestamp_lookup_keys + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> FeatureLookup: + """Deserializes the FeatureLookup from a dictionary.""" + return cls( + default_values=_repeated_dict(d, "default_values", DefaultValue), + feature_names=d.get("feature_names", None), + lookback_window=_duration(d, "lookback_window"), + lookup_keys=d.get("lookup_keys", None), + rename_outputs=_repeated_dict(d, "rename_outputs", RenameOutput), + table=d.get("table", None), + timestamp_lookup_keys=d.get("timestamp_lookup_keys", None), + ) + + @dataclass class FeatureReference: """A reference to one registered upstream Feature. A message rather than a bare name so an upstream @@ -2713,6 +2930,62 @@ def from_dict(cls, d: Dict[str, Any]) -> FeatureReference: return cls(feature=d.get("feature", None)) +@dataclass +class FeatureSpec: + """A named set of features to serve, stored as a Unity Catalog function.""" + + name: str + """Three-part name (catalog.schema.feature_spec) of the feature spec.""" + + exclude_columns: Optional[List[str]] = None + """Output names of columns to leave out of the feature spec.""" + + extra_pip_requirements: Optional[PipRequirements] = None + """Extra pip requirements to install in the serving container, in addition to the defaults and + those resolved from the features.""" + + features: Optional[List[Features]] = None + """The Features (either a declarative one, or a lookup from a feature table, or a feature function) + to serve. A declarative feature (``DeclarativeFeature``) cannot be combined with lookups + (``FeatureLookup``) and functions (``FeatureFunction``) in one feature spec.""" + + def as_dict(self) -> dict: + """Serializes the FeatureSpec into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.exclude_columns: + body["exclude_columns"] = [v for v in self.exclude_columns] + if self.extra_pip_requirements: + body["extra_pip_requirements"] = self.extra_pip_requirements.as_dict() + if self.features: + body["features"] = [v.as_dict() for v in self.features] + if self.name is not None: + body["name"] = self.name + return body + + def as_shallow_dict(self) -> dict: + """Serializes the FeatureSpec into a shallow dictionary of its immediate attributes.""" + body = {} + if self.exclude_columns: + body["exclude_columns"] = self.exclude_columns + if self.extra_pip_requirements: + body["extra_pip_requirements"] = self.extra_pip_requirements + if self.features: + body["features"] = self.features + if self.name is not None: + body["name"] = self.name + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> FeatureSpec: + """Deserializes the FeatureSpec from a dictionary.""" + return cls( + exclude_columns=d.get("exclude_columns", None), + extra_pip_requirements=_from_dict(d, "extra_pip_requirements", PipRequirements), + features=_repeated_dict(d, "features", Features), + name=d.get("name", None), + ) + + @dataclass class FeatureTag: """Represents a tag on a feature in a feature table.""" @@ -2772,6 +3045,51 @@ def from_dict(cls, d: Dict[str, Any]) -> FeatureViewSource: return cls(feature_references=_repeated_dict(d, "feature_references", FeatureReference)) +@dataclass +class Features: + """One entry of a feature spec: a Feature, a feature table lookup, or a feature function.""" + + declarative_feature: Optional[DeclarativeFeature] = None + """A Feature, referenced by its three-part Unity Catalog name.""" + + feature_function: Optional[FeatureFunction] = None + """A feature function.""" + + feature_lookup: Optional[FeatureLookup] = None + """A feature table lookup.""" + + def as_dict(self) -> dict: + """Serializes the Features into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.declarative_feature: + body["declarative_feature"] = self.declarative_feature.as_dict() + if self.feature_function: + body["feature_function"] = self.feature_function.as_dict() + if self.feature_lookup: + body["feature_lookup"] = self.feature_lookup.as_dict() + return body + + def as_shallow_dict(self) -> dict: + """Serializes the Features into a shallow dictionary of its immediate attributes.""" + body = {} + if self.declarative_feature: + body["declarative_feature"] = self.declarative_feature + if self.feature_function: + body["feature_function"] = self.feature_function + if self.feature_lookup: + body["feature_lookup"] = self.feature_lookup + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> Features: + """Deserializes the Features from a dictionary.""" + return cls( + declarative_feature=_from_dict(d, "declarative_feature", DeclarativeFeature), + feature_function=_from_dict(d, "feature_function", FeatureFunction), + feature_lookup=_from_dict(d, "feature_lookup", FeatureLookup), + ) + + @dataclass class FieldDefinition: """A single field definition within a FlatSchema, specifying the field name and its scalar data @@ -5527,14 +5845,13 @@ def from_dict(cls, d: Dict[str, Any]) -> MaterializationSetDeletionProgress: @dataclass class MaterializationSetOperationMetadata: - """Progress of a materialization set operation. Shared by create, update and delete, because - databricks.longrunning freezes a method's metadata type once it ships.""" + """Progress of a materialization set operation.""" create_time: Optional[Timestamp] = None """Time at which the operation was requested.""" deletion_progress: Optional[MaterializationSetDeletionProgress] = None - """Set by a DeleteMaterializationSet operation.""" + """Set by a ``DeleteMaterializationSet`` operation.""" feature_full_names: Optional[List[str]] = None """Full names of the features the set materializes, one entry per feature.""" @@ -5546,7 +5863,7 @@ class MaterializationSetOperationMetadata: """Current state of the operation.""" update_progress: Optional[MaterializationSetUpdateProgress] = None - """Set by an UpdateMaterializationSet operation.""" + """Set by an ``UpdateMaterializationSet`` operation.""" def as_dict(self) -> dict: """Serializes the MaterializationSetOperationMetadata into a dictionary suitable for use as a JSON request body.""" @@ -7046,6 +7363,33 @@ class PermissionLevel(Enum): CAN_READ = "CAN_READ" +@dataclass +class PipRequirements: + """Pip requirements to install in a serving container.""" + + requirements: List[str] + """Pip requirement strings (PEP 508) or UC Volume whl paths (/Volumes/...).""" + + def as_dict(self) -> dict: + """Serializes the PipRequirements into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.requirements: + body["requirements"] = [v for v in self.requirements] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the PipRequirements into a shallow dictionary of its immediate attributes.""" + body = {} + if self.requirements: + body["requirements"] = self.requirements + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> PipRequirements: + """Deserializes the PipRequirements from a dictionary.""" + return cls(requirements=d.get("requirements", None)) + + @dataclass class ProtoSchemaSpec: """A Protocol Buffer schema paired with the name of the message within it that describes the Kafka @@ -7281,7 +7625,11 @@ class PurgeFeatureEntitiesMetadataState(Enum): CANCELLED = "CANCELLED" FAILED = "FAILED" PENDING = "PENDING" + PENDING_CLEANUP = "PENDING_CLEANUP" + PENDING_PURGE = "PENDING_PURGE" RUNNING = "RUNNING" + RUNNING_CLEANUP = "RUNNING_CLEANUP" + RUNNING_PURGE = "RUNNING_PURGE" SUCCEEDED = "SUCCEEDED" @@ -7829,6 +8177,40 @@ def from_dict(cls, d: Dict[str, Any]) -> RenameModelResponse: return cls(registered_model=_from_dict(d, "registered_model", Model)) +@dataclass +class RenameOutput: + """An output name for a looked-up feature.""" + + feature: str + """Name of the feature in the feature table.""" + + output: str + """Output name for the feature.""" + + def as_dict(self) -> dict: + """Serializes the RenameOutput into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.feature is not None: + body["feature"] = self.feature + if self.output is not None: + body["output"] = self.output + return body + + def as_shallow_dict(self) -> dict: + """Serializes the RenameOutput into a shallow dictionary of its immediate attributes.""" + body = {} + if self.feature is not None: + body["feature"] = self.feature + if self.output is not None: + body["output"] = self.output + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> RenameOutput: + """Deserializes the RenameOutput from a dictionary.""" + return cls(feature=d.get("feature", None), output=d.get("output", None)) + + @dataclass class RequestSource: """A request-time data source whose value is provided at inference time: offline batch scoring or @@ -12022,6 +12404,29 @@ def create_feature(self, feature: Feature) -> Feature: res = self._api.do("POST", "/api/2.0/feature-engineering/features", body=body, headers=headers) return Feature.from_dict(res) + def create_feature_spec(self, feature_spec: FeatureSpec) -> FeatureSpec: + """Create a feature spec. + + :param feature_spec: :class:`FeatureSpec` + The feature spec to create. + + :returns: :class:`FeatureSpec` + """ + + body = feature_spec.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", "/api/2.0/feature-engineering/feature-specs", body=body, headers=headers) + return FeatureSpec.from_dict(res) + def create_kafka_config(self, kafka_config: KafkaConfig) -> KafkaConfig: """Create a Kafka config. During PrPr, Kafka configs can be read and used when creating features under the entire metastore. Only the creator of the Kafka config can delete it. @@ -12148,6 +12553,25 @@ def delete_feature(self, full_name: str): self._api.do("DELETE", f"/api/2.0/feature-engineering/features/{full_name}", headers=headers) + def delete_feature_spec(self, name: str): + """Delete a feature spec. + + :param name: str + Three-part name (catalog.schema.feature_spec) of the feature spec to delete. + + + """ + + headers = { + "Accept": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + self._api.do("DELETE", f"/api/2.0/feature-engineering/feature-specs/{name}", headers=headers) + def delete_kafka_config(self, name: str): """Delete a Kafka config. During PrPr, Kafka configs can be read and used when creating features under the entire metastore. Only the creator of the Kafka config can delete it. @@ -12273,6 +12697,26 @@ def get_feature(self, full_name: str) -> Feature: res = self._api.do("GET", f"/api/2.0/feature-engineering/features/{full_name}", headers=headers) return Feature.from_dict(res) + def get_feature_spec(self, name: str) -> FeatureSpec: + """Get a feature spec. + + :param name: str + Three-part name (catalog.schema.feature_spec) of the feature spec to get. + + :returns: :class:`FeatureSpec` + """ + + headers = { + "Accept": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do("GET", f"/api/2.0/feature-engineering/feature-specs/{name}", headers=headers) + return FeatureSpec.from_dict(res) + def get_kafka_config(self, name: str) -> KafkaConfig: """Get a Kafka config. During PrPr, Kafka configs can be read and used when creating features under the entire metastore. Only the creator of the Kafka config can delete it. diff --git a/databricks/sdk/version.py b/databricks/sdk/version.py index 36947cdb8..2b316f594 100644 --- a/databricks/sdk/version.py +++ b/databricks/sdk/version.py @@ -1 +1 @@ -__version__ = "0.160.0" +__version__ = "0.161.0" diff --git a/docs/dbdataclasses/aifunctions.rst b/docs/dbdataclasses/aifunctions.rst index 48bb9b4b0..e3307e8bd 100644 --- a/docs/dbdataclasses/aifunctions.rst +++ b/docs/dbdataclasses/aifunctions.rst @@ -67,3 +67,47 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. autoclass:: AiParseDocumentResponseMetadata :members: :undoc-members: + +.. autoclass:: Batch + :members: + :undoc-members: + +.. autoclass:: BatchRequestCounts + :members: + :undoc-members: + +.. py:class:: BatchStatus + + The status of a batch. + + .. py:attribute:: CANCELLED + :value: "CANCELLED" + + .. py:attribute:: CANCELLING + :value: "CANCELLING" + + .. py:attribute:: COMPLETED + :value: "COMPLETED" + + .. py:attribute:: EXPIRED + :value: "EXPIRED" + + .. py:attribute:: FAILED + :value: "FAILED" + + .. py:attribute:: FINALIZING + :value: "FINALIZING" + + .. py:attribute:: IN_PROGRESS + :value: "IN_PROGRESS" + + .. py:attribute:: PENDING + :value: "PENDING" + +.. autoclass:: CancelBatchResponse + :members: + :undoc-members: + +.. autoclass:: ListBatchesResponse + :members: + :undoc-members: diff --git a/docs/dbdataclasses/index.rst b/docs/dbdataclasses/index.rst index f77d8ee59..444f60b28 100644 --- a/docs/dbdataclasses/index.rst +++ b/docs/dbdataclasses/index.rst @@ -29,6 +29,7 @@ Dataclasses iamv2 jobs knowledgeassistants + lakewatch marketplace ml networking diff --git a/docs/dbdataclasses/jobs.rst b/docs/dbdataclasses/jobs.rst index 801f11f82..abca5bcb2 100644 --- a/docs/dbdataclasses/jobs.rst +++ b/docs/dbdataclasses/jobs.rst @@ -64,9 +64,6 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: BEST_EFFORT :value: "BEST_EFFORT" - .. py:attribute:: CRITICAL - :value: "CRITICAL" - .. py:attribute:: NORMAL :value: "NORMAL" diff --git a/docs/dbdataclasses/lakewatch.rst b/docs/dbdataclasses/lakewatch.rst new file mode 100644 index 000000000..e23c5acf7 --- /dev/null +++ b/docs/dbdataclasses/lakewatch.rst @@ -0,0 +1,264 @@ +Lakewatch +========= + +These dataclasses are used in the SDK to represent API requests and responses for services in the ``databricks.sdk.service.lakewatch`` module. + +.. py:currentmodule:: databricks.sdk.service.lakewatch +.. autoclass:: AlertConfig + :members: + :undoc-members: + +.. autoclass:: AutoLoader + :members: + :undoc-members: + +.. py:class:: AutoLoaderFormat + + File format of the source files that Auto Loader reads. + + .. py:attribute:: AUTOLOADER_FORMAT_CSV + :value: "AUTOLOADER_FORMAT_CSV" + + .. py:attribute:: AUTOLOADER_FORMAT_JSON + :value: "AUTOLOADER_FORMAT_JSON" + + .. py:attribute:: AUTOLOADER_FORMAT_JSONL + :value: "AUTOLOADER_FORMAT_JSONL" + + .. py:attribute:: AUTOLOADER_FORMAT_PARQUET + :value: "AUTOLOADER_FORMAT_PARQUET" + + .. py:attribute:: AUTOLOADER_FORMAT_XML + :value: "AUTOLOADER_FORMAT_XML" + +.. autoclass:: CloudWatchLogsEnvelope + :members: + :undoc-members: + +.. autoclass:: CommonConfig + :members: + :undoc-members: + +.. autoclass:: ContentPackReference + :members: + :undoc-members: + +.. autoclass:: Datasource + :members: + :undoc-members: + +.. py:class:: DatasourceComputeMode + + Compute mode for the datasource's ingestion job. + + .. py:attribute:: DS_COMPUTE_PERFORMANCE_OPTIMIZED + :value: "DS_COMPUTE_PERFORMANCE_OPTIMIZED" + + .. py:attribute:: DS_COMPUTE_STANDARD + :value: "DS_COMPUTE_STANDARD" + +.. autoclass:: DatasourceConfig + :members: + :undoc-members: + +.. autoclass:: DatasourceStatus + :members: + :undoc-members: + +.. py:class:: DatasourceView + + Which fields list responses return for each datasource. + + .. py:attribute:: DATASOURCE_VIEW_BASIC + :value: "DATASOURCE_VIEW_BASIC" + + .. py:attribute:: DATASOURCE_VIEW_FULL + :value: "DATASOURCE_VIEW_FULL" + +.. autoclass:: Envelope + :members: + :undoc-members: + +.. autoclass:: ExistingTable + :members: + :undoc-members: + +.. autoclass:: ExportConfig + :members: + :undoc-members: + +.. autoclass:: JiraExportConfig + :members: + :undoc-members: + +.. autoclass:: JsonArrayEnvelope + :members: + :undoc-members: + +.. autoclass:: ListDatasourcesResponse + :members: + :undoc-members: + +.. autoclass:: ListRulesResponse + :members: + :undoc-members: + +.. autoclass:: ParserRef + :members: + :undoc-members: + +.. autoclass:: Principal + :members: + :undoc-members: + +.. py:class:: PrincipalType + + Type of a principal (user / SP / system). + + .. py:attribute:: REGULAR_USER + :value: "REGULAR_USER" + + .. py:attribute:: SERVICE_PRINCIPAL + :value: "SERVICE_PRINCIPAL" + + .. py:attribute:: SYSTEM_USER + :value: "SYSTEM_USER" + +.. autoclass:: ResourceStatus + :members: + :undoc-members: + +.. py:class:: ResourceStatusLifecycle + + Configured scheduling / existence of the job or pipeline. + + .. py:attribute:: DELETED + :value: "DELETED" + + .. py:attribute:: SCHEDULED + :value: "SCHEDULED" + + .. py:attribute:: STOPPED + :value: "STOPPED" + + .. py:attribute:: UNSCHEDULED + :value: "UNSCHEDULED" + +.. py:class:: ResourceStatusRunState + + Outcome / progress of the most recent (or in-flight) run. + + .. py:attribute:: CANCELED + :value: "CANCELED" + + .. py:attribute:: FAILED + :value: "FAILED" + + .. py:attribute:: RUNNING + :value: "RUNNING" + + .. py:attribute:: SUCCEEDED + :value: "SUCCEEDED" + +.. autoclass:: Rule + :members: + :undoc-members: + +.. autoclass:: RuleConfig + :members: + :undoc-members: + +.. py:class:: RuleDeduplication + + Signal deduplication mode for a detection rule. + + .. py:attribute:: DEDUPLICATION_DEFAULT + :value: "DEDUPLICATION_DEFAULT" + + .. py:attribute:: DEDUPLICATION_DISABLED + :value: "DEDUPLICATION_DISABLED" + + .. py:attribute:: DEDUPLICATION_ENABLED + :value: "DEDUPLICATION_ENABLED" + +.. py:class:: RuleFidelity + + Maturity of a detection rule. + + .. py:attribute:: FIDELITY_DEVELOPMENT + :value: "FIDELITY_DEVELOPMENT" + + .. py:attribute:: FIDELITY_HIGH + :value: "FIDELITY_HIGH" + + .. py:attribute:: FIDELITY_INVESTIGATIVE + :value: "FIDELITY_INVESTIGATIVE" + + .. py:attribute:: FIDELITY_MEDIUM + :value: "FIDELITY_MEDIUM" + +.. autoclass:: RuleInputBatch + :members: + :undoc-members: + +.. autoclass:: RuleInputBatchSql + :members: + :undoc-members: + +.. autoclass:: RuleInputCustomNotebook + :members: + :undoc-members: + +.. autoclass:: RuleMitreMapping + :members: + :undoc-members: + +.. autoclass:: RuleOutput + :members: + :undoc-members: + +.. autoclass:: RuleSchedule + :members: + :undoc-members: + +.. autoclass:: RuleTimeRange + :members: + :undoc-members: + +.. py:class:: RuleView + + Selects how much of each rule a list response returns. + + .. py:attribute:: RULE_VIEW_BASIC + :value: "RULE_VIEW_BASIC" + + .. py:attribute:: RULE_VIEW_FULL + :value: "RULE_VIEW_FULL" + +.. autoclass:: Schedule + :members: + :undoc-members: + +.. autoclass:: ServiceNowExportConfig + :members: + :undoc-members: + +.. autoclass:: SlackExportConfig + :members: + :undoc-members: + +.. autoclass:: TableName + :members: + :undoc-members: + +.. autoclass:: TinesExportConfig + :members: + :undoc-members: + +.. autoclass:: UnityCatalogConnection + :members: + :undoc-members: + +.. autoclass:: WebhookExportConfig + :members: + :undoc-members: diff --git a/docs/dbdataclasses/ml.rst b/docs/dbdataclasses/ml.rst index b2ed0db8f..eeaf497c9 100644 --- a/docs/dbdataclasses/ml.rst +++ b/docs/dbdataclasses/ml.rst @@ -237,6 +237,14 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: DeclarativeFeature + :members: + :undoc-members: + +.. autoclass:: DefaultValue + :members: + :undoc-members: + .. autoclass:: DeleteCommentResponse :members: :undoc-members: @@ -601,6 +609,14 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: FeatureFunction + :members: + :undoc-members: + +.. autoclass:: FeatureFunctionInputBinding + :members: + :undoc-members: + .. autoclass:: FeatureLineage :members: :undoc-members: @@ -621,10 +637,18 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: FeatureLookup + :members: + :undoc-members: + .. autoclass:: FeatureReference :members: :undoc-members: +.. autoclass:: FeatureSpec + :members: + :undoc-members: + .. autoclass:: FeatureTag :members: :undoc-members: @@ -633,6 +657,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: Features + :members: + :undoc-members: + .. autoclass:: FieldDefinition :members: :undoc-members: @@ -1225,6 +1253,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: CAN_READ :value: "CAN_READ" +.. autoclass:: PipRequirements + :members: + :undoc-members: + .. autoclass:: ProtoSchemaSpec :members: :undoc-members: @@ -1265,9 +1297,21 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: PENDING :value: "PENDING" + .. py:attribute:: PENDING_CLEANUP + :value: "PENDING_CLEANUP" + + .. py:attribute:: PENDING_PURGE + :value: "PENDING_PURGE" + .. py:attribute:: RUNNING :value: "RUNNING" + .. py:attribute:: RUNNING_CLEANUP + :value: "RUNNING_CLEANUP" + + .. py:attribute:: RUNNING_PURGE + :value: "RUNNING_PURGE" + .. py:attribute:: SUCCEEDED :value: "SUCCEEDED" @@ -1412,6 +1456,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: RenameOutput + :members: + :undoc-members: + .. autoclass:: RequestSource :members: :undoc-members: diff --git a/docs/packages.py b/docs/packages.py index 2f03c7f80..1213c13ba 100644 --- a/docs/packages.py +++ b/docs/packages.py @@ -172,6 +172,11 @@ class Package: "Knowledge Assistants (legacy)", "Manage legacy Knowledge Assistants and related resources.", ), + Package( + "lakewatch", + "Lakewatch", + "Manage Lakewatch: security-analytics data sources, detection rules, and per-catalog configuration.", + ), Package("marketplace", "", ""), Package("ml", "", ""), Package("networking", "", ""), diff --git a/docs/workspace/aifunctions/ai_functions.rst b/docs/workspace/aifunctions/ai_functions.rst index 8fb3319cc..0365e5892 100644 --- a/docs/workspace/aifunctions/ai_functions.rst +++ b/docs/workspace/aifunctions/ai_functions.rst @@ -152,4 +152,48 @@ Function options. Omitted fields fall back to their documented defaults. :returns: :class:`AiParseDocumentResponse` + + + .. py:method:: cancel_batch(name: str) -> CancelBatchResponse + + Cancel a batch. Cancellation is asynchronous: requests in flight drain first, so read the batch to + observe the ``CANCELLING`` and ``CANCELLED`` statuses. + + :param name: str + ID of the batch, in the form ``batch_``. + + :returns: :class:`CancelBatchResponse` + + + .. py:method:: create_batch(batch: Batch) -> Batch + + Run an AI function over every request in a file on a Unity Catalog volume. The batch runs + asynchronously; read it to follow its progress. + + :param batch: :class:`Batch` + The batch to create. + + :returns: :class:`Batch` + + + .. py:method:: get_batch(name: str) -> Batch + + Read one batch, with its status and progress. + + :param name: str + ID of the batch, in the form ``batch_``. + + :returns: :class:`Batch` + + + .. py:method:: list_batches( [, page_size: Optional[int], page_token: Optional[str]]) -> Iterator[Batch] + + List the batches the caller created in this workspace. + + :param page_size: int (optional) + Maximum number of batches to return in one page. + :param page_token: str (optional) + Page token from an earlier response. Omit it to read the first page. + + :returns: Iterator over :class:`Batch` \ No newline at end of file diff --git a/docs/workspace/compute/instance_pools.rst b/docs/workspace/compute/instance_pools.rst index 7743ea56d..3b5ced2af 100644 --- a/docs/workspace/compute/instance_pools.rst +++ b/docs/workspace/compute/instance_pools.rst @@ -49,7 +49,7 @@ this cluster. For example, the Spark nodes can be provisioned and optimized for memory or compute intensive workloads. A list of available node types can be retrieved by using the `clusters/listNodeTypes `__ API - call. + call. The node type cannot be changed after the pool is created. :param aws_attributes: :class:`InstancePoolAwsAttributes` (optional) Attributes related to instance pools running on Amazon Web Services. If not specified at pool creation, a set of default values will be used. @@ -149,7 +149,7 @@ this cluster. For example, the Spark nodes can be provisioned and optimized for memory or compute intensive workloads. A list of available node types can be retrieved by using the `clusters/listNodeTypes `__ API - call. + call. The node type cannot be changed after the pool is created. :param custom_tags: Dict[str,str] (optional) Additional tags for pool resources. Databricks will tag all pool resources (e.g., AWS instances and EBS volumes) with these tags in addition to ``default_tags``. Notes: diff --git a/docs/workspace/index.rst b/docs/workspace/index.rst index 5fab702c1..07ac0b2d9 100644 --- a/docs/workspace/index.rst +++ b/docs/workspace/index.rst @@ -29,6 +29,7 @@ These APIs are available from WorkspaceClient iamv2/index jobs/index knowledgeassistants/index + lakewatch/index marketplace/index ml/index oauth2/index diff --git a/docs/workspace/lakewatch/index.rst b/docs/workspace/lakewatch/index.rst new file mode 100644 index 000000000..966d22065 --- /dev/null +++ b/docs/workspace/lakewatch/index.rst @@ -0,0 +1,10 @@ + +Lakewatch +========= + +Manage Lakewatch: security-analytics data sources, detection rules, and per-catalog configuration. + +.. toctree:: + :maxdepth: 1 + + lakewatch \ No newline at end of file diff --git a/docs/workspace/lakewatch/lakewatch.rst b/docs/workspace/lakewatch/lakewatch.rst new file mode 100644 index 000000000..5ed49c297 --- /dev/null +++ b/docs/workspace/lakewatch/lakewatch.rst @@ -0,0 +1,292 @@ +``w.lakewatch``: Lakewatch.v1 +============================= +.. currentmodule:: databricks.sdk.service.lakewatch + +.. py:class:: LakewatchAPI + + Manage Lakewatch: security-analytics data sources, detection rules, and per-catalog configuration. + + .. py:method:: create_datasource(datasource: Datasource) -> Datasource + + Creates a new datasource. + + :param datasource: :class:`Datasource` + + :returns: :class:`Datasource` + + + .. py:method:: create_rule(rule: Rule) -> Rule + + Creates a new rule. + + :param rule: :class:`Rule` + The rule to create. Server-assigned fields (``name``, ``rule_id``, timestamps, ``etag``, …) are + ignored on the request. + + :returns: :class:`Rule` + + + .. py:method:: delete_datasource(name: str) + + Deletes a datasource by its resource name. + + :param name: str + + + + + .. py:method:: delete_rule(name: str) + + Deletes a rule by its resource name. + + :param name: str + Resource name of the rule to delete. Format: ``rules/{rule_id}``. + + + + + .. py:method:: get_alert_config(name: str) -> AlertConfig + + Returns the alert delivery configuration for a catalog. + + :param name: str + + :returns: :class:`AlertConfig` + + + .. py:method:: get_common_config(name: str) -> CommonConfig + + Returns the common Lakewatch configuration for a catalog. + + :param name: str + + :returns: :class:`CommonConfig` + + + .. py:method:: get_datasource(name: str) -> Datasource + + Returns a single datasource by its resource name. + + :param name: str + + :returns: :class:`Datasource` + + + .. py:method:: get_datasource_config(name: str) -> DatasourceConfig + + Returns the datasource ingestion configuration for a catalog. + + :param name: str + + :returns: :class:`DatasourceConfig` + + + .. py:method:: get_rule(name: str) -> Rule + + Returns a single rule by its resource name. + + :param name: str + Full resource name of the detection rule, in the form ``rules/{rule_id}``. + + :returns: :class:`Rule` + + + .. py:method:: get_rule_config(name: str) -> RuleConfig + + Returns the detection-rule configuration for a catalog. + + :param name: str + + :returns: :class:`RuleConfig` + + + .. py:method:: list_datasources(catalog_name: str [, filter: Optional[str], page_size: Optional[int], page_token: Optional[str], schema_name: Optional[str], view: Optional[DatasourceView]]) -> Iterator[Datasource] + + Lists the datasources in a Unity Catalog catalog, optionally narrowed to one schema or filtered by + name. Results are paginated. Datasources that are still being created or deleted are not listed. + + :param catalog_name: str + Name of the Unity Catalog catalog to list datasources from. Required; each request lists the + datasources of a single catalog. + :param filter: str (optional) + Filter expression that restricts the results to datasources with a given name, for example + ``datasource_name = "my_source"``. Only a single exact-match comparison on ``datasource_name`` is + supported, and it matches the datasource's current Unity Catalog name, which reads report in + ``effective_datasource_name``. Filtered results are returned in a single page. + :param page_size: int (optional) + Maximum number of results to return in one page. This is a best-effort upper bound, not a guarantee: + the server may return fewer items than requested (even zero) while more results still exist. If + unset or 0, the server picks a default. + :param page_token: str (optional) + :param schema_name: str (optional) + Name of a Unity Catalog schema in ``catalog_name``. Optional; if set, only the datasources in this + schema are returned, and if omitted, the datasources in all schemas of the catalog are returned. + :param view: :class:`DatasourceView` (optional) + Which fields to return for each datasource. Defaults to ``DATASOURCE_VIEW_FULL``. + ``DATASOURCE_VIEW_BASIC`` returns only the identity and Unity Catalog location fields. + + :returns: Iterator over :class:`Datasource` + + + .. py:method:: list_rules(catalog_name: str [, filter: Optional[str], page_size: Optional[int], page_token: Optional[str], schema_name: Optional[str], view: Optional[RuleView]]) -> Iterator[Rule] + + Lists rules, with optional filtering and pagination. + + :param catalog_name: str + Name of the Unity Catalog catalog to list rules from. Required. + :param filter: str (optional) + Filter expression that restricts the results, such as ``rule_name = "my_rule"`` to return only the + rule with that name. The name is matched against the current UC name, which reads report in + ``effective_rule_name``. + :param page_size: int (optional) + Maximum number of rules to return. The server may return fewer. If unset or 0, the server picks a + default page size. + :param page_token: str (optional) + Page token from a previous ``ListRules`` response, used to fetch the next page of results. + :param schema_name: str (optional) + Name of a Unity Catalog schema in ``catalog_name``. If set, only rules in this schema are returned. + If unset, rules in all schemas of the catalog are returned. + :param view: :class:`RuleView` (optional) + How much of each rule to return. Defaults to ``RULE_VIEW_FULL``. ``RULE_VIEW_BASIC`` returns only + each rule's identity and Unity Catalog location. + + :returns: Iterator over :class:`Rule` + + + .. py:method:: update_alert_config(name: str, alert_config: AlertConfig, update_mask: FieldMask [, etag: Optional[str]]) -> AlertConfig + + Updates the alert delivery configuration for a catalog. + + :param name: str + :param alert_config: :class:`AlertConfig` + :param update_mask: FieldMask + The field mask must be a single string, with multiple fields separated by commas (no spaces). The + field path is relative to the resource object, using a dot (``.``) to navigate sub-fields (e.g., + ``author.given_name``). Specification of elements in sequence or map fields is not allowed, as only + the entire collection field can be specified. Field names must exactly match the resource field + names. + + A field mask of ``*`` indicates full replacement. It’s recommended to always explicitly list the + fields being updated and avoid using ``*`` wildcards, as it can lead to unintended results if the + API changes in the future. + :param etag: str (optional) + Optimistic-concurrency precondition: if set and it does not match the resource's current ``etag``, + the update is rejected as a conflict (HTTP 409). + + :returns: :class:`AlertConfig` + + + .. py:method:: update_common_config(name: str, common_config: CommonConfig, update_mask: FieldMask [, etag: Optional[str]]) -> CommonConfig + + Updates the common Lakewatch configuration for a catalog. + + :param name: str + :param common_config: :class:`CommonConfig` + :param update_mask: FieldMask + The field mask must be a single string, with multiple fields separated by commas (no spaces). The + field path is relative to the resource object, using a dot (``.``) to navigate sub-fields (e.g., + ``author.given_name``). Specification of elements in sequence or map fields is not allowed, as only + the entire collection field can be specified. Field names must exactly match the resource field + names. + + A field mask of ``*`` indicates full replacement. It’s recommended to always explicitly list the + fields being updated and avoid using ``*`` wildcards, as it can lead to unintended results if the + API changes in the future. + :param etag: str (optional) + Optimistic-concurrency precondition: if set and it does not match the resource's current ``etag``, + the update is rejected as a conflict (HTTP 409). + + :returns: :class:`CommonConfig` + + + .. py:method:: update_datasource(name: str, datasource: Datasource, update_mask: FieldMask [, etag: Optional[str]]) -> Datasource + + Updates an existing datasource. + + :param name: str + Resource name of the datasource, in the format ``datasources/{datasource_id}``. Unlike the + datasource's Unity Catalog location (``catalog_name``, ``schema_name``, and ``datasource_name``), + the name does not change when the datasource or its catalog or schema is renamed in Unity Catalog. + :param datasource: :class:`Datasource` + :param update_mask: FieldMask + The field mask must be a single string, with multiple fields separated by commas (no spaces). The + field path is relative to the resource object, using a dot (``.``) to navigate sub-fields (e.g., + ``author.given_name``). Specification of elements in sequence or map fields is not allowed, as only + the entire collection field can be specified. Field names must exactly match the resource field + names. + + A field mask of ``*`` indicates full replacement. It’s recommended to always explicitly list the + fields being updated and avoid using ``*`` wildcards, as it can lead to unintended results if the + API changes in the future. + :param etag: str (optional) + Optimistic-concurrency precondition: if set to a non-empty value that does not match the resource's + current ``etag``, the update is rejected as a conflict (HTTP 409). Omit this field or send an empty + string to update unconditionally. + + :returns: :class:`Datasource` + + + .. py:method:: update_datasource_config(name: str, datasource_config: DatasourceConfig, update_mask: FieldMask [, etag: Optional[str]]) -> DatasourceConfig + + Updates the datasource ingestion configuration for a catalog. + + :param name: str + :param datasource_config: :class:`DatasourceConfig` + :param update_mask: FieldMask + The field mask must be a single string, with multiple fields separated by commas (no spaces). The + field path is relative to the resource object, using a dot (``.``) to navigate sub-fields (e.g., + ``author.given_name``). Specification of elements in sequence or map fields is not allowed, as only + the entire collection field can be specified. Field names must exactly match the resource field + names. + + A field mask of ``*`` indicates full replacement. It’s recommended to always explicitly list the + fields being updated and avoid using ``*`` wildcards, as it can lead to unintended results if the + API changes in the future. + :param etag: str (optional) + Optimistic-concurrency precondition: if set and it does not match the resource's current ``etag``, + the update is rejected as a conflict (HTTP 409). + + :returns: :class:`DatasourceConfig` + + + .. py:method:: update_rule(name: str, rule: Rule, update_mask: FieldMask [, etag: Optional[str]]) -> Rule + + Updates an existing rule. + + :param name: str + Resource name, in the form ``rules/{rule_id}``. Assigned by the server on create; required on + update. + :param rule: :class:`Rule` + The rule with its new field values. ``rule.name`` identifies the rule to update. + :param update_mask: FieldMask + The set of fields to update. Only fields named in the mask are modified. + :param etag: str (optional) + Optimistic-concurrency token from the latest read. A non-empty token that does not match the current + ``etag`` rejects the update as a conflict (HTTP 409). Omit this field or provide an empty string to + update unconditionally. + + :returns: :class:`Rule` + + + .. py:method:: update_rule_config(name: str, rule_config: RuleConfig, update_mask: FieldMask [, etag: Optional[str]]) -> RuleConfig + + Updates the detection-rule configuration for a catalog. + + :param name: str + :param rule_config: :class:`RuleConfig` + :param update_mask: FieldMask + The field mask must be a single string, with multiple fields separated by commas (no spaces). The + field path is relative to the resource object, using a dot (``.``) to navigate sub-fields (e.g., + ``author.given_name``). Specification of elements in sequence or map fields is not allowed, as only + the entire collection field can be specified. Field names must exactly match the resource field + names. + + A field mask of ``*`` indicates full replacement. It’s recommended to always explicitly list the + fields being updated and avoid using ``*`` wildcards, as it can lead to unintended results if the + API changes in the future. + :param etag: str (optional) + Optimistic-concurrency precondition: if set and it does not match the resource's current ``etag``, + the update is rejected as a conflict (HTTP 409). + + :returns: :class:`RuleConfig` + \ No newline at end of file diff --git a/docs/workspace/ml/feature_engineering.rst b/docs/workspace/ml/feature_engineering.rst index ef00b0610..7fbc8dc20 100644 --- a/docs/workspace/ml/feature_engineering.rst +++ b/docs/workspace/ml/feature_engineering.rst @@ -62,6 +62,16 @@ :returns: :class:`Feature` + .. py:method:: create_feature_spec(feature_spec: FeatureSpec) -> FeatureSpec + + Create a feature spec. + + :param feature_spec: :class:`FeatureSpec` + The feature spec to create. + + :returns: :class:`FeatureSpec` + + .. py:method:: create_kafka_config(kafka_config: KafkaConfig) -> KafkaConfig Create a Kafka config. During PrPr, Kafka configs can be read and used when creating features under @@ -118,6 +128,16 @@ + .. py:method:: delete_feature_spec(name: str) + + Delete a feature spec. + + :param name: str + Three-part name (catalog.schema.feature_spec) of the feature spec to delete. + + + + .. py:method:: delete_kafka_config(name: str) Delete a Kafka config. During PrPr, Kafka configs can be read and used when creating features under @@ -178,6 +198,16 @@ :returns: :class:`Feature` + .. py:method:: get_feature_spec(name: str) -> FeatureSpec + + Get a feature spec. + + :param name: str + Three-part name (catalog.schema.feature_spec) of the feature spec to get. + + :returns: :class:`FeatureSpec` + + .. py:method:: get_kafka_config(name: str) -> KafkaConfig Get a Kafka config. During PrPr, Kafka configs can be read and used when creating features under the diff --git a/tests/test_files.py b/tests/test_files.py index d280b5e23..6e170ce68 100644 --- a/tests/test_files.py +++ b/tests/test_files.py @@ -822,6 +822,33 @@ def get_header(self, request: requests.Request) -> requests.Response: return resp +def _storage_proxy_endpoints_response(request: requests.Request) -> requests.Response: + """GET /api/2.0/fs/endpoints as a data-plane workspace answers it: the storage proxy, then the workspace.""" + endpoints = [ + { + "api": "FILES_API", + "api_version": "2.0", + "scheme": "http", + "hostname": "storage-proxy.databricks.com", + "multipart_upload_supported": True, + }, + { + "api": "FILES_API", + "api_version": "2.0", + "scheme": "http", + "hostname": "localhost", + "multipart_upload_supported": False, + }, + ] + resp = requests.Response() + resp.status_code = 200 + resp._content = json.dumps({"endpoints": endpoints}).encode() + resp.headers["Content-Type"] = "application/json" + resp.request = request + resp.url = request.url + return resp + + class PresignedUrlDownloadTestCase: _FILE_PATH = "/testfile/remote/path" # A fake path for the remote location of the file to be downloaded presigned_url_disabled_response = """ @@ -972,34 +999,14 @@ def match_request_to_response( request_url = urlparse(request.url) request_query = parse_qs(request_url.query) - # Storage proxy: workspace ID resolution via SCIM Me. + # Storage proxy: the workspace lists it first. if ( self.use_storage_proxy and request_url.hostname == "localhost" - and request_url.path == "/api/2.0/preview/scim/v2/Me" + and request_url.path == "/api/2.0/fs/endpoints" and request.method == "GET" ): - resp = requests.Response() - resp.status_code = 200 - resp._content = b"{}" - resp.headers["X-Databricks-Org-Id"] = "12345" - resp.request = request - resp.url = request.url - return resp - - # Storage proxy: probe. - elif ( - self.use_storage_proxy - and request_url.hostname == self._expected_hostname - and request_url.path == "/api/2.0/fs/files/DatabricksInternal/Probe/fullstack/wis" - and request.method == "HEAD" - ): - resp = requests.Response() - resp.status_code = 200 - resp._content = b"" - resp.request = request - resp.url = request.url - return resp + return _storage_proxy_endpoints_response(request) # Storage proxy: direct download. elif ( @@ -1083,7 +1090,7 @@ def run_one_case(self, config: Config, monkeypatch, download_mode: DownloadMode, if self.parallel_upload_part_size is not None: config.files_ext_parallel_upload_part_size = self.parallel_upload_part_size if self.use_storage_proxy: - config.experimental_files_ext_enable_storage_proxy = True + config.experimental_files_ext_enable_endpoints_api = True w = WorkspaceClient(config=config) state = PresignedUrlDownloadServerState(self.file_size, self.last_modified) @@ -1671,7 +1678,7 @@ def run_one_case(self, config: Config, use_parallel: bool, source_type: "UploadS self.customize_config(config) if self.use_storage_proxy: - config.experimental_files_ext_enable_storage_proxy = True + config.experimental_files_ext_enable_endpoints_api = True multipart_server_state = self.create_multipart_upload_server_state() single_shot_server_state = SingleShotUploadServerState() @@ -1686,37 +1693,13 @@ def run_one_case(self, config: Config, use_parallel: bool, source_type: "UploadS def custom_matcher(request: requests.Request) -> Optional[requests.Response]: parsed_url = urlparse(request.url) - if self.use_storage_proxy: - # Handle workspace ID resolution via SCIM Me. - if ( - parsed_url.hostname == "localhost" - and parsed_url.path == "/api/2.0/preview/scim/v2/Me" - and request.method == "GET" - ): - resp = requests.Response() - resp.status_code = 200 - resp._content = b"{}" - resp.headers["X-Databricks-Org-Id"] = "12345" - resp.request = request - resp.url = request.url - return resp - - # Handle storage proxy probe with workspace ID. - if ( - parsed_url.hostname == self._expected_hostname - and parsed_url.path == "/api/2.0/fs/files/DatabricksInternal/Probe/fullstack/wis" - and request.method == "HEAD" - ): - probe_query = parse_qs(parsed_url.query) - assert probe_query.get("ew") == [ - "12345" - ], f"Expected ew=12345 in probe URL, got: {probe_query}" - resp = requests.Response() - resp.status_code = 200 - resp._content = b"" - resp.request = request - resp.url = request.url - return resp + if ( + self.use_storage_proxy + and parsed_url.hostname == "localhost" + and parsed_url.path == "/api/2.0/fs/endpoints" + and request.method == "GET" + ): + return _storage_proxy_endpoints_response(request) # Try to match single-shot upload. if ( @@ -2783,6 +2766,9 @@ def __str__(self) -> str: _PARSED_FILES_V2_DP = _Endpoint( _ResolvedEndpoint.FILES_API_V2, "http://storage-proxy.example.com", multipart_upload_supported=True ) +_PARSED_FILES_V3_RIG = _Endpoint( + _ResolvedEndpoint.FILES_API_V3, "https://123.storage.us-west-2.example.com", multipart_upload_supported=True +) @pytest.mark.parametrize( @@ -2794,10 +2780,21 @@ def __str__(self) -> str: expected=[_PARSED_FILES_V2_DP, _PARSED_IO_OFFLOAD_WS, _PARSED_FILES_V2_GIG], ), ResolveEndpointsTestCase( - # This SDK version does not speak Files API 3.0, so RIG is skipped and the rest keeps its order. - name="Files API 3.0 is skipped", + name="Files API 3.0 is recognized", response={"endpoints": [_IO_OFFLOAD_WS, _FILES_V3_RIG, _FILES_V2_GIG]}, - expected=[_PARSED_IO_OFFLOAD_WS, _PARSED_FILES_V2_GIG], + expected=[_PARSED_IO_OFFLOAD_WS, _PARSED_FILES_V3_RIG, _PARSED_FILES_V2_GIG], + ), + ResolveEndpointsTestCase( + # Presigned URLs are coordinated by the workspace; plain HTTP to another host would expose credentials. + name="IO_OFFLOAD over plain HTTP to another host is skipped", + response={"endpoints": [{**_IO_OFFLOAD_WS, "scheme": "http"}, _FILES_V2_GIG]}, + expected=[_PARSED_FILES_V2_GIG], + ), + ResolveEndpointsTestCase( + # Workspace credentials go to a Files API 3.0 host, so it must use TLS. + name="Files API 3.0 over plain HTTP is skipped", + response={"endpoints": [{**_FILES_V3_RIG, "scheme": "http"}, _FILES_V2_GIG]}, + expected=[_PARSED_FILES_V2_GIG], ), ResolveEndpointsTestCase( name="unknown and unspecified api values are skipped", @@ -2864,7 +2861,9 @@ def __str__(self) -> str: ), ResolveEndpointsTestCase( name="only unrecognized entries falls back to the default", - response={"endpoints": [_FILES_V3_RIG, {**_IO_OFFLOAD_WS, "api": "SOMETHING_NEW"}]}, + response={ + "endpoints": [{**_FILES_V2_GIG, "api_version": "9.9"}, {**_IO_OFFLOAD_WS, "api": "SOMETHING_NEW"}] + }, expected=_DEFAULT_ENDPOINTS, ), ResolveEndpointsTestCase( @@ -3062,28 +3061,17 @@ def test_default_endpoints_are_offload_then_the_workspace(config: Config): ] -def test_default_endpoints_start_with_the_storage_proxy_when_probe_succeeds(config: Config): +def test_storage_proxy_flag_does_not_change_the_built_in_order(config: Config, requests_mock): config = config.copy() config.experimental_files_ext_enable_storage_proxy = True + config.experimental_files_ext_enable_endpoints_api = False w = WorkspaceClient(config=config) - # Skip the network probe; this test only covers the selection that follows a successful probe. - w.files._dp_hostname_available = True - assert w.files._default_endpoints() == [ - _Endpoint(_ResolvedEndpoint.FILES_API_V2, w.files._STORAGE_PROXY_HOSTNAME, multipart_upload_supported=True), + assert w.files._resolve_endpoints() == [ + _Endpoint(_ResolvedEndpoint.IO_OFFLOAD, config.host, multipart_upload_supported=True), _Endpoint(_ResolvedEndpoint.FILES_API_V2, config.host, multipart_upload_supported=False), ] - - -def test_default_endpoints_start_with_offload_when_probe_fails(config: Config): - config = config.copy() - config.experimental_files_ext_enable_storage_proxy = True - w = WorkspaceClient(config=config) - w.files._dp_hostname_available = False - - assert w.files._default_endpoints()[0] == _Endpoint( - _ResolvedEndpoint.IO_OFFLOAD, config.host, multipart_upload_supported=True - ) + assert requests_mock.call_count == 0 def test_request_builder_for_offload_is_presigned(config: Config): @@ -3106,3 +3094,93 @@ def test_request_builder_for_files_api_is_storage_proxy(config: Config): assert isinstance(builder, _StorageProxyRequestBuilder) assert builder._hostname == "http://storage-proxy.example.com" + + +def _storage_proxy_client(config: Config, endpoints_api: bool = False) -> WorkspaceClient: + config = config.copy() + config.experimental_files_ext_enable_storage_proxy = True + config.experimental_files_ext_enable_endpoints_api = endpoints_api + return WorkspaceClient(config=config) + + +def test_resolve_endpoints_skips_plain_http_to_a_tls_workspace(config: Config, requests_mock, caplog): + # Workspace credentials must not leave over plain HTTP when the SDK reaches the workspace over TLS. + config = config.copy() + config.host = "https://ws.example.com" + requests_mock.get( + "https://ws.example.com/api/2.0/fs/endpoints", + json={"endpoints": [{**_FILES_V2_GIG, "scheme": "http"}, {**_IO_OFFLOAD_WS, "scheme": "http"}, _FILES_V2_GIG]}, + ) + w = _endpoints_client(config) + + with caplog.at_level(logging.INFO, logger="databricks.sdk.mixins.files"): + assert w.files._resolve_endpoints() == [_PARSED_FILES_V2_GIG] + + assert [rec.levelname for rec in caplog.records if "Skipping" in rec.message] == ["WARNING", "WARNING"] + + +@pytest.mark.parametrize( + "hostname", + [ + "storage-proxy.example.com:abc", + "storage-proxy.example.com:70000", + "[::1", + "*.example.com", + ".example.com", + "localhost?x", + "localhost#x", + "localhost/x", + "storage-proxy.example.com@localhost", + "%6cocalhost", + "localhost\\", + ], +) +@pytest.mark.parametrize("scheme", ["http", "https"]) +def test_resolve_endpoints_skips_malformed_address(config: Config, requests_mock, caplog, scheme, hostname): + # Transfers route by each endpoint's parsed host, so an address that does not parse, or that names more than a + # host and port, must not reach them. + requests_mock.get( + _ENDPOINTS_URL, + json={"endpoints": [{**_FILES_V2_DP, "scheme": scheme, "hostname": hostname}, _FILES_V2_GIG]}, + ) + w = _endpoints_client(config) + + with caplog.at_level(logging.WARNING, logger="databricks.sdk.mixins.files"): + assert w.files._resolve_endpoints() == [_PARSED_FILES_V2_GIG] + + assert any(rec.levelname == "WARNING" and repr(hostname) in rec.message for rec in caplog.records) + + +@pytest.mark.parametrize("hostname", ["storage-proxy.example.com:8443", "[::1]", "[::1]:8443"]) +def test_resolve_endpoints_accepts_host_with_port(config: Config, requests_mock, hostname): + requests_mock.get(_ENDPOINTS_URL, json={"endpoints": [{**_FILES_V2_DP, "scheme": "https", "hostname": hostname}]}) + w = _endpoints_client(config) + + assert w.files._resolve_endpoints() == [_Endpoint(_ResolvedEndpoint.FILES_API_V2, f"https://{hostname}", True)] + + +def test_resolve_endpoints_ignores_the_storage_proxy_flag(config: Config, requests_mock): + requests_mock.get(_ENDPOINTS_URL, json={"endpoints": [_FILES_V2_GIG]}) + w = _storage_proxy_client(config, endpoints_api=True) + + assert w.files._resolve_endpoints() == [_PARSED_FILES_V2_GIG] + assert requests_mock.call_count == 1 + + +def test_resolve_endpoints_failure_falls_back_to_the_built_in_order_with_the_storage_proxy_flag( + config: Config, requests_mock +): + requests_mock.get(_ENDPOINTS_URL, status_code=404, json={"error_code": "NOT_FOUND", "message": "no"}) + w = _storage_proxy_client(config, endpoints_api=True) + + assert w.files._resolve_endpoints() == _DEFAULT_ENDPOINTS + assert requests_mock.call_count == 1 + + +def test_uses_api_client_for_the_workspace_host_however_it_is_spelled(config: Config): + w = WorkspaceClient(config=config) + + assert w.files._uses_api_client(_Endpoint(_ResolvedEndpoint.FILES_API_V2, "HTTP://LocalHost:80", False)) + assert not w.files._uses_api_client(_Endpoint(_ResolvedEndpoint.FILES_API_V2, "http://localhost:8080", False)) + assert not w.files._uses_api_client(_Endpoint(_ResolvedEndpoint.FILES_API_V2, "http://other", False)) + assert not w.files._uses_api_client(_Endpoint(_ResolvedEndpoint.IO_OFFLOAD, "http://localhost", True)) diff --git a/tests/test_files_walk.py b/tests/test_files_walk.py index 3e8317a19..42ad57b7a 100644 --- a/tests/test_files_walk.py +++ b/tests/test_files_walk.py @@ -1,8 +1,11 @@ """Which endpoints FilesExt transfers files through, and what a caller sees when they fail. +With experimental_files_ext_enable_endpoints_api, the endpoint list comes from GET /api/2.0/fs/endpoints. Otherwise, +or when the workspace does not answer with a usable list, it is the built-in order. + FakeFilesServer stands in for the hosts a transfer can reach: the workspace (the Files API and the presigned URL -coordination API), the data-plane storage proxy, and cloud storage behind presigned URLs. It records every -operation per endpoint. +coordination API), the data-plane storage proxy, a Files API 3.0 host (RIG), and cloud storage behind presigned URLs. +It records every operation per endpoint. The scenario matrix runs every combination of endpoint configuration, endpoint availability and transfer input, and checks only what a caller sees: the file stored or downloaded, or the error raised. The outcome each scenario @@ -28,18 +31,46 @@ from databricks.sdk import WorkspaceClient from databricks.sdk.core import Config +from databricks.sdk.errors.platform import AlreadyExists, InternalError, PermissionDenied from databricks.sdk.mixins import files as files_mixin -from databricks.sdk.errors.platform import AlreadyExists from .test_files_utils import NonSeekableBuffer WORKSPACE = "localhost" STORAGE_PROXY = "storage-proxy.databricks.com" +RIG = "rig.example.com" CLOUD = "cloud.example.com" FILE_PATH = "/Volumes/catalog/schema/volume/file.bin" LAST_MODIFIED = "Thu, 28 Nov 2024 16:39:14 GMT" PART_SIZE = 1024 +DP_ENTRY = { + "api": "FILES_API", + "api_version": "2.0", + "scheme": "http", + "hostname": STORAGE_PROXY, + "multipart_upload_supported": True, +} +RIG_ENTRY = { + "api": "FILES_API", + "api_version": "3.0", + "scheme": "https", + "hostname": RIG, + "multipart_upload_supported": True, +} +OFFLOAD_ENTRY = {"api": "IO_OFFLOAD", "scheme": "http", "hostname": WORKSPACE, "multipart_upload_supported": True} +FLOOR_ENTRY = { + "api": "FILES_API", + "api_version": "2.0", + "scheme": "http", + "hostname": WORKSPACE, + "multipart_upload_supported": False, +} +# What a data-plane caller gets: storage proxy, then presigned URLs, then the single-shot floor. +DATA_PLANE = [DP_ENTRY, OFFLOAD_ENTRY, FLOOR_ENTRY] +# What a caller that can reach RIG gets: Files API 3.0 there, then presigned URLs, then the single-shot floor. +RIG_PLAN = [RIG_ENTRY, OFFLOAD_ENTRY, FLOOR_ENTRY] + _ERROR_CODES = { 400: "INVALID_PARAMETER_VALUE", 403: "PERMISSION_DENIED", @@ -51,6 +82,9 @@ # Operations on presigned cloud URLs. They must never carry workspace credentials. _PRESIGNED_OPS = {"offload.part", "offload.get"} +# Requests that pick the endpoint list or read file metadata rather than transfer data. +_NOT_ENDPOINTS = {"endpoints", "scim", "probe", "workspace"} + class BreakAfter: """A download fault: the connection drops after `n` bytes of the body.""" @@ -96,8 +130,16 @@ def always(fault: Fault) -> Callable[[int], Optional[Fault]]: return lambda call: fault +def on_calls(fault: Fault, *calls: int) -> Callable[[int], Optional[Fault]]: + return lambda call: fault if call in calls else None + + class FakeFilesServer: - def __init__(self, faults: Optional[Dict[str, Callable[[int], Optional[Fault]]]] = None): + def __init__( + self, faults: Optional[Dict[str, Callable[[int], Optional[Fault]]]] = None, plan: Optional[List[dict]] = None + ): + # The list GET /api/2.0/fs/endpoints answers with; None means the workspace does not serve it. + self.plan = plan self.faults = faults or {} self.download_content = b"" self.calls: List[str] = [] @@ -113,6 +155,18 @@ def __init__(self, faults: Optional[Dict[str, Callable[[int], Optional[Fault]]]] self.violations: List[str] = [] self._lock = Lock() + def ops(self, prefix: str) -> List[str]: + return [op for op in self.calls if op.startswith(prefix)] + + def endpoints_touched(self) -> List[str]: + """Endpoints in the order they were first used, ignoring the probe and metadata requests.""" + seen = [] + for op in self.calls: + endpoint = op.split(".")[0] + if endpoint not in _NOT_ENDPOINTS and endpoint not in seen: + seen.append(endpoint) + return seen + def stored_file(self) -> Optional[bytes]: """The file as a caller would read it back, whichever endpoint stored it.""" assert len(self.stored) <= 1, f"the file was stored by more than one endpoint: {sorted(self.stored)}" @@ -129,10 +183,24 @@ def _operation(self, request: requests.Request) -> str: url = urlparse(request.url) query = parse_qs(url.query) host, method, path = url.hostname, request.method, url.path + if host == WORKSPACE and path == "/api/2.0/fs/endpoints" and method == "GET": + return "endpoints" if host == WORKSPACE and path == "/api/2.0/preview/scim/v2/Me" and method == "GET": return "scim" if host == STORAGE_PROXY and path == "/api/2.0/fs/files/DatabricksInternal/Probe/fullstack/wis": return "probe" + if path == f"/api/3.0/fs/files{FILE_PATH}": + assert host == RIG, f"Files API 3.0 request to {host}" + assert not {"upload_type", "session_token"} & query.keys(), f"2.0 parameters on {request.url}" + action = query.get("action", [None])[0] + if method == "POST" and action in ("initiate-upload", "complete-upload", "abort-upload"): + return f"rig.{action.split('-')[0]}" + if method == "PUT" and "part_number" in query: + return "rig.part" + if method == "PUT": + return "rig.single" + if method == "GET": + return "rig.get" if path == f"/api/2.0/fs/files{FILE_PATH}": assert host in (WORKSPACE, STORAGE_PROXY), f"Files API 2.0 request to {host}" endpoint = "dp" if host == STORAGE_PROXY else None @@ -180,6 +248,11 @@ def handle(self, request: requests.Request) -> requests.Response: return self._respond(op, request, fault) def _respond(self, op: str, request: requests.Request, fault: Optional[Fault]) -> requests.Response: + if op == "endpoints": + if self.plan is None: + body = {"error_code": "NOT_FOUND", "message": "no endpoints API"} + return requests_mock.create_response(request, status_code=404, json=body) + return requests_mock.create_response(request, json={"endpoints": self.plan}) if op == "scim": return requests_mock.create_response(request, json={}, headers={"X-Databricks-Org-Id": "12345"}) if op == "probe": @@ -187,6 +260,12 @@ def _respond(self, op: str, request: requests.Request, fault: Optional[Fault]) - endpoint, action = op.split(".", 1) url = urlparse(request.url) query = parse_qs(url.query) + if endpoint == "rig" and action in ("part", "complete", "abort"): + assert query.get("upload_id") == ["upload-rig"], f"{op} without the upload_id from initiate-upload" + if endpoint == "rig" and action in ("part", "single"): + assert request.headers.get("Content-Length") is not None, f"Files API 3.0 requires a Content-Length on {op}" + if op == "rig.initiate": + return requests_mock.create_response(request, json={"upload_id": "upload-rig"}) if action == "initiate": return requests_mock.create_response( request, json={"multipart_upload": {"session_token": f"token-{endpoint}"}} @@ -214,7 +293,8 @@ def _respond(self, op: str, request: requests.Request, fault: Optional[Fault]) - self.stored[endpoint] = b"".join(self.parts[endpoint][n] for n in numbers) return requests_mock.create_response(request, json={}) if action == "abort": - assert query.get("session_token") == [f"token-{endpoint}"] + if endpoint != "rig": + assert query.get("session_token") == [f"token-{endpoint}"] self.aborted.add(endpoint) return requests_mock.create_response(request, json={}) if action == "single": @@ -262,9 +342,10 @@ def _read_body(request: requests.Request) -> bytes: return body.encode() if isinstance(body, str) else bytes(body) -def _client(config: Config, storage_proxy: bool = True) -> WorkspaceClient: +def _client(config: Config, storage_proxy: bool = False, endpoints_api: bool = False) -> WorkspaceClient: config = config.copy() config.experimental_files_ext_enable_storage_proxy = storage_proxy + config.experimental_files_ext_enable_endpoints_api = endpoints_api config._header_factory = lambda: {"Authorization": "Bearer workspace-token"} config.files_ext_multipart_upload_min_stream_size = PART_SIZE config.files_ext_parallel_download_min_file_size = PART_SIZE @@ -321,11 +402,17 @@ class EndpointList: # The endpoints a transfer can use, best first. endpoints: Tuple[str, ...] storage_proxy: bool = False + # Served by GET /api/2.0/fs/endpoints; None leaves the client on its built-in order. + plan: Optional[List[dict]] = None ENDPOINT_LISTS = [ EndpointList("built-in", ("presigned", "workspace")), - EndpointList("built-in with the storage proxy", ("proxy", "presigned", "workspace"), storage_proxy=True), + # The storage-proxy flag is deprecated and has no effect. + EndpointList("built-in with the storage-proxy flag", ("presigned", "workspace"), storage_proxy=True), + EndpointList("server: proxy, presigned, workspace", ("proxy", "presigned", "workspace"), plan=DATA_PLANE), + EndpointList("server: rig, presigned, workspace", ("rig", "presigned", "workspace"), plan=RIG_PLAN), + EndpointList("server: presigned only", ("presigned",), plan=[OFFLOAD_ENTRY]), ] # up: the endpoint works. down: every request to it fails, as when it can't be reached. rejects data: its probe and @@ -344,9 +431,10 @@ class EndpointList: "proxy": ["probe", "dp.initiate", "dp.part", "dp.complete", "dp.abort", "dp.single", "dp.get"], "presigned": ["offload.part", "offload.get"], "workspace": ["floor.single", "floor.get"], + "rig": ["rig.initiate", "rig.part", "rig.complete", "rig.abort", "rig.single", "rig.get"], } # The fake's name for each endpoint that can come first in a list. -_FAKE_NAMES = {"proxy": "dp", "presigned": "offload"} +_FAKE_NAMES = {"proxy": "dp", "presigned": "offload", "rig": "rig"} SUCCEEDS = "succeeds" @@ -418,20 +506,24 @@ def _scenarios(transfer: str, modes: List[str], sources: List[Optional[str]]) -> DOWNLOAD_SCENARIOS = _scenarios("download", DOWNLOAD_MODES, [None]) -def _first_endpoint(scenario: Scenario) -> str: - """The endpoint a transfer starts on: the storage proxy when it is enabled and its probe reaches it, otherwise - presigned URLs.""" - if scenario.endpoint_list.storage_proxy and dict(scenario.availability)["proxy"] != DOWN: - return "proxy" - return "presigned" +def _walk_order(scenario: Scenario) -> List[str]: + """The endpoints a transfer tries, in order. The built-in order is presigned URLs, then the workspace. A list from + the server is used as it is.""" + endpoints = scenario.endpoint_list.endpoints + if scenario.endpoint_list.plan is not None: + # An upload with only presigned URLs listed gets the workspace after them, for small files. + if "workspace" not in endpoints and scenario.transfer == "upload": + return list(endpoints) + ["workspace"] + return list(endpoints) + return ["presigned", "workspace"] # Retries that run out raise TimeoutError from the API client's retry budget, which covers the workspace, and -# RuntimeError from the bounded retries of transfers to presigned URLs and the storage proxy. +# RuntimeError from the bounded retries of transfers to other hosts. def expected_upload(scenario: Scenario): """The outcome an upload has. - - Endpoints are tried in order: the first endpoint, then the workspace, which takes single-shot uploads only. + - Endpoints are tried in order, ending with the workspace, which takes single-shot uploads only. - A file under one part is sent in a single shot, which presigned URLs can't take, so they are skipped. - A non-seekable stream's size is known only once it is read, so it is read up to one part first: a short one is sent in a single shot as above. @@ -443,44 +535,51 @@ def expected_upload(scenario: Scenario): states = dict(scenario.availability) size = SIZES[scenario.size] non_seekable = scenario.source == "non-seekable" - first = _first_endpoint(scenario) - if size < PART_SIZE: - if first == "proxy" and states[first] == UP: - return SUCCEEDS - else: - if states[first] == UP: + for endpoint in _walk_order(scenario): + state = states[endpoint] + if endpoint == "workspace": + if state == UP: + return SUCCEEDS + return requests.ConnectionError if non_seekable and size >= PART_SIZE else TimeoutError + if size < PART_SIZE and endpoint == "presigned": + continue + if state == UP: return SUCCEEDS - if states[first] == BREAKS: + if state == BREAKS: return RuntimeError - if states["workspace"] == UP: - return SUCCEEDS - return requests.ConnectionError if non_seekable and size >= PART_SIZE else TimeoutError + raise AssertionError("every walk ends with the workspace") def expected_download(scenario: Scenario): """The outcome a download has. - - Endpoints are tried in order: the first endpoint, then the workspace. - - A failure before any bytes arrived moves the download to the next endpoint. + - Endpoints are tried in order. + - A failure before any bytes arrived moves the download to the next endpoint. When no endpoint could serve it, + the last failure is raised. - After bytes arrived, the download stays on that endpoint. A stream resumes there and, if that fails too, raises its broken connection; a parallel download raises the failure of a later range. """ states = dict(scenario.availability) - state = states[_first_endpoint(scenario)] - if state == UP: - return SUCCEEDS - if state == BREAKS: - return RuntimeError if scenario.mode == "parallel" else requests.exceptions.ChunkedEncodingError - return SUCCEEDS if states["workspace"] == UP else TimeoutError + for endpoint in _walk_order(scenario): + state = states[endpoint] + if state == UP: + return SUCCEEDS + if state == BREAKS: + return RuntimeError if scenario.mode == "parallel" else requests.exceptions.ChunkedEncodingError + if endpoint == "workspace": + return TimeoutError + # The last endpoint retries a dropped connection until its budget runs out, but not a rejected transfer. + return PermissionDenied if state == REJECTS else RuntimeError def _scenario_client(config: Config, scenario: Scenario) -> WorkspaceClient: - return _client(config, storage_proxy=scenario.endpoint_list.storage_proxy) + endpoint_list = scenario.endpoint_list + return _client(config, storage_proxy=endpoint_list.storage_proxy, endpoints_api=endpoint_list.plan is not None) @pytest.mark.parametrize("scenario", UPLOAD_SCENARIOS, ids=str) def test_upload_scenario(config, tmp_path, scenario): - server = FakeFilesServer(scenario.faults()) + server = FakeFilesServer(scenario.faults(), scenario.endpoint_list.plan) w = _scenario_client(config, scenario) data = _DATA[scenario.size] expected = expected_upload(scenario) @@ -496,6 +595,7 @@ def upload(): _run(server, upload) assert server.stored_file() is None server.assert_credentials() + assert "probe" not in server.calls and "scim" not in server.calls def _download(w: WorkspaceClient) -> bytes: @@ -520,7 +620,7 @@ def _download_to(w: WorkspaceClient, use_parallel: bool, parallelism: Optional[i @pytest.mark.parametrize("scenario", DOWNLOAD_SCENARIOS, ids=str) def test_download_scenario(config, scenario): - server = FakeFilesServer(scenario.faults()) + server = FakeFilesServer(scenario.faults(), scenario.endpoint_list.plan) server.download_content = _DATA[scenario.size] w = _scenario_client(config, scenario) download = { @@ -536,6 +636,7 @@ def test_download_scenario(config, scenario): with pytest.raises(expected): _run(server, download) server.assert_credentials() + assert "probe" not in server.calls and "scim" not in server.calls # ======================== Uploads before commit ======================== @@ -687,10 +788,10 @@ def test_an_empty_upload_is_sent_with_a_content_length(config, tmp_path, source, @pytest.mark.parametrize("source", ["file", "seekable"]) def test_storage_proxy_single_shot_onto_an_existing_file_is_already_exists(config, tmp_path, source): - server = FakeFilesServer({"dp.single": always(412)}) + server = FakeFilesServer({"dp.single": always(412)}, DATA_PLANE) config = config.copy() config.workspace_id = "1234" - w = _client(config) + w = _client(config, endpoints_api=True) with pytest.raises(AlreadyExists): _run(server, lambda: _upload(w, source, SMALL, False, tmp_path, overwrite=False)) @@ -698,3 +799,313 @@ def test_storage_proxy_single_shot_onto_an_existing_file_is_already_exists(confi # The conflict is the caller's, so the upload does not move on to the workspace. assert "floor.single" not in server.calls assert server.request_headers["dp.single"][0]["X-Databricks-Workspace-Id"] == "1234" + + +# ======================== Files API 3.0 (white box) ======================== + + +def _download_server(faults=None, plan=None) -> FakeFilesServer: + server = FakeFilesServer(faults, plan) + server.download_content = LARGE + return server + + +@pytest.mark.parametrize("use_parallel", [False, True], ids=["sequential", "parallel"]) +@pytest.mark.parametrize("source", SOURCES) +def test_rig_upload_uses_the_files_api_3_protocol(config, tmp_path, source, use_parallel): + server = FakeFilesServer(plan=RIG_PLAN) + w = _client(config, endpoints_api=True) + + _run(server, lambda: _upload(w, source, LARGE, use_parallel, tmp_path)) + + assert server.endpoints_touched() == ["rig"] + assert server.ops("rig") == ["rig.initiate"] + ["rig.part"] * 5 + ["rig.complete"] + assert server.stored == {"rig": LARGE} + server.assert_credentials() + + +@pytest.mark.parametrize("source", ["file", "seekable"]) +def test_rig_small_upload_uses_single_shot(config, tmp_path, source): + server = FakeFilesServer(plan=RIG_PLAN) + w = _client(config, endpoints_api=True) + + _run(server, lambda: _upload(w, source, SMALL, False, tmp_path)) + + assert server.calls == ["endpoints", "rig.single"] + assert server.stored == {"rig": SMALL} + server.assert_credentials() + + +@pytest.mark.parametrize("use_parallel", [False, True], ids=["sequential", "parallel"]) +@pytest.mark.parametrize("source", SOURCES) +def test_rig_first_part_failure_aborts_on_rig_and_moves_to_presigned(config, tmp_path, source, use_parallel): + server = FakeFilesServer({"rig.part": always(403)}, RIG_PLAN) + w = _client(config, endpoints_api=True) + + _run(server, lambda: _upload(w, source, LARGE, use_parallel, tmp_path)) + + assert server.endpoints_touched() == ["rig", "offload"] + assert server.counts["rig.part"] == 1 + assert server.aborted == {"rig"} + assert server.stored == {"offload": LARGE} + server.assert_credentials() + + +@pytest.mark.parametrize("use_parallel", [False, True], ids=["sequential", "parallel"]) +@pytest.mark.parametrize("source", SOURCES) +def test_rig_later_part_failure_raises_without_walking(config, tmp_path, source, use_parallel): + server = FakeFilesServer({"rig.part": on_calls(500, 2)}, RIG_PLAN) + w = _client(config, endpoints_api=True) + + with pytest.raises(InternalError): + _run(server, lambda: _upload(w, source, LARGE, use_parallel, tmp_path)) + + assert server.endpoints_touched() == ["rig"] + assert server.aborted == {"rig"} + assert server.stored == {} + + +@pytest.mark.parametrize("source", ["file", "seekable"]) +def test_rig_single_shot_onto_an_existing_file_is_already_exists(config, tmp_path, source): + server = FakeFilesServer({"rig.single": always(409)}, RIG_PLAN) + w = _client(config, endpoints_api=True) + + with pytest.raises(AlreadyExists): + _run(server, lambda: _upload(w, source, SMALL, False, tmp_path, overwrite=False)) + + assert server.endpoints_touched() == ["rig"] + assert parse_qs(urlparse(server.urls["rig.single"][0]).query) == {"overwrite": ["false"]} + + +@pytest.mark.parametrize("use_parallel", [False, True], ids=["sequential", "parallel"]) +@pytest.mark.parametrize("source", SOURCES) +def test_rig_complete_onto_an_existing_file_is_already_exists(config, tmp_path, source, use_parallel): + server = FakeFilesServer({"rig.complete": always(409)}, RIG_PLAN) + w = _client(config, endpoints_api=True) + + with pytest.raises(AlreadyExists): + _run(server, lambda: _upload(w, source, LARGE, use_parallel, tmp_path, overwrite=False)) + + # The parts landed, so the upload stays on RIG and aborts there. + assert server.endpoints_touched() == ["rig"] + assert server.aborted == {"rig"} + + +def test_rig_initiate_failure_moves_to_presigned(config, tmp_path): + server = FakeFilesServer({"rig.initiate": always(500)}, RIG_PLAN) + w = _client(config, endpoints_api=True) + + _run(server, lambda: _upload(w, "file", LARGE, False, tmp_path)) + + assert server.endpoints_touched() == ["rig", "offload"] + assert server.stored == {"offload": LARGE} + + +def test_rig_parts_are_capped_at_the_files_api_3_part_size_limit(config, tmp_path, monkeypatch): + monkeypatch.setattr(files_mixin, "_FILES_API_V3_MAX_PART_SIZE", PART_SIZE) + server = FakeFilesServer(plan=RIG_PLAN) + w = _client(config, endpoints_api=True) + local = tmp_path / "source.bin" + local.write_bytes(LARGE) + + _run(server, lambda: w.files.upload_from(FILE_PATH, str(local), part_size=2 * PART_SIZE, use_parallel=False)) + + assert server.stored == {"rig": LARGE} + assert max(len(part) for part in server.parts["rig"].values()) == PART_SIZE + + +@pytest.mark.parametrize("download", ["stream", "to-file"]) +def test_rig_download(config, download): + server = _download_server(plan=RIG_PLAN) + w = _client(config, endpoints_api=True) + result = {} + + _run(server, lambda: result.update(data=_download(w) if download == "stream" else _download_to(w, False))) + + assert result["data"] == LARGE + assert server.endpoints_touched() == ["rig"] + server.assert_credentials() + + +def test_rig_parallel_download_uses_rig_for_every_range(config): + server = _download_server(plan=RIG_PLAN) + w = _client(config, endpoints_api=True) + result = {} + + _run(server, lambda: result.update(data=_download_to(w, True, parallelism=2))) + + assert result["data"] == LARGE + assert server.endpoints_touched() == ["rig"] + assert len(server.ranges["rig"]) == 5 + server.assert_credentials() + + +@pytest.mark.parametrize("download", ["stream", "to-file", "parallel"]) +def test_rig_download_failure_before_any_bytes_moves_to_presigned(config, download): + server = _download_server({"rig.get": always(500)}, RIG_PLAN) + w = _client(config, endpoints_api=True) + result = {} + + def run(): + if download == "stream": + result.update(data=_download(w)) + else: + result.update(data=_download_to(w, download == "parallel", parallelism=2)) + + _run(server, run) + + assert result["data"] == LARGE + assert server.endpoints_touched() == ["rig", "offload"] + server.assert_credentials() + + +def test_rig_download_reopens_on_rig_after_a_dropped_connection(config): + server = _download_server({"rig.get": on_calls(BreakAfter(3 * 512), 1)}, RIG_PLAN) + w = _client(config, endpoints_api=True) + result = {} + + _run(server, lambda: result.update(data=_download(w))) + + assert result["data"] == LARGE + assert server.endpoints_touched() == ["rig"] + assert server.counts["rig.get"] == 2 + + +# ======================== Endpoints API: the workspace's list (white box) ======================== + + +@pytest.mark.parametrize("source", SOURCES) +def test_upload_follows_the_workspace_list(config, tmp_path, source): + server = FakeFilesServer({"dp.part": always(403)}, DATA_PLANE) + w = _client(config, storage_proxy=False, endpoints_api=True) + + _run(server, lambda: _upload(w, source, LARGE, False, tmp_path)) + + assert server.endpoints_touched() == ["dp", "offload"] + assert server.stored == {"offload": LARGE} + server.assert_credentials() + + +def test_server_list_takes_precedence_over_the_storage_proxy_probe(config, tmp_path): + server = FakeFilesServer(plan=[OFFLOAD_ENTRY, FLOOR_ENTRY]) + w = _client(config, storage_proxy=True, endpoints_api=True) + + _run(server, lambda: _upload(w, "file", LARGE, False, tmp_path)) + + assert "scim" not in server.calls and "probe" not in server.calls + assert server.endpoints_touched() == ["offload"] + + +@pytest.mark.parametrize("storage_proxy", [False, True], ids=["default", "storage-proxy-flag"]) +def test_unusable_endpoints_api_uses_the_built_in_order(config, tmp_path, storage_proxy): + server = FakeFilesServer() + w = _client(config, storage_proxy=storage_proxy, endpoints_api=True) + + _run(server, lambda: _upload(w, "file", LARGE, False, tmp_path)) + + assert server.calls[0] == "endpoints" + assert "probe" not in server.calls + assert server.endpoints_touched() == ["offload"] + + +def test_endpoints_are_resolved_once_per_client(config, tmp_path): + server = _download_server(plan=DATA_PLANE) + w = _client(config, storage_proxy=False, endpoints_api=True) + + def transfers(): + _upload(w, "file", LARGE, False, tmp_path) + _upload(w, "seekable", SMALL, False, tmp_path) + _download(w) + + _run(server, transfers) + + assert server.counts["endpoints"] == 1 + + +@pytest.mark.parametrize( + "source, data, use_parallel", + [ + ("file", SMALL, False), + # Their size is only known once read, so they reach the presigned path before they turn out to need a + # single-shot upload. + ("non-seekable", SMALL, False), + ("non-seekable", b"", False), + ("non-seekable", b"", True), + ], +) +def test_small_upload_with_only_presigned_listed_uses_the_workspace(config, tmp_path, source, data, use_parallel): + server = FakeFilesServer(plan=[OFFLOAD_ENTRY]) + w = _client(config, storage_proxy=False, endpoints_api=True) + + _run(server, lambda: _upload(w, source, data, use_parallel, tmp_path)) + + assert server.calls[0] == "endpoints" + assert server.stored == {"floor": data} + + +def test_upload_walks_the_non_data_plane_order_from_presigned_to_rig(config, tmp_path): + server = FakeFilesServer({"offload.part": always(403)}, [OFFLOAD_ENTRY, RIG_ENTRY, FLOOR_ENTRY]) + w = _client(config, storage_proxy=False, endpoints_api=True) + + _run(server, lambda: _upload(w, "file", LARGE, False, tmp_path)) + + assert server.endpoints_touched() == ["offload", "rig"] + assert server.aborted == {"offload"} + assert server.stored == {"rig": LARGE} + server.assert_credentials() + + +@pytest.mark.parametrize("use_parallel", [False, True], ids=["sequential", "parallel"]) +@pytest.mark.parametrize("source", SOURCES) +def test_large_upload_on_single_shot_only_endpoint_uses_single_shot(config, tmp_path, source, use_parallel): + server = FakeFilesServer(plan=[FLOOR_ENTRY]) + w = _client(config, storage_proxy=False, endpoints_api=True) + + _run(server, lambda: _upload(w, source, LARGE, use_parallel, tmp_path)) + + assert server.calls == ["endpoints", "floor.single"] + assert server.stored == {"floor": LARGE} + + +@pytest.mark.parametrize("source", ["file", "seekable"]) +def test_failed_single_shot_is_replayed_from_a_rereadable_source(config, tmp_path, source): + single_shot_dp = {**DP_ENTRY, "multipart_upload_supported": False} + server = FakeFilesServer({"dp.single": always(500)}, [single_shot_dp, OFFLOAD_ENTRY]) + w = _client(config, storage_proxy=False, endpoints_api=True) + + _run(server, lambda: _upload(w, source, LARGE, False, tmp_path)) + + assert server.endpoints_touched() == ["dp", "offload"] + assert server.stored == {"offload": LARGE} + + +def test_failed_single_shot_of_a_non_seekable_stream_cannot_be_replayed(config, tmp_path): + # The whole stream was handed to the single-shot PUT, so nothing is left to send to the next endpoint. + single_shot_dp = {**DP_ENTRY, "multipart_upload_supported": False} + server = FakeFilesServer({"dp.single": always(500)}, [single_shot_dp, OFFLOAD_ENTRY]) + w = _client(config, storage_proxy=False, endpoints_api=True) + + with pytest.raises(InternalError): + _run(server, lambda: _upload(w, "non-seekable", LARGE, False, tmp_path)) + + assert server.endpoints_touched() == ["dp"] + + +@pytest.mark.parametrize("download", ["stream", "to-file", "parallel"]) +def test_download_follows_the_workspace_list(config, download): + server = _download_server({"offload.get": always(403)}, [OFFLOAD_ENTRY, FLOOR_ENTRY]) + w = _client(config, storage_proxy=True, endpoints_api=True) + result = {} + + def run(): + if download == "stream": + result.update(data=_download(w)) + else: + result.update(data=_download_to(w, download == "parallel", parallelism=2)) + + _run(server, run) + + assert result["data"] == LARGE + assert "probe" not in server.calls + assert server.endpoints_touched() == ["offload", "floor"] + server.assert_credentials()