databricks-sdk 0.28.0__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
@@ -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,
@@ -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/config.py

@@ -6,15 +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 .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 .oauth import OidcEndpoints, Token
 from .version import __version__
 
 logger = logging.getLogger('databricks.sdk')
@@ -44,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')
@@ -66,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')
@@ -81,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']
@@ -107,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:
@@ -220,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:
@@ -436,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
 
@@ -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,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,

databricks/sdk/credentials_provider.py

@@ -22,12 +22,26 @@ from .azure import add_sp_management_token, add_workspace_id_header
 from .oauth import (ClientCredentials, OAuthClient, Refreshable, Token,
                     TokenCache, TokenSource)
 
-HeaderFactory = Callable[[], Dict[str, str]]
+CredentialsProvider = Callable[[], Dict[str, str]]
 
 logger = logging.getLogger('databricks.sdk')
 
 
-class CredentialsProvider(abc.ABC):
+class OAuthCredentialsProvider:
+    """ OAuthCredentialsProvider is a type of CredentialsProvider which exposes OAuth tokens. """
+
+    def __init__(self, credentials_provider: CredentialsProvider, token_provider: Callable[[], Token]):
+        self._credentials_provider = credentials_provider
+        self._token_provider = token_provider
+
+    def __call__(self) -> Dict[str, str]:
+        return self._credentials_provider()
+
+    def oauth_token(self) -> Token:
+        return self._token_provider()
+
+
+class CredentialsStrategy(abc.ABC):
     """ CredentialsProvider is the protocol (call-side interface)
      for authenticating requests to Databricks REST APIs"""
 
@@ -36,20 +50,39 @@ class CredentialsProvider(abc.ABC):
         ...
 
     @abc.abstractmethod
-    def __call__(self, cfg: 'Config') -> HeaderFactory:
+    def __call__(self, cfg: 'Config') -> CredentialsProvider:
         ...
 
 
-def credentials_provider(name: str, require: List[str]):
+class OauthCredentialsStrategy(CredentialsStrategy):
+    """ OauthCredentialsProvider is a CredentialsProvider which
+    supports Oauth tokens"""
+
+    def __init__(self, auth_type: str, headers_provider: Callable[['Config'], OAuthCredentialsProvider]):
+        self._headers_provider = headers_provider
+        self._auth_type = auth_type
+
+    def auth_type(self) -> str:
+        return self._auth_type
+
+    def __call__(self, cfg: 'Config') -> OAuthCredentialsProvider:
+        return self._headers_provider(cfg)
+
+    def oauth_token(self, cfg: 'Config') -> Token:
+        return self._headers_provider(cfg).oauth_token()
+
+
+def credentials_strategy(name: str, require: List[str]):
     """ Given the function that receives a Config and returns RequestVisitor,
     create CredentialsProvider with a given name and required configuration
     attribute names to be present for this function to be called. """
 
-    def inner(func: Callable[['Config'], HeaderFactory]) -> CredentialsProvider:
+    def inner(func: Callable[['Config'], CredentialsProvider]) -> CredentialsStrategy:
 
         @functools.wraps(func)
-        def wrapper(cfg: 'Config') -> Optional[HeaderFactory]:
+        def wrapper(cfg: 'Config') -> Optional[CredentialsProvider]:
             for attr in require:
+                getattr(cfg, attr)
                 if not getattr(cfg, attr):
                     return None
             return func(cfg)
@@ -60,8 +93,27 @@ def credentials_provider(name: str, require: List[str]):
     return inner
 
 
-@credentials_provider('basic', ['host', 'username', 'password'])
-def basic_auth(cfg: 'Config') -> HeaderFactory:
+def oauth_credentials_strategy(name: str, require: List[str]):
+    """ Given the function that receives a Config and returns an OauthHeaderFactory,
+    create an OauthCredentialsProvider with a given name and required configuration
+    attribute names to be present for this function to be called. """
+
+    def inner(func: Callable[['Config'], OAuthCredentialsProvider]) -> OauthCredentialsStrategy:
+
+        @functools.wraps(func)
+        def wrapper(cfg: 'Config') -> Optional[OAuthCredentialsProvider]:
+            for attr in require:
+                if not getattr(cfg, attr):
+                    return None
+            return func(cfg)
+
+        return OauthCredentialsStrategy(name, wrapper)
+
+    return inner
+
+
+@credentials_strategy('basic', ['host', 'username', 'password'])
+def basic_auth(cfg: 'Config') -> CredentialsProvider:
     """ Given username and password, add base64-encoded Basic credentials """
     encoded = base64.b64encode(f'{cfg.username}:{cfg.password}'.encode()).decode()
     static_credentials = {'Authorization': f'Basic {encoded}'}
@@ -72,8 +124,8 @@ def basic_auth(cfg: 'Config') -> HeaderFactory:
     return inner
 
 
-@credentials_provider('pat', ['host', 'token'])
-def pat_auth(cfg: 'Config') -> HeaderFactory:
+@credentials_strategy('pat', ['host', 'token'])
+def pat_auth(cfg: 'Config') -> CredentialsProvider:
     """ Adds Databricks Personal Access Token to every request """
     static_credentials = {'Authorization': f'Bearer {cfg.token}'}
 
@@ -83,8 +135,8 @@ def pat_auth(cfg: 'Config') -> HeaderFactory:
     return inner
 
 
-@credentials_provider('runtime', [])
-def runtime_native_auth(cfg: 'Config') -> Optional[HeaderFactory]:
+@credentials_strategy('runtime', [])
+def runtime_native_auth(cfg: 'Config') -> Optional[CredentialsProvider]:
     if 'DATABRICKS_RUNTIME_VERSION' not in os.environ:
         return None
 
@@ -107,8 +159,8 @@ def runtime_native_auth(cfg: 'Config') -> Optional[HeaderFactory]:
     return None
 
 
-@credentials_provider('oauth-m2m', ['host', 'client_id', 'client_secret'])
-def oauth_service_principal(cfg: 'Config') -> Optional[HeaderFactory]:
+@oauth_credentials_strategy('oauth-m2m', ['host', 'client_id', 'client_secret'])
+def oauth_service_principal(cfg: 'Config') -> Optional[CredentialsProvider]:
     """ Adds refreshed Databricks machine-to-machine OAuth Bearer token to every request,
     if /oidc/.well-known/oauth-authorization-server is available on the given host. """
     oidc = cfg.oidc_endpoints
@@ -124,11 +176,14 @@ def oauth_service_principal(cfg: 'Config') -> Optional[HeaderFactory]:
         token = token_source.token()
         return {'Authorization': f'{token.token_type} {token.access_token}'}
 
-    return inner
+    def token() -> Token:
+        return token_source.token()
+
+    return OAuthCredentialsProvider(inner, token)
 
 
-@credentials_provider('external-browser', ['host', 'auth_type'])
-def external_browser(cfg: 'Config') -> Optional[HeaderFactory]:
+@credentials_strategy('external-browser', ['host', 'auth_type'])
+def external_browser(cfg: 'Config') -> Optional[CredentialsProvider]:
     if cfg.auth_type != 'external-browser':
         return None
     if cfg.client_id:
@@ -178,9 +233,9 @@ def _ensure_host_present(cfg: 'Config', token_source_for: Callable[[str], TokenS
     cfg.host = f"https://{resp.json()['properties']['workspaceUrl']}"
 
 
-@credentials_provider('azure-client-secret',
-                      ['is_azure', 'azure_client_id', 'azure_client_secret', 'azure_tenant_id'])
-def azure_service_principal(cfg: 'Config') -> HeaderFactory:
+@oauth_credentials_strategy('azure-client-secret',
+                            ['is_azure', 'azure_client_id', 'azure_client_secret', 'azure_tenant_id'])
+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. """
 
@@ -203,11 +258,14 @@ def azure_service_principal(cfg: 'Config') -> HeaderFactory:
         add_sp_management_token(cloud, headers)
         return headers
 
-    return refreshed_headers
+    def token() -> Token:
+        return inner.token()
+
+    return OAuthCredentialsProvider(refreshed_headers, token)
 
 
-@credentials_provider('github-oidc-azure', ['host', 'azure_client_id'])
-def github_oidc_azure(cfg: 'Config') -> Optional[HeaderFactory]:
+@oauth_credentials_strategy('github-oidc-azure', ['host', 'azure_client_id'])
+def github_oidc_azure(cfg: 'Config') -> Optional[CredentialsProvider]:
     if 'ACTIONS_ID_TOKEN_REQUEST_TOKEN' not in os.environ:
         # not in GitHub actions
         return None
@@ -250,14 +308,17 @@ def github_oidc_azure(cfg: 'Config') -> Optional[HeaderFactory]:
         token = inner.token()
         return {'Authorization': f'{token.token_type} {token.access_token}'}
 
-    return refreshed_headers
+    def token() -> Token:
+        return inner.token()
+
+    return OAuthCredentialsProvider(refreshed_headers, token)
 
 
 GcpScopes = ["https://www.googleapis.com/auth/cloud-platform", "https://www.googleapis.com/auth/compute"]
 
 
-@credentials_provider('google-credentials', ['host', 'google_credentials'])
-def google_credentials(cfg: 'Config') -> Optional[HeaderFactory]:
+@oauth_credentials_strategy('google-credentials', ['host', 'google_credentials'])
+def google_credentials(cfg: 'Config') -> Optional[CredentialsProvider]:
     if not cfg.is_gcp:
         return None
     # Reads credentials as JSON. Credentials can be either a path to JSON file, or actual JSON string.
@@ -277,6 +338,10 @@ def google_credentials(cfg: 'Config') -> Optional[HeaderFactory]:
     gcp_credentials = service_account.Credentials.from_service_account_info(info=account_info,
                                                                             scopes=GcpScopes)
 
+    def token() -> Token:
+        credentials.refresh(request)
+        return credentials.token
+
     def refreshed_headers() -> Dict[str, str]:
         credentials.refresh(request)
         headers = {'Authorization': f'Bearer {credentials.token}'}
@@ -285,11 +350,11 @@ def google_credentials(cfg: 'Config') -> Optional[HeaderFactory]:
             headers["X-Databricks-GCP-SA-Access-Token"] = gcp_credentials.token
         return headers
 
-    return refreshed_headers
+    return OAuthCredentialsProvider(refreshed_headers, token)
 
 
-@credentials_provider('google-id', ['host', 'google_service_account'])
-def google_id(cfg: 'Config') -> Optional[HeaderFactory]:
+@oauth_credentials_strategy('google-id', ['host', 'google_service_account'])
+def google_id(cfg: 'Config') -> Optional[CredentialsProvider]:
     if not cfg.is_gcp:
         return None
     credentials, _project_id = google.auth.default()
@@ -309,6 +374,10 @@ def google_id(cfg: 'Config') -> Optional[HeaderFactory]:
 
     request = Request()
 
+    def token() -> Token:
+        id_creds.refresh(request)
+        return id_creds.token
+
     def refreshed_headers() -> Dict[str, str]:
         id_creds.refresh(request)
         headers = {'Authorization': f'Bearer {id_creds.token}'}
@@ -317,7 +386,7 @@ def google_id(cfg: 'Config') -> Optional[HeaderFactory]:
             headers["X-Databricks-GCP-SA-Access-Token"] = gcp_impersonated_credentials.token
         return headers
 
-    return refreshed_headers
+    return OAuthCredentialsProvider(refreshed_headers, token)
 
 
 class CliTokenSource(Refreshable):
@@ -422,8 +491,8 @@ class AzureCliTokenSource(CliTokenSource):
         return components[2]
 
 
-@credentials_provider('azure-cli', ['is_azure'])
-def azure_cli(cfg: 'Config') -> Optional[HeaderFactory]:
+@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. """
     token_source = None
     mgmt_token_source = None
@@ -516,8 +585,8 @@ class DatabricksCliTokenSource(CliTokenSource):
         raise err
 
 
-@credentials_provider('databricks-cli', ['host'])
-def databricks_cli(cfg: 'Config') -> Optional[HeaderFactory]:
+@oauth_credentials_strategy('databricks-cli', ['host'])
+def databricks_cli(cfg: 'Config') -> Optional[CredentialsProvider]:
     try:
         token_source = DatabricksCliTokenSource(cfg)
     except FileNotFoundError as e:
@@ -538,7 +607,7 @@ def databricks_cli(cfg: 'Config') -> Optional[HeaderFactory]:
         token = token_source.token()
         return {'Authorization': f'{token.token_type} {token.access_token}'}
 
-    return inner
+    return OAuthCredentialsProvider(inner, token_source.token)
 
 
 class MetadataServiceTokenSource(Refreshable):
@@ -577,8 +646,8 @@ class MetadataServiceTokenSource(Refreshable):
         return Token(access_token=access_token, token_type=token_type, expiry=expiry)
 
 
-@credentials_provider('metadata-service', ['host', 'metadata_service_url'])
-def metadata_service(cfg: 'Config') -> Optional[HeaderFactory]:
+@credentials_strategy('metadata-service', ['host', 'metadata_service_url'])
+def metadata_service(cfg: 'Config') -> Optional[CredentialsProvider]:
     """ Adds refreshed token granted by Databricks Metadata Service to every request. """
 
     token_source = MetadataServiceTokenSource(cfg)
@@ -597,17 +666,25 @@ class DefaultCredentials:
 
     def __init__(self) -> None:
         self._auth_type = 'default'
-
-    def auth_type(self) -> str:
-        return self._auth_type
-
-    def __call__(self, cfg: 'Config') -> HeaderFactory:
-        auth_providers = [
+        self._auth_providers = [
             pat_auth, basic_auth, metadata_service, oauth_service_principal, azure_service_principal,
             github_oidc_azure, azure_cli, external_browser, databricks_cli, runtime_native_auth,
             google_credentials, google_id
         ]
-        for provider in auth_providers:
+
+    def auth_type(self) -> str:
+        return self._auth_type
+
+    def oauth_token(self, cfg: 'Config') -> Token:
+        for provider in self._auth_providers:
+            auth_type = provider.auth_type()
+            if auth_type != self._auth_type:
+                # ignore other auth types if they don't match the selected one
+                continue
+            return provider.oauth_token(cfg)
+
+    def __call__(self, cfg: 'Config') -> CredentialsProvider:
+        for provider in self._auth_providers:
             auth_type = provider.auth_type()
             if cfg.auth_type and auth_type != cfg.auth_type:
                 # ignore other auth types if one is explicitly enforced

databricks/sdk/dbutils.py

@@ -1,10 +1,11 @@
 import base64
 import json
 import logging
-import os.path
+import os
 import threading
 from collections import namedtuple
-from typing import Callable, Dict, List
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional
 
 from .core import ApiClient, Config, DatabricksError
 from .mixins import compute as compute_ext
@@ -240,6 +241,76 @@ class RemoteDbUtils:
                           name=util)
 
 
+@dataclass
+class OverrideResult:
+    result: Any
+
+
+def get_local_notebook_path():
+    value = os.getenv("DATABRICKS_SOURCE_FILE")
+    if value is None:
+        raise ValueError(
+            "Getting the current notebook path is only supported when running a notebook using the `Databricks Connect: Run as File` or `Databricks Connect: Debug as File` commands in the Databricks extension for VS Code. To bypass this error, set environment variable `DATABRICKS_SOURCE_FILE` to the desired notebook path."
+        )
+
+    return value
+
+
+class _OverrideProxyUtil:
+
+    @classmethod
+    def new(cls, path: str):
+        if len(cls.__get_matching_overrides(path)) > 0:
+            return _OverrideProxyUtil(path)
+        return None
+
+    def __init__(self, name: str):
+        self._name = name
+
+    # These are the paths that we want to override and not send to remote dbutils. NOTE, for each of these paths, no prefixes
+    # are sent to remote either. This could lead to unintentional breakage.
+    # Our current proxy implementation (which sends everything to remote dbutils) uses `{util}.{method}(*args, **kwargs)` ONLY.
+    # This means, it is completely safe to override paths starting with `{util}.{attribute}.<other_parts>`, since none of the prefixes
+    # are being proxied to remote dbutils currently.
+    proxy_override_paths = {
+        'notebook.entry_point.getDbutils().notebook().getContext().notebookPath().get()':
+        get_local_notebook_path,
+    }
+
+    @classmethod
+    def __get_matching_overrides(cls, path: str):
+        return [x for x in cls.proxy_override_paths.keys() if x.startswith(path)]
+
+    def __run_override(self, path: str) -> Optional[OverrideResult]:
+        overrides = self.__get_matching_overrides(path)
+        if len(overrides) == 1 and overrides[0] == path:
+            return OverrideResult(self.proxy_override_paths[overrides[0]]())
+
+        if len(overrides) > 0:
+            return OverrideResult(_OverrideProxyUtil(name=path))
+
+        return None
+
+    def __call__(self, *args, **kwds) -> Any:
+        if len(args) != 0 or len(kwds) != 0:
+            raise TypeError(
+                f"Arguments are not supported for overridden method {self._name}. Invoke as: {self._name}()")
+
+        callable_path = f"{self._name}()"
+        result = self.__run_override(callable_path)
+        if result:
+            return result.result
+
+        raise TypeError(f"{self._name} is not callable")
+
+    def __getattr__(self, method: str) -> Any:
+        result = self.__run_override(f"{self._name}.{method}")
+        if result:
+            return result.result
+
+        raise AttributeError(f"module {self._name} has no attribute {method}")
+
+
 class _ProxyUtil:
     """Enables temporary workaround to call remote in-REPL dbutils without having to re-implement them"""
 
@@ -250,7 +321,14 @@ class _ProxyUtil:
         self._context_factory = context_factory
         self._name = name
 
-    def __getattr__(self, method: str) -> '_ProxyCall':
+    def __call__(self):
+        raise NotImplementedError(f"dbutils.{self._name} is not callable")
+
+    def __getattr__(self, method: str) -> '_ProxyCall | _ProxyUtil | _OverrideProxyUtil':
+        override = _OverrideProxyUtil.new(f"{self._name}.{method}")
+        if override:
+            return override
+
         return _ProxyCall(command_execution=self._commands,
                           cluster_id=self._cluster_id,
                           context_factory=self._context_factory,