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

databricks/sdk/__init__.py

@@ -5,8 +5,9 @@ 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
+from databricks.sdk.service.apps import AppsAPI
 from databricks.sdk.service.billing import (BillableUsageAPI, BudgetsAPI,
-                                            LogDeliveryAPI)
+                                            LogDeliveryAPI, UsageDashboardsAPI)
 from databricks.sdk.service.catalog import (AccountMetastoreAssignmentsAPI,
                                             AccountMetastoresAPI,
                                             AccountStorageCredentialsAPI,
@@ -27,7 +28,7 @@ from databricks.sdk.service.compute import (ClusterPoliciesAPI, ClustersAPI,
                                             InstancePoolsAPI,
                                             InstanceProfilesAPI, LibrariesAPI,
                                             PolicyFamiliesAPI)
-from databricks.sdk.service.dashboards import LakeviewAPI
+from databricks.sdk.service.dashboards import GenieAPI, LakeviewAPI
 from databricks.sdk.service.files import DbfsAPI, FilesAPI
 from databricks.sdk.service.iam import (AccountAccessControlAPI,
                                         AccountAccessControlProxyAPI,
@@ -55,7 +56,8 @@ from databricks.sdk.service.provisioning import (CredentialsAPI,
                                                  NetworksAPI, PrivateAccessAPI,
                                                  StorageAPI, VpcEndpointsAPI,
                                                  Workspace, WorkspacesAPI)
-from databricks.sdk.service.serving import AppsAPI, ServingEndpointsAPI
+from databricks.sdk.service.serving import (ServingEndpointsAPI,
+                                            ServingEndpointsDataPlaneAPI)
 from databricks.sdk.service.settings import (AccountIpAccessListsAPI,
                                              AccountSettingsAPI,
                                              AutomaticClusterUpdateAPI,
@@ -67,6 +69,7 @@ from databricks.sdk.service.settings import (AccountIpAccessListsAPI,
                                              EsmEnablementAccountAPI,
                                              IpAccessListsAPI,
                                              NetworkConnectivityAPI,
+                                             NotificationDestinationsAPI,
                                              PersonalComputeAPI,
                                              RestrictWorkspaceAdminsAPI,
                                              SettingsAPI, TokenManagementAPI,
@@ -74,11 +77,13 @@ from databricks.sdk.service.settings import (AccountIpAccessListsAPI,
 from databricks.sdk.service.sharing import (CleanRoomsAPI, ProvidersAPI,
                                             RecipientActivationAPI,
                                             RecipientsAPI, SharesAPI)
-from databricks.sdk.service.sql import (AlertsAPI, DashboardsAPI,
-                                        DashboardWidgetsAPI, DataSourcesAPI,
-                                        DbsqlPermissionsAPI, QueriesAPI,
+from databricks.sdk.service.sql import (AlertsAPI, AlertsLegacyAPI,
+                                        DashboardsAPI, DashboardWidgetsAPI,
+                                        DataSourcesAPI, DbsqlPermissionsAPI,
+                                        QueriesAPI, QueriesLegacyAPI,
                                         QueryHistoryAPI,
                                         QueryVisualizationsAPI,
+                                        QueryVisualizationsLegacyAPI,
                                         StatementExecutionAPI, WarehousesAPI)
 from databricks.sdk.service.vectorsearch import (VectorSearchEndpointsAPI,
                                                  VectorSearchIndexesAPI)
@@ -162,8 +167,10 @@ 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)
         self._account_access_control_proxy = AccountAccessControlProxyAPI(self._api_client)
         self._alerts = AlertsAPI(self._api_client)
+        self._alerts_legacy = AlertsLegacyAPI(self._api_client)
         self._apps = AppsAPI(self._api_client)
         self._artifact_allowlists = ArtifactAllowlistsAPI(self._api_client)
         self._catalogs = CatalogsAPI(self._api_client)
@@ -188,6 +195,7 @@ class WorkspaceClient:
         self._external_locations = ExternalLocationsAPI(self._api_client)
         self._files = FilesAPI(self._api_client)
         self._functions = FunctionsAPI(self._api_client)
+        self._genie = GenieAPI(self._api_client)
         self._git_credentials = GitCredentialsAPI(self._api_client)
         self._global_init_scripts = GlobalInitScriptsAPI(self._api_client)
         self._grants = GrantsAPI(self._api_client)
@@ -201,6 +209,7 @@ class WorkspaceClient:
         self._metastores = MetastoresAPI(self._api_client)
         self._model_registry = ModelRegistryAPI(self._api_client)
         self._model_versions = ModelVersionsAPI(self._api_client)
+        self._notification_destinations = NotificationDestinationsAPI(self._api_client)
         self._online_tables = OnlineTablesAPI(self._api_client)
         self._permission_migration = PermissionMigrationAPI(self._api_client)
         self._permissions = PermissionsAPI(self._api_client)
@@ -217,8 +226,10 @@ class WorkspaceClient:
         self._providers = ProvidersAPI(self._api_client)
         self._quality_monitors = QualityMonitorsAPI(self._api_client)
         self._queries = QueriesAPI(self._api_client)
+        self._queries_legacy = QueriesLegacyAPI(self._api_client)
         self._query_history = QueryHistoryAPI(self._api_client)
         self._query_visualizations = QueryVisualizationsAPI(self._api_client)
+        self._query_visualizations_legacy = QueryVisualizationsLegacyAPI(self._api_client)
         self._recipient_activation = RecipientActivationAPI(self._api_client)
         self._recipients = RecipientsAPI(self._api_client)
         self._registered_models = RegisteredModelsAPI(self._api_client)
@@ -226,7 +237,8 @@ class WorkspaceClient:
         self._schemas = SchemasAPI(self._api_client)
         self._secrets = SecretsAPI(self._api_client)
         self._service_principals = ServicePrincipalsAPI(self._api_client)
-        self._serving_endpoints = ServingEndpointsAPI(self._api_client)
+        self._serving_endpoints = serving_endpoints
+        self._serving_endpoints_data_plane = ServingEndpointsDataPlaneAPI(self._api_client, serving_endpoints)
         self._settings = SettingsAPI(self._api_client)
         self._shares = SharesAPI(self._api_client)
         self._statement_execution = StatementExecutionAPI(self._api_client)
@@ -267,6 +279,11 @@ class WorkspaceClient:
         """The alerts API can be used to perform CRUD operations on alerts."""
         return self._alerts
 
+    @property
+    def alerts_legacy(self) -> AlertsLegacyAPI:
+        """The alerts API can be used to perform CRUD operations on alerts."""
+        return self._alerts_legacy
+
     @property
     def apps(self) -> AppsAPI:
         """Apps run directly on a customer’s Databricks instance, integrate with their data, use and extend Databricks services, and enable users to interact through single sign-on."""
@@ -387,6 +404,11 @@ class WorkspaceClient:
         """Functions implement User-Defined Functions (UDFs) in Unity Catalog."""
         return self._functions
 
+    @property
+    def genie(self) -> GenieAPI:
+        """Genie provides a no-code experience for business users, powered by AI/BI."""
+        return self._genie
+
     @property
     def git_credentials(self) -> GitCredentialsAPI:
         """Registers personal access token for Databricks to do operations on behalf of the user."""
@@ -452,6 +474,11 @@ class WorkspaceClient:
         """Databricks provides a hosted version of MLflow Model Registry in Unity Catalog."""
         return self._model_versions
 
+    @property
+    def notification_destinations(self) -> NotificationDestinationsAPI:
+        """The notification destinations API lets you programmatically manage a workspace's notification destinations."""
+        return self._notification_destinations
+
     @property
     def online_tables(self) -> OnlineTablesAPI:
         """Online tables provide lower latency and higher QPS access to data from Delta tables."""
@@ -459,7 +486,7 @@ class WorkspaceClient:
 
     @property
     def permission_migration(self) -> PermissionMigrationAPI:
-        """This spec contains undocumented permission migration APIs used in https://github.com/databrickslabs/ucx."""
+        """APIs for migrating acl permissions, used only by the ucx tool: https://github.com/databrickslabs/ucx."""
         return self._permission_migration
 
     @property
@@ -524,19 +551,29 @@ class WorkspaceClient:
 
     @property
     def queries(self) -> QueriesAPI:
-        """These endpoints are used for CRUD operations on query definitions."""
+        """The queries API can be used to perform CRUD operations on queries."""
         return self._queries
 
+    @property
+    def queries_legacy(self) -> QueriesLegacyAPI:
+        """These endpoints are used for CRUD operations on query definitions."""
+        return self._queries_legacy
+
     @property
     def query_history(self) -> QueryHistoryAPI:
-        """Access the history of queries through SQL warehouses."""
+        """A service responsible for storing and retrieving the list of queries run against SQL endpoints, serverless compute, and DLT."""
         return self._query_history
 
     @property
     def query_visualizations(self) -> QueryVisualizationsAPI:
-        """This is an evolving API that facilitates the addition and removal of vizualisations from existing queries within the Databricks Workspace."""
+        """This is an evolving API that facilitates the addition and removal of visualizations from existing queries in the Databricks Workspace."""
         return self._query_visualizations
 
+    @property
+    def query_visualizations_legacy(self) -> QueryVisualizationsLegacyAPI:
+        """This is an evolving API that facilitates the addition and removal of vizualisations from existing queries within the Databricks Workspace."""
+        return self._query_visualizations_legacy
+
     @property
     def recipient_activation(self) -> RecipientActivationAPI:
         """The Recipient Activation API is only applicable in the open sharing model where the recipient object has the authentication type of `TOKEN`."""
@@ -577,6 +614,11 @@ class WorkspaceClient:
         """The Serving Endpoints API allows you to create, update, and delete model serving endpoints."""
         return self._serving_endpoints
 
+    @property
+    def serving_endpoints_data_plane(self) -> ServingEndpointsDataPlaneAPI:
+        """Serving endpoints DataPlane provides a set of operations to interact with data plane endpoints for Serving endpoints service."""
+        return self._serving_endpoints_data_plane
+
     @property
     def settings(self) -> SettingsAPI:
         """Workspace Settings API allows users to manage settings at the workspace level."""
@@ -734,7 +776,6 @@ class AccountClient:
         self._api_client = client.ApiClient(self._config)
         self._access_control = AccountAccessControlAPI(self._api_client)
         self._billable_usage = BillableUsageAPI(self._api_client)
-        self._budgets = BudgetsAPI(self._api_client)
         self._credentials = CredentialsAPI(self._api_client)
         self._custom_app_integration = CustomAppIntegrationAPI(self._api_client)
         self._encryption_keys = EncryptionKeysAPI(self._api_client)
@@ -753,10 +794,12 @@ class AccountClient:
         self._settings = AccountSettingsAPI(self._api_client)
         self._storage = StorageAPI(self._api_client)
         self._storage_credentials = AccountStorageCredentialsAPI(self._api_client)
+        self._usage_dashboards = UsageDashboardsAPI(self._api_client)
         self._users = AccountUsersAPI(self._api_client)
         self._vpc_endpoints = VpcEndpointsAPI(self._api_client)
         self._workspace_assignment = WorkspaceAssignmentAPI(self._api_client)
         self._workspaces = WorkspacesAPI(self._api_client)
+        self._budgets = BudgetsAPI(self._api_client)
 
     @property
     def config(self) -> client.Config:
@@ -776,11 +819,6 @@ class AccountClient:
         """This API allows you to download billable usage logs for the specified account and date range."""
         return self._billable_usage
 
-    @property
-    def budgets(self) -> BudgetsAPI:
-        """These APIs manage budget configuration including notifications for exceeding a budget for a period."""
-        return self._budgets
-
     @property
     def credentials(self) -> CredentialsAPI:
         """These APIs manage credential configurations for this workspace."""
@@ -788,7 +826,7 @@ class AccountClient:
 
     @property
     def custom_app_integration(self) -> CustomAppIntegrationAPI:
-        """These APIs enable administrators to manage custom oauth app integrations, which is required for adding/using Custom OAuth App Integration like Tableau Cloud for Databricks in AWS cloud."""
+        """These APIs enable administrators to manage custom OAuth app integrations, which is required for adding/using Custom OAuth App Integration like Tableau Cloud for Databricks in AWS cloud."""
         return self._custom_app_integration
 
     @property
@@ -843,7 +881,7 @@ class AccountClient:
 
     @property
     def published_app_integration(self) -> PublishedAppIntegrationAPI:
-        """These APIs enable administrators to manage published oauth app integrations, which is required for adding/using Published OAuth App Integration like Tableau Desktop for Databricks in AWS cloud."""
+        """These APIs enable administrators to manage published OAuth app integrations, which is required for adding/using Published OAuth App Integration like Tableau Desktop for Databricks in AWS cloud."""
         return self._published_app_integration
 
     @property
@@ -871,6 +909,11 @@ class AccountClient:
         """These APIs manage storage credentials for a particular metastore."""
         return self._storage_credentials
 
+    @property
+    def usage_dashboards(self) -> UsageDashboardsAPI:
+        """These APIs manage usage dashboards for this account."""
+        return self._usage_dashboards
+
     @property
     def users(self) -> AccountUsersAPI:
         """User identities recognized by Databricks and represented by email addresses."""
@@ -891,6 +934,11 @@ class AccountClient:
         """These APIs manage workspaces for this account."""
         return self._workspaces
 
+    @property
+    def budgets(self) -> BudgetsAPI:
+        """These APIs manage budget configurations for this account."""
+        return self._budgets
+
     def get_workspace_client(self, workspace: Workspace) -> WorkspaceClient:
         """Constructs a ``WorkspaceClient`` for the given workspace.
 

databricks/sdk/config.py

@@ -3,19 +3,18 @@ import copy
 import logging
 import os
 import pathlib
-import platform
 import sys
 import urllib.parse
-from typing import Dict, Iterable, List, Optional, Tuple
+from typing import Dict, Iterable, Optional
 
 import requests
 
+from . import useragent
 from .clock import Clock, RealClock
 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')
 
@@ -44,30 +43,14 @@ 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
+    useragent.with_product(product, product_version)
 
 
 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
+    useragent.with_extra(key, value)
 
 
 class Config:
@@ -111,21 +94,12 @@ class Config:
                  # Deprecated. Use credentials_strategy instead.
                  credentials_provider: CredentialsStrategy = None,
                  credentials_strategy: CredentialsStrategy = None,
-                 product=_DEFAULT_PRODUCT_NAME,
-                 product_version=_DEFAULT_PRODUCT_VERSION,
+                 product=None,
+                 product_version=None,
                  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 = []
         if credentials_strategy and credentials_provider:
             raise ValueError(
@@ -147,8 +121,7 @@ class Config:
             self._fix_host_if_needed()
             self._validate()
             self.init_auth()
-            self._product = product
-            self._product_version = product_version
+            self._init_product(product, product_version)
         except ValueError as e:
             message = self.wrap_debug_info(str(e))
             raise ValueError(message) from e
@@ -260,47 +233,19 @@ class Config:
     @property
     def user_agent(self):
         """ Returns User-Agent header used by this SDK """
-        py_version = platform.python_version()
-        os_name = platform.uname().system.lower()
-
-        ua = [
-            f"{self._product}/{self._product_version}", f"databricks-sdk-py/{__version__}",
-            f"python/{py_version}", f"os/{os_name}", f"auth/{self.auth_type}",
-        ]
-        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:
-            runtime_version = os.environ['DATABRICKS_RUNTIME_VERSION']
-            if runtime_version != '':
-                runtime_version = self._sanitize_header_value(runtime_version)
-                ua.append(f'runtime/{runtime_version}')
-
-        return ' '.join(ua)
 
-    @staticmethod
-    def _sanitize_header_value(value: str) -> str:
-        value = value.replace(' ', '-')
-        value = value.replace('/', '-')
-        return value
+        # global user agent includes SDK version, product name & version, platform info,
+        # and global extra info. Config can have specific extra info associated with it,
+        # such as an override product, auth type, and other user-defined information.
+        return useragent.to_string(self._product_info,
+                                   [("auth", self.auth_type)] + self._user_agent_other_info)
 
     @property
     def _upstream_user_agent(self) -> str:
-        product = os.environ.get('DATABRICKS_SDK_UPSTREAM', None)
-        product_version = os.environ.get('DATABRICKS_SDK_UPSTREAM_VERSION', None)
-        if product is not None and product_version is not None:
-            return f"upstream/{product} upstream-version/{product_version}"
-        return ""
+        return " ".join(f"{k}/{v}" for k, v in useragent._get_upstream_user_agent_info())
 
     def with_user_agent_extra(self, key: str, value: str) -> 'Config':
-        self._user_agent_other_info.append(f"{key}/{value}")
+        self._user_agent_other_info.append((key, value))
         return self
 
     @property
@@ -403,13 +348,47 @@ class Config:
     def _fix_host_if_needed(self):
         if not self.host:
             return
-        # fix url to remove trailing slash
+
+        # Add a default scheme if it's missing
+        if '://' not in self.host:
+            self.host = 'https://' + self.host
+
         o = urllib.parse.urlparse(self.host)
-        if not o.hostname:
-            # only hostname is specified
-            self.host = f"https://{self.host}"
-        else:
-            self.host = f"{o.scheme}://{o.netloc}"
+        # remove trailing slash
+        path = o.path.rstrip('/')
+        # remove port if 443
+        netloc = o.netloc
+        if o.port == 443:
+            netloc = netloc.split(':')[0]
+
+        self.host = urllib.parse.urlunparse((o.scheme, netloc, path, o.params, o.query, o.fragment))
+
+    def load_azure_tenant_id(self):
+        """[Internal] Load the Azure tenant ID from the Azure Databricks login page.
+
+        If the tenant ID is already set, this method does nothing."""
+        if not self.is_azure or self.azure_tenant_id is not None or self.host is None:
+            return
+        login_url = f'{self.host}/aad/auth'
+        logger.debug(f'Loading tenant ID from {login_url}')
+        resp = requests.get(login_url, allow_redirects=False)
+        if resp.status_code // 100 != 3:
+            logger.debug(
+                f'Failed to get tenant ID from {login_url}: expected status code 3xx, got {resp.status_code}')
+            return
+        entra_id_endpoint = resp.headers.get('Location')
+        if entra_id_endpoint is None:
+            logger.debug(f'No Location header in response from {login_url}')
+            return
+        # The Location header has the following form: https://login.microsoftonline.com/<tenant-id>/oauth2/authorize?...
+        # The domain may change depending on the Azure cloud (e.g. login.microsoftonline.us for US Government cloud).
+        url = urllib.parse.urlparse(entra_id_endpoint)
+        path_segments = url.path.split('/')
+        if len(path_segments) < 2:
+            logger.debug(f'Invalid path in Location header: {url.path}')
+            return
+        self.azure_tenant_id = path_segments[1]
+        logger.debug(f'Loaded tenant ID: {self.azure_tenant_id}')
 
     def _set_inner_config(self, keyword_args: Dict[str, any]):
         for attr in self.attributes():
@@ -498,6 +477,13 @@ class Config:
         except ValueError as e:
             raise ValueError(f'{self._credentials_strategy.auth_type()} auth: {e}') from e
 
+    def _init_product(self, product, product_version):
+        if product is not None or product_version is not None:
+            default_product, default_version = useragent.product()
+            self._product_info = (product or default_product, product_version or default_version)
+        else:
+            self._product_info = None
+
     def __repr__(self):
         return f'<{self.debug_string()}>'
 

databricks/sdk/core.py

@@ -133,31 +133,36 @@ class ApiClient:
 
     def do(self,
            method: str,
-           path: 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, 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)
         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)
         response = retryable(self._perform)(method,
-                                            path,
+                                            url,
                                             query=query,
                                             headers=headers,
                                             body=body,
                                             raw=raw,
                                             files=files,
-                                            data=data)
+                                            data=data,
+                                            auth=auth)
 
         resp = dict()
         for header in response_headers if response_headers else []:
@@ -239,20 +244,22 @@ class ApiClient:
 
     def _perform(self,
                  method: str,
-                 path: str,
+                 url: str,
                  query: dict = None,
                  headers: dict = None,
                  body: dict = None,
                  raw: bool = False,
                  files=None,
-                 data=None):
+                 data=None,
+                 auth: Callable[[requests.PreparedRequest], requests.PreparedRequest] = None):
         response = self._session.request(method,
-                                         f"{self._cfg.host}{path}",
+                                         url,
                                          params=self._fix_query_string(query),
                                          json=body,
                                          headers=headers,
                                          files=files,
                                          data=data,
+                                         auth=auth,
                                          stream=raw,
                                          timeout=self._http_timeout_seconds)
         try:

databricks/sdk/credentials_provider.py

@@ -233,8 +233,7 @@ def _ensure_host_present(cfg: 'Config', token_source_for: Callable[[str], TokenS
     cfg.host = f"https://{resp.json()['properties']['workspaceUrl']}"
 
 
-@oauth_credentials_strategy('azure-client-secret',
-                            ['is_azure', 'azure_client_id', 'azure_client_secret', 'azure_tenant_id'])
+@oauth_credentials_strategy('azure-client-secret', ['is_azure', 'azure_client_id', 'azure_client_secret'])
 def azure_service_principal(cfg: 'Config') -> CredentialsProvider:
     """ Adds refreshed Azure Active Directory (AAD) Service Principal OAuth tokens
     to every request, while automatically resolving different Azure environment endpoints. """
@@ -248,6 +247,7 @@ def azure_service_principal(cfg: 'Config') -> CredentialsProvider:
                                  use_params=True)
 
     _ensure_host_present(cfg, token_source_for)
+    cfg.load_azure_tenant_id()
     logger.info("Configured AAD token for Service Principal (%s)", cfg.azure_client_id)
     inner = token_source_for(cfg.effective_azure_login_app_id)
     cloud = token_source_for(cfg.arm_environment.service_management_endpoint)
@@ -432,11 +432,13 @@ class CliTokenSource(Refreshable):
 class AzureCliTokenSource(CliTokenSource):
     """ Obtain the token granted by `az login` CLI command """
 
-    def __init__(self, resource: str, subscription: str = ""):
+    def __init__(self, resource: str, subscription: Optional[str] = None, tenant: Optional[str] = None):
         cmd = ["az", "account", "get-access-token", "--resource", resource, "--output", "json"]
-        if subscription != "":
+        if subscription is not None:
             cmd.append("--subscription")
             cmd.append(subscription)
+        if tenant:
+            cmd.extend(["--tenant", tenant])
         super().__init__(cmd=cmd,
                          token_type_field='tokenType',
                          access_token_field='accessToken',
@@ -464,8 +466,10 @@ class AzureCliTokenSource(CliTokenSource):
     @staticmethod
     def for_resource(cfg: 'Config', resource: str) -> 'AzureCliTokenSource':
         subscription = AzureCliTokenSource.get_subscription(cfg)
-        if subscription != "":
-            token_source = AzureCliTokenSource(resource, subscription)
+        if subscription is not None:
+            token_source = AzureCliTokenSource(resource,
+                                               subscription=subscription,
+                                               tenant=cfg.azure_tenant_id)
             try:
                 # This will fail if the user has access to the workspace, but not to the subscription
                 # itself.
@@ -475,25 +479,26 @@ class AzureCliTokenSource(CliTokenSource):
             except OSError:
                 logger.warning("Failed to get token for subscription. Using resource only token.")
 
-        token_source = AzureCliTokenSource(resource)
+        token_source = AzureCliTokenSource(resource, subscription=None, tenant=cfg.azure_tenant_id)
         token_source.token()
         return token_source
 
     @staticmethod
-    def get_subscription(cfg: 'Config') -> str:
+    def get_subscription(cfg: 'Config') -> Optional[str]:
         resource = cfg.azure_workspace_resource_id
         if resource is None or resource == "":
-            return ""
+            return None
         components = resource.split('/')
         if len(components) < 3:
             logger.warning("Invalid azure workspace resource ID")
-            return ""
+            return None
         return components[2]
 
 
 @credentials_strategy('azure-cli', ['is_azure'])
 def azure_cli(cfg: 'Config') -> Optional[CredentialsProvider]:
     """ Adds refreshed OAuth token granted by `az login` command to every request. """
+    cfg.load_azure_tenant_id()
     token_source = None
     mgmt_token_source = None
     try:
@@ -517,11 +522,6 @@ def azure_cli(cfg: 'Config') -> Optional[CredentialsProvider]:
 
     _ensure_host_present(cfg, lambda resource: AzureCliTokenSource.for_resource(cfg, resource))
     logger.info("Using Azure CLI authentication with AAD tokens")
-    if not cfg.is_account_client and AzureCliTokenSource.get_subscription(cfg) == "":
-        logger.warning(
-            "azure_workspace_resource_id field not provided. "
-            "It is recommended to specify this field in the Databricks configuration to avoid authentication errors."
-        )
 
     def inner() -> Dict[str, str]:
         token = token_source.token()