From 3e3fa29d5cdd01ddd450ffc353a3f9c793e151f4 Mon Sep 17 00:00:00 2001 From: "databricks-ci-ghec-1[bot]" <184311507+databricks-ci-ghec-1[bot]@users.noreply.github.com> Date: Tue, 29 Sep 2026 03:18:36 +0000 Subject: [PATCH] Release databricks-sdk-py --- .codegen/_last_sha | 2 +- CHANGELOG.md | 16 ++ databricks/sdk/__init__.py | 5 +- databricks/sdk/_sandbox_grpc/__init__.py | 28 +++ .../sdk/_sandbox_grpc/sandbox_service_pb2.py | 74 ++++++ .../_sandbox_grpc/sandbox_service_pb2_grpc.py | 233 ++++++++++++++++++ databricks/sdk/mixins/sandbox.py | 187 ++++++++++++++ databricks/sdk/service/jobs.py | 185 +++++++++++++- databricks/sdk/service/ml.py | 5 +- databricks/sdk/version.py | 2 +- docs/dbdataclasses/jobs.rst | 8 + docs/workspace/jobs/jobs.rst | 16 +- docs/workspace/sandbox/sandbox.rst | 51 +++- pyproject.toml | 24 +- tests/test_sandbox_mixin.py | 216 ++++++++++++++++ 15 files changed, 1029 insertions(+), 23 deletions(-) create mode 100644 databricks/sdk/_sandbox_grpc/__init__.py create mode 100644 databricks/sdk/_sandbox_grpc/sandbox_service_pb2.py create mode 100644 databricks/sdk/_sandbox_grpc/sandbox_service_pb2_grpc.py create mode 100644 databricks/sdk/mixins/sandbox.py create mode 100644 tests/test_sandbox_mixin.py diff --git a/.codegen/_last_sha b/.codegen/_last_sha index 6f2f696f1..c9ebb7284 100644 --- a/.codegen/_last_sha +++ b/.codegen/_last_sha @@ -1 +1 @@ -231772f37bb5c01154bfc3d1401dc241ff1b7409 \ No newline at end of file +80beeb9b732df6fa8eb43abd5c73a3a0207c1584 \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md index 11f5283f5..a8c2176f1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,21 @@ # Version changelog +## Release v0.144.0 (2026-09-29) + +### API Changes +* Add `environment_variables` field for `databricks.sdk.service.jobs.BaseRun`. +* Add `environment_variables` field for `databricks.sdk.service.jobs.CreateJob`. +* Add `environment_variables` field for `databricks.sdk.service.jobs.JobSettings`. +* Add `environment_variables` field for `databricks.sdk.service.jobs.Run`. +* Add `environment_variables_key` field for `databricks.sdk.service.jobs.RunTask`. +* Add `environment_variables` field for `databricks.sdk.service.jobs.SubmitRun`. +* Add `environment_variables_key` field for `databricks.sdk.service.jobs.SubmitTask`. +* Add `environment_variables_key` field for `databricks.sdk.service.jobs.Task`. +* Change `command_path` field for `databricks.sdk.service.jobs.DeploymentSpec` to no longer be required. +* [Breaking] Change `command_path` field for `databricks.sdk.service.jobs.DeploymentSpec` to no longer be required. +* Change `api_secret_ref` field for `databricks.sdk.service.ml.SchemaRegistryConfig` to no longer be required. +* [Breaking] Change `api_secret_ref` field for `databricks.sdk.service.ml.SchemaRegistryConfig` to no longer be required. + ## Release v0.143.0 (2026-09-26) ### API Changes diff --git a/databricks/sdk/__init__.py b/databricks/sdk/__init__.py index 5e670ff81..ea3367399 100644 --- a/databricks/sdk/__init__.py +++ b/databricks/sdk/__init__.py @@ -55,6 +55,7 @@ from databricks.sdk.mixins.workspace import WorkspaceExt from databricks.sdk.mixins.open_ai_client import ServingEndpointsExt from databricks.sdk.mixins.jobs import JobsExt +from databricks.sdk.mixins.sandbox import SandboxExt from databricks.sdk.oauth import AuthorizationDetail from databricks.sdk.service.iam import AccessControlAPI from databricks.sdk.service.iam import AccountAccessControlAPI @@ -468,7 +469,7 @@ def __init__( self._repos = pkg_workspace.ReposAPI(self._api_client) self._resource_quotas = pkg_catalog.ResourceQuotasAPI(self._api_client) self._rfa = pkg_catalog.RfaAPI(self._api_client) - self._sandbox = pkg_sandbox.SandboxAPI(self._api_client) + self._sandbox = SandboxExt(self._api_client) self._schemas = pkg_catalog.SchemasAPI(self._api_client) self._secrets = pkg_workspace.SecretsAPI(self._api_client) self._secrets_uc = pkg_catalog.SecretsUcAPI(self._api_client) @@ -1031,7 +1032,7 @@ def rfa(self) -> pkg_catalog.RfaAPI: return self._rfa @property - def sandbox(self) -> pkg_sandbox.SandboxAPI: + def sandbox(self) -> SandboxExt: """Create, manage, and control the lifecycle of sandboxes -- isolated, pre-configured, low-latency Serverless compute environments for running code.""" return self._sandbox diff --git a/databricks/sdk/_sandbox_grpc/__init__.py b/databricks/sdk/_sandbox_grpc/__init__.py new file mode 100644 index 000000000..e75a6d8ba --- /dev/null +++ b/databricks/sdk/_sandbox_grpc/__init__.py @@ -0,0 +1,28 @@ +"""Committed gRPC stubs for the Sandbox command-execution service. + +sandbox_service_pb2 / sandbox_service_pb2_grpc are generated from the shared +client proto (kept out of this package) at +compute-fabric/sandbox-sdk/proto/databricks/sandbox/sandbox_service.proto, a +dependency-stripped copy of +compute-fabric/sandbox-daemon/proto/sandbox_service.proto. The mixin uses only the +streaming RPCs (ExecuteCommand, AttachCommand, StreamInput); the stub also carries +the unary ListCommands, which the mixin does not use. Stubs are committed, not +built: the daemon proto pulls in internal imports that cannot ship publicly. + +SDK codegen can't emit streaming clients yet. TODO(XTA-19715): migrate the mixin to +native streaming codegen and drop this package. + +Regenerate after the shared proto changes, with a grpcio-tools whose bundled protoc +matches the pinned protobuf major (the sandbox extra in pyproject.toml). Stage the +proto at this package's path so the module path stays databricks.sdk._sandbox_grpc: + + root=$(mktemp -d) + mkdir -p "$root/databricks/sdk/_sandbox_grpc" + cp /compute-fabric/sandbox-sdk/proto/databricks/sandbox/sandbox_service.proto \\ + "$root/databricks/sdk/_sandbox_grpc/sandbox_service.proto" + python -m grpc_tools.protoc -I "$root" \\ + --python_out="$root" --grpc_python_out="$root" \\ + databricks/sdk/_sandbox_grpc/sandbox_service.proto + cp "$root"/databricks/sdk/_sandbox_grpc/sandbox_service_pb2*.py \\ + databricks/sdk/_sandbox_grpc/ +""" diff --git a/databricks/sdk/_sandbox_grpc/sandbox_service_pb2.py b/databricks/sdk/_sandbox_grpc/sandbox_service_pb2.py new file mode 100644 index 000000000..6d39fcfd4 --- /dev/null +++ b/databricks/sdk/_sandbox_grpc/sandbox_service_pb2.py @@ -0,0 +1,74 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# NO CHECKED-IN PROTOBUF GENCODE +# source: databricks/sdk/_sandbox_grpc/sandbox_service.proto +# Protobuf Python Version: 6.33.5 +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import runtime_version as _runtime_version +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder +_runtime_version.ValidateProtobufRuntimeVersion( + _runtime_version.Domain.PUBLIC, + 6, + 33, + 5, + '', + 'databricks/sdk/_sandbox_grpc/sandbox_service.proto' +) +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n2databricks/sdk/_sandbox_grpc/sandbox_service.proto\x12\x12\x64\x61tabricks.sandbox\"\xe2\x01\n\x15\x45xecuteCommandRequest\x12\x0b\n\x03\x63md\x18\x01 \x01(\t\x12\x0c\n\x04\x61rgs\x18\x02 \x03(\t\x12*\n\x03pty\x18\x03 \x01(\x0b\x32\x1d.databricks.sandbox.PtyConfig\x12\x41\n\x04\x65nvs\x18\x04 \x03(\x0b\x32\x33.databricks.sandbox.ExecuteCommandRequest.EnvsEntry\x12\x12\n\ncommand_id\x18\x05 \x01(\t\x1a+\n\tEnvsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"H\n\x14\x41ttachCommandRequest\x12\x12\n\ncommand_id\x18\x01 \x01(\t\x12\x1c\n\x14last_sequence_number\x18\x02 \x01(\x03\"\xc0\x01\n\x12StreamInputRequest\x12\x35\n\x05start\x18\x01 \x01(\x0b\x32$.databricks.sandbox.StreamInputStartH\x00\x12\x33\n\x04\x64\x61ta\x18\x02 \x01(\x0b\x32#.databricks.sandbox.StreamInputDataH\x00\x12\x35\n\x05\x63lose\x18\x03 \x01(\x0b\x32$.databricks.sandbox.StreamInputCloseH\x00\x42\x07\n\x05\x65vent\"&\n\x10StreamInputStart\x12\x12\n\ncommand_id\x18\x01 \x01(\t\"\x1f\n\x0fStreamInputData\x12\x0c\n\x04\x64\x61ta\x18\x01 \x01(\x0c\"\x12\n\x10StreamInputClose\"\x15\n\x13StreamInputResponse\"<\n\x13ListCommandsRequest\x12\x11\n\tpage_size\x18\x01 \x01(\x05\x12\x12\n\npage_token\x18\x02 \x01(\t\"^\n\x14ListCommandsResponse\x12-\n\x08\x63ommands\x18\x01 \x03(\x0b\x32\x1b.databricks.sandbox.Command\x12\x17\n\x0fnext_page_token\x18\x02 \x01(\t\"j\n\x07\x43ommand\x12\x12\n\ncommand_id\x18\x01 \x01(\t\x12\x0b\n\x03\x63md\x18\x02 \x01(\t\x12\x0c\n\x04\x61rgs\x18\x03 \x03(\t\x12\x0b\n\x03pid\x18\x04 \x01(\x03\x12\x10\n\x08\x66inished\x18\x05 \x01(\x08\x12\x11\n\texit_code\x18\x06 \x01(\x05\"i\n\x15\x41ttachCommandResponse\x12\x37\n\rprocess_event\x18\x01 \x01(\x0b\x32 .databricks.sandbox.ProcessEvent\x12\x17\n\x0fsequence_number\x18\x02 \x01(\x03\"j\n\x16\x45xecuteCommandResponse\x12\x37\n\rprocess_event\x18\x01 \x01(\x0b\x32 .databricks.sandbox.ProcessEvent\x12\x17\n\x0fsequence_number\x18\x02 \x01(\x03\"\xa4\x01\n\x0cProcessEvent\x12/\n\x05start\x18\x01 \x01(\x0b\x32\x1e.databricks.sandbox.StartEventH\x00\x12-\n\x04\x64\x61ta\x18\x02 \x01(\x0b\x32\x1d.databricks.sandbox.DataEventH\x00\x12+\n\x03\x65nd\x18\x03 \x01(\x0b\x32\x1c.databricks.sandbox.EndEventH\x00\x42\x07\n\x05\x65vent\"-\n\nStartEvent\x12\x0b\n\x03pid\x18\x01 \x01(\x03\x12\x12\n\ncommand_id\x18\x02 \x01(\t\"H\n\tDataEvent\x12\x10\n\x06stdout\x18\x01 \x01(\x0cH\x00\x12\x10\n\x06stderr\x18\x02 \x01(\x0cH\x00\x12\r\n\x03pty\x18\x03 \x01(\x0cH\x00\x42\x08\n\x06output\"\x1d\n\x08\x45ndEvent\x12\x11\n\texit_code\x18\x01 \x01(\x05\"\'\n\tPtyConfig\x12\x0c\n\x04\x63ols\x18\x01 \x01(\x05\x12\x0c\n\x04rows\x18\x02 \x01(\x05\x32\xb0\x03\n\x0eSandboxService\x12k\n\x0e\x45xecuteCommand\x12).databricks.sandbox.ExecuteCommandRequest\x1a*.databricks.sandbox.ExecuteCommandResponse\"\x00\x30\x01\x12h\n\rAttachCommand\x12(.databricks.sandbox.AttachCommandRequest\x1a).databricks.sandbox.AttachCommandResponse\"\x00\x30\x01\x12\x62\n\x0bStreamInput\x12&.databricks.sandbox.StreamInputRequest\x1a\'.databricks.sandbox.StreamInputResponse\"\x00(\x01\x12\x63\n\x0cListCommands\x12\'.databricks.sandbox.ListCommandsRequest\x1a(.databricks.sandbox.ListCommandsResponse\"\x00') + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'databricks.sdk._sandbox_grpc.sandbox_service_pb2', _globals) +if not _descriptor._USE_C_DESCRIPTORS: + DESCRIPTOR._loaded_options = None + _globals['_EXECUTECOMMANDREQUEST_ENVSENTRY']._loaded_options = None + _globals['_EXECUTECOMMANDREQUEST_ENVSENTRY']._serialized_options = b'8\001' + _globals['_EXECUTECOMMANDREQUEST']._serialized_start=75 + _globals['_EXECUTECOMMANDREQUEST']._serialized_end=301 + _globals['_EXECUTECOMMANDREQUEST_ENVSENTRY']._serialized_start=258 + _globals['_EXECUTECOMMANDREQUEST_ENVSENTRY']._serialized_end=301 + _globals['_ATTACHCOMMANDREQUEST']._serialized_start=303 + _globals['_ATTACHCOMMANDREQUEST']._serialized_end=375 + _globals['_STREAMINPUTREQUEST']._serialized_start=378 + _globals['_STREAMINPUTREQUEST']._serialized_end=570 + _globals['_STREAMINPUTSTART']._serialized_start=572 + _globals['_STREAMINPUTSTART']._serialized_end=610 + _globals['_STREAMINPUTDATA']._serialized_start=612 + _globals['_STREAMINPUTDATA']._serialized_end=643 + _globals['_STREAMINPUTCLOSE']._serialized_start=645 + _globals['_STREAMINPUTCLOSE']._serialized_end=663 + _globals['_STREAMINPUTRESPONSE']._serialized_start=665 + _globals['_STREAMINPUTRESPONSE']._serialized_end=686 + _globals['_LISTCOMMANDSREQUEST']._serialized_start=688 + _globals['_LISTCOMMANDSREQUEST']._serialized_end=748 + _globals['_LISTCOMMANDSRESPONSE']._serialized_start=750 + _globals['_LISTCOMMANDSRESPONSE']._serialized_end=844 + _globals['_COMMAND']._serialized_start=846 + _globals['_COMMAND']._serialized_end=952 + _globals['_ATTACHCOMMANDRESPONSE']._serialized_start=954 + _globals['_ATTACHCOMMANDRESPONSE']._serialized_end=1059 + _globals['_EXECUTECOMMANDRESPONSE']._serialized_start=1061 + _globals['_EXECUTECOMMANDRESPONSE']._serialized_end=1167 + _globals['_PROCESSEVENT']._serialized_start=1170 + _globals['_PROCESSEVENT']._serialized_end=1334 + _globals['_STARTEVENT']._serialized_start=1336 + _globals['_STARTEVENT']._serialized_end=1381 + _globals['_DATAEVENT']._serialized_start=1383 + _globals['_DATAEVENT']._serialized_end=1455 + _globals['_ENDEVENT']._serialized_start=1457 + _globals['_ENDEVENT']._serialized_end=1486 + _globals['_PTYCONFIG']._serialized_start=1488 + _globals['_PTYCONFIG']._serialized_end=1527 + _globals['_SANDBOXSERVICE']._serialized_start=1530 + _globals['_SANDBOXSERVICE']._serialized_end=1962 +# @@protoc_insertion_point(module_scope) diff --git a/databricks/sdk/_sandbox_grpc/sandbox_service_pb2_grpc.py b/databricks/sdk/_sandbox_grpc/sandbox_service_pb2_grpc.py new file mode 100644 index 000000000..2889cc8fe --- /dev/null +++ b/databricks/sdk/_sandbox_grpc/sandbox_service_pb2_grpc.py @@ -0,0 +1,233 @@ +# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! +"""Client and server classes corresponding to protobuf-defined services.""" +import grpc +import warnings + +from databricks.sdk._sandbox_grpc import sandbox_service_pb2 as databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2 + +GRPC_GENERATED_VERSION = '1.81.1' +GRPC_VERSION = grpc.__version__ +_version_not_supported = False + +try: + from grpc._utilities import first_version_is_lower + _version_not_supported = first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION) +except ImportError: + _version_not_supported = True + +if _version_not_supported: + raise RuntimeError( + f'The grpc package installed is at version {GRPC_VERSION},' + + ' but the generated code in databricks/sdk/_sandbox_grpc/sandbox_service_pb2_grpc.py depends on' + + f' grpcio>={GRPC_GENERATED_VERSION}.' + + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}' + + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.' + ) + + +class SandboxServiceStub: + """Service for executing commands inside running sandboxes. + """ + + def __init__(self, channel): + """Constructor. + + Args: + channel: A grpc.Channel. + """ + self.ExecuteCommand = channel.unary_stream( + '/databricks.sandbox.SandboxService/ExecuteCommand', + request_serializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.ExecuteCommandRequest.SerializeToString, + response_deserializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.ExecuteCommandResponse.FromString, + _registered_method=True) + self.AttachCommand = channel.unary_stream( + '/databricks.sandbox.SandboxService/AttachCommand', + request_serializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.AttachCommandRequest.SerializeToString, + response_deserializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.AttachCommandResponse.FromString, + _registered_method=True) + self.StreamInput = channel.stream_unary( + '/databricks.sandbox.SandboxService/StreamInput', + request_serializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.StreamInputRequest.SerializeToString, + response_deserializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.StreamInputResponse.FromString, + _registered_method=True) + self.ListCommands = channel.unary_unary( + '/databricks.sandbox.SandboxService/ListCommands', + request_serializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.ListCommandsRequest.SerializeToString, + response_deserializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.ListCommandsResponse.FromString, + _registered_method=True) + + +class SandboxServiceServicer: + """Service for executing commands inside running sandboxes. + """ + + def ExecuteCommand(self, request, context): + """Executes a command inside the sandbox. Streams process lifecycle events back to the caller. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def AttachCommand(self, request, context): + """Attaches to a previously started command execution, streaming its output. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def StreamInput(self, request_iterator, context): + """Streams stdin data to a running command. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def ListCommands(self, request, context): + """Lists all tracked command executions (running and completed). + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + +def add_SandboxServiceServicer_to_server(servicer, server): + rpc_method_handlers = { + 'ExecuteCommand': grpc.unary_stream_rpc_method_handler( + servicer.ExecuteCommand, + request_deserializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.ExecuteCommandRequest.FromString, + response_serializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.ExecuteCommandResponse.SerializeToString, + ), + 'AttachCommand': grpc.unary_stream_rpc_method_handler( + servicer.AttachCommand, + request_deserializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.AttachCommandRequest.FromString, + response_serializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.AttachCommandResponse.SerializeToString, + ), + 'StreamInput': grpc.stream_unary_rpc_method_handler( + servicer.StreamInput, + request_deserializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.StreamInputRequest.FromString, + response_serializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.StreamInputResponse.SerializeToString, + ), + 'ListCommands': grpc.unary_unary_rpc_method_handler( + servicer.ListCommands, + request_deserializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.ListCommandsRequest.FromString, + response_serializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.ListCommandsResponse.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler( + 'databricks.sandbox.SandboxService', rpc_method_handlers) + server.add_generic_rpc_handlers((generic_handler,)) + server.add_registered_method_handlers('databricks.sandbox.SandboxService', rpc_method_handlers) + + + # This class is part of an EXPERIMENTAL API. +class SandboxService: + """Service for executing commands inside running sandboxes. + """ + + @staticmethod + def ExecuteCommand(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_stream( + request, + target, + '/databricks.sandbox.SandboxService/ExecuteCommand', + databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.ExecuteCommandRequest.SerializeToString, + databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.ExecuteCommandResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def AttachCommand(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_stream( + request, + target, + '/databricks.sandbox.SandboxService/AttachCommand', + databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.AttachCommandRequest.SerializeToString, + databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.AttachCommandResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def StreamInput(request_iterator, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.stream_unary( + request_iterator, + target, + '/databricks.sandbox.SandboxService/StreamInput', + databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.StreamInputRequest.SerializeToString, + databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.StreamInputResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def ListCommands(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/databricks.sandbox.SandboxService/ListCommands', + databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.ListCommandsRequest.SerializeToString, + databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.ListCommandsResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) diff --git a/databricks/sdk/mixins/sandbox.py b/databricks/sdk/mixins/sandbox.py new file mode 100644 index 000000000..8b047a72a --- /dev/null +++ b/databricks/sdk/mixins/sandbox.py @@ -0,0 +1,187 @@ +"""Sandbox command-execution mixin. + +Extends the generated :class:`~databricks.sdk.service.sandbox.SandboxAPI` with +the streaming command RPCs that codegen does not produce a client for: +``execute_command_streaming``, ``attach_command`` (both gRPC server-streaming) +and ``stream_input`` (gRPC client-streaming). They run on the generic gRPC +transport in :mod:`databricks.sdk.mixins._grpc_transport`. + +The wire types come from the committed stubs in +:mod:`databricks.sdk._sandbox_grpc` (generated from a slim copy of the sandbox +daemon proto; see that package for provenance). grpcio and those stubs are +imported lazily inside the methods, so this module imports without the optional +``sandbox`` extra installed. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Callable, Dict, Iterable, Iterator, Optional, Tuple + +from databricks.sdk.mixins._grpc_transport import ( + _DEFAULT_MAX_ATTEMPTS, + _DEFAULT_RETRY_INTERVAL_SECONDS, + open_channel, + open_stream_with_retry, +) +from databricks.sdk.service.sandbox import SandboxAPI + +if TYPE_CHECKING: + from databricks.sdk._sandbox_grpc import sandbox_service_pb2 as _pb + from databricks.sdk._sandbox_grpc import sandbox_service_pb2_grpc as _pb_grpc + +# gRPC metadata key that routes a command RPC to a specific sandbox. +_SANDBOX_NAME_METADATA = "x-databricks-sandbox-name" + +_Metadata = Tuple[Tuple[str, str], ...] + + +def _never_retry(_error: Exception) -> bool: + """Retry predicate for :meth:`SandboxExt.execute_command_streaming`: never + retry. ExecuteCommand starts a process and is not idempotent (its + ``command_id`` is a stream-resume token, not a dedup key), so re-issuing a + failed open could run the command twice.""" + return False + + +class SandboxExt(SandboxAPI): + __doc__ = SandboxAPI.__doc__ + + def execute_command_streaming( + self, + name: str, + cmd: str, + *, + args: Optional[list[str]] = None, + envs: Optional[Dict[str, str]] = None, + is_retryable: Optional[Callable[[Exception], bool]] = None, + max_attempts: int = _DEFAULT_MAX_ATTEMPTS, + retry_interval_seconds: float = _DEFAULT_RETRY_INTERVAL_SECONDS, + ) -> Iterator["_pb.ExecuteCommandResponse"]: + """Run ``cmd`` in sandbox ``name`` and stream its process events. + + Streaming counterpart of :meth:`SandboxAPI.execute_command_sync`: the same + ``name`` (``sandboxes/{sandbox_id}``), ``cmd``, ``args``, and ``envs``, + but the process output is streamed instead of returned in one response. + Yields ``ExecuteCommandResponse`` messages, each carrying one + ``ProcessEvent`` (a start event, then stdout/stderr data chunks, then an + end event with the exit code) and a monotonic ``sequence_number``. + + Retries are off by default (``is_retryable`` defaults to + :func:`_never_retry`) because ExecuteCommand is not idempotent: re-issuing + a failed open could run the command twice. Pass an ``is_retryable`` + predicate - e.g. one matching ``UNAVAILABLE`` - to retry the stream open + when the command is safe to run more than once, optionally tuning + ``max_attempts`` and ``retry_interval_seconds``. Only the stream open is + retried; a failure once events are flowing is always raised. + """ + from databricks.sdk._sandbox_grpc import sandbox_service_pb2 as _pb # noqa: FlagLocalImports - lazy; stubs ship only with the optional "sandbox" extra + + request = _pb.ExecuteCommandRequest(cmd=cmd, args=list(args or []), envs=dict(envs or {})) + yield from self._server_stream( + name, + lambda stub, metadata: stub.ExecuteCommand(request, metadata=metadata), + is_retryable=_never_retry if is_retryable is None else is_retryable, + max_attempts=max_attempts, + retry_interval_seconds=retry_interval_seconds, + ) + + def attach_command( + self, + name: str, + command_id: str, + *, + last_sequence_number: Optional[int] = None, + is_retryable: Optional[Callable[[Exception], bool]] = None, + max_attempts: int = _DEFAULT_MAX_ATTEMPTS, + retry_interval_seconds: float = _DEFAULT_RETRY_INTERVAL_SECONDS, + ) -> Iterator["_pb.AttachCommandResponse"]: + """Attach to a running command in sandbox ``name`` and stream its output. + + ``command_id`` comes from the ``StartEvent`` of a prior + :meth:`execute_command_streaming`. The server replays buffered events with + ``sequence_number`` greater than ``last_sequence_number`` (0 replays from + the beginning; omitted tails live output only), then streams live events. + + Unlike :meth:`execute_command_streaming`, retries default to the + transport's ``UNAVAILABLE`` policy: attach is a resumable replay, so + re-opening from ``last_sequence_number`` is safe. Pass ``is_retryable`` to + override. + """ + from databricks.sdk._sandbox_grpc import sandbox_service_pb2 as _pb # noqa: FlagLocalImports - lazy; stubs ship only with the optional "sandbox" extra + + request = _pb.AttachCommandRequest(command_id=command_id) + if last_sequence_number is not None: + request.last_sequence_number = last_sequence_number + yield from self._server_stream( + name, + lambda stub, metadata: stub.AttachCommand(request, metadata=metadata), + is_retryable=is_retryable, + max_attempts=max_attempts, + retry_interval_seconds=retry_interval_seconds, + ) + + def stream_input( + self, + name: str, + command_id: str, + inputs: Iterable[bytes], + *, + close_stdin: bool = True, + ) -> "_pb.StreamInputResponse": + """Forward stdin ``inputs`` to a running command in sandbox ``name``. + + ``command_id`` comes from the ``StartEvent`` of a prior + :meth:`execute_command_streaming`. Sends a start message, one data message + per chunk of ``inputs``, then (unless ``close_stdin`` is false) a close + message that sends EOF to the process's stdin. Returns the server's + acknowledgement once the input stream is fully sent. + + Not retried: forwarding input is stateful and side-effecting, so a + transport failure is raised to the caller. + """ + from databricks.sdk._sandbox_grpc import sandbox_service_pb2 as _pb # noqa: FlagLocalImports - lazy; stubs ship only with the optional "sandbox" extra + from databricks.sdk._sandbox_grpc import sandbox_service_pb2_grpc as _pb_grpc # noqa: FlagLocalImports - lazy; stubs ship only with the optional "sandbox" extra + + def requests() -> Iterator["_pb.StreamInputRequest"]: + yield _pb.StreamInputRequest(start=_pb.StreamInputStart(command_id=command_id)) + for chunk in inputs: + yield _pb.StreamInputRequest(data=_pb.StreamInputData(data=chunk)) + if close_stdin: + yield _pb.StreamInputRequest(close=_pb.StreamInputClose()) + + channel = open_channel(self._api._cfg) + try: + stub = _pb_grpc.SandboxServiceStub(channel) + metadata: _Metadata = ((_SANDBOX_NAME_METADATA, name),) + return stub.StreamInput(requests(), metadata=metadata) + finally: + channel.close() + + def _server_stream( + self, + name: str, + make_stream: Callable[["_pb_grpc.SandboxServiceStub", _Metadata], Iterator], + *, + is_retryable: Optional[Callable[[Exception], bool]], + max_attempts: int, + retry_interval_seconds: float, + ) -> Iterator: + """Open a per-call authenticated channel, run ``make_stream`` on the stub + under :func:`open_stream_with_retry`, and yield its items, closing the + channel once iteration ends or is abandoned.""" + # Lazy import: grpcio and the generated stubs ship only with the optional + # "sandbox" extra, so importing this module must not require them. + from databricks.sdk._sandbox_grpc import sandbox_service_pb2_grpc as _pb_grpc # noqa: FlagLocalImports - lazy; stubs ship only with the optional "sandbox" extra + + channel = open_channel(self._api._cfg) + try: + stub = _pb_grpc.SandboxServiceStub(channel) + metadata: _Metadata = ((_SANDBOX_NAME_METADATA, name),) + yield from open_stream_with_retry( + lambda: make_stream(stub, metadata), + is_retryable=is_retryable, + max_attempts=max_attempts, + retry_interval_seconds=retry_interval_seconds, + ) + finally: + channel.close() diff --git a/databricks/sdk/service/jobs.py b/databricks/sdk/service/jobs.py index e1fba90ad..d299b0a4a 100644 --- a/databricks/sdk/service/jobs.py +++ b/databricks/sdk/service/jobs.py @@ -513,6 +513,11 @@ class BaseRun: """The time at which this run ended in epoch milliseconds (milliseconds since 1/1/1970 UTC). This field is set to 0 if the job is still running.""" + environment_variables: Optional[List[JobEnvironmentVariables]] = None + """Snapshot of ``JobSettings.environment_variables`` as it was at run launch — the full list of + named environment-variable entries the job defined. To find which entry a given task ran with, + look at ``RunTaskSettings.environment_variables_key``.""" + execution_duration: Optional[int] = None """The time in milliseconds it took to execute the commands in the JAR or notebook until they completed, failed, timed out, were cancelled, or encountered an unexpected error. The duration @@ -640,6 +645,8 @@ def as_dict(self) -> dict: body["effective_usage_policy_id"] = self.effective_usage_policy_id if self.end_time is not None: body["end_time"] = self.end_time + if self.environment_variables: + body["environment_variables"] = [v.as_dict() for v in self.environment_variables] if self.execution_duration is not None: body["execution_duration"] = self.execution_duration if self.git_source: @@ -717,6 +724,8 @@ def as_shallow_dict(self) -> dict: body["effective_usage_policy_id"] = self.effective_usage_policy_id if self.end_time is not None: body["end_time"] = self.end_time + if self.environment_variables: + body["environment_variables"] = self.environment_variables if self.execution_duration is not None: body["execution_duration"] = self.execution_duration if self.git_source: @@ -785,6 +794,7 @@ def from_dict(cls, d: Dict[str, Any]) -> BaseRun: effective_performance_target=_enum(d, "effective_performance_target", PerformanceTarget), effective_usage_policy_id=d.get("effective_usage_policy_id", None), end_time=_int64(d, "end_time"), + environment_variables=_repeated_dict(d, "environment_variables", JobEnvironmentVariables), execution_duration=_int64(d, "execution_duration"), git_source=_from_dict(d, "git_source", GitSource), has_more=d.get("has_more", None), @@ -2103,7 +2113,10 @@ class DeploymentSpec: same command; role-split workloads (driver + worker, parameter server, separate eval node, etc.) use multiple entries.""" - command_path: str + compute: ComputeSpec + """Compute resources allocated to each node in this deployment.""" + + command_path: Optional[str] = None """Workspace path of the script to run on each node in this deployment. Upload the script to this path and supply the path here. When the task runs, the file at this path is run on each node; if it fails, the task fails with its exit code. @@ -2121,9 +2134,6 @@ class DeploymentSpec: # Distributed via torchrun: torchrun --nproc_per_node=8 train.py""" - compute: ComputeSpec - """Compute resources allocated to each node in this deployment.""" - name: Optional[str] = None """Optional human-readable name for this deployment (for example, ``driver``, ``worker``, ``param_server``). Used for log and UI display. Distinct names are recommended so deployments @@ -3380,6 +3390,100 @@ def from_dict(cls, d: Dict[str, Any]) -> JobEnvironment: return cls(environment_key=d.get("environment_key", None), spec=_from_dict(d, "spec", compute.Environment)) +@dataclass +class JobEnvironmentVariables: + """A named environment-variable entry, defined once at the job level and referenced by key from one + or more tasks. Entries live on ``JobSettings.environment_variables``, and tasks select one via + ``TaskSettings.environment_variables_key``.""" + + environment_variables_key: Optional[str] = None + """Identifier for this entry. Must be unique within ``JobSettings.environment_variables``. Tasks + reference it from ``TaskSettings.environment_variables_key``.""" + + spec: Optional[JobEnvironmentVariablesSpec] = None + """The environment variable specification.""" + + def as_dict(self) -> dict: + """Serializes the JobEnvironmentVariables into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.environment_variables_key is not None: + body["environment_variables_key"] = self.environment_variables_key + if self.spec: + body["spec"] = self.spec.as_dict() + return body + + def as_shallow_dict(self) -> dict: + """Serializes the JobEnvironmentVariables into a shallow dictionary of its immediate attributes.""" + body = {} + if self.environment_variables_key is not None: + body["environment_variables_key"] = self.environment_variables_key + if self.spec: + body["spec"] = self.spec + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> JobEnvironmentVariables: + """Deserializes the JobEnvironmentVariables from a dictionary.""" + return cls( + environment_variables_key=d.get("environment_variables_key", None), + spec=_from_dict(d, "spec", JobEnvironmentVariablesSpec), + ) + + +@dataclass +class JobEnvironmentVariablesSpec: + """The environment variables and files associated with a job environment variable entry. Runtime + environment variables override inline ``variables``, which override values from ``files``, on + duplicate keys.""" + + files: Optional[List[str]] = None + """Workspace (``/Workspace/...``) or UC Volumes (``/Volumes/...``) paths to ``.env`` files. Maximum + 5 files. Files are read, parsed, and merged at task execution time, not at job creation or + update API call time. + + File format: each line containing a variable must be exactly ``KEY=VALUE``. Empty and + whitespace-only lines, and lines beginning with ``#``, are ignored. Keys must match the same + regex as inlined variable names (``^[A-Za-z_][A-Za-z0-9_]*$``); the value continues to the end + of the line. No other syntax is supported — no inline comments, no quoted values, no escape + sequences, no variable interpolation. Any other line that does not match the ``KEY=VALUE`` shape + fails the run. + + Size limits: maximum 32,768 bytes (32 KiB) per file on disk; maximum 1,024 bytes (1 KiB) per + ``KEY=VALUE`` line combined. Files or lines exceeding these limits fail the run. + + On a duplicate key, the later file wins; ``variables`` override values from any file.""" + + variables: Optional[Dict[str, str]] = None + """Environment variables specified directly as key/value pairs. Maximum 20 entries. + + Each key must be 1 to 256 characters and match ``^[A-Za-z_][A-Za-z0-9_]*$``: it must start with + an ASCII letter or underscore and contain only ASCII letters, digits, and underscores. Each + value can be any Unicode string of up to 512 characters, including an empty string.""" + + def as_dict(self) -> dict: + """Serializes the JobEnvironmentVariablesSpec into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.files: + body["files"] = [v for v in self.files] + if self.variables: + body["variables"] = self.variables + return body + + def as_shallow_dict(self) -> dict: + """Serializes the JobEnvironmentVariablesSpec into a shallow dictionary of its immediate attributes.""" + body = {} + if self.files: + body["files"] = self.files + if self.variables: + body["variables"] = self.variables + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> JobEnvironmentVariablesSpec: + """Deserializes the JobEnvironmentVariablesSpec from a dictionary.""" + return cls(files=d.get("files", None), variables=d.get("variables", None)) + + @dataclass class JobNotificationSettings: no_alert_for_canceled_runs: Optional[bool] = None @@ -3693,6 +3797,12 @@ class JobSettings: """An optional set of email addresses that is notified when runs of this job begin or complete as well as when this job is deleted.""" + environment_variables: Optional[List[JobEnvironmentVariables]] = None + """Named environment-variable entries that tasks can reference by key from + ``TaskSettings.environment_variables_key``. Each entry's ``spec`` holds inline ``variables`` and + optional ``.env`` ``files``. Maximum 10 entries per job. A task can reference at most one entry + from this list.""" + environments: Optional[List[JobEnvironment]] = None """A list of task execution environment specifications that can be referenced by tasks that use serverless compute or a compute resource that uses Environments mode. @@ -3819,6 +3929,8 @@ def as_dict(self) -> dict: body["edit_mode"] = self.edit_mode.value if self.email_notifications: body["email_notifications"] = self.email_notifications.as_dict() + if self.environment_variables: + body["environment_variables"] = [v.as_dict() for v in self.environment_variables] if self.environments: body["environments"] = [v.as_dict() for v in self.environments] if self.format is not None: @@ -3878,6 +3990,8 @@ def as_shallow_dict(self) -> dict: body["edit_mode"] = self.edit_mode if self.email_notifications: body["email_notifications"] = self.email_notifications + if self.environment_variables: + body["environment_variables"] = self.environment_variables if self.environments: body["environments"] = self.environments if self.format is not None: @@ -3932,6 +4046,7 @@ def from_dict(cls, d: Dict[str, Any]) -> JobSettings: description=d.get("description", None), edit_mode=_enum(d, "edit_mode", JobEditMode), email_notifications=_from_dict(d, "email_notifications", JobEmailNotifications), + environment_variables=_repeated_dict(d, "environment_variables", JobEnvironmentVariables), environments=_repeated_dict(d, "environments", JobEnvironment), format=_enum(d, "format", Format), git_source=_from_dict(d, "git_source", GitSource), @@ -5738,6 +5853,11 @@ class Run: """The time at which this run ended in epoch milliseconds (milliseconds since 1/1/1970 UTC). This field is set to 0 if the job is still running.""" + environment_variables: Optional[List[JobEnvironmentVariables]] = None + """Snapshot of ``JobSettings.environment_variables`` as it was at run launch — the full list of + named environment-variable entries the job defined. To find which entry a given task ran with, + look at ``RunTaskSettings.environment_variables_key``.""" + execution_duration: Optional[int] = None """The time in milliseconds it took to execute the commands in the JAR or notebook until they completed, failed, timed out, were cancelled, or encountered an unexpected error. The duration @@ -5871,6 +5991,8 @@ def as_dict(self) -> dict: body["effective_usage_policy_id"] = self.effective_usage_policy_id if self.end_time is not None: body["end_time"] = self.end_time + if self.environment_variables: + body["environment_variables"] = [v.as_dict() for v in self.environment_variables] if self.execution_duration is not None: body["execution_duration"] = self.execution_duration if self.git_source: @@ -5952,6 +6074,8 @@ def as_shallow_dict(self) -> dict: body["effective_usage_policy_id"] = self.effective_usage_policy_id if self.end_time is not None: body["end_time"] = self.end_time + if self.environment_variables: + body["environment_variables"] = self.environment_variables if self.execution_duration is not None: body["execution_duration"] = self.execution_duration if self.git_source: @@ -6024,6 +6148,7 @@ def from_dict(cls, d: Dict[str, Any]) -> Run: effective_performance_target=_enum(d, "effective_performance_target", PerformanceTarget), effective_usage_policy_id=d.get("effective_usage_policy_id", None), end_time=_int64(d, "end_time"), + environment_variables=_repeated_dict(d, "environment_variables", JobEnvironmentVariables), execution_duration=_int64(d, "execution_duration"), git_source=_from_dict(d, "git_source", GitSource), has_more=d.get("has_more", None), @@ -6991,6 +7116,13 @@ class RunTask: Python wheel and dbt tasks when using serverless compute or a compute resource that uses Environments mode.""" + environment_variables_key: Optional[str] = None + """Reference to a ``JobEnvironmentVariables`` entry defined in + ``RunSettings.environment_variables`` for one-time runs or preserved in + ``Run.environment_variables`` for run snapshots. The selected entry's variables are applied to + this task at execution time. This field supports serverless tasks using environment version 5 or + later.""" + execution_duration: Optional[int] = None """The time in milliseconds it took to execute the commands in the JAR or notebook until they completed, failed, timed out, were cancelled, or encountered an unexpected error. The duration @@ -7169,6 +7301,8 @@ def as_dict(self) -> dict: body["end_time"] = self.end_time if self.environment_key is not None: body["environment_key"] = self.environment_key + if self.environment_variables_key is not None: + body["environment_variables_key"] = self.environment_variables_key if self.execution_duration is not None: body["execution_duration"] = self.execution_duration if self.existing_cluster_id is not None: @@ -7286,6 +7420,8 @@ def as_shallow_dict(self) -> dict: body["end_time"] = self.end_time if self.environment_key is not None: body["environment_key"] = self.environment_key + if self.environment_variables_key is not None: + body["environment_variables_key"] = self.environment_variables_key if self.execution_duration is not None: body["execution_duration"] = self.execution_duration if self.existing_cluster_id is not None: @@ -7383,6 +7519,7 @@ def from_dict(cls, d: Dict[str, Any]) -> RunTask: email_notifications=_from_dict(d, "email_notifications", JobEmailNotifications), end_time=_int64(d, "end_time"), environment_key=d.get("environment_key", None), + environment_variables_key=d.get("environment_variables_key", None), execution_duration=_int64(d, "execution_duration"), existing_cluster_id=d.get("existing_cluster_id", None), for_each_task=_from_dict(d, "for_each_task", RunForEachTask), @@ -8493,6 +8630,13 @@ class SubmitTask: Python wheel and dbt tasks when using serverless compute or a compute resource that uses Environments mode.""" + environment_variables_key: Optional[str] = None + """Reference to a ``JobEnvironmentVariables`` entry defined in + ``RunSettings.environment_variables`` for one-time runs or preserved in + ``Run.environment_variables`` for run snapshots. The selected entry's variables are applied to + this task at execution time. This field supports serverless tasks using environment version 5 or + later.""" + existing_cluster_id: Optional[str] = None """If existing_cluster_id, the ID of an existing cluster that is used for all runs. When running jobs or tasks on an existing cluster, you may need to manually restart the cluster if it stops @@ -8611,6 +8755,8 @@ def as_dict(self) -> dict: body["email_notifications"] = self.email_notifications.as_dict() if self.environment_key is not None: body["environment_key"] = self.environment_key + if self.environment_variables_key is not None: + body["environment_variables_key"] = self.environment_variables_key if self.existing_cluster_id is not None: body["existing_cluster_id"] = self.existing_cluster_id if self.for_each_task: @@ -8694,6 +8840,8 @@ def as_shallow_dict(self) -> dict: body["email_notifications"] = self.email_notifications if self.environment_key is not None: body["environment_key"] = self.environment_key + if self.environment_variables_key is not None: + body["environment_variables_key"] = self.environment_variables_key if self.existing_cluster_id is not None: body["existing_cluster_id"] = self.existing_cluster_id if self.for_each_task: @@ -8763,6 +8911,7 @@ def from_dict(cls, d: Dict[str, Any]) -> SubmitTask: disabled=d.get("disabled", None), email_notifications=_from_dict(d, "email_notifications", JobEmailNotifications), environment_key=d.get("environment_key", None), + environment_variables_key=d.get("environment_variables_key", None), existing_cluster_id=d.get("existing_cluster_id", None), for_each_task=_from_dict(d, "for_each_task", ForEachTask), gen_ai_compute_task=_from_dict(d, "gen_ai_compute_task", GenAiComputeTask), @@ -9053,6 +9202,11 @@ class Task: Python wheel and dbt tasks when using serverless compute or a compute resource that uses Environments mode.""" + environment_variables_key: Optional[str] = None + """Reference to a ``JobEnvironmentVariables`` entry defined in + ``JobSettings.environment_variables``. The selected entry's variables are applied to this task + at execution time. This field supports serverless tasks using environment version 5 or later.""" + existing_cluster_id: Optional[str] = None """If existing_cluster_id, the ID of an existing cluster that is used for all runs. When running jobs or tasks on an existing cluster, you may need to manually restart the cluster if it stops @@ -9180,6 +9334,8 @@ def as_dict(self) -> dict: body["email_notifications"] = self.email_notifications.as_dict() if self.environment_key is not None: body["environment_key"] = self.environment_key + if self.environment_variables_key is not None: + body["environment_variables_key"] = self.environment_variables_key if self.existing_cluster_id is not None: body["existing_cluster_id"] = self.existing_cluster_id if self.for_each_task: @@ -9265,6 +9421,8 @@ def as_shallow_dict(self) -> dict: body["email_notifications"] = self.email_notifications if self.environment_key is not None: body["environment_key"] = self.environment_key + if self.environment_variables_key is not None: + body["environment_variables_key"] = self.environment_variables_key if self.existing_cluster_id is not None: body["existing_cluster_id"] = self.existing_cluster_id if self.for_each_task: @@ -9336,6 +9494,7 @@ def from_dict(cls, d: Dict[str, Any]) -> Task: disabled=d.get("disabled", None), email_notifications=_from_dict(d, "email_notifications", TaskEmailNotifications), environment_key=d.get("environment_key", None), + environment_variables_key=d.get("environment_variables_key", None), existing_cluster_id=d.get("existing_cluster_id", None), for_each_task=_from_dict(d, "for_each_task", ForEachTask), gen_ai_compute_task=_from_dict(d, "gen_ai_compute_task", GenAiComputeTask), @@ -10406,6 +10565,7 @@ def create( description: Optional[str] = None, edit_mode: Optional[JobEditMode] = None, email_notifications: Optional[JobEmailNotifications] = None, + environment_variables: Optional[List[JobEnvironmentVariables]] = None, environments: Optional[List[JobEnvironment]] = None, format: Optional[Format] = None, git_source: Optional[GitSource] = None, @@ -10454,6 +10614,11 @@ def create( :param email_notifications: :class:`JobEmailNotifications` (optional) An optional set of email addresses that is notified when runs of this job begin or complete as well as when this job is deleted. + :param environment_variables: List[:class:`JobEnvironmentVariables`] (optional) + Named environment-variable entries that tasks can reference by key from + ``TaskSettings.environment_variables_key``. Each entry's ``spec`` holds inline ``variables`` and + optional ``.env`` ``files``. Maximum 10 entries per job. A task can reference at most one entry from + this list. :param environments: List[:class:`JobEnvironment`] (optional) A list of task execution environment specifications that can be referenced by tasks that use serverless compute or a compute resource that uses Environments mode. @@ -10562,6 +10727,8 @@ def create( body["edit_mode"] = edit_mode.value if email_notifications is not None: body["email_notifications"] = email_notifications.as_dict() + if environment_variables is not None: + body["environment_variables"] = [v.as_dict() for v in environment_variables] if environments is not None: body["environments"] = [v.as_dict() for v in environments] if format is not None: @@ -11480,6 +11647,7 @@ def submit( access_control_list: Optional[List[JobAccessControlRequest]] = None, budget_policy_id: Optional[str] = None, email_notifications: Optional[JobEmailNotifications] = None, + environment_variables: Optional[List[JobEnvironmentVariables]] = None, environments: Optional[List[JobEnvironment]] = None, git_source: Optional[GitSource] = None, health: Optional[JobsHealthRules] = None, @@ -11511,6 +11679,11 @@ def submit( will be not be attributed to any budget policy. :param email_notifications: :class:`JobEmailNotifications` (optional) An optional set of email addresses notified when the run begins or completes. + :param environment_variables: List[:class:`JobEnvironmentVariables`] (optional) + Named environment-variable entries that tasks of this one-time run can reference by key from + ``RunTaskSettings.environment_variables_key``. Each entry's ``spec`` holds inline ``variables`` and + optional ``.env`` ``files``. Handled identically to ``JobSettings.environment_variables``. Maximum + 10 entries. Entries are independent of one another — there is no cross-entry merging. :param environments: List[:class:`JobEnvironment`] (optional) A list of task execution environment specifications that can be referenced by tasks of this run. :param git_source: :class:`GitSource` (optional) @@ -11573,6 +11746,8 @@ def submit( body["budget_policy_id"] = budget_policy_id if email_notifications is not None: body["email_notifications"] = email_notifications.as_dict() + if environment_variables is not None: + body["environment_variables"] = [v.as_dict() for v in environment_variables] if environments is not None: body["environments"] = [v.as_dict() for v in environments] if git_source is not None: @@ -11621,6 +11796,7 @@ def submit_and_wait( access_control_list: Optional[List[JobAccessControlRequest]] = None, budget_policy_id: Optional[str] = None, email_notifications: Optional[JobEmailNotifications] = None, + environment_variables: Optional[List[JobEnvironmentVariables]] = None, environments: Optional[List[JobEnvironment]] = None, git_source: Optional[GitSource] = None, health: Optional[JobsHealthRules] = None, @@ -11640,6 +11816,7 @@ def submit_and_wait( access_control_list=access_control_list, budget_policy_id=budget_policy_id, email_notifications=email_notifications, + environment_variables=environment_variables, environments=environments, git_source=git_source, health=health, diff --git a/databricks/sdk/service/ml.py b/databricks/sdk/service/ml.py index a55c4982e..0ea1b331e 100644 --- a/databricks/sdk/service/ml.py +++ b/databricks/sdk/service/ml.py @@ -7713,8 +7713,9 @@ class SchemaLocatorFormat(Enum): class SchemaRegistryConfig: """Configuration for resolving a Stream's schema from an external schema registry (e.g. Confluent).""" - api_secret_ref: SecretScopeReference - """Reference to the schema registry API secret in a Databricks secret scope.""" + api_secret_ref: Optional[SecretScopeReference] = None + """Reference to the schema registry API secret in a Databricks secret scope. Set this only if + required for authentication for the schema registry.""" key_schema_locator: Optional[SchemaLocator] = None """Schema locator for the message key. Only used for Kafka streams. At least one of diff --git a/databricks/sdk/version.py b/databricks/sdk/version.py index 5a6bb17fb..3b5229b1c 100644 --- a/databricks/sdk/version.py +++ b/databricks/sdk/version.py @@ -1 +1 @@ -__version__ = "0.143.0" +__version__ = "0.144.0" diff --git a/docs/dbdataclasses/jobs.rst b/docs/dbdataclasses/jobs.rst index 6402a8ad6..6f6bccc03 100644 --- a/docs/dbdataclasses/jobs.rst +++ b/docs/dbdataclasses/jobs.rst @@ -491,6 +491,14 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: JobEnvironmentVariables + :members: + :undoc-members: + +.. autoclass:: JobEnvironmentVariablesSpec + :members: + :undoc-members: + .. autoclass:: JobNotificationSettings :members: :undoc-members: diff --git a/docs/workspace/jobs/jobs.rst b/docs/workspace/jobs/jobs.rst index c4a51fbbe..10a217e02 100644 --- a/docs/workspace/jobs/jobs.rst +++ b/docs/workspace/jobs/jobs.rst @@ -125,7 +125,7 @@ .. py:method:: cancel_run_and_wait(run_id: int, timeout: datetime.timedelta = 0:20:00) -> Run - .. py:method:: create( [, access_control_list: Optional[List[JobAccessControlRequest]], budget_policy_id: Optional[str], continuous: Optional[Continuous], deployment: Optional[JobDeployment], description: Optional[str], edit_mode: Optional[JobEditMode], email_notifications: Optional[JobEmailNotifications], environments: Optional[List[JobEnvironment]], format: Optional[Format], git_source: Optional[GitSource], health: Optional[JobsHealthRules], job_clusters: Optional[List[JobCluster]], max_concurrent_runs: Optional[int], name: Optional[str], notification_settings: Optional[JobNotificationSettings], parameters: Optional[List[JobParameterDefinition]], parent_path: Optional[str], performance_target: Optional[PerformanceTarget], queue: Optional[QueueSettings], run_as: Optional[JobRunAs], schedule: Optional[CronSchedule], tags: Optional[Dict[str, str]], tasks: Optional[List[Task]], timeout_seconds: Optional[int], trigger: Optional[TriggerSettings], triggers: Optional[List[TriggerConfiguration]], usage_policy_id: Optional[str], webhook_notifications: Optional[WebhookNotifications]]) -> CreateResponse + .. py:method:: create( [, access_control_list: Optional[List[JobAccessControlRequest]], budget_policy_id: Optional[str], continuous: Optional[Continuous], deployment: Optional[JobDeployment], description: Optional[str], edit_mode: Optional[JobEditMode], email_notifications: Optional[JobEmailNotifications], environment_variables: Optional[List[JobEnvironmentVariables]], environments: Optional[List[JobEnvironment]], format: Optional[Format], git_source: Optional[GitSource], health: Optional[JobsHealthRules], job_clusters: Optional[List[JobCluster]], max_concurrent_runs: Optional[int], name: Optional[str], notification_settings: Optional[JobNotificationSettings], parameters: Optional[List[JobParameterDefinition]], parent_path: Optional[str], performance_target: Optional[PerformanceTarget], queue: Optional[QueueSettings], run_as: Optional[JobRunAs], schedule: Optional[CronSchedule], tags: Optional[Dict[str, str]], tasks: Optional[List[Task]], timeout_seconds: Optional[int], trigger: Optional[TriggerSettings], triggers: Optional[List[TriggerConfiguration]], usage_policy_id: Optional[str], webhook_notifications: Optional[WebhookNotifications]]) -> CreateResponse Usage: @@ -188,6 +188,11 @@ :param email_notifications: :class:`JobEmailNotifications` (optional) An optional set of email addresses that is notified when runs of this job begin or complete as well as when this job is deleted. + :param environment_variables: List[:class:`JobEnvironmentVariables`] (optional) + Named environment-variable entries that tasks can reference by key from + ``TaskSettings.environment_variables_key``. Each entry's ``spec`` holds inline ``variables`` and + optional ``.env`` ``files``. Maximum 10 entries per job. A task can reference at most one entry from + this list. :param environments: List[:class:`JobEnvironment`] (optional) A list of task execution environment specifications that can be referenced by tasks that use serverless compute or a compute resource that uses Environments mode. @@ -1030,7 +1035,7 @@ :returns: :class:`JobPermissions` - .. py:method:: submit( [, access_control_list: Optional[List[JobAccessControlRequest]], budget_policy_id: Optional[str], email_notifications: Optional[JobEmailNotifications], environments: Optional[List[JobEnvironment]], git_source: Optional[GitSource], health: Optional[JobsHealthRules], idempotency_token: Optional[str], notification_settings: Optional[JobNotificationSettings], performance_target: Optional[PerformanceTarget], queue: Optional[QueueSettings], run_as: Optional[JobRunAs], run_name: Optional[str], tasks: Optional[List[SubmitTask]], timeout_seconds: Optional[int], usage_policy_id: Optional[str], webhook_notifications: Optional[WebhookNotifications]]) -> Wait[Run] + .. py:method:: submit( [, access_control_list: Optional[List[JobAccessControlRequest]], budget_policy_id: Optional[str], email_notifications: Optional[JobEmailNotifications], environment_variables: Optional[List[JobEnvironmentVariables]], environments: Optional[List[JobEnvironment]], git_source: Optional[GitSource], health: Optional[JobsHealthRules], idempotency_token: Optional[str], notification_settings: Optional[JobNotificationSettings], performance_target: Optional[PerformanceTarget], queue: Optional[QueueSettings], run_as: Optional[JobRunAs], run_name: Optional[str], tasks: Optional[List[SubmitTask]], timeout_seconds: Optional[int], usage_policy_id: Optional[str], webhook_notifications: Optional[WebhookNotifications]]) -> Wait[Run] Usage: @@ -1082,6 +1087,11 @@ will be not be attributed to any budget policy. :param email_notifications: :class:`JobEmailNotifications` (optional) An optional set of email addresses notified when the run begins or completes. + :param environment_variables: List[:class:`JobEnvironmentVariables`] (optional) + Named environment-variable entries that tasks of this one-time run can reference by key from + ``RunTaskSettings.environment_variables_key``. Each entry's ``spec`` holds inline ``variables`` and + optional ``.env`` ``files``. Handled identically to ``JobSettings.environment_variables``. Maximum + 10 entries. Entries are independent of one another — there is no cross-entry merging. :param environments: List[:class:`JobEnvironment`] (optional) A list of task execution environment specifications that can be referenced by tasks of this run. :param git_source: :class:`GitSource` (optional) @@ -1137,7 +1147,7 @@ See :method:wait_get_run_job_terminated_or_skipped for more details. - .. py:method:: submit_and_wait( [, access_control_list: Optional[List[JobAccessControlRequest]], budget_policy_id: Optional[str], email_notifications: Optional[JobEmailNotifications], environments: Optional[List[JobEnvironment]], git_source: Optional[GitSource], health: Optional[JobsHealthRules], idempotency_token: Optional[str], notification_settings: Optional[JobNotificationSettings], performance_target: Optional[PerformanceTarget], queue: Optional[QueueSettings], run_as: Optional[JobRunAs], run_name: Optional[str], tasks: Optional[List[SubmitTask]], timeout_seconds: Optional[int], usage_policy_id: Optional[str], webhook_notifications: Optional[WebhookNotifications], timeout: datetime.timedelta = 0:20:00]) -> Run + .. py:method:: submit_and_wait( [, access_control_list: Optional[List[JobAccessControlRequest]], budget_policy_id: Optional[str], email_notifications: Optional[JobEmailNotifications], environment_variables: Optional[List[JobEnvironmentVariables]], environments: Optional[List[JobEnvironment]], git_source: Optional[GitSource], health: Optional[JobsHealthRules], idempotency_token: Optional[str], notification_settings: Optional[JobNotificationSettings], performance_target: Optional[PerformanceTarget], queue: Optional[QueueSettings], run_as: Optional[JobRunAs], run_name: Optional[str], tasks: Optional[List[SubmitTask]], timeout_seconds: Optional[int], usage_policy_id: Optional[str], webhook_notifications: Optional[WebhookNotifications], timeout: datetime.timedelta = 0:20:00]) -> Run .. py:method:: update(job_id: int [, fields_to_remove: Optional[List[str]], new_settings: Optional[JobSettings]]) diff --git a/docs/workspace/sandbox/sandbox.rst b/docs/workspace/sandbox/sandbox.rst index f686eedd2..00216a60c 100644 --- a/docs/workspace/sandbox/sandbox.rst +++ b/docs/workspace/sandbox/sandbox.rst @@ -2,11 +2,26 @@ ========================= .. currentmodule:: databricks.sdk.service.sandbox -.. py:class:: SandboxAPI +.. py:class:: SandboxExt Create, manage, and control the lifecycle of sandboxes -- isolated, pre-configured, low-latency Serverless compute environments for running code. + .. py:method:: attach_command(name: str, command_id: str [, last_sequence_number: Optional[int], is_retryable: Optional[Callable[[Exception], bool]], max_attempts: int = 30, retry_interval_seconds: float = 3.0]) -> Iterator['_pb.AttachCommandResponse'] + + Attach to a running command in sandbox ``name`` and stream its output. + + ``command_id`` comes from the ``StartEvent`` of a prior + :meth:`execute_command_streaming`. The server replays buffered events with + ``sequence_number`` greater than ``last_sequence_number`` (0 replays from + the beginning; omitted tails live output only), then streams live events. + + Unlike :meth:`execute_command_streaming`, retries default to the + transport's ``UNAVAILABLE`` policy: attach is a resumable replay, so + re-opening from ``last_sequence_number`` is safe. Pass ``is_retryable`` to + override. + + .. py:method:: create_sandbox(sandbox: Sandbox, sandbox_id: str) -> Sandbox Creates a new Sandbox. @@ -28,6 +43,26 @@ + .. py:method:: execute_command_streaming(name: str, cmd: str [, args: Optional[list[str]], envs: Optional[Dict[str, str]], is_retryable: Optional[Callable[[Exception], bool]], max_attempts: int = 30, retry_interval_seconds: float = 3.0]) -> Iterator['_pb.ExecuteCommandResponse'] + + Run ``cmd`` in sandbox ``name`` and stream its process events. + + Streaming counterpart of :meth:`SandboxAPI.execute_command_sync`: the same + ``name`` (``sandboxes/{sandbox_id}``), ``cmd``, ``args``, and ``envs``, + but the process output is streamed instead of returned in one response. + Yields ``ExecuteCommandResponse`` messages, each carrying one + ``ProcessEvent`` (a start event, then stdout/stderr data chunks, then an + end event with the exit code) and a monotonic ``sequence_number``. + + Retries are off by default (``is_retryable`` defaults to + :func:`_never_retry`) because ExecuteCommand is not idempotent: re-issuing + a failed open could run the command twice. Pass an ``is_retryable`` + predicate - e.g. one matching ``UNAVAILABLE`` - to retry the stream open + when the command is safe to run more than once, optionally tuning + ``max_attempts`` and ``retry_interval_seconds``. Only the stream open is + retried; a failure once events are flowing is always raised. + + .. py:method:: execute_command_sync(name: str, cmd: str [, args: Optional[List[str]], envs: Optional[Dict[str, str]], execution_timeout: Optional[Duration]]) -> ExecuteCommandSyncResponse Runs a command in the sandbox and blocks until it exits, returning the captured stdout, stderr and @@ -91,6 +126,20 @@ :returns: :class:`Sandbox` + .. py:method:: stream_input(name: str, command_id: str, inputs: Iterable[bytes] [, close_stdin: bool = True]) -> '_pb.StreamInputResponse' + + Forward stdin ``inputs`` to a running command in sandbox ``name``. + + ``command_id`` comes from the ``StartEvent`` of a prior + :meth:`execute_command_streaming`. Sends a start message, one data message + per chunk of ``inputs``, then (unless ``close_stdin`` is false) a close + message that sends EOF to the process's stdin. Returns the server's + acknowledgement once the input stream is fully sent. + + Not retried: forwarding input is stateful and side-effecting, so a + transport failure is raised to the caller. + + .. py:method:: update_sandbox(name: str, sandbox: Sandbox, update_mask: FieldMask) -> Sandbox Updates mutable fields on an existing Sandbox. Allowlisted update_mask paths today: display_name, diff --git a/pyproject.toml b/pyproject.toml index edcbe610f..8f4a1090a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,7 +51,8 @@ dev = [ "openai", "langchain-openai", "httpx", - "grpcio<2", # the sandbox extra's dep; also needed to run the gRPC transport unit tests + "grpcio>=1.81.1,<2", # the sandbox extra's dep; also needed to run the gRPC transport unit tests + "build", # some integration tests depend on the databricks-sdk-py wheel "check-manifest", ] notebook = [ @@ -64,15 +65,20 @@ openai = [ "httpx", ] sandbox = [ - # gRPC transport for the sandbox command-execution mixin. Optional so - # users who don't use the sandbox APIs don't pull the large native grpcio - # wheel; imported lazily in databricks/sdk/mixins/_grpc_transport.py. + # gRPC transport and committed stubs for the sandbox command-execution + # mixin. Optional so users who don't use the sandbox APIs don't pull the + # large native grpcio wheel; grpcio and the stubs are imported lazily (see + # databricks/sdk/mixins/_grpc_transport.py and mixins/sandbox.py). # - # `<2` caps the major version to guard against a breaking grpcio 2.x. No - # lower bound is needed: the transport uses only long-stable grpc APIs, and - # pip installs the newest grpcio with wheels for the user's Python (the SDK - # requires >=3.10), so a too-old release is never selected. - "grpcio<2", + # The floors match the version checks baked into the committed stubs: the + # generated sandbox_service_pb2.py validates protobuf >= 6.33.5 and + # sandbox_service_pb2_grpc.py validates grpcio >= 1.81.1 at import, so an + # older runtime that satisfied only the core protobuf floor / an unbounded + # grpcio would raise on import. `<2` caps grpcio's major to guard against a + # breaking 2.x. protobuf's core dependency (above) supplies the `<7` cap and + # the CVE exclusions. + "protobuf>=6.33.5", + "grpcio>=1.81.1,<2", ] [tool.setuptools.dynamic] diff --git a/tests/test_sandbox_mixin.py b/tests/test_sandbox_mixin.py new file mode 100644 index 000000000..fffec0750 --- /dev/null +++ b/tests/test_sandbox_mixin.py @@ -0,0 +1,216 @@ +import pytest + +from databricks.sdk._sandbox_grpc import sandbox_service_pb2 as pb +from databricks.sdk._sandbox_grpc import sandbox_service_pb2_grpc as pb_grpc +from databricks.sdk.mixins import sandbox as sandbox_mod +from databricks.sdk.mixins.sandbox import SandboxExt + +# The fakes below stand in for the authenticated gRPC channel and stub so these +# unit tests exercise execute_command's routing, streaming, retry, and channel +# lifecycle without opening a real network connection. + + +class _FakeApiClient: + def __init__(self): + self._cfg = object() + + +class _FakeChannel: + def __init__(self): + self.closed = 0 + + def close(self): + self.closed += 1 + + +def _response(seq, **event): + return pb.ExecuteCommandResponse(process_event=pb.ProcessEvent(**event), sequence_number=seq) + + +def test_execute_command_streams_responses_and_closes_channel(monkeypatch): + channel = _FakeChannel() + captured = {} + + class _FakeStub: + def __init__(self, ch): + captured["channel"] = ch + + def ExecuteCommand(self, request, metadata=None): + captured["request"] = request + captured["metadata"] = metadata + return iter( + [ + _response(1, start=pb.StartEvent(pid=7)), + _response(2, data=pb.DataEvent(stdout=b"hi\n")), + _response(3, end=pb.EndEvent(exit_code=0)), + ] + ) + + monkeypatch.setattr(sandbox_mod, "open_channel", lambda cfg: channel) + monkeypatch.setattr(pb_grpc, "SandboxServiceStub", _FakeStub) + + api = SandboxExt(_FakeApiClient()) + events = list(api.execute_command_streaming("sandboxes/s1", "/bin/echo", args=["hi"], envs={"K": "v"})) + + assert [e.sequence_number for e in events] == [1, 2, 3] + assert events[0].process_event.start.pid == 7 + assert events[2].process_event.end.exit_code == 0 + assert captured["channel"] is channel + assert captured["request"].cmd == "/bin/echo" + assert list(captured["request"].args) == ["hi"] + assert dict(captured["request"].envs) == {"K": "v"} + assert captured["metadata"] == (("x-databricks-sandbox-name", "sandboxes/s1"),) + assert channel.closed == 1 + + +def test_execute_command_closes_channel_when_abandoned(monkeypatch): + channel = _FakeChannel() + + class _FakeStub: + def __init__(self, ch): + pass + + def ExecuteCommand(self, request, metadata=None): + return iter([_response(1, start=pb.StartEvent(pid=1))]) + + monkeypatch.setattr(sandbox_mod, "open_channel", lambda cfg: channel) + monkeypatch.setattr(pb_grpc, "SandboxServiceStub", _FakeStub) + + stream = SandboxExt(_FakeApiClient()).execute_command_streaming("s", "cmd") + assert next(stream).sequence_number == 1 + stream.close() + assert channel.closed == 1 + + +def test_execute_command_does_not_retry_by_default(monkeypatch): + channel = _FakeChannel() + calls = {"count": 0} + + class _FailingStub: + def __init__(self, ch): + pass + + def ExecuteCommand(self, request, metadata=None): + calls["count"] += 1 + raise RuntimeError("transport failed") + + monkeypatch.setattr(sandbox_mod, "open_channel", lambda cfg: channel) + monkeypatch.setattr(pb_grpc, "SandboxServiceStub", _FailingStub) + + stream = SandboxExt(_FakeApiClient()).execute_command_streaming("s", "cmd") + with pytest.raises(RuntimeError, match="transport failed"): + next(stream) + + assert calls["count"] == 1 + assert channel.closed == 1 + + +def test_execute_command_retries_stream_open_when_opted_in(monkeypatch): + channel = _FakeChannel() + calls = {"count": 0} + + class _FlakyStub: + def __init__(self, ch): + pass + + def ExecuteCommand(self, request, metadata=None): + calls["count"] += 1 + if calls["count"] == 1: + raise RuntimeError("transient") + return iter([_response(1, end=pb.EndEvent(exit_code=0))]) + + monkeypatch.setattr(sandbox_mod, "open_channel", lambda cfg: channel) + monkeypatch.setattr(pb_grpc, "SandboxServiceStub", _FlakyStub) + + events = list( + SandboxExt(_FakeApiClient()).execute_command_streaming( + "s", + "cmd", + is_retryable=lambda _exc: True, + max_attempts=2, + retry_interval_seconds=0.0, + ) + ) + + assert calls["count"] == 2 + assert events[0].process_event.end.exit_code == 0 + assert channel.closed == 1 + + +def test_attach_command_streams_and_closes_channel(monkeypatch): + channel = _FakeChannel() + captured = {} + + class _FakeStub: + def __init__(self, ch): + pass + + def AttachCommand(self, request, metadata=None): + captured["request"] = request + captured["metadata"] = metadata + return iter( + [ + pb.AttachCommandResponse( + process_event=pb.ProcessEvent(end=pb.EndEvent(exit_code=0)), sequence_number=5 + ) + ] + ) + + monkeypatch.setattr(sandbox_mod, "open_channel", lambda cfg: channel) + monkeypatch.setattr(pb_grpc, "SandboxServiceStub", _FakeStub) + + events = list(SandboxExt(_FakeApiClient()).attach_command("sandboxes/s1", "cmd-1", last_sequence_number=4)) + + assert [e.sequence_number for e in events] == [5] + assert captured["request"].command_id == "cmd-1" + assert captured["request"].last_sequence_number == 4 + assert captured["metadata"] == (("x-databricks-sandbox-name", "sandboxes/s1"),) + assert channel.closed == 1 + + +def test_stream_input_sends_start_data_close_and_closes_channel(monkeypatch): + channel = _FakeChannel() + captured = {} + + class _FakeStub: + def __init__(self, ch): + pass + + def StreamInput(self, request_iterator, metadata=None): + captured["requests"] = list(request_iterator) + captured["metadata"] = metadata + return pb.StreamInputResponse() + + monkeypatch.setattr(sandbox_mod, "open_channel", lambda cfg: channel) + monkeypatch.setattr(pb_grpc, "SandboxServiceStub", _FakeStub) + + resp = SandboxExt(_FakeApiClient()).stream_input("sandboxes/s1", "cmd-1", [b"a", b"bc"]) + + assert isinstance(resp, pb.StreamInputResponse) + reqs = captured["requests"] + assert [r.WhichOneof("event") for r in reqs] == ["start", "data", "data", "close"] + assert reqs[0].start.command_id == "cmd-1" + assert [reqs[1].data.data, reqs[2].data.data] == [b"a", b"bc"] + assert captured["metadata"] == (("x-databricks-sandbox-name", "sandboxes/s1"),) + assert channel.closed == 1 + + +def test_stream_input_omits_close_when_disabled(monkeypatch): + channel = _FakeChannel() + captured = {} + + class _FakeStub: + def __init__(self, ch): + pass + + def StreamInput(self, request_iterator, metadata=None): + captured["requests"] = list(request_iterator) + return pb.StreamInputResponse() + + monkeypatch.setattr(sandbox_mod, "open_channel", lambda cfg: channel) + monkeypatch.setattr(pb_grpc, "SandboxServiceStub", _FakeStub) + + SandboxExt(_FakeApiClient()).stream_input("sandboxes/s1", "cmd-1", [b"x"], close_stdin=False) + + assert [r.WhichOneof("event") for r in captured["requests"]] == ["start", "data"] + assert channel.closed == 1