databricks-sdk 0.33.0__tar.gz → 0.35.0__tar.gz

{databricks_sdk-0.33.0 → databricks_sdk-0.35.0}/NOTICE

@@ -12,8 +12,22 @@ googleapis/google-auth-library-python - https://github.com/googleapis/google-aut
 Copyright google-auth-library-python authors
 License - https://github.com/googleapis/google-auth-library-python/blob/main/LICENSE
 
+openai/openai-python - https://github.com/openai/openai-python
+Copyright 2024 OpenAI
+License - https://github.com/openai/openai-python/blob/main/LICENSE
+
 This software contains code from the following open source projects, licensed under the BSD (3-clause) license.
 
 x/oauth2 - https://cs.opensource.google/go/x/oauth2/+/master:oauth2.go
 Copyright 2014 The Go Authors. All rights reserved.
 License - https://cs.opensource.google/go/x/oauth2/+/master:LICENSE
+
+encode/httpx - https://github.com/encode/httpx
+Copyright 2019, Encode OSS Ltd
+License - https://github.com/encode/httpx/blob/master/LICENSE.md
+
+This software contains code from the following open source projects, licensed under the MIT license:
+
+langchain-ai/langchain - https://github.com/langchain-ai/langchain/blob/master/libs/partners/openai
+Copyright 2023 LangChain, Inc.
+License - https://github.com/langchain-ai/langchain/blob/master/libs/partners/openai/LICENSE

