Add new credential provider

Author: aravind-segu

databricks/sdk/__init__.py

@@ -113,6 +113,8 @@ class Config:
     disable_experimental_files_api_client: bool = ConfigAttribute(
         env="DATABRICKS_DISABLE_EXPERIMENTAL_FILES_API_CLIENT"
     )
+    scopes: str = ConfigAttribute()
+    authorization_details: str = ConfigAttribute()
 
     files_ext_client_download_streaming_chunk_size: int = 2 * 1024 * 1024  # 2 MiB
 

databricks/sdk/config.py

@@ -176,6 +176,42 @@ def runtime_native_auth(cfg: "Config") -> Optional[CredentialsProvider]:
     return None
 
 
+@oauth_credentials_strategy("runtime-oauth", ["scopes"])
+def runtime_oauth(cfg: "Config") -> Optional[CredentialsProvider]:
+    if "DATABRICKS_RUNTIME_VERSION" not in os.environ:
+        return None
+
+    def get_notebook_pat_token() -> Optional[str]:
+        native_auth = runtime_native_auth(cfg)
+        if native_auth is None:
+            return None
+        notebook_pat_token = None
+        notebook_pat_authorization = native_auth().get("Authorization", "").strip()
+        if notebook_pat_authorization.lower().startswith("bearer "):
+            notebook_pat_token = notebook_pat_authorization[len("bearer ") :].strip()
+        return notebook_pat_token
+
+    notebook_pat_token = get_notebook_pat_token()
+    if notebook_pat_token is None:
+        return None
+
+    token_source = oauth.PATOAuthTokenExchange(
+        get_original_token=get_notebook_pat_token,
+        host=cfg.host,
+        scopes=cfg.scopes,
+        authorization_details=cfg.authorization_details,
+    )
+
+    def inner() -> Dict[str, str]:
+        token = token_source.token()
+        return {"Authorization": f"{token.token_type} {token.access_token}"}
+
+    def token() -> oauth.Token:
+        return token_source.token()
+
+    return OAuthCredentialsProvider(inner, token)
+
+
 @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,
@@ -189,9 +225,10 @@ def oauth_service_principal(cfg: "Config") -> Optional[CredentialsProvider]:
         client_id=cfg.client_id,
         client_secret=cfg.client_secret,
         token_url=oidc.token_endpoint,
-        scopes=["all-apis"],
+        scopes=cfg.scopes or "all-apis",
         use_header=True,
         disable_async=cfg.disable_async_token_refresh,
+        authorization_details=cfg.authorization_details,
     )
 
     def inner() -> Dict[str, str]:
@@ -292,6 +329,8 @@ def token_source_for(resource: str) -> oauth.TokenSource:
             endpoint_params={"resource": resource},
             use_params=True,
             disable_async=cfg.disable_async_token_refresh,
+            scopes=cfg.scopes,
+            authorization_details=cfg.authorization_details,
         )
 
     _ensure_host_present(cfg, token_source_for)
@@ -411,9 +450,10 @@ def token_source_for(audience: str) -> oauth.TokenSource:
                 "subject_token": id_token,
                 "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
             },
-            scopes=["all-apis"],
+            scopes=cfg.scopes or "all-apis",
             use_params=True,
             disable_async=cfg.disable_async_token_refresh,
+            authorization_details=cfg.authorization_details,
         )
 
     def refreshed_headers() -> Dict[str, str]:
@@ -493,6 +533,8 @@ def github_oidc_azure(cfg: "Config") -> Optional[CredentialsProvider]:
         },
         use_params=True,
         disable_async=cfg.disable_async_token_refresh,
+        scopes=cfg.scopes,
+        authorization_details=cfg.authorization_details,
     )
 
     def refreshed_headers() -> Dict[str, str]:
@@ -1070,6 +1112,7 @@ def __init__(self) -> None:
             azure_devops_oidc,
             external_browser,
             databricks_cli,
+            runtime_oauth,
             runtime_native_auth,
             google_credentials,
             google_id,

