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

databricks/sdk/__init__.py

@@ -15,6 +15,7 @@ from databricks.sdk.service.catalog import (AccountMetastoreAssignmentsAPI,
                                             ExternalLocationsAPI, FunctionsAPI,
                                             GrantsAPI, LakehouseMonitorsAPI,
                                             MetastoresAPI, ModelVersionsAPI,
+                                            OnlineTablesAPI,
                                             RegisteredModelsAPI, SchemasAPI,
                                             StorageCredentialsAPI,
                                             SystemSchemasAPI,
@@ -180,6 +181,7 @@ class WorkspaceClient:
         self._metastores = MetastoresAPI(self._api_client)
         self._model_registry = ModelRegistryAPI(self._api_client)
         self._model_versions = ModelVersionsAPI(self._api_client)
+        self._online_tables = OnlineTablesAPI(self._api_client)
         self._permissions = PermissionsAPI(self._api_client)
         self._pipelines = PipelinesAPI(self._api_client)
         self._policy_families = PolicyFamiliesAPI(self._api_client)
@@ -322,7 +324,7 @@ class WorkspaceClient:
 
     @property
     def files(self) -> FilesAPI:
-        """The Files API allows you to read, write, and delete files and directories in Unity Catalog volumes."""
+        """The Files API allows you to read, write, list, and delete files and directories."""
         return self._files
 
     @property
@@ -400,6 +402,11 @@ class WorkspaceClient:
         """Databricks provides a hosted version of MLflow Model Registry in Unity Catalog."""
         return self._model_versions
 
+    @property
+    def online_tables(self) -> OnlineTablesAPI:
+        """Online tables provide lower latency and higher QPS access to data from Delta tables."""
+        return self._online_tables
+
     @property
     def permissions(self) -> PermissionsAPI:
         """Permissions API are used to create read, write, edit, update and manage access for various users on different objects and endpoints."""
@@ -560,6 +567,13 @@ class WorkspaceClient:
         """This API allows updating known workspace settings for advanced users."""
         return self._workspace_conf
 
+    def get_workspace_id(self) -> int:
+        """Get the workspace ID of the workspace that this client is connected to."""
+        response = self._api_client.do("GET",
+                                       "/api/2.0/preview/scim/v2/Me",
+                                       response_headers=['X-Databricks-Org-Id'])
+        return int(response["X-Databricks-Org-Id"])
+
     def __repr__(self):
         return f"WorkspaceClient(host='{self._config.host}', auth_type='{self._config.auth_type}', ...)"
 
@@ -802,7 +816,7 @@ class AccountClient:
         :param workspace: The workspace to construct a client for.
         :return: A ``WorkspaceClient`` for the given workspace.
         """
-        config = self._config.copy()
+        config = self._config.deep_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

databricks/sdk/casing.py

@@ -0,0 +1,38 @@
+class _Name(object):
+    """Parses a name in camelCase, PascalCase, snake_case, or kebab-case into its segments."""
+
+    def __init__(self, raw_name: str):
+        #
+        self._segments = []
+        segment = []
+        for ch in raw_name:
+            if ch.isupper():
+                if segment:
+                    self._segments.append(''.join(segment))
+                segment = [ch.lower()]
+            elif ch.islower():
+                segment.append(ch)
+            else:
+                if segment:
+                    self._segments.append(''.join(segment))
+                segment = []
+        if segment:
+            self._segments.append(''.join(segment))
+
+    def to_snake_case(self) -> str:
+        return '_'.join(self._segments)
+
+    def to_header_case(self) -> str:
+        return '-'.join([s.capitalize() for s in self._segments])
+
+
+class Casing(object):
+
+    @staticmethod
+    def to_header_case(name: str) -> str:
+        """
+        Convert a name from camelCase, PascalCase, snake_case, or kebab-case to header-case.
+        :param name:
+        :return:
+        """
+        return _Name(name).to_header_case()

databricks/sdk/config.py

@@ -457,3 +457,8 @@ class Config:
         cpy: Config = copy.copy(self)
         cpy._user_agent_other_info = copy.deepcopy(self._user_agent_other_info)
         return cpy
+
+    def deep_copy(self):
+        """Creates a deep copy of the config object.
+        """
+        return copy.deepcopy(self)

databricks/sdk/core.py

@@ -7,6 +7,7 @@ from typing import Any, BinaryIO, Iterator, Type
 
 from requests.adapters import HTTPAdapter
 
+from .casing import Casing
 from .config import *
 # To preserve backwards compatibility (as these definitions were previously in this module)
 from .credentials_provider import *
@@ -115,7 +116,8 @@ class ApiClient:
            body: dict = None,
            raw: bool = False,
            files=None,
-           data=None) -> Union[dict, BinaryIO]:
+           data=None,
+           response_headers: List[str] = None) -> Union[dict, BinaryIO]:
         # 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)
@@ -125,14 +127,29 @@ class ApiClient:
         retryable = retried(timeout=timedelta(seconds=self._retry_timeout_seconds),
                             is_retryable=self._is_retryable,
                             clock=self._cfg.clock)
-        return retryable(self._perform)(method,
-                                        path,
-                                        query=query,
-                                        headers=headers,
-                                        body=body,
-                                        raw=raw,
-                                        files=files,
-                                        data=data)
+        response = retryable(self._perform)(method,
+                                            path,
+                                            query=query,
+                                            headers=headers,
+                                            body=body,
+                                            raw=raw,
+                                            files=files,
+                                            data=data)
+
+        resp = dict()
+        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)
+            return resp
+        if not len(response.content):
+            return resp
+
+        json = response.json()
+        if isinstance(json, list):
+            return json
+
+        return {**resp, **json}
 
     @staticmethod
     def _is_retryable(err: BaseException) -> Optional[str]:
@@ -219,11 +236,7 @@ class ApiClient:
                 # See https://stackoverflow.com/a/58821552/277035
                 payload = response.json()
                 raise self._make_nicer_error(response=response, **payload) from None
-            if raw:
-                return StreamingResponse(response)
-            if not len(response.content):
-                return {}
-            return response.json()
+            return response
         except requests.exceptions.JSONDecodeError:
             message = self._make_sense_from_html(response.text)
             if not message:

databricks/sdk/mixins/workspace.py

@@ -102,4 +102,5 @@ class WorkspaceExt(WorkspaceAPI):
         """
         query = {'path': path, 'direct_download': 'true'}
         if format: query['format'] = format.value
-        return self._api.do('GET', '/api/2.0/workspace/export', query=query, raw=True)
+        response = self._api.do('GET', '/api/2.0/workspace/export', query=query, raw=True)
+        return response["contents"]

databricks/sdk/service/billing.py

@@ -335,6 +335,17 @@ class DeliveryStatus(Enum):
 class DownloadResponse:
     contents: Optional[BinaryIO] = None
 
+    def as_dict(self) -> dict:
+        """Serializes the DownloadResponse into a dictionary suitable for use as a JSON request body."""
+        body = {}
+        if self.contents: body['contents'] = self.contents
+        return body
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, any]) -> DownloadResponse:
+        """Deserializes the DownloadResponse from a dictionary."""
+        return cls(contents=d.get('contents', None))
+
 
 class LogDeliveryConfigStatus(Enum):
     """Status of log delivery configuration. Set to `ENABLED` (enabled) or `DISABLED` (disabled).
@@ -704,12 +715,13 @@ class BillableUsageAPI:
         if personal_data is not None: query['personal_data'] = personal_data
         if start_month is not None: query['start_month'] = start_month
         headers = {'Accept': 'text/plain', }
+
         res = self._api.do('GET',
                            f'/api/2.0/accounts/{self._api.account_id}/usage/download',
                            query=query,
                            headers=headers,
                            raw=True)
-        return DownloadResponse(contents=res)
+        return DownloadResponse.from_dict(res)
 
 
 class BudgetsAPI:
@@ -732,6 +744,7 @@ class BudgetsAPI:
         body = {}
         if budget is not None: body['budget'] = budget.as_dict()
         headers = {'Accept': 'application/json', 'Content-Type': 'application/json', }
+
         res = self._api.do('POST',
                            f'/api/2.0/accounts/{self._api.account_id}/budget',
                            body=body,
@@ -750,6 +763,7 @@ class BudgetsAPI:
         """
 
         headers = {'Accept': 'application/json', }
