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

databricks/sdk/environments.py

@@ -0,0 +1,72 @@
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional
+
+from .azure import ARM_DATABRICKS_RESOURCE_ID, ENVIRONMENTS, AzureEnvironment
+
+
+class Cloud(Enum):
+    AWS = "AWS"
+    AZURE = "AZURE"
+    GCP = "GCP"
+
+
+@dataclass
+class DatabricksEnvironment:
+    cloud: Cloud
+    dns_zone: str
+    azure_application_id: Optional[str] = None
+    azure_environment: Optional[AzureEnvironment] = None
+
+    def deployment_url(self, name: str) -> str:
+        return f"https://{name}{self.dns_zone}"
+
+    @property
+    def azure_service_management_endpoint(self) -> Optional[str]:
+        if self.azure_environment is None:
+            return None
+        return self.azure_environment.service_management_endpoint
+
+    @property
+    def azure_resource_manager_endpoint(self) -> Optional[str]:
+        if self.azure_environment is None:
+            return None
+        return self.azure_environment.resource_manager_endpoint
+
+    @property
+    def azure_active_directory_endpoint(self) -> Optional[str]:
+        if self.azure_environment is None:
+            return None
+        return self.azure_environment.active_directory_endpoint
+
+
+DEFAULT_ENVIRONMENT = DatabricksEnvironment(Cloud.AWS, ".cloud.databricks.com")
+
+ALL_ENVS = [
+    DatabricksEnvironment(Cloud.AWS, ".dev.databricks.com"),
+    DatabricksEnvironment(Cloud.AWS, ".staging.cloud.databricks.com"),
+    DatabricksEnvironment(Cloud.AWS, ".cloud.databricks.us"), DEFAULT_ENVIRONMENT,
+    DatabricksEnvironment(Cloud.AZURE,
+                          ".dev.azuredatabricks.net",
+                          azure_application_id="62a912ac-b58e-4c1d-89ea-b2dbfc7358fc",
+                          azure_environment=ENVIRONMENTS["PUBLIC"]),
+    DatabricksEnvironment(Cloud.AZURE,
+                          ".staging.azuredatabricks.net",
+                          azure_application_id="4a67d088-db5c-48f1-9ff2-0aace800ae68",
+                          azure_environment=ENVIRONMENTS["PUBLIC"]),
+    DatabricksEnvironment(Cloud.AZURE,
+                          ".azuredatabricks.net",
+                          azure_application_id=ARM_DATABRICKS_RESOURCE_ID,
+                          azure_environment=ENVIRONMENTS["PUBLIC"]),
+    DatabricksEnvironment(Cloud.AZURE,
+                          ".databricks.azure.us",
+                          azure_application_id=ARM_DATABRICKS_RESOURCE_ID,
+                          azure_environment=ENVIRONMENTS["USGOVERNMENT"]),
+    DatabricksEnvironment(Cloud.AZURE,
+                          ".databricks.azure.cn",
+                          azure_application_id=ARM_DATABRICKS_RESOURCE_ID,
+                          azure_environment=ENVIRONMENTS["CHINA"]),
+    DatabricksEnvironment(Cloud.GCP, ".dev.gcp.databricks.com"),
+    DatabricksEnvironment(Cloud.GCP, ".staging.gcp.databricks.com"),
+    DatabricksEnvironment(Cloud.GCP, ".gcp.databricks.com")
+]

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/oauth.py

@@ -357,7 +357,8 @@ class OAuthClient:
                  scopes: List[str] = None,
                  client_secret: str = None):
         # TODO: is it a circular dependency?..
-        from .core import Config, credentials_provider
+        from .core import Config
+        from .credentials_provider import credentials_provider
 
         @credentials_provider('noop', [])
         def noop_credentials(_: any):

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
 

databricks/sdk/service/_internal.py

@@ -10,7 +10,7 @@ def _from_dict(d: Dict[str, any], field: str, cls: Type) -> any:
 
 def _repeated_dict(d: Dict[str, any], field: str, cls: Type) -> any:
     if field not in d or not d[field]:
-        return None
+        return []
     from_dict = getattr(cls, 'from_dict')
     return [from_dict(v) for v in d[field]]