formatted

Author: madhav-db

src/databricks/sql/auth/auth.py

@@ -15,6 +15,7 @@ class AuthType(Enum):
     AZURE_OAUTH = "azure-oauth"
     CLIENT_CREDENTIALS = "client-credentials"
 
+
 class ClientContext:
     def __init__(
         self,
@@ -48,31 +49,42 @@ def __init__(
         self.oauth_client_secret = oauth_client_secret
         self.tenant_id = tenant_id
 
-def _create_azure_client_credentials_provider(cfg: ClientContext) -> ClientCredentialsProvider:
+
+def _create_azure_client_credentials_provider(
+    cfg: ClientContext,
+) -> ClientCredentialsProvider:
     """Create an Azure client credentials provider."""
     if not cfg.oauth_client_id or not cfg.oauth_client_secret or not cfg.tenant_id:
-        raise ValueError("Azure client credentials flow requires oauth_client_id, oauth_client_secret, and tenant_id")
-    
-    token_endpoint = "https://login.microsoftonline.com/{}/oauth2/v2.0/token".format(cfg.tenant_id)
+        raise ValueError(
+            "Azure client credentials flow requires oauth_client_id, oauth_client_secret, and tenant_id"
+        )
+
+    token_endpoint = "https://login.microsoftonline.com/{}/oauth2/v2.0/token".format(
+        cfg.tenant_id
+    )
     return ClientCredentialsProvider(
         client_id=cfg.oauth_client_id,
         client_secret=cfg.oauth_client_secret,
         token_endpoint=token_endpoint,
-        auth_type_value="azure-client-credentials"
+        auth_type_value="azure-client-credentials",
     )
 
 
-def _create_databricks_client_credentials_provider(cfg: ClientContext) -> ClientCredentialsProvider:
+def _create_databricks_client_credentials_provider(
+    cfg: ClientContext,
+) -> ClientCredentialsProvider:
     """Create a Databricks client credentials provider for service principals."""
     if not cfg.oauth_client_id or not cfg.oauth_client_secret:
-        raise ValueError("Databricks client credentials flow requires oauth_client_id and oauth_client_secret")
-    
+        raise ValueError(
+            "Databricks client credentials flow requires oauth_client_id and oauth_client_secret"
+        )
+
     token_endpoint = "{}oidc/v1/token".format(cfg.hostname)
     return ClientCredentialsProvider(
         client_id=cfg.oauth_client_id,
         client_secret=cfg.oauth_client_secret,
         token_endpoint=token_endpoint,
-        auth_type_value="client-credentials"
+        auth_type_value="client-credentials",
     )
 
 
@@ -102,9 +114,9 @@ def get_auth_provider(cfg: ClientContext):
         RuntimeError: If no valid authentication settings are provided
     """
     from databricks.sql.auth.token_federation import DatabricksTokenFederationProvider
-    
+
     base_provider = None
-    
+
     if cfg.credentials_provider:
         base_provider = ExternalAuthProvider(cfg.credentials_provider)
     elif cfg.access_token is not None:

src/databricks/sql/auth/authenticators.py

@@ -11,7 +11,11 @@
 # Private API: this is an evolving interface and it will change in the future.
 # Please must not depend on it in your applications.
 from databricks.sql.experimental.oauth_persistence import OAuthToken, OAuthPersistence
-from databricks.sql.auth.endpoint import AzureOAuthEndpointCollection, InHouseOAuthEndpointCollection
+from databricks.sql.auth.endpoint import (
+    AzureOAuthEndpointCollection,
+    InHouseOAuthEndpointCollection,
+)
+
 
 class AuthProvider:
     def add_headers(self, request_headers: Dict[str, str]):
@@ -56,8 +60,10 @@ def auth_type(self) -> str:
     def __call__(self, *args, **kwargs) -> HeaderFactory:
         def get_headers():
             return {"Authorization": self.__authorization_header_value}
+
         return get_headers
 
+
 # Private API: this is an evolving interface and it will change in the future.
 # Please must not depend on it in your applications.
 class DatabricksOAuthProvider(AuthProvider, CredentialsProvider):
@@ -81,11 +87,8 @@ def __init__(
 
         idp_endpoint = get_oauth_endpoints(hostname, auth_type == "azure-oauth")
         if not idp_endpoint:
-            raise NotImplementedError(
-                    f"OAuth is not supported for host ${hostname}"
-                )
+            raise NotImplementedError(f"OAuth is not supported for host ${hostname}")
 
-        
         cloud_scopes = idp_endpoint.get_scopes_mapping(scopes)
         self._scopes_as_str = self.SCOPE_DELIM.join(cloud_scopes)
 
@@ -107,6 +110,7 @@ def __call__(self, *args, **kwargs) -> HeaderFactory:
         def get_headers():
             self._update_token_if_expired()
             return {"Authorization": "Bearer {}".format(self._access_token)}
+
         return get_headers
 
     def _initial_get_token(self):
@@ -170,25 +174,24 @@ def __init__(
         client_id: str,
         client_secret: str,
         token_endpoint: str,
-        auth_type_value: str = "client-credentials"
+        auth_type_value: str = "client-credentials",
     ):
         """
         Initialize a ClientCredentialsProvider.
-        
+
         Args:
             client_id: OAuth client ID
-            client_secret: OAuth client secret  
+            client_secret: OAuth client secret
             token_endpoint: OAuth token endpoint URL
             auth_type_value: Auth type identifier
         """
         self.client_id = client_id
         self.client_secret = client_secret
         self.token_endpoint = token_endpoint
         self.auth_type_value = auth_type_value
-        
+
         self._cached_token = None
         self._token_expires_at = None
-        
 
     def auth_type(self) -> str:
         return self.auth_type_value
@@ -197,50 +200,54 @@ def __call__(self, *args, **kwargs) -> HeaderFactory:
         def get_headers() -> Dict[str, str]:
             token = self._get_access_token()
             return {"Authorization": "Bearer {}".format(token)}
+
         return get_headers
-    
+
     def add_headers(self, request_headers: Dict[str, str]):
         token = self._get_access_token()
         request_headers["Authorization"] = "Bearer {}".format(token)
 
     def _get_access_token(self) -> str:
         """Get a valid access token using client credentials flow, with caching."""
         # Check if we have a valid cached token (with 40 second buffer since azure doesn't respect a token with less than 30s expiry)
-        if (self._cached_token and self._token_expires_at and 
-            time.time() < self._token_expires_at - 40):
+        if (
+            self._cached_token
+            and self._token_expires_at
+            and time.time() < self._token_expires_at - 40
+        ):
             return self._cached_token
-        
+
         # Get new token using client credentials flow
         token_data = self._request_token()
-        
-        self._cached_token = token_data['access_token']
+
+        self._cached_token = token_data["access_token"]
         # expires_in is in seconds, convert to absolute time
-        self._token_expires_at = time.time() + token_data.get('expires_in', 3600)
-        
+        self._token_expires_at = time.time() + token_data.get("expires_in", 3600)
+
         return self._cached_token
 
     def _request_token(self) -> dict:
         """Request a new token using OAuth client credentials flow."""
         data = {
-            'grant_type': 'client_credentials',
-            'client_id': self.client_id,
-            'client_secret': self.client_secret,
-            'scope': self.AZURE_DATABRICKS_SCOPE,
+            "grant_type": "client_credentials",
+            "client_id": self.client_id,
+            "client_secret": self.client_secret,
+            "scope": self.AZURE_DATABRICKS_SCOPE,
         }
-        
-        headers = {'Content-Type': 'application/x-www-form-urlencoded'}
-        
+
+        headers = {"Content-Type": "application/x-www-form-urlencoded"}
+
         try:
             response = requests.post(self.token_endpoint, data=data, headers=headers)
             response.raise_for_status()
-            
+
             token_data = response.json()
-            
-            if 'access_token' not in token_data:
+
+            if "access_token" not in token_data:
                 raise ValueError("No access_token in response: {}".format(token_data))
-                
+
             return token_data
-            
+
         except requests.exceptions.RequestException as e:
             raise RuntimeError("Token request failed: {}".format(e)) from e
         except ValueError as e:

src/databricks/sql/auth/token_federation.py

@@ -19,7 +19,7 @@
 class DatabricksTokenFederationProvider(CredentialsProvider):
     """
     Token federation provider that exchanges external tokens for Databricks tokens.
-    
+
     This implementation follows the JDBC pattern:
     1. Try token exchange without HTTP Basic authentication (per RFC 8693)
     2. Fall back to using external token directly if exchange fails
@@ -118,7 +118,9 @@ def _parse_jwt_claims(self, token: str) -> Dict[str, Any]:
             Dict[str, Any]: Parsed JWT claims
         """
         try:
-            return jwt.decode(token, options={"verify_signature": False, "verify_aud": False})
+            return jwt.decode(
+                token, options={"verify_signature": False, "verify_aud": False}
+            )
         except Exception as e:
             logger.debug("Failed to parse JWT: %s", str(e))
             return {}
@@ -143,7 +145,6 @@ def _get_expiry_from_jwt(self, token: str) -> Optional[datetime]:
             except (ValueError, TypeError) as e:
                 logger.warning("Invalid JWT expiry value: %s", e)
 
-
     def refresh_token(self) -> Token:
         """
         Refresh the token and return the new Token object.
@@ -184,7 +185,9 @@ def refresh_token(self) -> Token:
                 self.current_token = new_token
                 return new_token
             except Exception as e:
-                logger.debug("Token exchange failed: %s. Using external token as fallback.", e)
+                logger.debug(
+                    "Token exchange failed: %s. Using external token as fallback.", e
+                )
                 expiry = self._get_expiry_from_jwt(access_token)
                 fallback_token = Token(access_token, token_type, "", expiry)
                 self.current_token = fallback_token
@@ -223,7 +226,9 @@ def get_auth_headers(self) -> Dict[str, str]:
             # Always get the latest headers from the credentials provider
             header_factory = self.credentials_provider()
             headers = dict(header_factory()) if header_factory else {}
-            headers["Authorization"] = "{} {}".format(token.token_type, token.access_token)
+            headers["Authorization"] = "{} {}".format(
+                token.token_type, token.access_token
+            )
             return headers
         except Exception as e:
             return dict(self.external_headers) if self.external_headers else {}
@@ -233,7 +238,7 @@ def _send_token_exchange_request(
     ) -> Dict[str, Any]:
         """
         Send the token exchange request to the token endpoint.
-        
+
         For M2M flows, this should include HTTP Basic authentication using client credentials.
         For U2M flows, token exchange is validated purely based on the JWT token and federation policies.
 
@@ -250,23 +255,29 @@ def _send_token_exchange_request(
             raise ValueError("Token endpoint not initialized")
 
         auth = None
-        if hasattr(self.credentials_provider, 'client_id') and hasattr(self.credentials_provider, 'client_secret'):
+        if hasattr(self.credentials_provider, "client_id") and hasattr(
+            self.credentials_provider, "client_secret"
+        ):
             client_id = self.credentials_provider.client_id
             client_secret = self.credentials_provider.client_secret
             auth = (client_id, client_secret)
         else:
-            logger.debug("No client credentials available, sending request without authentication")
-        
+            logger.debug(
+                "No client credentials available, sending request without authentication"
+            )
+
         response = requests.post(
-            self.token_endpoint, 
-            data=token_exchange_data, 
+            self.token_endpoint,
+            data=token_exchange_data,
             headers=self.EXCHANGE_HEADERS,
-            auth=auth
+            auth=auth,
         )
 
         if response.status_code != 200:
             raise requests.HTTPError(
-                "Token exchange failed with status code {}: {}".format(response.status_code, response.text),
+                "Token exchange failed with status code {}: {}".format(
+                    response.status_code, response.text
+                ),
                 response=response,
             )
 
@@ -289,7 +300,7 @@ def _exchange_token(self, access_token: str) -> Token:
             self.token_endpoint = OIDCDiscoveryUtil.discover_token_endpoint(
                 self.hostname
             )
-            
+
         # Prepare the request data according to RFC 8693
         token_exchange_data = dict(self.TOKEN_EXCHANGE_PARAMS)
         token_exchange_data["subject_token"] = access_token
@@ -319,4 +330,4 @@ def add_headers(self, request_headers: Dict[str, str]):
         """
         headers = self.get_auth_headers()
         for k, v in headers.items():
-            request_headers[k] = v
+            request_headers[k] = v

tests/unit/test_token_federation.py

@@ -4,9 +4,7 @@
 import jwt
 
 from databricks.sql.auth.token import Token
-from databricks.sql.auth.token_federation import (
-    DatabricksTokenFederationProvider
-)
+from databricks.sql.auth.token_federation import DatabricksTokenFederationProvider
 from databricks.sql.auth.oidc_utils import OIDCDiscoveryUtil