+
         self._api.do('DELETE',
                      f'/api/2.0/accounts/{self._api.account_id}/budget/{budget_id}',
                      headers=headers)
@@ -767,6 +781,7 @@ class BudgetsAPI:
         """
 
         headers = {'Accept': 'application/json', }
+
         res = self._api.do('GET',
                            f'/api/2.0/accounts/{self._api.account_id}/budget/{budget_id}',
                            headers=headers)
@@ -782,6 +797,7 @@ class BudgetsAPI:
         """
 
         headers = {'Accept': 'application/json', }
+
         json = self._api.do('GET', f'/api/2.0/accounts/{self._api.account_id}/budget', headers=headers)
         parsed = BudgetList.from_dict(json).budgets
         return parsed if parsed is not None else []
@@ -801,6 +817,7 @@ class BudgetsAPI:
         body = {}
         if budget is not None: body['budget'] = budget.as_dict()
         headers = {'Accept': 'application/json', 'Content-Type': 'application/json', }
+
         self._api.do('PATCH',
                      f'/api/2.0/accounts/{self._api.account_id}/budget/{budget_id}',
                      body=body,
@@ -894,6 +911,7 @@ class LogDeliveryAPI:
         if log_delivery_configuration is not None:
             body['log_delivery_configuration'] = log_delivery_configuration.as_dict()
         headers = {'Accept': 'application/json', 'Content-Type': 'application/json', }
+
         res = self._api.do('POST',
                            f'/api/2.0/accounts/{self._api.account_id}/log-delivery',
                            body=body,
@@ -912,6 +930,7 @@ class LogDeliveryAPI:
         """
 
         headers = {'Accept': 'application/json', }
+
         res = self._api.do(
             'GET',
             f'/api/2.0/accounts/{self._api.account_id}/log-delivery/{log_delivery_configuration_id}',
@@ -942,6 +961,7 @@ class LogDeliveryAPI:
         if status is not None: query['status'] = status.value
         if storage_configuration_id is not None: query['storage_configuration_id'] = storage_configuration_id
         headers = {'Accept': 'application/json', }
+
         json = self._api.do('GET',
                             f'/api/2.0/accounts/{self._api.account_id}/log-delivery',
                             query=query,
@@ -970,6 +990,7 @@ class LogDeliveryAPI:
         body = {}
         if status is not None: body['status'] = status.value
         headers = {'Accept': 'application/json', 'Content-Type': 'application/json', }
+
         self._api.do('PATCH',
                      f'/api/2.0/accounts/{self._api.account_id}/log-delivery/{log_delivery_configuration_id}',
                      body=body,