databricks/sdk/credentials_provider.py

@@ -14,7 +14,7 @@
 from datetime import datetime, timedelta
 from enum import Enum
 from http.server import BaseHTTPRequestHandler, HTTPServer
-from typing import Any, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional
 
 import requests
 import requests.auth
@@ -32,6 +32,30 @@
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class AuthorizationDetail:
+    type: str
+    object_type: str
+    object_path: str
+    actions: list[str]
+
+    def as_dict(self) -> dict:
+        return {
+            "type": self.type,
+            "object_type": self.object_type,
+            "object_path": self.object_path,
+            "actions": self.actions,
+        }
+
+    def from_dict(self, d: dict) -> "AuthorizationDetail":
+        return AuthorizationDetail(
+            type=d.get("type"),
+            object_type=d.get("object_type"),
+            object_path=d.get("object_path"),
+            actions=d.get("actions"),
+        )
+
+
 class IgnoreNetrcAuth(requests.auth.AuthBase):
     """This auth method is a no-op.
 
@@ -706,18 +730,21 @@ class ClientCredentials(Refreshable):
     client_secret: str
     token_url: str
     endpoint_params: dict = None
-    scopes: List[str] = None
+    scopes: str = None
     use_params: bool = False
     use_header: bool = False
     disable_async: bool = True
+    authorization_details: str = None
 
     def __post_init__(self):
         super().__init__(disable_async=self.disable_async)
 
     def refresh(self) -> Token:
         params = {"grant_type": "client_credentials"}
         if self.scopes:
-            params["scope"] = " ".join(self.scopes)
+            params["scope"] = self.scopes
+        if self.authorization_details:
+            params["authorization_details"] = self.authorization_details
         if self.endpoint_params:
             for k, v in self.endpoint_params.items():
                 params[k] = v
@@ -731,6 +758,51 @@ def refresh(self) -> Token:
         )
 
 
+@dataclass
+class PATOAuthTokenExchange(Refreshable):
+    get_original_token: Callable[[], Optional[str]]
+    host: str
+    scopes: str
+    authorization_details: str = None
+    disable_async: bool = True
+
+    def __post_init__(self):
+        super().__init__(disable_async=self.disable_async)
+
+    def refresh(self) -> Token:
+        token_exchange_url = f"{self.host}/oidc/v1/token"
+        params = {
+            "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
+            "subject_token": self.get_original_token(),
+            "subject_token_type": "urn:databricks:params:oauth:token-type:personal-access-token",
+            "requested_token_type": "urn:ietf:params:oauth:token-type:access_token",
+            "scope": self.scopes,
+        }
+        if self.authorization_details:
+            params["authorization_details"] = self.authorization_details
+
+        resp = requests.post(token_exchange_url, params)
+        if not resp.ok:
+            if resp.headers["Content-Type"].startswith("application/json"):
+                err = resp.json()
+                code = err.get("errorCode", err.get("error", "unknown"))
+                summary = err.get("errorSummary", err.get("error_description", "unknown"))
+                summary = summary.replace("\r\n", " ")
+                raise ValueError(f"{code}: {summary}")
+            raise ValueError(resp.content)
+        try:
+            j = resp.json()
+            expires_in = int(j["expires_in"])
+            expiry = datetime.now() + timedelta(seconds=expires_in)
+            return Token(
+                access_token=j["access_token"],
+                expiry=expiry,
+                token_type=j["token_type"],
+            )
+        except Exception as e:
+            raise ValueError(f"Failed to exchange PAT for OAuth token: {e}")
+
+
 class TokenCache:
     BASE_PATH = "~/.config/databricks-sdk-py/oauth"
 

databricks/sdk/oauth.py

@@ -202,7 +202,7 @@ def _exchange_id_token(self, id_token: IdToken) -> oauth.Token:
                 "subject_token": id_token.jwt,
                 "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
             },
-            scopes=["all-apis"],
+            scopes="all-apis",
             use_params=True,
             disable_async=self._disable_async,
         )