databricks-sdk 0.44.1__py3-none-any.whl → 0.46.0__py3-none-any.whl

databricks/sdk/retries.py

@@ -9,16 +9,18 @@ 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),
-            clock: Clock = None,
-            before_retry: Callable = None):
+def retried(
+    *,
+    on: Optional[Sequence[Type[BaseException]]] = None,
+    is_retryable: Optional[Callable[[BaseException], Optional[str]]] = None,
+    timeout=timedelta(minutes=20),
+    clock: Optional[Clock] = None,
+    before_retry: Optional[Callable] = 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')
+        raise SyntaxError("either on=[Exception] or callback=lambda x: .. is required")
     if clock is None:
         clock = RealClock()
 
@@ -37,30 +39,30 @@ def retried(*,
                     retry_reason = None
                     # sleep 10s max per attempt, unless it's HTTP 429 or 503
                     sleep = min(10, attempt)
-                    retry_after_secs = getattr(err, 'retry_after_secs', None)
+                    retry_after_secs = getattr(err, "retry_after_secs", None)
                     if retry_after_secs is not None:
                         # cannot depend on DatabricksError directly because of circular dependency
                         sleep = retry_after_secs
-                        retry_reason = 'throttled by platform'
+                        retry_reason = "throttled by platform"
                     elif is_retryable is not None:
                         retry_reason = is_retryable(err)
                     elif on is not None:
                         for err_type in on:
                             if not isinstance(err, err_type):
                                 continue
-                            retry_reason = f'{type(err).__name__} is allowed to retry'
+                            retry_reason = f"{type(err).__name__} is allowed to retry"
 
                     if retry_reason is None:
                         # raise if exception is not retryable
                         raise err
 
-                    logger.debug(f'Retrying: {retry_reason} (sleeping ~{sleep}s)')
+                    logger.debug(f"Retrying: {retry_reason} (sleeping ~{sleep}s)")
                     if before_retry:
                         before_retry()
 
                     clock.sleep(sleep + random())
                     attempt += 1
-            raise TimeoutError(f'Timed out after {timeout}') from last_err
+            raise TimeoutError(f"Timed out after {timeout}") from last_err
 
         return wrapper
 

databricks/sdk/runtime/__init__.py

@@ -3,14 +3,23 @@ from __future__ import annotations
 import logging
 from typing import Dict, Optional, Union, cast
 
-logger = logging.getLogger('databricks.sdk')
+logger = logging.getLogger("databricks.sdk")
 is_local_implementation = True
 
 # All objects that are injected into the Notebook's user namespace should also be made
 # available to be imported from databricks.sdk.runtime.globals. This import can be used
 # in Python modules so users can access these objects from Files more easily.
 dbruntime_objects = [
-    "display", "displayHTML", "dbutils", "table", "sql", "udf", "getArgument", "sc", "sqlContext", "spark"
+    "display",
+    "displayHTML",
+    "dbutils",
+    "table",
+    "sql",
+    "udf",
+    "getArgument",
+    "sc",
+    "sqlContext",
+    "spark",
 ]
 
 # DO NOT MOVE THE TRY-CATCH BLOCK BELOW AND DO NOT ADD THINGS BEFORE IT! WILL MAKE TEST FAIL.
@@ -18,7 +27,8 @@ try:
     # We don't want to expose additional entity to user namespace, so
     # a workaround here for exposing required information in notebook environment
     from dbruntime.sdk_credential_provider import init_runtime_native_auth
-    logger.debug('runtime SDK credential provider available')
+
+    logger.debug("runtime SDK credential provider available")
     dbruntime_objects.append("init_runtime_native_auth")
 except ImportError:
     init_runtime_native_auth = None
@@ -29,18 +39,19 @@ globals()["init_runtime_native_auth"] = init_runtime_native_auth
 def init_runtime_repl_auth():
     try:
         from dbruntime.databricks_repl_context import get_context
+
         ctx = get_context()
         if ctx is None:
-            logger.debug('Empty REPL context returned, skipping runtime auth')
+            logger.debug("Empty REPL context returned, skipping runtime auth")
             return None, None
         if ctx.workspaceUrl is None:
-            logger.debug('Workspace URL is not available, skipping runtime auth')
+            logger.debug("Workspace URL is not available, skipping runtime auth")
             return None, None
-        host = f'https://{ctx.workspaceUrl}'
+        host = f"https://{ctx.workspaceUrl}"
 
         def inner() -> Dict[str, str]:
             ctx = get_context()
-            return {'Authorization': f'Bearer {ctx.apiToken}'}
+            return {"Authorization": f"Bearer {ctx.apiToken}"}
 
         return host, inner
     except ImportError:
@@ -50,11 +61,12 @@ def init_runtime_repl_auth():
 def init_runtime_legacy_auth():
     try:
         import IPython
+
         ip_shell = IPython.get_ipython()
         if ip_shell is None:
             return None, None
         global_ns = ip_shell.ns_table["user_global"]
-        if 'dbutils' not in global_ns:
+        if "dbutils" not in global_ns:
             return None, None
         dbutils = global_ns["dbutils"].notebook.entry_point.getDbutils()
         if dbutils is None:
@@ -62,11 +74,11 @@ def init_runtime_legacy_auth():
         ctx = dbutils.notebook().getContext()
         if ctx is None:
             return None, None
-        host = getattr(ctx, 'apiUrl')().get()
+        host = getattr(ctx, "apiUrl")().get()
 
         def inner() -> Dict[str, str]:
             ctx = dbutils.notebook().getContext()
-            return {'Authorization': f'Bearer {getattr(ctx, "apiToken")().get()}'}
+            return {"Authorization": f'Bearer {getattr(ctx, "apiToken")().get()}'}
 
         return host, inner
     except ImportError:
@@ -97,7 +109,8 @@ except ImportError:
     try:
         # We expect this to fail and only do this for providing types
         from pyspark.sql.context import SQLContext
-        sqlContext: SQLContext = None # type: ignore
+
+        sqlContext: SQLContext = None  # type: ignore
         table = sqlContext.table
     except Exception as e:
         logging.debug(f"Failed to initialize globals 'sqlContext' and 'table', continuing. Cause: {e}")
@@ -109,8 +122,9 @@ except ImportError:
 
     try:
         from databricks.connect import DatabricksSession  # type: ignore
+
         spark = DatabricksSession.builder.getOrCreate()
-        sql = spark.sql # type: ignore
+        sql = spark.sql  # type: ignore
     except Exception as e:
         # We are ignoring all failures here because user might want to initialize
         # spark session themselves and we don't want to interfere with that
@@ -118,11 +132,11 @@ except ImportError:
 
     try:
         # We expect this to fail locally since dbconnect does not support sparkcontext. This is just for typing
-        sc = spark.sparkContext
+        sc = spark.sparkContext  # type: ignore
     except Exception as e:
         logging.debug(f"Failed to initialize global 'sc', continuing. Cause: {e}")
 
-    def display(input=None, *args, **kwargs) -> None: # type: ignore
+    def display(input=None, *args, **kwargs) -> None:  # type: ignore
         """
         Display plots or data.
         Display plot:
@@ -144,9 +158,10 @@ except ImportError:
         """
         # Import inside the function so that imports are only triggered on usage.
         from IPython import display as IPDisplay
-        return IPDisplay.display(input, *args, **kwargs) # type: ignore
 
-    def displayHTML(html) -> None: # type: ignore
+        return IPDisplay.display(input, *args, **kwargs)  # type: ignore
+
+    def displayHTML(html) -> None:  # type: ignore
         """
         Display HTML data.
         Parameters
@@ -160,13 +175,15 @@ except ImportError:
         """
         # Import inside the function so that imports are only triggered on usage.
         from IPython import display as IPDisplay
-        return IPDisplay.display_html(html, raw=True) # type: ignore
+
+        return IPDisplay.display_html(html, raw=True)  # type: ignore
 
     # We want to propagate the error in initialising dbutils because this is a core
     # functionality of the sdk
     from databricks.sdk.dbutils import RemoteDbUtils
 
     from . import dbutils_stub
+
     dbutils_type = Union[dbutils_stub.dbutils, RemoteDbUtils]
 
     dbutils = RemoteDbUtils()

databricks/sdk/runtime/dbutils_stub.py

@@ -2,21 +2,21 @@ import typing
 from collections import namedtuple
 
 
-class FileInfo(namedtuple('FileInfo', ['path', 'name', 'size', "modificationTime"])):
+class FileInfo(namedtuple("FileInfo", ["path", "name", "size", "modificationTime"])):
     pass
 
 
-class MountInfo(namedtuple('MountInfo', ['mountPoint', 'source', 'encryptionType'])):
+class MountInfo(namedtuple("MountInfo", ["mountPoint", "source", "encryptionType"])):
     pass
 
 
-class SecretScope(namedtuple('SecretScope', ['name'])):
+class SecretScope(namedtuple("SecretScope", ["name"])):
 
     def getName(self):
         return self.name
 
 
-class SecretMetadata(namedtuple('SecretMetadata', ['key'])):
+class SecretMetadata(namedtuple("SecretMetadata", ["key"])):
     pass
 
 
@@ -49,8 +49,7 @@ class dbutils:
             ...
 
         @staticmethod
-        def getCurrentCredentials() -> typing.Mapping[str, str]:
-            ...
+        def getCurrentCredentials() -> typing.Mapping[str, str]: ...
 
     class data:
         """
@@ -129,40 +128,38 @@ class dbutils:
             ...
 
         @staticmethod
-        def cacheFiles(*files):
-            ...
+        def cacheFiles(*files): ...
 
         @staticmethod
-        def cacheTable(name: str):
-            ...
+        def cacheTable(name: str): ...
 
         @staticmethod
-        def uncacheFiles(*files):
-            ...
+        def uncacheFiles(*files): ...
 
         @staticmethod
-        def uncacheTable(name: str):
-            ...
+        def uncacheTable(name: str): ...
 
         @staticmethod
-        def mount(source: str,
-                  mount_point: str,
-                  encryption_type: str = "",
-                  owner: typing.Optional[str] = None,
-                  extra_configs: typing.Mapping[str, str] = {},
-                  ) -> bool:
+        def mount(
+            source: str,
+            mount_point: str,
+            encryption_type: str = "",
+            owner: typing.Optional[str] = None,
+            extra_configs: typing.Mapping[str, str] = {},
+        ) -> bool:
             """
             Mounts the given source directory into DBFS at the given mount point
             """
             ...
 
         @staticmethod
-        def updateMount(source: str,
-                        mount_point: str,
-                        encryption_type: str = "",
-                        owner: typing.Optional[str] = None,
-                        extra_configs: typing.Mapping[str, str] = {},
-                        ) -> bool:
+        def updateMount(
+            source: str,
+            mount_point: str,
+            encryption_type: str = "",
+            owner: typing.Optional[str] = None,
+            extra_configs: typing.Mapping[str, str] = {},
+        ) -> bool:
             """
             Similar to mount(), but updates an existing mount point (if present) instead of creating a new one
             """
@@ -200,7 +197,12 @@ class dbutils:
             """
 
             @staticmethod
-            def get(taskKey: str, key: str, default: any = None, debugValue: any = None) -> None:
+            def get(
+                taskKey: str,
+                key: str,
+                default: any = None,
+                debugValue: any = None,
+            ) -> None:
                 """
                 Returns the latest task value that belongs to the current job run
                 """
@@ -238,7 +240,11 @@ class dbutils:
             ...
 
         @staticmethod
-        def run(path: str, timeout_seconds: int, arguments: typing.Mapping[str, str]) -> str:
+        def run(
+            path: str,
+            timeout_seconds: int,
+            arguments: typing.Mapping[str, str],
+        ) -> str:
             """
             This method runs a notebook and returns its exit value
             """
@@ -307,7 +313,12 @@ class dbutils:
             ...
 
         @staticmethod
-        def dropdown(name: str, defaultValue: str, choices: typing.List[str], label: str = None):
+        def dropdown(
+            name: str,
+            defaultValue: str,
+            choices: typing.List[str],
+            label: str = None,
+        ):
             """Creates a dropdown input widget with given specification.
             :param name: Name of argument associated with the new input widget
             :param defaultValue: Default value of the input widget (must be one of choices)
@@ -317,11 +328,12 @@ class dbutils:
             ...
 
         @staticmethod
-        def combobox(name: str,
-                     defaultValue: str,
-                     choices: typing.List[str],
-                     label: typing.Optional[str] = None,
-                     ):
+        def combobox(
+            name: str,
+            defaultValue: str,
+            choices: typing.List[str],
+            label: typing.Optional[str] = None,
+        ):
             """Creates a combobox input widget with given specification.
             :param name: Name of argument associated with the new input widget
             :param defaultValue: Default value of the input widget
@@ -331,11 +343,12 @@ class dbutils:
             ...
 
         @staticmethod
-        def multiselect(name: str,
-                        defaultValue: str,
-                        choices: typing.List[str],
-                        label: typing.Optional[str] = None,
-                        ):
+        def multiselect(
+            name: str,
+            defaultValue: str,
+            choices: typing.List[str],
+            label: typing.Optional[str] = None,
+        ):
             """Creates a multiselect input widget with given specification.
             :param name: Name of argument associated with the new input widget
             :param defaultValue: Default value of the input widget (must be one of choices)

databricks/sdk/service/_internal.py

@@ -6,18 +6,21 @@ from typing import Callable, Dict, Generic, Optional, Type, TypeVar
 def _from_dict(d: Dict[str, any], field: str, cls: Type) -> any:
     if field not in d or d[field] is None:
         return None
-    return getattr(cls, 'from_dict')(d[field])
+    return getattr(cls, "from_dict")(d[field])
 
 
 def _repeated_dict(d: Dict[str, any], field: str, cls: Type) -> any:
     if field not in d or not d[field]:
         return []
-    from_dict = getattr(cls, 'from_dict')
+    from_dict = getattr(cls, "from_dict")
     return [from_dict(v) for v in d[field]]
 
 
 def _get_enum_value(cls: Type, value: str) -> Optional[Type]:
-    return next((v for v in getattr(cls, '__members__').values() if v.value == value), None)
+    return next(
+        (v for v in getattr(cls, "__members__").values() if v.value == value),
+        None,
+    )
 
 
 def _enum(d: Dict[str, any], field: str, cls: Type) -> any:
@@ -43,7 +46,7 @@ def _escape_multi_segment_path_parameter(param: str) -> str:
     return urllib.parse.quote(param)
 
 
-ReturnType = TypeVar('ReturnType')
+ReturnType = TypeVar("ReturnType")
 
 
 class Wait(Generic[ReturnType]):
@@ -60,8 +63,10 @@ class Wait(Generic[ReturnType]):
     def bind(self) -> dict:
         return self._bind
 
-    def result(self,
-               timeout: datetime.timedelta = datetime.timedelta(minutes=20),
-               callback: Callable[[ReturnType], None] = None) -> ReturnType:
+    def result(
+        self,
+        timeout: datetime.timedelta = datetime.timedelta(minutes=20),
+        callback: Callable[[ReturnType], None] = None,
+    ) -> ReturnType:
         kwargs = self._bind.copy()
         return self._waiter(callback=callback, timeout=timeout, **kwargs)