databricks-sdk 0.27.0__tar.gz → 0.28.0__tar.gz

{databricks-sdk-0.27.0/databricks_sdk.egg-info → databricks-sdk-0.28.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: databricks-sdk
-Version: 0.27.0
+Version: 0.28.0
 Summary: Databricks SDK for Python (Beta)
 Home-page: https://databricks-sdk-py.readthedocs.io
 Author: Serge Smertin

{databricks-sdk-0.27.0 → databricks-sdk-0.28.0}/databricks/sdk/__init__.py

@@ -13,9 +13,9 @@ from databricks.sdk.service.catalog import (AccountMetastoreAssignmentsAPI,
                                             ArtifactAllowlistsAPI, CatalogsAPI,
                                             ConnectionsAPI,
                                             ExternalLocationsAPI, FunctionsAPI,
-                                            GrantsAPI, LakehouseMonitorsAPI,
-                                            MetastoresAPI, ModelVersionsAPI,
-                                            OnlineTablesAPI,
+                                            GrantsAPI, MetastoresAPI,
+                                            ModelVersionsAPI, OnlineTablesAPI,
+                                            QualityMonitorsAPI,
                                             RegisteredModelsAPI, SchemasAPI,
                                             StorageCredentialsAPI,
                                             SystemSchemasAPI,
@@ -194,7 +194,6 @@ class WorkspaceClient:
         self._instance_profiles = InstanceProfilesAPI(self._api_client)
         self._ip_access_lists = IpAccessListsAPI(self._api_client)
         self._jobs = JobsAPI(self._api_client)
-        self._lakehouse_monitors = LakehouseMonitorsAPI(self._api_client)
         self._lakeview = LakeviewAPI(self._api_client)
         self._libraries = LibrariesAPI(self._api_client)
         self._metastores = MetastoresAPI(self._api_client)
@@ -214,6 +213,7 @@ class WorkspaceClient:
             self._api_client)
         self._provider_providers = ProviderProvidersAPI(self._api_client)
         self._providers = ProvidersAPI(self._api_client)
+        self._quality_monitors = QualityMonitorsAPI(self._api_client)
         self._queries = QueriesAPI(self._api_client)
         self._query_history = QueryHistoryAPI(self._api_client)
         self._query_visualizations = QueryVisualizationsAPI(self._api_client)
@@ -425,11 +425,6 @@ class WorkspaceClient:
         """The Jobs API allows you to create, edit, and delete jobs."""
         return self._jobs
 
-    @property
-    def lakehouse_monitors(self) -> LakehouseMonitorsAPI:
-        """A monitor computes and monitors data or model quality metrics for a table over time."""
-        return self._lakehouse_monitors
-
     @property
     def lakeview(self) -> LakeviewAPI:
         """These APIs provide specific management operations for Lakeview dashboards."""
@@ -520,6 +515,11 @@ class WorkspaceClient:
         """A data provider is an object representing the organization in the real world who shares the data."""
         return self._providers
 
+    @property
+    def quality_monitors(self) -> QualityMonitorsAPI:
+        """A monitor computes and monitors data or model quality metrics for a table over time."""
+        return self._quality_monitors
+
     @property
     def queries(self) -> QueriesAPI:
         """These endpoints are used for CRUD operations on query definitions."""

databricks-sdk-0.28.0/databricks/sdk/azure.py

@@ -0,0 +1,27 @@
+from typing import Dict
+
+from .oauth import TokenSource
+from .service.provisioning import Workspace
+
+
+def add_workspace_id_header(cfg: 'Config', headers: Dict[str, str]):
+    if cfg.azure_workspace_resource_id:
+        headers["X-Databricks-Azure-Workspace-Resource-Id"] = cfg.azure_workspace_resource_id
+
+
+def add_sp_management_token(token_source: 'TokenSource', headers: Dict[str, str]):
+    mgmt_token = token_source.token()
+    headers['X-Databricks-Azure-SP-Management-Token'] = mgmt_token.access_token
+
+
+def get_azure_resource_id(workspace: Workspace):
+    """
+    Returns the Azure Resource ID for the given workspace, if it is an Azure workspace.
+    :param workspace:
+    :return:
+    """
+    if workspace.azure_workspace_info is None:
+        return None
+    return (f'/subscriptions/{workspace.azure_workspace_info.subscription_id}'
+            f'/resourceGroups/{workspace.azure_workspace_info.resource_group}'
+            f'/providers/Microsoft.Databricks/workspaces/{workspace.workspace_name}')

{databricks-sdk-0.27.0 → databricks-sdk-0.28.0}/databricks/sdk/config.py

@@ -10,11 +10,10 @@ from typing import Dict, Iterable, Optional
 
 import requests
 
-from .azure import AzureEnvironment
 from .clock import Clock, RealClock
 from .credentials_provider import CredentialsProvider, DefaultCredentials
-from .environments import (ALL_ENVS, DEFAULT_ENVIRONMENT, Cloud,
-                           DatabricksEnvironment)
+from .environments import (ALL_ENVS, AzureEnvironment, Cloud,
+                           DatabricksEnvironment, get_environment_for_hostname)
 from .oauth import OidcEndpoints
 from .version import __version__
 
@@ -154,11 +153,7 @@ class Config:
         """Returns the environment based on configuration."""
         if self.databricks_environment:
             return self.databricks_environment
-        if self.host:
-            for environment in ALL_ENVS:
-                if self.host.endswith(environment.dns_zone):
-                    return environment
-        if self.azure_workspace_resource_id:
+        if not self.host and self.azure_workspace_resource_id:
             azure_env = self._get_azure_environment_name()
             for environment in ALL_ENVS:
                 if environment.cloud != Cloud.AZURE:
@@ -168,10 +163,12 @@ class Config:
                 if environment.dns_zone.startswith(".dev") or environment.dns_zone.startswith(".staging"):
                     continue
                 return environment
-        return DEFAULT_ENVIRONMENT
+        return get_environment_for_hostname(self.host)
 
     @property
     def is_azure(self) -> bool:
+        if self.azure_workspace_resource_id:
+            return True
         return self.environment.cloud == Cloud.AZURE
 
     @property

{databricks-sdk-0.27.0 → databricks-sdk-0.28.0}/databricks/sdk/core.py

@@ -12,6 +12,7 @@ from .config import *
 # To preserve backwards compatibility (as these definitions were previously in this module)
 from .credentials_provider import *
 from .errors import DatabricksError, error_mapper
+from .errors.private_link import _is_private_link_redirect
 from .retries import retried
 
 __all__ = ['Config', 'DatabricksError']
@@ -145,11 +146,14 @@ class ApiClient:
         if not len(response.content):
             return resp
 
-        json = response.json()
-        if isinstance(json, list):
-            return json
+        jsonResponse = response.json()
+        if jsonResponse is None:
+            return resp
+
+        if isinstance(jsonResponse, list):
+            return jsonResponse
 
-        return {**resp, **json}
+        return {**resp, **jsonResponse}
 
     @staticmethod
     def _is_retryable(err: BaseException) -> Optional[str]:
@@ -236,6 +240,10 @@ class ApiClient:
                 # See https://stackoverflow.com/a/58821552/277035
                 payload = response.json()
                 raise self._make_nicer_error(response=response, **payload) from None
+            # Private link failures happen via a redirect to the login page. From a requests-perspective, the request
+            # is successful, but the response is not what we expect. We need to handle this case separately.
+            if _is_private_link_redirect(response):
+                raise self._make_nicer_error(response=response) from None
             return response
         except requests.exceptions.JSONDecodeError:
             message = self._make_sense_from_html(response.text)

{databricks-sdk-0.27.0 → databricks-sdk-0.28.0}/databricks/sdk/environments.py

@@ -2,7 +2,31 @@ from dataclasses import dataclass
 from enum import Enum
 from typing import Optional
 
-from .azure import ARM_DATABRICKS_RESOURCE_ID, ENVIRONMENTS, AzureEnvironment
+
+@dataclass
+class AzureEnvironment:
+    name: str
+    service_management_endpoint: str
+    resource_manager_endpoint: str
+    active_directory_endpoint: str
+
+
+ARM_DATABRICKS_RESOURCE_ID = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d"
+
+ENVIRONMENTS = dict(
+    PUBLIC=AzureEnvironment(name="PUBLIC",
+                            service_management_endpoint="https://management.core.windows.net/",
+                            resource_manager_endpoint="https://management.azure.com/",
+                            active_directory_endpoint="https://login.microsoftonline.com/"),
+    USGOVERNMENT=AzureEnvironment(name="USGOVERNMENT",
+                                  service_management_endpoint="https://management.core.usgovcloudapi.net/",
+                                  resource_manager_endpoint="https://management.usgovcloudapi.net/",
+                                  active_directory_endpoint="https://login.microsoftonline.us/"),
+    CHINA=AzureEnvironment(name="CHINA",
+                           service_management_endpoint="https://management.core.chinacloudapi.cn/",
+                           resource_manager_endpoint="https://management.chinacloudapi.cn/",
+                           active_directory_endpoint="https://login.chinacloudapi.cn/"),
+)
 
 
 class Cloud(Enum):
@@ -70,3 +94,12 @@ ALL_ENVS = [
     DatabricksEnvironment(Cloud.GCP, ".staging.gcp.databricks.com"),
     DatabricksEnvironment(Cloud.GCP, ".gcp.databricks.com")
 ]
+
+
+def get_environment_for_hostname(hostname: str) -> DatabricksEnvironment:
+    if not hostname:
+        return DEFAULT_ENVIRONMENT
+    for env in ALL_ENVS:
+        if hostname.endswith(env.dns_zone):
+            return env
+    return DEFAULT_ENVIRONMENT

{databricks-sdk-0.27.0 → databricks-sdk-0.28.0}/databricks/sdk/errors/__init__.py

@@ -1,4 +1,5 @@
 from .base import DatabricksError, ErrorDetail
 from .mapper import error_mapper
 from .platform import *
+from .private_link import PrivateLinkValidationError
 from .sdk import *

{databricks-sdk-0.27.0 → databricks-sdk-0.28.0}/databricks/sdk/errors/mapper.py

@@ -4,6 +4,8 @@ from databricks.sdk.errors import platform
 from databricks.sdk.errors.base import DatabricksError
 
 from .overrides import _ALL_OVERRIDES
+from .private_link import (_get_private_link_validation_error,
+                           _is_private_link_redirect)
 
 
 def error_mapper(response: requests.Response, raw: dict) -> DatabricksError:
@@ -21,6 +23,8 @@ def error_mapper(response: requests.Response, raw: dict) -> DatabricksError:
         # 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 platform.STATUS_CODE_MAPPING[status_code](**raw)
+    if _is_private_link_redirect(response):
+        return _get_private_link_validation_error(response.url)
 
     # backwards-compatible error creation for cases like using older versions of
     # the SDK on way never releases of the platform.

databricks-sdk-0.28.0/databricks/sdk/errors/private_link.py

@@ -0,0 +1,60 @@
+from dataclasses import dataclass
+from urllib import parse
+
+import requests
+
+from ..environments import Cloud, get_environment_for_hostname
+from .platform import PermissionDenied
+
+
+@dataclass
+class _PrivateLinkInfo:
+    serviceName: str
+    endpointName: str
+    referencePage: str
+
+    def error_message(self):
+        return (
+            f'The requested workspace has {self.serviceName} enabled and is not accessible from the current network. '
+            f'Ensure that {self.serviceName} is properly configured and that your device has access to the '
+            f'{self.endpointName}. For more information, see {self.referencePage}.')
+
+
+_private_link_info_map = {
+    Cloud.AWS:
+    _PrivateLinkInfo(serviceName='AWS PrivateLink',
+                     endpointName='AWS VPC endpoint',
+                     referencePage='https://docs.databricks.com/en/security/network/classic/privatelink.html',
+                     ),
+    Cloud.AZURE:
+    _PrivateLinkInfo(
+        serviceName='Azure Private Link',
+        endpointName='Azure Private Link endpoint',
+        referencePage='https://learn.microsoft.com/en-us/azure/databricks/security/network/classic/private-link-standard#authentication-troubleshooting',
+    ),
+    Cloud.GCP:
+    _PrivateLinkInfo(
+        serviceName='Private Service Connect',
+        endpointName='GCP VPC endpoint',
+        referencePage='https://docs.gcp.databricks.com/en/security/network/classic/private-service-connect.html',
+    )
+}
+
+
+class PrivateLinkValidationError(PermissionDenied):
+    """Raised when a user tries to access a Private Link-enabled workspace, but the user's network does not have access
+    to the workspace."""
+
+
+def _is_private_link_redirect(resp: requests.Response) -> bool:
+    parsed = parse.urlparse(resp.url)
+    return parsed.path == '/login.html' and 'error=private-link-validation-error' in parsed.query
+
+
+def _get_private_link_validation_error(url: str) -> _PrivateLinkInfo:
+    parsed = parse.urlparse(url)
+    env = get_environment_for_hostname(parsed.hostname)
+    return PrivateLinkValidationError(message=_private_link_info_map[env.cloud].error_message(),
+                                      error_code='PRIVATE_LINK_VALIDATION_ERROR',
+                                      status_code=403,
+                                      )