diff --git a/databricks/sdk/core.py b/databricks/sdk/core.py index f8bd89eb0..f72addd9f 100644 --- a/databricks/sdk/core.py +++ b/databricks/sdk/core.py @@ -87,6 +87,17 @@ def inner() -> Dict[str, str]: return inner +@credentials_provider('runtime', []) +def runtime_native_auth(cfg: 'Config') -> Optional[HeaderFactory]: + from databricks.sdk.runtime import init_runtime_native_auth + try: + host, inner = init_runtime_native_auth() + cfg.host = host + return inner + except NotImplementedError: + return None + + @credentials_provider('oauth-m2m', ['is_aws', 'host', 'client_id', 'client_secret']) def oauth_service_principal(cfg: 'Config') -> Optional[HeaderFactory]: """ Adds refreshed Databricks machine-to-machine OAuth Bearer token to every request, @@ -309,7 +320,7 @@ def auth_type(self) -> str: def __call__(self, cfg: 'Config') -> HeaderFactory: auth_providers = [ pat_auth, basic_auth, oauth_service_principal, azure_service_principal, azure_cli, - external_browser, bricks_cli + external_browser, bricks_cli, runtime_native_auth ] for provider in auth_providers: auth_type = provider.auth_type() @@ -754,7 +765,11 @@ def _authenticate(self, r: requests.PreparedRequest) -> requests.PreparedRequest def do(self, method: str, path: str, query: dict = None, body: dict = None) -> dict: headers = {'Accept': 'application/json', 'User-Agent': self._user_agent_base} - response = self._session.request(method, f"{self._cfg.host}{path}", params=query, json=body, headers=headers) + response = self._session.request(method, + f"{self._cfg.host}{path}", + params=query, + json=body, + headers=headers) try: self._record_request_log(response) if not response.ok: diff --git a/databricks/sdk/runtime/__init__.py b/databricks/sdk/runtime/__init__.py index 2bfcc4939..e83c9d65c 100644 --- a/databricks/sdk/runtime/__init__.py +++ b/databricks/sdk/runtime/__init__.py @@ -1,3 +1,5 @@ +from typing import Callable, Dict, Tuple + is_local_implementation = True # All objects that are injected into the Notebook's user namespace should also be made @@ -7,6 +9,13 @@ "display", "displayHTML", "dbutils", "table", "sql", "udf", "getArgument", "sc", "sqlContext", "spark" ] +RuntimeAuth = Tuple[str, Callable[[], Dict[str, str]]] + + +def init_runtime_native_auth() -> RuntimeAuth: + raise NotImplementedError + + try: # Internal implementation from dbruntime import UserNamespaceInitializer