From ecbef18f2609c25e561dc6a4f05db71b355d25ec Mon Sep 17 00:00:00 2001 From: "databricks-ci-ghec-1[bot]" <184311507+databricks-ci-ghec-1[bot]@users.noreply.github.com> Date: Wed, 16 Sep 2026 03:16:49 +0000 Subject: [PATCH] Release databricks-sdk-py --- .codegen.json | 2 +- .codegen/_last_sha | 2 +- CHANGELOG.md | 5 + databricks/sdk/service/aifunctions.py | 3 +- databricks/sdk/service/apps.py | 3 +- databricks/sdk/service/files.py | 2 +- databricks/sdk/service/jobs.py | 11 +- databricks/sdk/service/ml.py | 4 + databricks/sdk/service/pipelines.py | 15 +- databricks/sdk/version.py | 2 +- docs/workspace/aifunctions/ai_functions.rst | 3 +- docs/workspace/files/files.rst | 2 +- docs/workspace/ml/feature_store.rst | 4 + tests/test_compute_mixins.py | 170 +++++++++++++++++ tests/test_data_plane.py | 66 ++++++- tests/test_dbfs_mixins.py | 198 +++++++++++++------- tests/test_files.py | 47 +++++ tests/test_jobs_mixin.py | 59 +++++- tests/test_workspace.py | 92 +++++++++ 19 files changed, 597 insertions(+), 93 deletions(-) mode change 100755 => 100644 docs/workspace/files/files.rst mode change 100755 => 100644 docs/workspace/ml/feature_store.rst create mode 100644 tests/test_workspace.py diff --git a/.codegen.json b/.codegen.json index 4a145805b..07d77d1ce 100644 --- a/.codegen.json +++ b/.codegen.json @@ -10,7 +10,7 @@ "tagging.py", "tagging.py.lock" ], - "allow_from_source": [ + "allow_from_packages": [ "**/*.{md,py,rst,png,svg,gif,tmpl,ps1}", "**/py.typed", ".codegen/{_last_sha,_openapi_sha}", diff --git a/.codegen/_last_sha b/.codegen/_last_sha index d720f4024..ecc74edf8 100644 --- a/.codegen/_last_sha +++ b/.codegen/_last_sha @@ -1 +1 @@ -8758d3b71cae7ab16f4bfb80629e5c156885ccbb \ No newline at end of file +94a696f2e030bf6a63214df926bc6952522c12dd \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md index d47b1e399..63951e230 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,10 @@ # Version changelog +## Release v0.140.0 (2026-09-16) + +### API Changes +* Add `group_name` field for `databricks.sdk.service.pipelines.RunAs`. + ## Release v0.139.0 (2026-09-13) ### API Changes diff --git a/databricks/sdk/service/aifunctions.py b/databricks/sdk/service/aifunctions.py index 04ab3dc9a..bc7b25568 100644 --- a/databricks/sdk/service/aifunctions.py +++ b/databricks/sdk/service/aifunctions.py @@ -710,8 +710,7 @@ def ai_extract(self, content: any, schema: any, *, options: Optional[AiExtractOp def ai_parse_document( self, content: str, *, options: Optional[AiParseDocumentOptions] = None ) -> AiParseDocumentResponse: - """Parse structured content from unstructured documents. For REST API requests, the default rate limit is - 120 pages per minute per workspace. Contact your Databricks account team to request a higher limit. + """Parse structured content from unstructured documents. :param content: str The document to parse, given as a Unity Catalog volume path to the source file (the REST API accepts diff --git a/databricks/sdk/service/apps.py b/databricks/sdk/service/apps.py index b5bedd687..256795156 100644 --- a/databricks/sdk/service/apps.py +++ b/databricks/sdk/service/apps.py @@ -1878,8 +1878,7 @@ class ComputeState(Enum): @dataclass class ComputeStatus: active_instances: Optional[int] = None - """The number of compute instances currently serving requests for this application. An instance is - considered active if it is reachable and ready to handle requests.""" + """The number of compute instances used and billed for this application.""" message: Optional[str] = None """Compute status message""" diff --git a/databricks/sdk/service/files.py b/databricks/sdk/service/files.py index f897417cd..5e606f2e3 100644 --- a/databricks/sdk/service/files.py +++ b/databricks/sdk/service/files.py @@ -787,7 +787,7 @@ class FilesAPI: The API supports `Unity Catalog volumes `__, where files and directories to operate on are specified using their volume URI path, which follows the format - /Volumes/<catalog_name>/<schema_name>/<volume_name>/<path_to_file>. + /Volumes////. The Files API has two distinct endpoints, one for working with files (``/fs/files``) and another one for working with directories (``/fs/directories``). Both endpoints use the standard HTTP methods GET, HEAD, diff --git a/databricks/sdk/service/jobs.py b/databricks/sdk/service/jobs.py index 13c158081..8bb64d0e8 100644 --- a/databricks/sdk/service/jobs.py +++ b/databricks/sdk/service/jobs.py @@ -3593,15 +3593,16 @@ def from_dict(cls, d: Dict[str, Any]) -> JobPermissionsDescription: @dataclass class JobRunAs: - """Write-only setting. Specifies the user or service principal that the job runs as. If not + """Write-only setting. Specifies the user, service principal, or group that the job runs as. If not specified, the job runs as the user who created the job. - Either ``user_name`` or ``service_principal_name`` should be specified. If not, an error is - thrown.""" + One of ``user_name``, ``service_principal_name``, or ``group_name`` should be specified. If not, + an error is thrown.""" group_name: Optional[str] = None - """Group name of an account group assigned to the workspace. Setting this field requires being a - member of the group.""" + """Group name of an account group assigned to the workspace. When set, all tasks run as the group + and the group's permissions are used for data access. Setting this field requires being a member + of the group, or having the ``Assume`` permission on the group.""" service_principal_name: Optional[str] = None """Application ID of an active service principal. Setting this field requires the diff --git a/databricks/sdk/service/ml.py b/databricks/sdk/service/ml.py index 45373e198..6bb62ae26 100644 --- a/databricks/sdk/service/ml.py +++ b/databricks/sdk/service/ml.py @@ -11803,6 +11803,10 @@ def publish_table(self, source_table_name: str, publish_spec: PublishSpec) -> Pu def update_online_store(self, name: str, online_store: OnlineStore, update_mask: str) -> OnlineStore: """Update an Online Feature Store. + This update is not guaranteed to be atomic: when a request changes multiple fields, some may be + applied while others fail. On a failed response, treat the update as partially applied and retry until + it succeeds. + :param name: str The name of the online store. This is the unique identifier for the online store. :param online_store: :class:`OnlineStore` diff --git a/databricks/sdk/service/pipelines.py b/databricks/sdk/service/pipelines.py index 591e7882d..469fe2e3a 100644 --- a/databricks/sdk/service/pipelines.py +++ b/databricks/sdk/service/pipelines.py @@ -4465,6 +4465,11 @@ class RunAs: Only ``user_name`` or ``service_principal_name`` can be specified. If both are specified, an error is thrown.""" + group_name: Optional[str] = None + """Group name of an account group assigned to the workspace. When set, the pipeline runs as the + group and the group's permissions are used for data access. Setting this field requires being a + member of the group, or having the ``Assume`` permission on the group.""" + service_principal_name: Optional[str] = None """Application ID of an active service principal. Setting this field requires the ``servicePrincipal/user`` role.""" @@ -4475,6 +4480,8 @@ class RunAs: def as_dict(self) -> dict: """Serializes the RunAs into a dictionary suitable for use as a JSON request body.""" body = {} + if self.group_name is not None: + body["group_name"] = self.group_name if self.service_principal_name is not None: body["service_principal_name"] = self.service_principal_name if self.user_name is not None: @@ -4484,6 +4491,8 @@ def as_dict(self) -> dict: def as_shallow_dict(self) -> dict: """Serializes the RunAs into a shallow dictionary of its immediate attributes.""" body = {} + if self.group_name is not None: + body["group_name"] = self.group_name if self.service_principal_name is not None: body["service_principal_name"] = self.service_principal_name if self.user_name is not None: @@ -4493,7 +4502,11 @@ def as_shallow_dict(self) -> dict: @classmethod def from_dict(cls, d: Dict[str, Any]) -> RunAs: """Deserializes the RunAs from a dictionary.""" - return cls(service_principal_name=d.get("service_principal_name", None), user_name=d.get("user_name", None)) + return cls( + group_name=d.get("group_name", None), + service_principal_name=d.get("service_principal_name", None), + user_name=d.get("user_name", None), + ) @dataclass diff --git a/databricks/sdk/version.py b/databricks/sdk/version.py index 99ac4c403..120dcc32b 100644 --- a/databricks/sdk/version.py +++ b/databricks/sdk/version.py @@ -1 +1 @@ -__version__ = "0.139.0" +__version__ = "0.140.0" diff --git a/docs/workspace/aifunctions/ai_functions.rst b/docs/workspace/aifunctions/ai_functions.rst index d2384cd07..423f82ab4 100644 --- a/docs/workspace/aifunctions/ai_functions.rst +++ b/docs/workspace/aifunctions/ai_functions.rst @@ -48,8 +48,7 @@ .. py:method:: ai_parse_document(content: str [, options: Optional[AiParseDocumentOptions]]) -> AiParseDocumentResponse - Parse structured content from unstructured documents. For REST API requests, the default rate limit is - 120 pages per minute per workspace. Contact your Databricks account team to request a higher limit. + Parse structured content from unstructured documents. :param content: str The document to parse, given as a Unity Catalog volume path to the source file (the REST API accepts diff --git a/docs/workspace/files/files.rst b/docs/workspace/files/files.rst old mode 100755 new mode 100644 index f3e2f3aa6..8b6f2e2a2 --- a/docs/workspace/files/files.rst +++ b/docs/workspace/files/files.rst @@ -11,7 +11,7 @@ The API supports `Unity Catalog volumes `__, where files and directories to operate on are specified using their volume URI path, which follows the format - /Volumes/<catalog_name>/<schema_name>/<volume_name>/<path_to_file>. + /Volumes////. The Files API has two distinct endpoints, one for working with files (``/fs/files``) and another one for working with directories (``/fs/directories``). Both endpoints use the standard HTTP methods GET, HEAD, diff --git a/docs/workspace/ml/feature_store.rst b/docs/workspace/ml/feature_store.rst old mode 100755 new mode 100644 index 40ba8dc3b..54aef3b0e --- a/docs/workspace/ml/feature_store.rst +++ b/docs/workspace/ml/feature_store.rst @@ -79,6 +79,10 @@ Update an Online Feature Store. + This update is not guaranteed to be atomic: when a request changes multiple fields, some may be + applied while others fail. On a failed response, treat the update as partially applied and retry until + it succeeds. + :param name: str The name of the online store. This is the unique identifier for the online store. :param online_store: :class:`OnlineStore` diff --git a/tests/test_compute_mixins.py b/tests/test_compute_mixins.py index ec895b022..d0e4858b2 100644 --- a/tests/test_compute_mixins.py +++ b/tests/test_compute_mixins.py @@ -1,6 +1,8 @@ import pytest +from databricks.sdk import WorkspaceClient from databricks.sdk.mixins.compute import SemVer +from databricks.sdk.service import compute @pytest.mark.parametrize( @@ -40,3 +42,171 @@ def test_sorting_semver(): SemVer(1, 0, 0), SemVer(12, 0, 0), ] + + +def test_select_spark_version_filters_and_selects_latest_lts_from_http(config, requests_mock): + requests_mock.get( + "http://localhost/api/2.1/clusters/spark-versions", + json={ + "versions": [ + { + "key": "14.3.x-scala2.12", + "name": "14.3 LTS (includes Apache Spark 3.5.0, Scala 2.12)", + }, + { + "key": "15.4.x-scala2.12", + "name": "15.4 LTS (includes Apache Spark 3.5.0, Scala 2.12)", + }, + { + "key": "16.0.x-scala2.12", + "name": "16.0 Beta (includes Apache Spark 4.0.0, Scala 2.12)", + }, + { + "key": "15.4.x-photon-scala2.12", + "name": "15.4 LTS Photon (includes Apache Spark 3.5.0, Scala 2.12)", + }, + ] + }, + ) + workspace = WorkspaceClient(config=config) + + selected = workspace.clusters.select_spark_version(long_term_support=True) + + assert selected == "15.4.x-scala2.12" + assert requests_mock.last_request.method == "GET" + + +def test_select_node_type_filters_diskless_and_unavailable_nodes_from_http(config, requests_mock): + def node_type( + node_type_id, + *, + memory_mb, + num_cores, + local_disks, + local_disk_size_gb, + node_info=None, + ): + node = { + "category": "General Purpose", + "description": node_type_id, + "instance_type_id": node_type_id, + "is_deprecated": False, + "memory_mb": memory_mb, + "node_instance_type": { + "instance_type_id": node_type_id, + "local_disk_size_gb": local_disk_size_gb, + "local_disks": local_disks, + "local_nvme_disk_size_gb": 0, + "local_nvme_disks": 0, + }, + "node_type_id": node_type_id, + "num_cores": num_cores, + "num_gpus": 0, + } + if node_info is not None: + node["node_info"] = node_info + return node + + requests_mock.get( + "http://localhost/api/2.1/clusters/list-node-types", + json={ + "node_types": [ + node_type( + "unavailable-local", + memory_mb=4096, + num_cores=2, + local_disks=1, + local_disk_size_gb=100, + node_info={"status": ["NotAvailableInRegion"]}, + ), + node_type( + "diskless-small", + memory_mb=4096, + num_cores=2, + local_disks=0, + local_disk_size_gb=0, + ), + node_type( + "local-medium", + memory_mb=8192, + num_cores=4, + local_disks=1, + local_disk_size_gb=100, + ), + node_type( + "local-large", + memory_mb=16384, + num_cores=8, + local_disks=1, + local_disk_size_gb=200, + ), + ] + }, + ) + workspace = WorkspaceClient(config=config) + + selected = workspace.clusters.select_node_type(local_disk=True) + + assert selected == "local-medium" + assert requests_mock.last_request.method == "GET" + + +def test_ensure_cluster_is_running_waits_for_termination_then_starts(config, monkeypatch, requests_mock): + cluster_url = "http://localhost/api/2.1/clusters/get?cluster_id=cluster-1" + requests_mock.register_uri( + "GET", + cluster_url, + [ + {"json": {"cluster_id": "cluster-1", "cluster_name": "fixture", "state": "TERMINATING"}}, + {"json": {"cluster_id": "cluster-1", "cluster_name": "fixture", "state": "TERMINATED"}}, + {"json": {"cluster_id": "cluster-1", "cluster_name": "fixture", "state": "PENDING"}}, + {"json": {"cluster_id": "cluster-1", "cluster_name": "fixture", "state": "RUNNING"}}, + ], + ) + requests_mock.post("http://localhost/api/2.1/clusters/start", json={}) + monkeypatch.setattr(compute.time, "sleep", lambda _: None) + workspace = WorkspaceClient(config=config) + + workspace.clusters.ensure_cluster_is_running("cluster-1") + + assert [request.method for request in requests_mock.request_history] == ["GET", "GET", "POST", "GET", "GET"] + assert requests_mock.request_history[2].json() == {"cluster_id": "cluster-1"} + + +def test_cluster_events_reposts_next_page_request_at_http_boundary(config, requests_mock): + next_page = { + "cluster_id": "cluster-1", + "event_types": ["STARTING", "TERMINATING"], + "limit": 1, + "offset": 1, + } + requests_mock.register_uri( + "POST", + "http://localhost/api/2.1/clusters/events", + [ + { + "json": { + "events": [{"cluster_id": "cluster-1", "timestamp": 1000, "type": "STARTING"}], + "next_page": next_page, + } + }, + {"json": {"events": [{"cluster_id": "cluster-1", "timestamp": 2000, "type": "TERMINATING"}]}}, + ], + ) + workspace = WorkspaceClient(config=config) + + events = list( + workspace.clusters.events( + "cluster-1", + event_types=[compute.EventType.STARTING, compute.EventType.TERMINATING], + limit=1, + ) + ) + + assert [event.type for event in events] == [compute.EventType.STARTING, compute.EventType.TERMINATING] + assert requests_mock.request_history[0].json() == { + "cluster_id": "cluster-1", + "event_types": ["STARTING", "TERMINATING"], + "limit": 1, + } + assert requests_mock.request_history[1].json() == next_page diff --git a/tests/test_data_plane.py b/tests/test_data_plane.py index 1cade49d3..499bba649 100644 --- a/tests/test_data_plane.py +++ b/tests/test_data_plane.py @@ -2,7 +2,9 @@ from unittest.mock import patch from urllib import parse -from databricks.sdk import data_plane, oauth +from databricks.sdk import WorkspaceClient, data_plane, oauth +from databricks.sdk.config import Config +from databricks.sdk.credentials_provider import OAuthCredentialsProvider, OauthCredentialsStrategy from databricks.sdk.oauth import Token cp_token = Token(access_token="control plane token", token_type="type", expiry=datetime.now() + timedelta(hours=1)) @@ -67,3 +69,65 @@ def test_token_source_get_token_existing(config): retrieve_token.assert_not_called() assert result_token.access_token == another_token.access_token + + +def test_serving_endpoint_data_plane_discovers_exchanges_and_queries_over_http(requests_mock): + control_plane_token = Token(access_token="control-plane-token", token_type="Bearer") + credentials_strategy = OauthCredentialsStrategy( + "fixture-oauth", + lambda _: OAuthCredentialsProvider( + lambda: {"Authorization": "Bearer control-plane-token"}, + lambda: control_plane_token, + ), + ) + config = Config( + host="http://localhost", + workspace_id="12345", + credentials_strategy=credentials_strategy, + disable_async_token_refresh=True, + ) + endpoint_name = "route-optimized" + data_plane_url = "https://data-plane.example/invocations" + authorization_details = '{"type":"query","resource":"route-optimized"}' + requests_mock.get( + f"http://localhost/api/2.0/serving-endpoints/{endpoint_name}", + json={ + "name": endpoint_name, + "data_plane_info": { + "query_info": { + "endpoint_url": data_plane_url, + "authorization_details": authorization_details, + } + }, + }, + ) + requests_mock.post( + "http://localhost/oidc/v1/token", + json={"access_token": "data-plane-token", "token_type": "Bearer", "expires_in": 3600}, + ) + requests_mock.post( + data_plane_url, + json={"predictions": [0.75]}, + headers={"served-model-name": "model-v1"}, + ) + + response = WorkspaceClient(config=config).serving_endpoints_data_plane.query( + name=endpoint_name, + dataframe_records=[{"col": 1.0}], + ) + + discovery_request, exchange_request, query_request = requests_mock.request_history + assert response.predictions == [0.75] + assert response.served_model_name == "model-v1" + assert discovery_request.method == "GET" + assert discovery_request.headers["Authorization"] == "Bearer control-plane-token" + assert parse.parse_qs(exchange_request.text) == { + "assertion": ["control-plane-token"], + "authorization_details": [authorization_details], + "grant_type": ["urn:ietf:params:oauth:grant-type:jwt-bearer"], + } + assert exchange_request.headers["Content-Type"] == "application/x-www-form-urlencoded" + assert query_request.url == data_plane_url + assert query_request.headers["Authorization"] == "Bearer data-plane-token" + assert query_request.headers["X-Databricks-Workspace-Id"] == "12345" + assert query_request.json() == {"dataframe_records": [{"col": 1.0}]} diff --git a/tests/test_dbfs_mixins.py b/tests/test_dbfs_mixins.py index 5d785585d..4b0d3a240 100644 --- a/tests/test_dbfs_mixins.py +++ b/tests/test_dbfs_mixins.py @@ -1,64 +1,121 @@ +import base64 +import io + import pytest -from databricks.sdk.errors import NotFound +from databricks.sdk import WorkspaceClient from databricks.sdk.mixins.files import DbfsExt, _DbfsPath, _LocalPath, _VolumesPath -def test_moving_dbfs_file_to_local_dir(config, tmp_path, mocker): - from databricks.sdk import WorkspaceClient - from databricks.sdk.service.files import FileInfo, ReadResponse - - get_status = mocker.patch( - "databricks.sdk.service.files.DbfsAPI.get_status", - return_value=FileInfo(path="a", is_dir=False, file_size=4), +def test_dbfs_upload_chunks_payload_through_http_api(config, requests_mock): + chunk_size = 1024 * 1024 + payload = b"a" * chunk_size + b"tail" + requests_mock.post("http://localhost/api/2.0/dbfs/create", json={"handle": 123}) + requests_mock.post("http://localhost/api/2.0/dbfs/add-block", json={}) + requests_mock.post("http://localhost/api/2.0/dbfs/close", json={}) + + WorkspaceClient(config=config).dbfs.upload("/tmp/payload.bin", io.BytesIO(payload), overwrite=True) + + create_request, first_block, second_block, close_request = requests_mock.request_history + assert create_request.json() == {"overwrite": True, "path": "/tmp/payload.bin"} + assert first_block.url == "http://localhost/api/2.0/dbfs/add-block" + assert second_block.url == "http://localhost/api/2.0/dbfs/add-block" + assert base64.b64decode(first_block.json()["data"]) == payload[:chunk_size] + assert base64.b64decode(second_block.json()["data"]) == payload[chunk_size:] + assert first_block.json()["handle"] == 123 + assert second_block.json()["handle"] == 123 + assert close_request.json() == {"handle": 123} + + +def test_copy_routes_dbfs_source_to_volumes_http_api(config, requests_mock): + source_path = "/tmp/source.txt" + target_path = "/Volumes/main/default/volume/copied.txt" + target_url = f"http://localhost/api/2.0/fs/files{target_path}" + target_directory_url = f"http://localhost/api/2.0/fs/directories{target_path}" + payload = b"hello" + not_found = {"error_code": "RESOURCE_DOES_NOT_EXIST", "message": "not found"} + requests_mock.head(target_url, status_code=404, json=not_found) + requests_mock.head(target_directory_url, status_code=404, json=not_found) + requests_mock.get( + "http://localhost/api/2.0/dbfs/get-status", + json={"path": source_path, "is_dir": False, "file_size": len(payload)}, ) - def fake_read(path: str, *, length: int = None, offset: int = None): - assert path == "a" - assert length == 1048576 - if not offset: - return ReadResponse(bytes_read=4, data="aGVsbG8=") - return ReadResponse(bytes_read=0) - - mocker.patch("databricks.sdk.service.files.DbfsAPI.read", wraps=fake_read) - delete = mocker.patch("databricks.sdk.service.files.DbfsAPI.delete") - - w = WorkspaceClient(config=config) - w.dbfs.move_("a", f"file:{tmp_path}", recursive=True) - - get_status.assert_called_with("a") - delete.assert_called_with("a", recursive=True) - - with (tmp_path / "a").open("rb") as f: - assert f.read() == b"hello" - + def read_chunk(request, _context): + offset = int(request.qs["offset"][0]) + length = int(request.qs["length"][0]) + chunk = payload[offset : offset + length] + return {"bytes_read": len(chunk), "data": base64.b64encode(chunk).decode()} + + requests_mock.get("http://localhost/api/2.0/dbfs/read", json=read_chunk) + uploaded = {} + + def capture_upload(request, _context): + body = request.body + uploaded["body"] = body.read() if hasattr(body, "read") else body + return "" + + requests_mock.put(target_url, text=capture_upload) + + WorkspaceClient(config=config).dbfs.copy(source_path, target_path) + + status_requests = [request for request in requests_mock.request_history if "/dbfs/get-status" in request.url] + read_requests = [request for request in requests_mock.request_history if "/dbfs/read" in request.url] + upload_request = requests_mock.request_history[-1] + assert len(status_requests) == 2 + assert all(request.qs == {"path": [source_path]} for request in status_requests) + assert [request.qs["offset"] for request in read_requests] == [["0"], [str(len(payload))]] + assert all(request.qs["length"] == [str(1024 * 1024)] for request in read_requests) + assert upload_request.method == "PUT" + assert upload_request.url == f"{target_url}?overwrite=false" + assert upload_request.headers["Content-Type"] == "application/octet-stream" + assert uploaded["body"] == payload + + +def test_move_routes_volume_to_local_and_deletes_remote_source(config, requests_mock, tmp_path): + source_path = "/Volumes/main/default/volume/source.txt" + source_url = f"http://localhost/api/2.0/fs/files{source_path}" + source_directory_url = f"http://localhost/api/2.0/fs/directories{source_path}" + payload = b"moved from a volume" + requests_mock.head( + source_directory_url, + status_code=404, + json={"error_code": "RESOURCE_DOES_NOT_EXIST", "message": "not a directory"}, + ) + requests_mock.get( + source_url, + content=payload, + headers={ + "Content-Length": str(len(payload)), + "Content-Type": "application/octet-stream", + "Last-Modified": "Wed, 10 Sep 2025 12:00:00 GMT", + }, + ) + requests_mock.delete(source_url, status_code=204) -def test_moving_local_dir_to_dbfs(config, tmp_path, mocker): - from databricks.sdk import WorkspaceClient - from databricks.sdk.service.files import CreateResponse + WorkspaceClient(config=config).dbfs.move_(source_path, f"file:{tmp_path}", recursive=True) - with (tmp_path / "a").open("wb") as f: - f.write(b"hello") + assert (tmp_path / "source.txt").read_bytes() == payload + assert [request.method for request in requests_mock.request_history] == [ + "HEAD", + "HEAD", + "GET", + "HEAD", + "DELETE", + ] + assert requests_mock.request_history[-1].url == source_url - mocker.patch( - "databricks.sdk.service.files.DbfsAPI.create", - return_value=CreateResponse(123), - ) - get_status = mocker.patch( - "databricks.sdk.service.files.DbfsAPI.get_status", - side_effect=NotFound(), - ) - add_block = mocker.patch("databricks.sdk.service.files.DbfsAPI.add_block") - close = mocker.patch("databricks.sdk.service.files.DbfsAPI.close") +def test_move_within_dbfs_uses_dbfs_move_http_api(config, requests_mock): + requests_mock.post("http://localhost/api/2.0/dbfs/move", json={}) - w = WorkspaceClient(config=config) - w.dbfs.move_(f"file:{tmp_path}", "a", recursive=True) + WorkspaceClient(config=config).dbfs.move_("/tmp/source", "/tmp/destination") - get_status.assert_called_with("a") - close.assert_called_with(123) - add_block.assert_called_with(123, "aGVsbG8=") - assert not (tmp_path / "a").exists() + request = requests_mock.last_request + assert request.json() == { + "destination_path": "/tmp/destination", + "source_path": "/tmp/source", + } @pytest.mark.parametrize( @@ -85,33 +142,36 @@ def test_fs_path_invalid(config): def test_dbfs_local_path_mkdir(config, tmp_path): - from databricks.sdk import WorkspaceClient - - w = WorkspaceClient(config=config) - w.dbfs._path(f"file:{tmp_path}/test_dir").mkdir() - assert w.dbfs.exists(f"file:{tmp_path}/test_dir") - + client = WorkspaceClient(config=config) + client.dbfs._path(f"file:{tmp_path}/test_dir").mkdir() + assert client.dbfs.exists(f"file:{tmp_path}/test_dir") -def test_dbfs_exists(config, mocker): - from databricks.sdk import WorkspaceClient - get_status = mocker.patch( - "databricks.sdk.service.files.DbfsAPI.get_status", - side_effect=NotFound(), +def test_dbfs_exists_maps_not_found_http_response_to_false(config, requests_mock): + requests_mock.get( + "http://localhost/api/2.0/dbfs/get-status", + status_code=404, + json={"error_code": "RESOURCE_DOES_NOT_EXIST", "message": "not found"}, ) - client = WorkspaceClient(config=config) - client.dbfs.exists("/abc/def/ghi") - - get_status.assert_called_with("/abc/def/ghi") + exists = WorkspaceClient(config=config).dbfs.exists("/abc/def/ghi") + assert exists is False + assert requests_mock.last_request.qs == {"path": ["/abc/def/ghi"]} -def test_volume_exists(config, mocker): - from databricks.sdk import WorkspaceClient - get_metadata = mocker.patch("databricks.sdk.service.files.FilesAPI.get_metadata") +def test_volume_exists_reads_file_metadata_over_http(config, requests_mock): + path = "/Volumes/abc/def/ghi" + requests_mock.head( + f"http://localhost/api/2.0/fs/files{path}", + headers={ + "Content-Length": "4", + "Content-Type": "application/octet-stream", + "Last-Modified": "Wed, 10 Sep 2025 12:00:00 GMT", + }, + ) - client = WorkspaceClient(config=config) - client.dbfs.exists("/Volumes/abc/def/ghi") + exists = WorkspaceClient(config=config).dbfs.exists(path) - get_metadata.assert_called_with("/Volumes/abc/def/ghi") + assert exists is True + assert requests_mock.last_request.method == "HEAD" diff --git a/tests/test_files.py b/tests/test_files.py index 25c7c9689..63ab60715 100644 --- a/tests/test_files.py +++ b/tests/test_files.py @@ -41,6 +41,53 @@ logger = logging.getLogger(__name__) +def test_files_http_wire_format_preserves_bytes_metadata_and_stream_lifecycle(config: Config, requests_mock): + config = config.copy() + config.disable_experimental_files_api_client = True + file_path = "/Volumes/main/default/volume/report ?#%.bin" + file_url = "http://localhost/api/2.0/fs/files/Volumes/main/default/volume/report%20%3F%23%25.bin" + payload = b"raw file contents\x00\xff" + uploaded = {} + + def capture_upload(request, _context): + body = request.body + uploaded["body"] = body.read() if hasattr(body, "read") else body + return "" + + requests_mock.put(file_url, text=capture_upload) + response_headers = { + "Content-Length": str(len(payload)), + "Content-Type": "application/octet-stream", + "Last-Modified": "Wed, 10 Sep 2025 12:00:00 GMT", + } + requests_mock.head(file_url, headers=response_headers) + requests_mock.get(file_url, content=payload, headers=response_headers) + files = WorkspaceClient(config=config).files + + files.upload(file_path, io.BytesIO(payload), overwrite=False) + metadata = files.get_metadata(file_path) + download = files.download(file_path) + with download.contents as stream: + assert stream.read() == payload + with pytest.raises(ValueError, match="I/O operation on closed file"): + download.contents.read() + + upload_request, metadata_request, download_request = requests_mock.request_history + assert uploaded["body"] == payload + assert upload_request.url == f"{file_url}?overwrite=false" + assert upload_request.headers["Content-Type"] == "application/octet-stream" + assert metadata.content_length == len(payload) + assert metadata.content_type == "application/octet-stream" + assert metadata.last_modified == "Wed, 10 Sep 2025 12:00:00 GMT" + assert metadata_request.method == "HEAD" + assert metadata_request.url == file_url + assert download.content_length == len(payload) + assert download.content_type == "application/octet-stream" + assert download.last_modified == "Wed, 10 Sep 2025 12:00:00 GMT" + assert download_request.method == "GET" + assert download_request.url == file_url + + class CustomResponse: """Custom response allows to override the "default" response generated by the server with the "custom" response to simulate failure error code, unexpected response body or diff --git a/tests/test_jobs_mixin.py b/tests/test_jobs_mixin.py index c14c7dd83..96339be3d 100644 --- a/tests/test_jobs_mixin.py +++ b/tests/test_jobs_mixin.py @@ -3,6 +3,7 @@ from typing import Optional, Pattern from databricks.sdk import WorkspaceClient +from databricks.sdk.service import compute, jobs def make_getrun_path_pattern(run_id: int, page_token: Optional[str] = None) -> Pattern[str]: @@ -53,8 +54,6 @@ def test_get_run_with_no_pagination(config, requests_mock): def test_get_run_pagination_with_tasks(config, requests_mock): - from databricks.sdk.service import compute, jobs - cluster_spec = compute.ClusterSpec( spark_version="11.3.x-scala2.12", custom_tags={"ResourceClass": "SingleNode"}, @@ -199,8 +198,6 @@ def test_get_job_forwards_include_trigger_state(config, requests_mock): def test_get_job_pagination_with_tasks(config, requests_mock): - from databricks.sdk.service import compute, jobs - cluster_spec = compute.ClusterSpec( spark_version="11.3.x-scala2.12", custom_tags={"ResourceClass": "SingleNode"}, @@ -374,8 +371,6 @@ def test_list_jobs_without_task_expansion(config, requests_mock): def test_list_jobs_with_many_tasks(config, requests_mock): - from databricks.sdk.service import compute, jobs - cluster_spec = compute.ClusterSpec( spark_version="11.3.x-scala2.12", custom_tags={"ResourceClass": "SingleNode"}, @@ -816,3 +811,55 @@ def test_list_runs(config, requests_mock): history = requests_mock.request_history assert all("300" not in request.qs.get("run_id", [""]) for request in history) assert all("200" not in request.qs.get("run_id", [""]) for request in history) + + +def test_submit_waiter_polls_real_jobs_service_until_terminal_state(config, monkeypatch, requests_mock): + requests_mock.post( + "http://localhost/api/2.2/jobs/runs/submit", + json={"run_id": 123}, + ) + requests_mock.register_uri( + "GET", + make_getrun_path_pattern(123), + [ + { + "json": { + "run_id": 123, + "run_name": "fixture-run", + "state": {"life_cycle_state": "PENDING", "state_message": "Waiting for compute"}, + } + }, + { + "json": { + "run_id": 123, + "run_name": "fixture-run", + "state": {"life_cycle_state": "RUNNING", "state_message": "Running"}, + } + }, + { + "json": { + "run_id": 123, + "run_name": "fixture-run", + "state": { + "life_cycle_state": "TERMINATED", + "result_state": "SUCCESS", + "state_message": "Completed", + }, + } + }, + ], + ) + monkeypatch.setattr(jobs.time, "sleep", lambda _: None) + workspace = WorkspaceClient(config=config) + observed_states = [] + + run = workspace.jobs.submit(run_name="fixture-run").result( + callback=lambda polled: observed_states.append(polled.state.life_cycle_state) + ) + + assert run.run_id == 123 + assert run.state.life_cycle_state == jobs.RunLifeCycleState.TERMINATED + assert run.state.result_state == jobs.RunResultState.SUCCESS + assert observed_states == [jobs.RunLifeCycleState.PENDING, jobs.RunLifeCycleState.RUNNING] + assert requests_mock.request_history[0].json() == {"run_name": "fixture-run"} + assert [request.method for request in requests_mock.request_history] == ["POST", "GET", "GET", "GET"] diff --git a/tests/test_workspace.py b/tests/test_workspace.py new file mode 100644 index 000000000..c7c71e479 --- /dev/null +++ b/tests/test_workspace.py @@ -0,0 +1,92 @@ +import io +from urllib.parse import parse_qs, urlsplit + +import pytest + +from databricks.sdk import WorkspaceClient +from databricks.sdk.service.workspace import ExportFormat, ObjectType + + +def test_recursive_list_traverses_directories_through_workspace_http_api(config, requests_mock): + requests_mock.get( + "http://localhost/api/2.0/workspace/list", + response_list=[ + { + "json": { + "objects": [ + {"object_type": "DIRECTORY", "path": "/root/nested"}, + {"object_type": "NOTEBOOK", "path": "/root/top.py"}, + ] + } + }, + { + "json": { + "objects": [ + {"object_type": "DIRECTORY", "path": "/root/nested/empty"}, + {"object_type": "FILE", "path": "/root/nested/data.bin"}, + ] + } + }, + {"json": {"objects": []}}, + ], + ) + workspace = WorkspaceClient(config=config).workspace + + objects = list(workspace.list("/root", notebooks_modified_after=123, recursive=True)) + + assert [(obj.path, obj.object_type) for obj in objects] == [ + ("/root/top.py", ObjectType.NOTEBOOK), + ("/root/nested/data.bin", ObjectType.FILE), + ] + assert [parse_qs(urlsplit(request.url).query) for request in requests_mock.request_history] == [ + {"notebooks_modified_after": ["123"], "path": ["/root"]}, + {"notebooks_modified_after": ["123"], "path": ["/root/nested"]}, + {"notebooks_modified_after": ["123"], "path": ["/root/nested/empty"]}, + ] + + +def test_upload_and_download_use_workspace_http_wire_format(config, requests_mock): + uploaded = {} + + def capture_upload(request, _context): + uploaded["body"] = request.body + return "" + + requests_mock.post("http://localhost/api/2.0/workspace/import", text=capture_upload) + requests_mock.get( + "http://localhost/api/2.0/workspace/export", + content=b"print('downloaded')\n", + ) + workspace = WorkspaceClient(config=config).workspace + + workspace.upload( + "/Users/sdk/example.py", + io.BytesIO(b"print('uploaded')\n"), + overwrite=True, + ) + download = workspace.download("/Users/sdk/example.py", format=ExportFormat.AUTO) + with download as stream: + assert stream.read() == b"print('downloaded')\n" + with pytest.raises(ValueError, match="I/O operation on closed file"): + download.read() + + upload_request, download_request = requests_mock.request_history + assert upload_request.method == "POST" + assert upload_request.url == "http://localhost/api/2.0/workspace/import" + assert upload_request.headers["Content-Type"].startswith("multipart/form-data; boundary=") + assert b'name="path"' in uploaded["body"] + assert b"/Users/sdk/example.py" in uploaded["body"] + assert b'name="language"' in uploaded["body"] + assert b"PYTHON" in uploaded["body"] + assert b'name="overwrite"' in uploaded["body"] + assert b"true" in uploaded["body"] + assert b'name="content"' in uploaded["body"] + assert b"print('uploaded')\n" in uploaded["body"] + + assert download_request.method == "GET" + assert download_request.url.startswith("http://localhost/api/2.0/workspace/export?") + assert parse_qs(urlsplit(download_request.url).query) == { + "direct_download": ["true"], + "format": ["AUTO"], + "path": ["/Users/sdk/example.py"], + }