{databricks_sdk-0.33.0/databricks_sdk.egg-info → databricks_sdk-0.35.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: databricks-sdk
-Version: 0.33.0
+Version: 0.35.0
 Summary: Databricks SDK for Python (Beta)
 Home-page: https://databricks-sdk-py.readthedocs.io
 Author: Serge Smertin
@@ -40,9 +40,16 @@ Requires-Dist: requests-mock; extra == "dev"
 Requires-Dist: pyfakefs; extra == "dev"
 Requires-Dist: databricks-connect; extra == "dev"
 Requires-Dist: pytest-rerunfailures; extra == "dev"
+Requires-Dist: openai; extra == "dev"
+Requires-Dist: langchain-openai; python_version > "3.7" and extra == "dev"
+Requires-Dist: httpx; extra == "dev"
 Provides-Extra: notebook
 Requires-Dist: ipython<9,>=8; extra == "notebook"
 Requires-Dist: ipywidgets<9,>=8; extra == "notebook"
+Provides-Extra: openai
+Requires-Dist: openai; extra == "openai"
+Requires-Dist: langchain-openai; python_version > "3.7" and extra == "openai"
+Requires-Dist: httpx; extra == "openai"
 
 # Databricks SDK for Python (Beta)
 

{databricks_sdk-0.33.0 → databricks_sdk-0.35.0}/databricks/sdk/__init__.py

@@ -6,6 +6,7 @@ from databricks.sdk import azure
 from databricks.sdk.credentials_provider import CredentialsStrategy
 from databricks.sdk.mixins.compute import ClustersExt
 from databricks.sdk.mixins.files import DbfsExt
+from databricks.sdk.mixins.open_ai_client import ServingEndpointsExt
 from databricks.sdk.mixins.workspace import WorkspaceExt
 from databricks.sdk.service.apps import AppsAPI
 from databricks.sdk.service.billing import (BillableUsageAPI, BudgetsAPI,
@@ -71,6 +72,7 @@ from databricks.sdk.service.settings import (AccountIpAccessListsAPI,
                                              CspEnablementAccountAPI,
                                              DefaultNamespaceAPI,
                                              DisableLegacyAccessAPI,
+                                             DisableLegacyDbfsAPI,
                                              DisableLegacyFeaturesAPI,
                                              EnhancedSecurityMonitoringAPI,
                                              EsmEnablementAccountAPI,
@@ -174,7 +176,7 @@ class WorkspaceClient:
         self._config = config.copy()
         self._dbutils = _make_dbutils(self._config)
         self._api_client = client.ApiClient(self._config)
-        serving_endpoints = ServingEndpointsAPI(self._api_client)
+        serving_endpoints = ServingEndpointsExt(self._api_client)
         self._account_access_control_proxy = AccountAccessControlProxyAPI(self._api_client)
         self._alerts = AlertsAPI(self._api_client)
         self._alerts_legacy = AlertsLegacyAPI(self._api_client)
@@ -636,7 +638,7 @@ class WorkspaceClient:
         return self._service_principals
 
     @property
-    def serving_endpoints(self) -> ServingEndpointsAPI:
+    def serving_endpoints(self) -> ServingEndpointsExt:
         """The Serving Endpoints API allows you to create, update, and delete model serving endpoints."""
         return self._serving_endpoints
 

databricks_sdk-0.33.0/databricks/sdk/core.py → databricks_sdk-0.35.0/databricks/sdk/_base_client.py

@@ -1,90 +1,82 @@
-import re
+import logging
 from datetime import timedelta
 from types import TracebackType
-from typing import Any, BinaryIO, Iterator, Type
-from urllib.parse import urlencode
+from typing import (Any, BinaryIO, Callable, Dict, Iterable, Iterator, List,
+                    Optional, Type, Union)
 
-from requests.adapters import HTTPAdapter
+import requests
+import requests.adapters
 
+from . import useragent
 from .casing import Casing
-from .config import *
-# To preserve backwards compatibility (as these definitions were previously in this module)
-from .credentials_provider import *
+from .clock import Clock, RealClock
 from .errors import DatabricksError, _ErrorCustomizer, _Parser
 from .logger import RoundTrip
-from .oauth import retrieve_token
 from .retries import retried
 
-__all__ = ['Config', 'DatabricksError']
-
 logger = logging.getLogger('databricks.sdk')
 
-URL_ENCODED_CONTENT_TYPE = "application/x-www-form-urlencoded"
-JWT_BEARER_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer"
-OIDC_TOKEN_PATH = "/oidc/v1/token"
-
-
-class ApiClient:
-    _cfg: Config
-    _RETRY_AFTER_DEFAULT: int = 1
-
-    def __init__(self, cfg: Config = None):
 
-        if cfg is None:
-            cfg = Config()
-
-        self._cfg = cfg
-        # See https://github.com/databricks/databricks-sdk-go/blob/main/client/client.go#L34-L35
-        self._debug_truncate_bytes = cfg.debug_truncate_bytes if cfg.debug_truncate_bytes else 96
-        self._retry_timeout_seconds = cfg.retry_timeout_seconds if cfg.retry_timeout_seconds else 300
-        self._user_agent_base = cfg.user_agent
+class _BaseClient:
+
+    def __init__(self,
+                 debug_truncate_bytes: int = None,
+                 retry_timeout_seconds: int = None,
+                 user_agent_base: str = None,
+                 header_factory: Callable[[], dict] = None,
+                 max_connection_pools: int = None,
+                 max_connections_per_pool: int = None,
+                 pool_block: bool = True,
+                 http_timeout_seconds: float = None,
+                 extra_error_customizers: List[_ErrorCustomizer] = None,
+                 debug_headers: bool = False,
+                 clock: Clock = None):
+        """
+        :param debug_truncate_bytes:
+        :param retry_timeout_seconds:
+        :param user_agent_base:
+        :param header_factory: A function that returns a dictionary of headers to include in the request.
+        :param max_connection_pools: Number of urllib3 connection pools to cache before discarding the least
+            recently used pool. Python requests default value is 10.
+        :param max_connections_per_pool: The maximum number of connections to save in the pool. Improves performance
+            in multithreaded situations. For now, we're setting it to the same value as connection_pool_size.
+        :param pool_block: If pool_block is False, then more connections will are created, but not saved after the
+            first use. Blocks when no free connections are available. urllib3 ensures that no more than
+            pool_maxsize connections are used at a time. Prevents platform from flooding. By default, requests library
+            doesn't block.
+        :param http_timeout_seconds:
+        :param extra_error_customizers:
+        :param debug_headers: Whether to include debug headers in the request log.
+        :param clock: Clock object to use for time-related operations.
+        """
+
+        self._debug_truncate_bytes = debug_truncate_bytes or 96
+        self._debug_headers = debug_headers
+        self._retry_timeout_seconds = retry_timeout_seconds or 300
+        self._user_agent_base = user_agent_base or useragent.to_string()
+        self._header_factory = header_factory
+        self._clock = clock or RealClock()
         self._session = requests.Session()
         self._session.auth = self._authenticate
 
-        # Number of urllib3 connection pools to cache before discarding the least
-        # recently used pool. Python requests default value is 10.
-        pool_connections = cfg.max_connection_pools
-        if pool_connections is None:
-            pool_connections = 20
-
-        # The maximum number of connections to save in the pool. Improves performance
-        # in multithreaded situations. For now, we're setting it to the same value
-        # as connection_pool_size.
-        pool_maxsize = cfg.max_connections_per_pool
-        if cfg.max_connections_per_pool is None:
-            pool_maxsize = pool_connections
-
-        # If pool_block is False, then more connections will are created,
-        # but not saved after the first use. Blocks when no free connections are available.
-        # urllib3 ensures that no more than pool_maxsize connections are used at a time.
-        # Prevents platform from flooding. By default, requests library doesn't block.
-        pool_block = True
-
         # We don't use `max_retries` from HTTPAdapter to align with a more production-ready
         # retry strategy established in the Databricks SDK for Go. See _is_retryable and
         # @retried for more details.
-        http_adapter = HTTPAdapter(pool_connections=pool_connections,
-                                   pool_maxsize=pool_maxsize,
-                                   pool_block=pool_block)
+        http_adapter = requests.adapters.HTTPAdapter(pool_connections=max_connections_per_pool or 20,
+                                                     pool_maxsize=max_connection_pools or 20,
+                                                     pool_block=pool_block)
         self._session.mount("https://", http_adapter)
 
         # Default to 60 seconds
-        self._http_timeout_seconds = cfg.http_timeout_seconds if cfg.http_timeout_seconds else 60
-
-        self._error_parser = _Parser(extra_error_customizers=[_AddDebugErrorCustomizer(cfg)])
+        self._http_timeout_seconds = http_timeout_seconds or 60
 
-    @property
-    def account_id(self) -> str:
-        return self._cfg.account_id
-
-    @property
-    def is_account_client(self) -> bool:
-        return self._cfg.is_account_client
+        self._error_parser = _Parser(extra_error_customizers=extra_error_customizers)
 
     def _authenticate(self, r: requests.PreparedRequest) -> requests.PreparedRequest:
-        headers = self._cfg.authenticate()
-        for k, v in headers.items():
-            r.headers[k] = v
+        if self._header_factory:
+            headers = self._header_factory()
+            for k, v in headers.items():
+                r.headers[k] = v
         return r
 
     @staticmethod
@@ -115,26 +107,9 @@ class ApiClient:
         flattened = dict(flatten_dict(with_fixed_bools))
         return flattened
 
-    def get_oauth_token(self, auth_details: str) -> Token:
-        if not self._cfg.auth_type:
-            self._cfg.authenticate()
-        original_token = self._cfg.oauth_token()
-        headers = {"Content-Type": URL_ENCODED_CONTENT_TYPE}
-        params = urlencode({
-            "grant_type": JWT_BEARER_GRANT_TYPE,
-            "authorization_details": auth_details,
-            "assertion": original_token.access_token
-        })
-        return retrieve_token(client_id=self._cfg.client_id,
-                              client_secret=self._cfg.client_secret,
-                              token_url=self._cfg.host + OIDC_TOKEN_PATH,
-                              params=params,
-                              headers=headers)
-
     def do(self,
            method: str,
-           path: str = None,
-           url: str = None,
+           url: str,
            query: dict = None,
            headers: dict = None,
            body: dict = None,
@@ -142,18 +117,13 @@ class ApiClient:
            files=None,
            data=None,
            auth: Callable[[requests.PreparedRequest], requests.PreparedRequest] = None,
-           response_headers: List[str] = None) -> Union[dict, BinaryIO]:
+           response_headers: List[str] = None) -> Union[dict, list, BinaryIO]:
         if headers is None:
             headers = {}
-        if url is None:
-            # Remove extra `/` from path for Files API
-            # Once we've fixed the OpenAPI spec, we can remove this
-            path = re.sub('^/api/2.0/fs/files//', '/api/2.0/fs/files/', path)
-            url = f"{self._cfg.host}{path}"
         headers['User-Agent'] = self._user_agent_base
         retryable = retried(timeout=timedelta(seconds=self._retry_timeout_seconds),
                             is_retryable=self._is_retryable,
-                            clock=self._cfg.clock)
+                            clock=self._clock)
         response = retryable(self._perform)(method,
                                             url,
                                             query=query,
@@ -168,19 +138,19 @@ class ApiClient:
         for header in response_headers if response_headers else []:
             resp[header] = response.headers.get(Casing.to_header_case(header))
         if raw:
-            resp["contents"] = StreamingResponse(response)
+            resp["contents"] = _StreamingResponse(response)
             return resp
         if not len(response.content):
             return resp
 
-        jsonResponse = response.json()
-        if jsonResponse is None:
+        json_response = response.json()
+        if json_response is None:
             return resp
 
-        if isinstance(jsonResponse, list):
-            return jsonResponse
+        if isinstance(json_response, list):
+            return json_response
 
-        return {**resp, **jsonResponse}
+        return {**resp, **json_response}
 
     @staticmethod
     def _is_retryable(err: BaseException) -> Optional[str]:
@@ -250,23 +220,10 @@ class ApiClient:
     def _record_request_log(self, response: requests.Response, raw: bool = False) -> None:
         if not logger.isEnabledFor(logging.DEBUG):
             return
-        logger.debug(RoundTrip(response, self._cfg.debug_headers, self._debug_truncate_bytes, raw).generate())
-
-
-class _AddDebugErrorCustomizer(_ErrorCustomizer):
-    """An error customizer that adds debug information about the configuration to unauthenticated and
-    unauthorized errors."""
-
-    def __init__(self, cfg: Config):
-        self._cfg = cfg
-
-    def customize_error(self, response: requests.Response, kwargs: dict):
-        if response.status_code in (401, 403):
-            message = kwargs.get('message', 'request failed')
-            kwargs['message'] = self._cfg.wrap_debug_info(message)
+        logger.debug(RoundTrip(response, self._debug_headers, self._debug_truncate_bytes, raw).generate())
 
 
-class StreamingResponse(BinaryIO):
+class _StreamingResponse(BinaryIO):
     _response: requests.Response
     _buffer: bytes
     _content: Union[Iterator[bytes], None]

databricks_sdk-0.35.0/databricks/sdk/core.py

@@ -0,0 +1,99 @@
+import re
+from typing import BinaryIO
+from urllib.parse import urlencode
+
+from ._base_client import _BaseClient
+from .config import *
+# To preserve backwards compatibility (as these definitions were previously in this module)
+from .credentials_provider import *
+from .errors import DatabricksError, _ErrorCustomizer
+from .oauth import retrieve_token
+
+__all__ = ['Config', 'DatabricksError']
+
+logger = logging.getLogger('databricks.sdk')
+
+URL_ENCODED_CONTENT_TYPE = "application/x-www-form-urlencoded"
+JWT_BEARER_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer"
+OIDC_TOKEN_PATH = "/oidc/v1/token"
+
+
+class ApiClient:
+
+    def __init__(self, cfg: Config):
+        self._cfg = cfg
+        self._api_client = _BaseClient(debug_truncate_bytes=cfg.debug_truncate_bytes,
+                                       retry_timeout_seconds=cfg.retry_timeout_seconds,
+                                       user_agent_base=cfg.user_agent,
+                                       header_factory=cfg.authenticate,
+                                       max_connection_pools=cfg.max_connection_pools,
+                                       max_connections_per_pool=cfg.max_connections_per_pool,
+                                       pool_block=True,
+                                       http_timeout_seconds=cfg.http_timeout_seconds,
+                                       extra_error_customizers=[_AddDebugErrorCustomizer(cfg)],
+                                       clock=cfg.clock)
+
+    @property
+    def account_id(self) -> str:
+        return self._cfg.account_id
+
+    @property
+    def is_account_client(self) -> bool:
+        return self._cfg.is_account_client
+
+    def get_oauth_token(self, auth_details: str) -> Token:
+        if not self._cfg.auth_type:
+            self._cfg.authenticate()
+        original_token = self._cfg.oauth_token()
+        headers = {"Content-Type": URL_ENCODED_CONTENT_TYPE}
+        params = urlencode({
+            "grant_type": JWT_BEARER_GRANT_TYPE,
+            "authorization_details": auth_details,
+            "assertion": original_token.access_token
+        })
+        return retrieve_token(client_id=self._cfg.client_id,
+                              client_secret=self._cfg.client_secret,
+                              token_url=self._cfg.host + OIDC_TOKEN_PATH,
+                              params=params,
+                              headers=headers)
+
+    def do(self,
+           method: str,
+           path: str = None,
+           url: str = None,
+           query: dict = None,
+           headers: dict = None,
+           body: dict = None,
+           raw: bool = False,
+           files=None,
+           data=None,
+           auth: Callable[[requests.PreparedRequest], requests.PreparedRequest] = None,
+           response_headers: List[str] = None) -> Union[dict, list, BinaryIO]:
+        if url is None:
+            # Remove extra `/` from path for Files API
+            # Once we've fixed the OpenAPI spec, we can remove this
+            path = re.sub('^/api/2.0/fs/files//', '/api/2.0/fs/files/', path)
+            url = f"{self._cfg.host}{path}"
+        return self._api_client.do(method=method,
+                                   url=url,
+                                   query=query,
+                                   headers=headers,
+                                   body=body,
+                                   raw=raw,
+                                   files=files,
+                                   data=data,
+                                   auth=auth,
+                                   response_headers=response_headers)
+
+
+class _AddDebugErrorCustomizer(_ErrorCustomizer):
+    """An error customizer that adds debug information about the configuration to unauthenticated and
+    unauthorized errors."""
+
+    def __init__(self, cfg: Config):
+        self._cfg = cfg
+
+    def customize_error(self, response: requests.Response, kwargs: dict):
+        if response.status_code in (401, 403):
+            message = kwargs.get('message', 'request failed')
+            kwargs['message'] = self._cfg.wrap_debug_info(message)

{databricks_sdk-0.33.0 → databricks_sdk-0.35.0}/databricks/sdk/mixins/files.py

@@ -167,7 +167,7 @@ class _DbfsIO(BinaryIO):
         return f"<_DbfsIO {self._path} {'read' if self.readable() else 'write'}=True>"
 
 
-class _FilesIO(BinaryIO):
+class _VolumesIO(BinaryIO):
 
     def __init__(self, api: files.FilesAPI, path: str, *, read: bool, write: bool, overwrite: bool):
         self._buffer = []
@@ -262,7 +262,7 @@ class _FilesIO(BinaryIO):
         self.close()
 
     def __repr__(self) -> str:
-        return f"<_FilesIO {self._path} {'read' if self.readable() else 'write'}=True>"
+        return f"<_VolumesIO {self._path} {'read' if self.readable() else 'write'}=True>"
 
 
 class _Path(ABC):
@@ -398,7 +398,7 @@ class _LocalPath(_Path):
         return f'<_LocalPath {self._path}>'
 
 
-class _FilesPath(_Path):
+class _VolumesPath(_Path):
 
     def __init__(self, api: files.FilesAPI, src: Union[str, pathlib.Path]):
         self._path = pathlib.PurePosixPath(str(src).replace('dbfs:', '').replace('file:', ''))
@@ -411,7 +411,7 @@ class _FilesPath(_Path):
         return False
 
     def child(self, path: str) -> Self:
-        return _FilesPath(self._api, str(self._path / path))
+        return _VolumesPath(self._api, str(self._path / path))
 
     def _is_dir(self) -> bool:
         try:
@@ -431,7 +431,7 @@ class _FilesPath(_Path):
             return self.is_dir
 
     def open(self, *, read=False, write=False, overwrite=False) -> BinaryIO:
-        return _FilesIO(self._api, self.as_string, read=read, write=write, overwrite=overwrite)
+        return _VolumesIO(self._api, self.as_string, read=read, write=write, overwrite=overwrite)
 
     def list(self, *, recursive=False) -> Generator[files.FileInfo, None, None]:
         if not self.is_dir:
@@ -458,13 +458,13 @@ class _FilesPath(_Path):
     def delete(self, *, recursive=False):
         if self.is_dir:
             for entry in self.list(recursive=False):
-                _FilesPath(self._api, entry.path).delete(recursive=True)
+                _VolumesPath(self._api, entry.path).delete(recursive=True)
             self._api.delete_directory(self.as_string)
         else:
             self._api.delete(self.as_string)
 
     def __repr__(self) -> str:
-        return f'<_FilesPath {self._path}>'
+        return f'<_VolumesPath {self._path}>'
 
 
 class _DbfsPath(_Path):
@@ -589,8 +589,8 @@ class DbfsExt(files.DbfsAPI):
                 'UC Volumes paths, not external locations or DBFS mount points.')
         if src.scheme == 'file':
             return _LocalPath(src.geturl())
-        if src.path.startswith(('/Volumes', '/Models')):
-            return _FilesPath(self._files_api, src.geturl())
+        if src.path.startswith('/Volumes'):
+            return _VolumesPath(self._files_api, src.geturl())
         return _DbfsPath(self._dbfs_api, src.geturl())
 
     def copy(self, src: str, dst: str, *, recursive=False, overwrite=False):

databricks_sdk-0.35.0/databricks/sdk/mixins/open_ai_client.py

@@ -0,0 +1,52 @@
+from databricks.sdk.service.serving import ServingEndpointsAPI
+
+
+class ServingEndpointsExt(ServingEndpointsAPI):
+
+    # Using the HTTP Client to pass in the databricks authorization
+    # This method will be called on every invocation, so when using with model serving will always get the refreshed token
+    def _get_authorized_http_client(self):
+        import httpx
+
+        class BearerAuth(httpx.Auth):
+
+            def __init__(self, get_headers_func):
+                self.get_headers_func = get_headers_func
+
+            def auth_flow(self, request: httpx.Request) -> httpx.Request:
+                auth_headers = self.get_headers_func()
+                request.headers["Authorization"] = auth_headers["Authorization"]
+                yield request
+
+        databricks_token_auth = BearerAuth(self._api._cfg.authenticate)
+
+        # Create an HTTP client with Bearer Token authentication
+        http_client = httpx.Client(auth=databricks_token_auth)
+        return http_client
+
+    def get_open_ai_client(self):
+        try:
+            from openai import OpenAI
+        except Exception:
+            raise ImportError(
+                "Open AI is not installed. Please install the Databricks SDK with the following command `pip isntall databricks-sdk[openai]`"
+            )
+
+        return OpenAI(
+            base_url=self._api._cfg.host + "/serving-endpoints",
+            api_key="no-token", # Passing in a placeholder to pass validations, this will not be used
+            http_client=self._get_authorized_http_client())
+
+    def get_langchain_chat_open_ai_client(self, model):
+        try:
+            from langchain_openai import ChatOpenAI
+        except Exception:
+            raise ImportError(
+                "Langchain Open AI is not installed. Please install the Databricks SDK with the following command `pip isntall databricks-sdk[openai]` and ensure you are using python>3.7"
+            )
+
+        return ChatOpenAI(
+            model=model,
+            openai_api_base=self._api._cfg.host + "/serving-endpoints",
+            api_key="no-token", # Passing in a placeholder to pass validations, this will not be used
+            http_client=self._get_authorized_http_client())