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

databricks/sdk/__init__.py

@@ -1,12 +1,13 @@
 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
+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)
@@ -131,7 +136,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 +158,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,
@@ -160,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)
@@ -186,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)
@@ -199,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)
@@ -215,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)
@@ -224,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)
@@ -265,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."""
@@ -385,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."""
@@ -450,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."""
@@ -457,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
@@ -522,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`."""
@@ -575,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."""
@@ -700,7 +744,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 +766,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,
@@ -730,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)
@@ -749,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:
@@ -772,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."""
@@ -784,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
@@ -839,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
@@ -867,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."""
@@ -887,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, Optional
 
 import requests
 
+from . import useragent
 from .clock import Clock, RealClock
-from .credentials_provider import CredentialsProvider, DefaultCredentials
+from .credentials_provider import CredentialsStrategy, DefaultCredentials
 from .environments import (ALL_ENVS, AzureEnvironment, Cloud,
                            DatabricksEnvironment, get_environment_for_hostname)
-from .oauth import OidcEndpoints
-from .version import __version__
+from .oauth import OidcEndpoints, Token
 
 logger = logging.getLogger('databricks.sdk')
 
@@ -44,6 +43,16 @@ class ConfigAttribute:
         return f"<ConfigAttribute '{self.name}' {self.transform.__name__}>"
 
 
+def with_product(product: str, product_version: str):
+    """[INTERNAL API] Change the product name and version used in the User-Agent header."""
+    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."""
+    useragent.with_extra(key, value)
+
+
 class Config:
     host: str = ConfigAttribute(env='DATABRICKS_HOST')
     account_id: str = ConfigAttribute(env='DATABRICKS_ACCOUNT_ID')
@@ -66,6 +75,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')
@@ -81,15 +91,25 @@ 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=None,
+                 product_version=None,
                  clock: Clock = None,
                  **kwargs):
         self._header_factory = None
         self._inner = {}
         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']
@@ -101,12 +121,14 @@ 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
 
+    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:
@@ -211,41 +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))
-        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
@@ -348,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():
@@ -436,12 +470,19 @@ 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 _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

@@ -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
 
@@ -13,12 +14,17 @@ from .config import *
 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
@@ -109,33 +115,54 @@ 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,
+           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 []:
@@ -217,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: