databricks-sdk 0.27.1__py3-none-any.whl → 0.29.0__py3-none-any.whl

databricks/sdk/__init__.py

@@ -1,7 +1,7 @@
 import databricks.sdk.core as client
 import databricks.sdk.dbutils as dbutils
 from databricks.sdk import azure
-from databricks.sdk.credentials_provider import CredentialsProvider
+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.workspace import WorkspaceExt
@@ -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,
@@ -131,7 +131,8 @@ class WorkspaceClient:
                  debug_headers: bool = None,
                  product="unknown",
                  product_version="0.0.0",
-                 credentials_provider: CredentialsProvider = None,
+                 credentials_strategy: CredentialsStrategy = None,
+                 credentials_provider: CredentialsStrategy = None,
                  config: client.Config = None):
         if not config:
             config = client.Config(host=host,
@@ -152,6 +153,7 @@ class WorkspaceClient:
                                    cluster_id=cluster_id,
                                    google_credentials=google_credentials,
                                    google_service_account=google_service_account,
+                                   credentials_strategy=credentials_strategy,
                                    credentials_provider=credentials_provider,
                                    debug_truncate_bytes=debug_truncate_bytes,
                                    debug_headers=debug_headers,
@@ -194,7 +196,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 +215,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 +427,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 +517,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."""
@@ -700,7 +702,8 @@ class AccountClient:
                  debug_headers: bool = None,
                  product="unknown",
                  product_version="0.0.0",
-                 credentials_provider: CredentialsProvider = None,
+                 credentials_strategy: CredentialsStrategy = None,
+                 credentials_provider: CredentialsStrategy = None,
                  config: client.Config = None):
         if not config:
             config = client.Config(host=host,
@@ -721,6 +724,7 @@ class AccountClient:
                                    cluster_id=cluster_id,
                                    google_credentials=google_credentials,
                                    google_service_account=google_service_account,
+                                   credentials_strategy=credentials_strategy,
                                    credentials_provider=credentials_provider,
                                    debug_truncate_bytes=debug_truncate_bytes,
                                    debug_headers=debug_headers,

databricks/sdk/azure.py

@@ -1,36 +1,9 @@
-from dataclasses import dataclass
 from typing import Dict
 
 from .oauth import TokenSource
 from .service.provisioning import Workspace
 
 
-@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/"),
-)
-
-
 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

databricks/sdk/config.py

@@ -6,16 +6,15 @@ import pathlib
 import platform
 import sys
 import urllib.parse
-from typing import Dict, Iterable, Optional
+from typing import Dict, Iterable, List, Optional, Tuple
 
 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 .oauth import OidcEndpoints
+from .credentials_provider import CredentialsStrategy, DefaultCredentials
+from .environments import (ALL_ENVS, AzureEnvironment, Cloud,
+                           DatabricksEnvironment, get_environment_for_hostname)
+from .oauth import OidcEndpoints, Token
 from .version import __version__
 
 logger = logging.getLogger('databricks.sdk')
@@ -45,6 +44,32 @@ class ConfigAttribute:
         return f"<ConfigAttribute '{self.name}' {self.transform.__name__}>"
 
 
+_DEFAULT_PRODUCT_NAME = 'unknown'
+_DEFAULT_PRODUCT_VERSION = '0.0.0'
+_STATIC_USER_AGENT: Tuple[str, str, List[str]] = (_DEFAULT_PRODUCT_NAME, _DEFAULT_PRODUCT_VERSION, [])
+
+
+def with_product(product: str, product_version: str):
+    """[INTERNAL API] Change the product name and version used in the User-Agent header."""
+    global _STATIC_USER_AGENT
+    prev_product, prev_version, prev_other_info = _STATIC_USER_AGENT
+    logger.debug(f'Changing product from {prev_product}/{prev_version} to {product}/{product_version}')
+    _STATIC_USER_AGENT = product, product_version, prev_other_info
+
+
+def with_user_agent_extra(key: str, value: str):
+    """[INTERNAL API] Add extra metadata to the User-Agent header when developing a library."""
+    global _STATIC_USER_AGENT
+    product_name, product_version, other_info = _STATIC_USER_AGENT
+    for item in other_info:
+        if item.startswith(f"{key}/"):
+            # ensure that we don't have duplicates
+            other_info.remove(item)
+            break
+    other_info.append(f"{key}/{value}")
+    _STATIC_USER_AGENT = product_name, product_version, other_info
+
+
 class Config:
     host: str = ConfigAttribute(env='DATABRICKS_HOST')
     account_id: str = ConfigAttribute(env='DATABRICKS_ACCOUNT_ID')
@@ -67,6 +92,7 @@ class Config:
     auth_type: str = ConfigAttribute(env='DATABRICKS_AUTH_TYPE')
     cluster_id: str = ConfigAttribute(env='DATABRICKS_CLUSTER_ID')
     warehouse_id: str = ConfigAttribute(env='DATABRICKS_WAREHOUSE_ID')
+    serverless_compute_id: str = ConfigAttribute(env='DATABRICKS_SERVERLESS_COMPUTE_ID')
     skip_verify: bool = ConfigAttribute()
     http_timeout_seconds: float = ConfigAttribute()
     debug_truncate_bytes: int = ConfigAttribute(env='DATABRICKS_DEBUG_TRUNCATE_BYTES')
@@ -82,15 +108,34 @@ class Config:
 
     def __init__(self,
                  *,
-                 credentials_provider: CredentialsProvider = None,
-                 product="unknown",
-                 product_version="0.0.0",
+                 # Deprecated. Use credentials_strategy instead.
+                 credentials_provider: CredentialsStrategy = None,
+                 credentials_strategy: CredentialsStrategy = None,
+                 product=_DEFAULT_PRODUCT_NAME,
+                 product_version=_DEFAULT_PRODUCT_VERSION,
                  clock: Clock = None,
                  **kwargs):
         self._header_factory = None
         self._inner = {}
+        # as in SDK for Go, pull information from global static user agent context,
+        # so that we can track additional metadata for mid-stream libraries, as well
+        # as for cases, when the downstream product is used as a library and is not
+        # configured with a proper product name and version.
+        static_product, static_version, _ = _STATIC_USER_AGENT
+        if product == _DEFAULT_PRODUCT_NAME:
+            product = static_product
+        if product_version == _DEFAULT_PRODUCT_VERSION:
+            product_version = static_version
         self._user_agent_other_info = []
-        self._credentials_provider = credentials_provider if credentials_provider else DefaultCredentials()
+        if credentials_strategy and credentials_provider:
+            raise ValueError(
+                "When providing `credentials_strategy` field, `credential_provider` cannot be specified.")
+        if credentials_provider:
+            logger.warning(
+                "parameter 'credentials_provider' is deprecated. Use 'credentials_strategy' instead.")
+        self._credentials_strategy = next(
+            s for s in [credentials_strategy, credentials_provider,
+                        DefaultCredentials()] if s is not None)
         if 'databricks_environment' in kwargs:
             self.databricks_environment = kwargs['databricks_environment']
             del kwargs['databricks_environment']
@@ -108,6 +153,9 @@ class Config:
             message = self.wrap_debug_info(str(e))
             raise ValueError(message) from e
 
+    def oauth_token(self) -> Token:
+        return self._credentials_strategy.oauth_token(self)
+
     def wrap_debug_info(self, message: str) -> str:
         debug_string = self.debug_string()
         if debug_string:
@@ -154,11 +202,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 +212,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
@@ -223,6 +269,12 @@ class Config:
         ]
         if len(self._user_agent_other_info) > 0:
             ua.append(' '.join(self._user_agent_other_info))
+        # as in SDK for Go, pull information from global static user agent context,
+        # so that we can track additional metadata for mid-stream libraries. this value
+        # is shared across all instances of Config objects intentionally.
+        _, _, static_info = _STATIC_USER_AGENT
+        if len(static_info) > 0:
+            ua.append(' '.join(static_info))
         if len(self._upstream_user_agent) > 0:
             ua.append(self._upstream_user_agent)
         if 'DATABRICKS_RUNTIME_VERSION' in os.environ:
@@ -439,12 +491,12 @@ class Config:
 
     def init_auth(self):
         try:
-            self._header_factory = self._credentials_provider(self)
-            self.auth_type = self._credentials_provider.auth_type()
+            self._header_factory = self._credentials_strategy(self)
+            self.auth_type = self._credentials_strategy.auth_type()
             if not self._header_factory:
                 raise ValueError('not configured')
         except ValueError as e:
-            raise ValueError(f'{self._credentials_provider.auth_type()} auth: {e}') from e
+            raise ValueError(f'{self._credentials_strategy.auth_type()} auth: {e}') from e
 
     def __repr__(self):
         return f'<{self.debug_string()}>'

databricks/sdk/core.py

@@ -4,6 +4,7 @@ from datetime import timedelta
 from json import JSONDecodeError
 from types import TracebackType
 from typing import Any, BinaryIO, Iterator, Type
+from urllib.parse import urlencode
 
 from requests.adapters import HTTPAdapter
 
@@ -12,12 +13,18 @@ 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 .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
@@ -108,6 +115,22 @@ 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,
@@ -239,6 +262,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)