databricks-sdk 0.29.0__py3-none-any.whl → 0.31.0__py3-none-any.whl

databricks/sdk/credentials_provider.py

@@ -233,8 +233,7 @@ def _ensure_host_present(cfg: 'Config', token_source_for: Callable[[str], TokenS
     cfg.host = f"https://{resp.json()['properties']['workspaceUrl']}"
 
 
-@oauth_credentials_strategy('azure-client-secret',
-                            ['is_azure', 'azure_client_id', 'azure_client_secret', 'azure_tenant_id'])
+@oauth_credentials_strategy('azure-client-secret', ['is_azure', 'azure_client_id', 'azure_client_secret'])
 def azure_service_principal(cfg: 'Config') -> CredentialsProvider:
     """ Adds refreshed Azure Active Directory (AAD) Service Principal OAuth tokens
     to every request, while automatically resolving different Azure environment endpoints. """
@@ -248,6 +247,7 @@ def azure_service_principal(cfg: 'Config') -> CredentialsProvider:
                                  use_params=True)
 
     _ensure_host_present(cfg, token_source_for)
+    cfg.load_azure_tenant_id()
     logger.info("Configured AAD token for Service Principal (%s)", cfg.azure_client_id)
     inner = token_source_for(cfg.effective_azure_login_app_id)
     cloud = token_source_for(cfg.arm_environment.service_management_endpoint)
@@ -432,11 +432,13 @@ class CliTokenSource(Refreshable):
 class AzureCliTokenSource(CliTokenSource):
     """ Obtain the token granted by `az login` CLI command """
 
-    def __init__(self, resource: str, subscription: str = ""):
+    def __init__(self, resource: str, subscription: Optional[str] = None, tenant: Optional[str] = None):
         cmd = ["az", "account", "get-access-token", "--resource", resource, "--output", "json"]
-        if subscription != "":
+        if subscription is not None:
             cmd.append("--subscription")
             cmd.append(subscription)
+        if tenant:
+            cmd.extend(["--tenant", tenant])
         super().__init__(cmd=cmd,
                          token_type_field='tokenType',
                          access_token_field='accessToken',
@@ -464,8 +466,10 @@ class AzureCliTokenSource(CliTokenSource):
     @staticmethod
     def for_resource(cfg: 'Config', resource: str) -> 'AzureCliTokenSource':
         subscription = AzureCliTokenSource.get_subscription(cfg)
-        if subscription != "":
-            token_source = AzureCliTokenSource(resource, subscription)
+        if subscription is not None:
+            token_source = AzureCliTokenSource(resource,
+                                               subscription=subscription,
+                                               tenant=cfg.azure_tenant_id)
             try:
                 # This will fail if the user has access to the workspace, but not to the subscription
                 # itself.
@@ -475,25 +479,26 @@ class AzureCliTokenSource(CliTokenSource):
             except OSError:
                 logger.warning("Failed to get token for subscription. Using resource only token.")
 
-        token_source = AzureCliTokenSource(resource)
+        token_source = AzureCliTokenSource(resource, subscription=None, tenant=cfg.azure_tenant_id)
         token_source.token()
         return token_source
 
     @staticmethod
-    def get_subscription(cfg: 'Config') -> str:
+    def get_subscription(cfg: 'Config') -> Optional[str]:
         resource = cfg.azure_workspace_resource_id
         if resource is None or resource == "":
-            return ""
+            return None
         components = resource.split('/')
         if len(components) < 3:
             logger.warning("Invalid azure workspace resource ID")
-            return ""
+            return None
         return components[2]
 
 
 @credentials_strategy('azure-cli', ['is_azure'])
 def azure_cli(cfg: 'Config') -> Optional[CredentialsProvider]:
     """ Adds refreshed OAuth token granted by `az login` command to every request. """
+    cfg.load_azure_tenant_id()
     token_source = None
     mgmt_token_source = None
     try:
@@ -517,11 +522,6 @@ def azure_cli(cfg: 'Config') -> Optional[CredentialsProvider]:
 
     _ensure_host_present(cfg, lambda resource: AzureCliTokenSource.for_resource(cfg, resource))
     logger.info("Using Azure CLI authentication with AAD tokens")
-    if not cfg.is_account_client and AzureCliTokenSource.get_subscription(cfg) == "":
-        logger.warning(
-            "azure_workspace_resource_id field not provided. "
-            "It is recommended to specify this field in the Databricks configuration to avoid authentication errors."
-        )
 
     def inner() -> Dict[str, str]:
         token = token_source.token()

databricks/sdk/data_plane.py

@@ -0,0 +1,65 @@
+import threading
+from dataclasses import dataclass
+from typing import Callable, List
+
+from databricks.sdk.oauth import Token
+from databricks.sdk.service.oauth2 import DataPlaneInfo
+
+
+@dataclass
+class DataPlaneDetails:
+    """
+    Contains details required to query a DataPlane endpoint.
+    """
+    endpoint_url: str
+    """URL used to query the endpoint through the DataPlane."""
+    token: Token
+    """Token to query the DataPlane endpoint."""
+
+
+class DataPlaneService:
+    """Helper class to fetch and manage DataPlane details."""
+
+    def __init__(self):
+        self._data_plane_info = {}
+        self._tokens = {}
+        self._lock = threading.Lock()
+
+    def get_data_plane_details(self, method: str, params: List[str], info_getter: Callable[[], DataPlaneInfo],
+                               refresh: Callable[[str], Token]):
+        """Get and cache information required to query a Data Plane endpoint using the provided methods.
+
+        Returns a cached DataPlaneDetails if the details have already been fetched previously and are still valid.
+        If not, it uses the provided functions to fetch the details.
+
+        :param method: method name. Used to construct a unique key for the cache.
+        :param params: path params used in the "get" operation which uniquely determine the object. Used to construct a unique key for the cache.
+        :param info_getter: function which returns the DataPlaneInfo. It will only be called if the information is not already present in the cache.
+        :param refresh: function to refresh the token. It will only be called if the token is missing or expired.
+        """
+        all_elements = params.copy()
+        all_elements.insert(0, method)
+        map_key = "/".join(all_elements)
+        info = self._data_plane_info.get(map_key)
+        if not info:
+            self._lock.acquire()
+            try:
+                info = self._data_plane_info.get(map_key)
+                if not info:
+                    info = info_getter()
+                    self._data_plane_info[map_key] = info
+            finally:
+                self._lock.release()
+
+        token = self._tokens.get(map_key)
+        if not token or not token.valid:
+            self._lock.acquire()
+            token = self._tokens.get(map_key)
+            try:
+                if not token or not token.valid:
+                    token = refresh(info.authorization_details)
+                    self._tokens[map_key] = token
+            finally:
+                self._lock.release()
+
+        return DataPlaneDetails(endpoint_url=info.endpoint_url, token=token)

databricks/sdk/errors/overrides.py

@@ -22,4 +22,12 @@ _ALL_OVERRIDES = [
                    message_matcher=re.compile(r'Job .* does not exist'),
                    custom_error=ResourceDoesNotExist,
                    ),
+    _ErrorOverride(debug_name="Job Runs InvalidParameterValue=>ResourceDoesNotExist",
+                   path_regex=re.compile(r'^/api/2\.\d/jobs/runs/get'),
+                   verb="GET",
+                   status_code_matcher=re.compile(r'^400$'),
+                   error_code_matcher=re.compile(r'INVALID_PARAMETER_VALUE'),
+                   message_matcher=re.compile(r'(Run .* does not exist|Run: .* in job: .* doesn\'t exist)'),
+                   custom_error=ResourceDoesNotExist,
+                   ),
 ]

databricks/sdk/errors/platform.py

@@ -47,6 +47,10 @@ class DeadlineExceeded(DatabricksError):
     """the deadline expired before the operation could complete"""
 
 
+class InvalidState(BadRequest):
+    """unexpected state"""
+
+
 class InvalidParameterValue(BadRequest):
     """supplied value for a parameter was invalid"""
 
@@ -99,6 +103,7 @@ STATUS_CODE_MAPPING = {
 }
 
 ERROR_CODE_MAPPING = {
+    'INVALID_STATE': InvalidState,
     'INVALID_PARAMETER_VALUE': InvalidParameterValue,
     'RESOURCE_DOES_NOT_EXIST': ResourceDoesNotExist,
     'ABORTED': Aborted,

databricks/sdk/mixins/files.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import base64
 import os
 import pathlib
+import platform
 import shutil
 import sys
 from abc import ABC, abstractmethod
@@ -266,8 +267,9 @@ class _VolumesIO(BinaryIO):
 
 class _Path(ABC):
 
-    def __init__(self, path: str):
-        self._path = pathlib.Path(str(path).replace('dbfs:', '').replace('file:', ''))
+    @abstractmethod
+    def __init__(self):
+        ...
 
     @property
     def is_local(self) -> bool:
@@ -327,6 +329,12 @@ class _Path(ABC):
 
 class _LocalPath(_Path):
 
+    def __init__(self, path: str):
+        if platform.system() == "Windows":
+            self._path = pathlib.Path(str(path).replace('file:///', '').replace('file:', ''))
+        else:
+            self._path = pathlib.Path(str(path).replace('file:', ''))
+
     def _is_local(self) -> bool:
         return True
 
@@ -393,7 +401,7 @@ class _LocalPath(_Path):
 class _VolumesPath(_Path):
 
     def __init__(self, api: files.FilesAPI, src: Union[str, pathlib.Path]):
-        super().__init__(src)
+        self._path = pathlib.PurePosixPath(str(src).replace('dbfs:', '').replace('file:', ''))
         self._api = api
 
     def _is_local(self) -> bool:
@@ -462,7 +470,7 @@ class _VolumesPath(_Path):
 class _DbfsPath(_Path):
 
     def __init__(self, api: files.DbfsAPI, src: str):
-        super().__init__(src)
+        self._path = pathlib.PurePosixPath(str(src).replace('dbfs:', '').replace('file:', ''))
         self._api = api
 
     def _is_local(self) -> bool: