databricks-sdk 0.18.0__py3-none-any.whl → 0.19.0__py3-none-any.whl

databricks/sdk/__init__.py

@@ -1,5 +1,6 @@
 import databricks.sdk.core as client
 import databricks.sdk.dbutils as dbutils
+from databricks.sdk import azure
 from databricks.sdk.credentials_provider import CredentialsProvider
 from databricks.sdk.mixins.compute import ClustersExt
 from databricks.sdk.mixins.files import DbfsExt
@@ -46,7 +47,7 @@ from databricks.sdk.service.provisioning import (CredentialsAPI,
                                                  EncryptionKeysAPI,
                                                  NetworksAPI, PrivateAccessAPI,
                                                  StorageAPI, VpcEndpointsAPI,
-                                                 WorkspacesAPI)
+                                                 Workspace, WorkspacesAPI)
 from databricks.sdk.service.serving import AppsAPI, ServingEndpointsAPI
 from databricks.sdk.service.settings import (AccountIpAccessListsAPI,
                                              AccountSettingsAPI,
@@ -779,5 +780,33 @@ class AccountClient:
         """These APIs manage workspaces for this account."""
         return self._workspaces
 
+    def get_workspace_client(self, workspace: Workspace) -> WorkspaceClient:
+        """Constructs a ``WorkspaceClient`` for the given workspace.
+
+        Returns a ``WorkspaceClient`` that is configured to use the same
+        credentials as this ``AccountClient``. The underlying config is
+        copied from this ``AccountClient``, but the ``host`` and
+        ``azure_workspace_resource_id`` are overridden to match the
+        given workspace, and the ``account_id`` field is cleared.
+
+        Usage:
+
+        .. code-block::
+
+            wss = list(a.workspaces.list())
+            if len(wss) == 0:
+                pytest.skip("no workspaces")
+            w = a.get_workspace_client(wss[0])
+            assert w.current_user.me().active
+
+        :param workspace: The workspace to construct a client for.
+        :return: A ``WorkspaceClient`` for the given workspace.
+        """
+        config = self._config.copy()
+        config.host = config.environment.deployment_url(workspace.deployment_name)
+        config.azure_workspace_resource_id = azure.get_azure_resource_id(workspace)
+        config.account_id = None
+        return WorkspaceClient(config=config)
+
     def __repr__(self):
         return f"AccountClient(account_id='{self._config.account_id}', auth_type='{self._config.auth_type}', ...)"

databricks/sdk/azure.py

@@ -2,6 +2,7 @@ from dataclasses import dataclass
 from typing import Dict
 
 from .oauth import TokenSource
+from .service.provisioning import Workspace
 
 
 @dataclass
@@ -38,3 +39,16 @@ def add_workspace_id_header(cfg: 'Config', headers: Dict[str, str]):
 def add_sp_management_token(token_source: 'TokenSource', headers: Dict[str, str]):
     mgmt_token = token_source.token()
     headers['X-Databricks-Azure-SP-Management-Token'] = mgmt_token.access_token
+
+
+def get_azure_resource_id(workspace: Workspace):
+    """
+    Returns the Azure Resource ID for the given workspace, if it is an Azure workspace.
+    :param workspace:
+    :return:
+    """
+    if workspace.azure_workspace_info is None:
+        return None
+    return (f'/subscriptions/{workspace.azure_workspace_info.subscription_id}'
+            f'/resourceGroups/{workspace.azure_workspace_info.resource_group}'
+            f'/providers/Microsoft.Databricks/workspaces/{workspace.workspace_name}')

databricks/sdk/clock.py

@@ -0,0 +1,49 @@
+import abc
+import time
+
+
+class Clock(metaclass=abc.ABCMeta):
+
+    @abc.abstractmethod
+    def time(self) -> float:
+        """
+        Return the current time in seconds since the Epoch.
+        Fractions of a second may be present if the system clock provides them.
+
+        :return: The current time in seconds since the Epoch.
+        """
+
+    @abc.abstractmethod
+    def sleep(self, seconds: float) -> None:
+        """
+        Delay execution for a given number of seconds.  The argument may be
+        a floating point number for subsecond precision.
+
+        :param seconds: The duration to sleep in seconds.
+        :return:
+        """
+
+
+class RealClock(Clock):
+    """
+    A real clock that uses the ``time`` module to get the current time and sleep.
+    """
+
+    def time(self) -> float:
+        """
+        Return the current time in seconds since the Epoch.
+        Fractions of a second may be present if the system clock provides them.
+
+        :return: The current time in seconds since the Epoch.
+        """
+        return time.time()
+
+    def sleep(self, seconds: float) -> None:
+        """
+        Delay execution for a given number of seconds.  The argument may be
+        a floating point number for subsecond precision.
+
+        :param seconds: The duration to sleep in seconds.
+        :return:
+        """
+        time.sleep(seconds)

databricks/sdk/config.py

@@ -11,6 +11,7 @@ from typing import Dict, Iterable, Optional
 import requests
 
 from .azure import AzureEnvironment
+from .clock import Clock, RealClock
 from .credentials_provider import CredentialsProvider, DefaultCredentials
 from .environments import (ALL_ENVS, DEFAULT_ENVIRONMENT, Cloud,
                            DatabricksEnvironment)
@@ -84,6 +85,7 @@ class Config:
                  credentials_provider: CredentialsProvider = None,
                  product="unknown",
                  product_version="0.0.0",
+                 clock: Clock = None,
                  **kwargs):
         self._inner = {}
         self._user_agent_other_info = []
