fmt

Author: madhav-db

src/databricks/sql/auth/token_federation.py

@@ -1,17 +1,15 @@
 import base64
 import json
 import logging
-import urllib.parse
 from datetime import datetime, timezone, timedelta
-from typing import Dict, Optional, Any, Tuple, List, Union
+from typing import Dict, Optional, Any, Tuple
 from urllib.parse import urlparse
 
 import requests
 from requests.exceptions import RequestException
 
 from databricks.sql.auth.authenticators import CredentialsProvider, HeaderFactory
 from databricks.sql.auth.endpoint import (
-    get_databricks_oidc_url,
     get_oauth_endpoints,
     infer_cloud_from_host,
 )
@@ -25,11 +23,7 @@
     "return_original_token_if_authenticated": "true",
 }
 
-# Special client IDs for different IdPs
-AZURE_AD_MULTI_TENANT_APP_ID = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d"
-
-# Buffer time in seconds before token expiry to trigger a refresh (5 minutes)
-TOKEN_REFRESH_BUFFER_SECONDS = 300
+TOKEN_REFRESH_BUFFER_SECONDS = 10
 
 
 class Token:
@@ -85,7 +79,6 @@ def __init__(
         self.hostname = hostname
         self.identity_federation_client_id = identity_federation_client_id
         self.external_provider_headers: Dict[str, str] = {}
-        self.token = None
         self.token_endpoint: Optional[str] = None
         self.idp_endpoints = None
         self.openid_config = None
@@ -123,9 +116,7 @@ def get_headers() -> Dict[str, str]:
             self.external_provider_headers = header_factory()
 
             # Extract the token from the headers
-            token_info = self._extract_token_info_from_header(
-                self.external_provider_headers
-            )
+            token_info = self._extract_token_info_from_header(self.external_provider_headers)
             token_type, access_token = token_info
 
             try:
@@ -148,10 +139,7 @@ def get_headers() -> Dict[str, str]:
                     return self.external_provider_headers
                 else:
                     # Token is from a different host, need to exchange
-                    return self._try_token_exchange_or_fallback(
-                        access_token, token_type
-                    )
-
+                    return self._try_token_exchange_or_fallback(access_token, token_type)
             except Exception as e:
                 logger.error(f"Failed to process token: {str(e)}")
                 # Fall back to original headers in case of error
@@ -171,10 +159,8 @@ def _init_oidc_discovery(self):
 
             if self.idp_endpoints:
                 # Get the OpenID configuration URL
-                openid_config_url = self.idp_endpoints.get_openid_config_url(
-                    self.hostname
-                )
-
+                openid_config_url = self.idp_endpoints.get_openid_config_url(self.hostname)
+                
                 # Fetch the OpenID configuration
                 response = requests.get(openid_config_url)
                 if response.status_code == 200:
@@ -189,33 +175,26 @@ def _init_oidc_discovery(self):
 
             # Fallback to default token endpoint if discovery fails
             if not self.token_endpoint:
-                # Make sure hostname has proper format with https:// prefix and trailing slash
-                hostname = self.hostname
-                if not hostname.startswith("https://"):
-                    hostname = f"https://{hostname}"
-                if not hostname.endswith("/"):
-                    hostname = f"{hostname}/"
+                hostname = self._format_hostname(self.hostname)
                 self.token_endpoint = f"{hostname}oidc/v1/token"
                 logger.info(f"Using default token endpoint: {self.token_endpoint}")
-
         except Exception as e:
             logger.warning(
                 f"OIDC discovery failed: {str(e)}. Using default token endpoint."
             )
-            # Make sure hostname has proper format with https:// prefix and trailing slash
-            hostname = self.hostname
-            if not hostname.startswith("https://"):
-                hostname = f"https://{hostname}"
-            if not hostname.endswith("/"):
-                hostname = f"{hostname}/"
+            hostname = self._format_hostname(self.hostname)
             self.token_endpoint = f"{hostname}oidc/v1/token"
-            logger.info(
-                f"Using default token endpoint after error: {self.token_endpoint}"
-            )
+            logger.info(f"Using default token endpoint after error: {self.token_endpoint}")
 
-    def _extract_token_info_from_header(
-        self, headers: Dict[str, str]
-    ) -> Tuple[str, str]:
+    def _format_hostname(self, hostname: str) -> str:
+        """Format hostname to ensure it has proper https:// prefix and trailing slash."""
+        if not hostname.startswith("https://"):
+            hostname = f"https://{hostname}"
+        if not hostname.endswith("/"):
+            hostname = f"{hostname}/"
+        return hostname
+
+    def _extract_token_info_from_header(self, headers: Dict[str, str]) -> Tuple[str, str]:
         """Extract token type and token value from authorization header."""
         auth_header = headers.get("Authorization")
         if not auth_header:
@@ -286,10 +265,6 @@ def _refresh_token(self, access_token: str, token_type: str) -> Dict[str, str]:
         Attempt to refresh an expired token by first getting a fresh external token
         and then exchanging it for a new Databricks token.
 
-        This implementation follows the JDBC driver approach by first requesting
-        a fresh token from the underlying credentials provider before performing
-        the token exchange.
-
         Args:
             access_token: The original external access token (will be replaced)
             token_type: The token type (Bearer, etc.)
@@ -300,7 +275,7 @@ def _refresh_token(self, access_token: str, token_type: str) -> Dict[str, str]:
         try:
             logger.info("Refreshing expired token by getting a new external token")
 
-            # ENHANCEMENT: Get a fresh token from the underlying credentials provider
+            # Get a fresh token from the underlying credentials provider
             # instead of reusing the same access_token
             fresh_headers = self.credentials_provider()()
 
@@ -333,20 +308,14 @@ def _refresh_token(self, access_token: str, token_type: str) -> Dict[str, str]:
 
             # Create new headers with the refreshed token
             headers = dict(fresh_headers)  # Use the fresh headers as base
-            headers[
-                "Authorization"
-            ] = f"{refreshed_token.token_type} {refreshed_token.access_token}"
+            headers["Authorization"] = f"{refreshed_token.token_type} {refreshed_token.access_token}"
             return headers
         except Exception as e:
-            logger.error(
-                f"Token refresh failed, falling back to original token: {str(e)}"
-            )
+            logger.error(f"Token refresh failed, falling back to original token: {str(e)}")
             # If refresh fails, fall back to the original headers
             return self.external_provider_headers
 
-    def _try_token_exchange_or_fallback(
-        self, access_token: str, token_type: str
-    ) -> Dict[str, str]:
+    def _try_token_exchange_or_fallback(self, access_token: str, token_type: str) -> Dict[str, str]:
         """Try to exchange the token or fall back to the original token."""
         try:
             # Parse the token to get claims for IdP-specific adjustments
@@ -362,14 +331,10 @@ def _try_token_exchange_or_fallback(
 
             # Create new headers with the exchanged token
             headers = dict(self.external_provider_headers)
-            headers[
-                "Authorization"
-            ] = f"{exchanged_token.token_type} {exchanged_token.access_token}"
+            headers["Authorization"] = f"{exchanged_token.token_type} {exchanged_token.access_token}"
             return headers
         except Exception as e:
-            logger.error(
-                f"Token exchange failed, falling back to using external token: {str(e)}"
-            )
+            logger.error(f"Token exchange failed, falling back to using external token: {str(e)}")
             # Fall back to original headers
             return self.external_provider_headers
 
@@ -431,14 +396,10 @@ def _exchange_token(self, access_token: str, idp_type: str = "unknown") -> Token
                 try:
                     # Calculate expiry by adding expires_in seconds to current time
                     expires_in_seconds = int(resp_data["expires_in"])
-                    token.expiry = datetime.now(tz=timezone.utc) + timedelta(
-                        seconds=expires_in_seconds
-                    )
+                    token.expiry = datetime.now(tz=timezone.utc) + timedelta(seconds=expires_in_seconds)
                     logger.debug(f"Token expiry set from expires_in: {token.expiry}")
                 except (ValueError, TypeError) as e:
-                    logger.warning(
-                        f"Could not parse expires_in from response: {str(e)}"
-                    )
+                    logger.warning(f"Could not parse expires_in from response: {str(e)}")
 
             # If expires_in wasn't available, try to parse expiry from the token JWT
             if token.expiry == datetime.now(tz=timezone.utc):
@@ -447,9 +408,7 @@ def _exchange_token(self, access_token: str, idp_type: str = "unknown") -> Token
                     exp_time = token_claims.get("exp")
                     if exp_time:
                         token.expiry = datetime.fromtimestamp(exp_time, tz=timezone.utc)
-                        logger.debug(
-                            f"Token expiry set from JWT exp claim: {token.expiry}"
-                        )
+                        logger.debug(f"Token expiry set from JWT exp claim: {token.expiry}")
                 except Exception as e:
                     logger.warning(f"Could not parse expiry from token: {str(e)}")