@@ -91,6 +93,7 @@ class Config:
         if 'databricks_environment' in kwargs:
             self.databricks_environment = kwargs['databricks_environment']
             del kwargs['databricks_environment']
+        self._clock = clock if clock is not None else RealClock()
         try:
             self._set_inner_config(kwargs)
             self._load_from_env()
@@ -317,6 +320,10 @@ class Config:
         if self.warehouse_id:
             return f'/sql/1.0/warehouses/{self.warehouse_id}'
 
+    @property
+    def clock(self) -> Clock:
+        return self._clock
+
     @classmethod
     def attributes(cls) -> Iterable[ConfigAttribute]:
         """ Returns a list of Databricks SDK configuration metadata """

databricks/sdk/core.py

@@ -123,7 +123,8 @@ class ApiClient:
             headers = {}
         headers['User-Agent'] = self._user_agent_base
         retryable = retried(timeout=timedelta(seconds=self._retry_timeout_seconds),
-                            is_retryable=self._is_retryable)
+                            is_retryable=self._is_retryable,
+                            clock=self._cfg.clock)
         return retryable(self._perform)(method,
                                         path,
                                         query=query,

databricks/sdk/credentials_provider.py

@@ -6,6 +6,7 @@ import json
 import logging
 import os
 import pathlib
+import platform
 import subprocess
 import sys
 from datetime import datetime
@@ -473,11 +474,21 @@ class DatabricksCliTokenSource(CliTokenSource):
             args += ['--account-id', cfg.account_id]
 
         cli_path = cfg.databricks_cli_path
+
+        # If the path is not specified look for "databricks" / "databricks.exe" in PATH.
         if not cli_path:
-            cli_path = 'databricks'
+            try:
+                # Try to find "databricks" in PATH
+                cli_path = self.__class__._find_executable("databricks")
+            except FileNotFoundError as e:
+                # If "databricks" is not found, try to find "databricks.exe" in PATH (Windows)
+                if platform.system() == "Windows":
+                    cli_path = self.__class__._find_executable("databricks.exe")
+                else:
+                    raise e
 
         # If the path is unqualified, look it up in PATH.
-        if cli_path.count("/") == 0:
+        elif cli_path.count("/") == 0:
             cli_path = self.__class__._find_executable(cli_path)
 
         super().__init__(cmd=[cli_path, *args],
@@ -505,7 +516,7 @@ class DatabricksCliTokenSource(CliTokenSource):
         raise err
 
 
-@credentials_provider('databricks-cli', ['host', 'is_aws'])
+@credentials_provider('databricks-cli', ['host'])
 def databricks_cli(cfg: 'Config') -> Optional[HeaderFactory]:
     try:
         token_source = DatabricksCliTokenSource(cfg)

databricks/sdk/environments.py

@@ -19,7 +19,7 @@ class DatabricksEnvironment:
     azure_environment: Optional[AzureEnvironment] = None
 
     def deployment_url(self, name: str) -> str:
-        return f"https://{name}.{self.dns_zone}"
+        return f"https://{name}{self.dns_zone}"
 
     @property
     def azure_service_management_endpoint(self) -> Optional[str]:

databricks/sdk/errors/__init__.py

@@ -1,4 +1,4 @@
 from .base import DatabricksError, ErrorDetail
 from .mapper import error_mapper
-from .mapping import *
+from .platform import *
 from .sdk import *

databricks/sdk/errors/mapper.py

@@ -1,18 +1,18 @@
-from databricks.sdk.errors import mapping
+from databricks.sdk.errors import platform
 from databricks.sdk.errors.base import DatabricksError
 
 
 def error_mapper(status_code: int, raw: dict) -> DatabricksError:
     error_code = raw.get('error_code', None)
-    if error_code in mapping.ERROR_CODE_MAPPING:
+    if error_code in platform.ERROR_CODE_MAPPING:
         # more specific error codes override more generic HTTP status codes
-        return mapping.ERROR_CODE_MAPPING[error_code](**raw)
+        return platform.ERROR_CODE_MAPPING[error_code](**raw)
 
-    if status_code in mapping.STATUS_CODE_MAPPING:
+    if status_code in platform.STATUS_CODE_MAPPING:
         # more generic HTTP status codes matched after more specific error codes,
         # where there's a default exception class per HTTP status code, and we do
         # rely on Databricks platform exception mapper to do the right thing.
-        return mapping.STATUS_CODE_MAPPING[status_code](**raw)
+        return platform.STATUS_CODE_MAPPING[status_code](**raw)
 
     # backwards-compatible error creation for cases like using older versions of
     # the SDK on way never releases of the platform.

databricks/sdk/mixins/workspace.py

@@ -1,4 +1,4 @@
-from typing import BinaryIO, Iterator, Optional
+from typing import BinaryIO, Iterator, Optional, Union
 
 from ..core import DatabricksError
 from ..service.workspace import (ExportFormat, ImportFormat, Language,
@@ -37,7 +37,7 @@ class WorkspaceExt(WorkspaceAPI):
 
     def upload(self,
                path: str,
-               content: BinaryIO,
+               content: Union[bytes, BinaryIO],
                *,
                format: Optional[ImportFormat] = None,
                language: Optional[Language] = None,
@@ -51,7 +51,7 @@ class WorkspaceExt(WorkspaceAPI):
          * `INVALID_PARAMETER_VALUE`: if `format` and `content` values are not compatible.
 
         :param path:     target location of the file on workspace.
-        :param content:  file-like `io.BinaryIO` of the `path` contents.
+        :param content:  the contents as either raw binary data `bytes` or a file-like the file-like `io.BinaryIO` of the `path` contents.
         :param format:   By default, `ImportFormat.SOURCE`. If using `ImportFormat.AUTO` the `path`
                          is imported or exported as either a workspace file or a notebook, depending
                          on an analysis of the `item`’s extension and the header content provided in

databricks/sdk/retries.py

@@ -1,30 +1,34 @@
 import functools
 import logging
-import time
 from datetime import timedelta
 from random import random
 from typing import Callable, Optional, Sequence, Type
 
+from .clock import Clock, RealClock
+
 logger = logging.getLogger(__name__)
 
 
 def retried(*,
             on: Sequence[Type[BaseException]] = None,
             is_retryable: Callable[[BaseException], Optional[str]] = None,
-            timeout=timedelta(minutes=20)):
+            timeout=timedelta(minutes=20),
+            clock: Clock = None):
     has_allowlist = on is not None
     has_callback = is_retryable is not None
     if not (has_allowlist or has_callback) or (has_allowlist and has_callback):
         raise SyntaxError('either on=[Exception] or callback=lambda x: .. is required')
+    if clock is None:
+        clock = RealClock()
 
     def decorator(func):
 
         @functools.wraps(func)
         def wrapper(*args, **kwargs):
-            deadline = time.time() + timeout.total_seconds()
+            deadline = clock.time() + timeout.total_seconds()
             attempt = 1
             last_err = None
-            while time.time() < deadline:
+            while clock.time() < deadline:
                 try:
                     return func(*args, **kwargs)
                 except Exception as err:
@@ -50,7 +54,7 @@ def retried(*,
                         raise err
 
                     logger.debug(f'Retrying: {retry_reason} (sleeping ~{sleep}s)')
-                    time.sleep(sleep + random())
+                    clock.sleep(sleep + random())
                     attempt += 1
             raise TimeoutError(f'Timed out after {timeout}') from last_err