fmt

Author: madhav-db

src/databricks/sql/auth/token_federation.py

@@ -150,9 +150,7 @@ def get_headers() -> Dict[str, str]:
                 else:
                     # Token is from a different host, need to exchange
                     logger.debug("Token from different host, attempting 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"Error processing token: {str(e)}")
                 # Fall back to original headers in case of error
@@ -172,9 +170,7 @@ def _init_oidc_discovery(self):
 
             if idp_endpoints:
                 # Get the OpenID configuration URL
-                openid_config_url = idp_endpoints.get_openid_config_url(
-                    self.hostname
-                )
+                openid_config_url = idp_endpoints.get_openid_config_url(self.hostname)
 
                 # Fetch the OpenID configuration
                 response = requests.get(openid_config_url)
@@ -185,7 +181,8 @@ def _init_oidc_discovery(self):
                     logger.info(f"Discovered token endpoint: {self.token_endpoint}")
                 else:
                     logger.warning(
-                        f"Failed to fetch OpenID configuration from {openid_config_url}: {response.status_code}"
+                        f"Failed to fetch OpenID configuration from {openid_config_url}: "
+                        f"{response.status_code}"
                     )
         except Exception as e:
             logger.warning(
@@ -282,9 +279,15 @@ def _refresh_token(self, access_token: str, token_type: str) -> Dict[str, str]:
             self.last_external_token = access_token
 
             # Update the headers with the new token
-            return {"Authorization": f"{exchanged_token.token_type} {exchanged_token.access_token}"}
+            return {
+                "Authorization": (
+                    f"{exchanged_token.token_type} {exchanged_token.access_token}"
+                )
+            }
         except Exception as e:
-            logger.error(f"Token refresh failed: {str(e)}, falling back to original token")
+            logger.error(
+                f"Token refresh failed: {str(e)}, falling back to original token"
+            )
             return self.external_provider_headers
 
     def _try_token_exchange_or_fallback(
@@ -305,12 +308,20 @@ def _try_token_exchange_or_fallback(
             self.last_exchanged_token = exchanged_token
             self.last_external_token = access_token
 
-            return {"Authorization": f"{exchanged_token.token_type} {exchanged_token.access_token}"}
+            return {
+                "Authorization": (
+                    f"{exchanged_token.token_type} {exchanged_token.access_token}"
+                )
+            }
         except Exception as e:
-            logger.warning(f"Token exchange failed: {str(e)}, falling back to original token")
+            logger.warning(
+                f"Token exchange failed: {str(e)}, falling back to original token"
+            )
             return self.external_provider_headers
 
-    def _send_token_exchange_request(self, token_exchange_data: Dict[str, str]) -> Dict[str, Any]:
+    def _send_token_exchange_request(
+        self, token_exchange_data: Dict[str, str]
+    ) -> Dict[str, Any]:
         """
         Send the token exchange request to the token endpoint.
         
@@ -325,20 +336,19 @@ def _send_token_exchange_request(self, token_exchange_data: Dict[str, str]) -> D
         """
         if not self.token_endpoint:
             raise ValueError("Token endpoint not initialized")
-            
+
         headers = {"Accept": "*/*", "Content-Type": "application/x-www-form-urlencoded"}
-        
+
         response = requests.post(
-            self.token_endpoint,
-            data=token_exchange_data,
-            headers=headers
+            self.token_endpoint, data=token_exchange_data, headers=headers
         )
-        
+
         if response.status_code != 200:
             raise ValueError(
-                f"Token exchange failed with status code {response.status_code}: {response.text}"
+                f"Token exchange failed with status code {response.status_code}: "
+                f"{response.text}"
             )
-            
+
         return response.json()
 
     def _exchange_token(self, access_token: str) -> Token:
@@ -365,26 +375,28 @@ def _exchange_token(self, access_token: str) -> Token:
         try:
             # Send the token exchange request
             resp_data = self._send_token_exchange_request(token_exchange_data)
-            
+
             # Extract token information
             new_access_token = resp_data.get("access_token")
             if not new_access_token:
                 raise ValueError("No access token in exchange response")
-                
+
             token_type = resp_data.get("token_type", "Bearer")
             refresh_token = resp_data.get("refresh_token", "")
-            
+
             # Parse expiry time from token claims if possible
             expiry = datetime.now(tz=timezone.utc)
-            
+
             # First try to get expiry from the response's expires_in field
             if "expires_in" in resp_data and resp_data["expires_in"]:
                 try:
                     expires_in = int(resp_data["expires_in"])
-                    expiry = datetime.now(tz=timezone.utc) + timedelta(seconds=expires_in)
+                    expiry = datetime.now(tz=timezone.utc) + timedelta(
+                        seconds=expires_in
+                    )
                 except (ValueError, TypeError) as e:
                     logger.warning(f"Invalid expires_in value: {str(e)}")
-            
+
             # If that didn't work, try to parse JWT claims for expiry
             if expiry == datetime.now(tz=timezone.utc):
                 token_claims = self._parse_jwt_claims(new_access_token)
@@ -394,9 +406,9 @@ def _exchange_token(self, access_token: str) -> Token:
                         expiry = datetime.fromtimestamp(exp_timestamp, tz=timezone.utc)
                     except (ValueError, TypeError) as e:
                         logger.warning(f"Invalid exp claim in token: {str(e)}")
-            
+
             return Token(new_access_token, token_type, refresh_token, expiry)
-            
+
         except Exception as e:
             logger.error(f"Token exchange failed: {str(e)}")
             raise

tests/token_federation/github_oidc_test.py

@@ -14,13 +14,17 @@
 import base64
 import logging
 from databricks import sql
-import jwt
 
+try:
+    import jwt
+
+    HAS_JWT_LIBRARY = True
+except ImportError:
+    HAS_JWT_LIBRARY = False
 
 
 logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s - %(levelname)s - %(message)s"
+    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
 )
 logger = logging.getLogger(__name__)
 
@@ -35,10 +39,29 @@ def decode_jwt(token):
     Returns:
         dict: The decoded token claims or None if decoding fails
     """
+    if HAS_JWT_LIBRARY:
+        try:
+            # Using PyJWT library (preferred method)
+            # Note: we're not verifying the signature as this is just for debugging
+            return jwt.decode(token, options={"verify_signature": False})
+        except Exception as e:
+            logger.error(f"Failed to decode token with PyJWT: {str(e)}")
+
+    # Fallback to manual decoding
     try:
-        return jwt.decode(token, options={"verify_signature": False})
+        parts = token.split(".")
+        if len(parts) != 3:
+            raise ValueError("Invalid JWT format")
+
+        payload = parts[1]
+        # Add padding if needed
+        padding = "=" * (4 - len(payload) % 4)
+        payload += padding
+
+        decoded = base64.b64decode(payload)
+        return json.loads(decoded)
     except Exception as e:
-        logger.error(f"Failed to decode token with PyJWT: {str(e)}")
+        logger.error(f"Failed to decode token: {str(e)}")
         return {}
 
 
@@ -53,7 +76,7 @@ def get_environment_variables():
     host = os.environ.get("DATABRICKS_HOST_FOR_TF")
     http_path = os.environ.get("DATABRICKS_HTTP_PATH_FOR_TF")
     identity_federation_client_id = os.environ.get("IDENTITY_FEDERATION_CLIENT_ID")
-    
+
     return github_token, host, http_path, identity_federation_client_id
 
 
@@ -62,7 +85,7 @@ def display_token_info(claims):
     if not claims:
         logger.warning("No token claims available to display")
         return
-        
+
     logger.info("=== GitHub OIDC Token Claims ===")
     logger.info(f"Token issuer: {claims.get('iss')}")
     logger.info(f"Token subject: {claims.get('sub')}")
@@ -74,7 +97,9 @@ def display_token_info(claims):
     logger.info("===============================")
 
 
-def test_databricks_connection(host, http_path, github_token, identity_federation_client_id):
+def test_databricks_connection(
+    host, http_path, github_token, identity_federation_client_id
+):
     """
     Test connection to Databricks using token federation.
     
@@ -90,30 +115,30 @@ def test_databricks_connection(host, http_path, github_token, identity_federatio
     logger.info("=== Testing Connection via Connector ===")
     logger.info(f"Connecting to Databricks at {host}{http_path}")
     logger.info(f"Using client ID: {identity_federation_client_id}")
-    
+
     connection_params = {
         "server_hostname": host,
         "http_path": http_path,
         "access_token": github_token,
         "auth_type": "token-federation",
         "identity_federation_client_id": identity_federation_client_id,
     }
-    
+
     try:
         with sql.connect(**connection_params) as connection:
             logger.info("Connection established successfully")
-            
+
             # Execute a simple query
             cursor = connection.cursor()
             cursor.execute("SELECT 1 + 1 as result")
             result = cursor.fetchall()
             logger.info(f"Query result: {result[0][0]}")
-            
+
             # Show current user
             cursor.execute("SELECT current_user() as user")
             result = cursor.fetchall()
             logger.info(f"Connected as user: {result[0][0]}")
-            
+
             logger.info("Token federation test successful!")
             return True
     except Exception as e:
@@ -125,29 +150,34 @@ def main():
     """Main entry point for the test script."""
     try:
         # Get environment variables
-        github_token, host, http_path, identity_federation_client_id = get_environment_variables()
-        
+        github_token, host, http_path, identity_federation_client_id = (
+            get_environment_variables()
+        )
+
         if not github_token:
             logger.error("Missing GitHub OIDC token (OIDC_TOKEN)")
             sys.exit(1)
-            
+
         if not host or not http_path:
-            logger.error("Missing Databricks connection parameters (DATABRICKS_HOST_FOR_TF, DATABRICKS_HTTP_PATH_FOR_TF)")
+            logger.error(
+                "Missing Databricks connection parameters "
+                "(DATABRICKS_HOST_FOR_TF, DATABRICKS_HTTP_PATH_FOR_TF)"
+            )
             sys.exit(1)
-        
+
         # Display token claims
         claims = decode_jwt(github_token)
         display_token_info(claims)
-        
+
         # Test Databricks connection
         success = test_databricks_connection(
             host, http_path, github_token, identity_federation_client_id
         )
-        
+
         if not success:
             logger.error("Token federation test failed")
             sys.exit(1)
-            
+
     except Exception as e:
         logger.error(f"Unexpected error: {str(e)}")
         sys.exit(1)

tests/unit/test_token_federation.py

@@ -13,7 +13,6 @@
     Token,
     DatabricksTokenFederationProvider,
     SimpleCredentialsProvider,
-    create_token_federation_provider,
     TOKEN_REFRESH_BUFFER_SECONDS,
 )
 
@@ -136,19 +135,25 @@ def test_init_oidc_discovery(mock_request_get, mock_get_oauth_endpoints):
 
 @pytest.fixture
 def mock_parse_jwt_claims():
-    with patch("databricks.sql.auth.token_federation.DatabricksTokenFederationProvider._parse_jwt_claims") as mock:
+    with patch(
+        "databricks.sql.auth.token_federation.DatabricksTokenFederationProvider._parse_jwt_claims"
+    ) as mock:
         yield mock
 
 
 @pytest.fixture
 def mock_exchange_token():
-    with patch("databricks.sql.auth.token_federation.DatabricksTokenFederationProvider._exchange_token") as mock:
+    with patch(
+        "databricks.sql.auth.token_federation.DatabricksTokenFederationProvider._exchange_token"
+    ) as mock:
         yield mock
 
 
 @pytest.fixture
 def mock_is_same_host():
-    with patch("databricks.sql.auth.token_federation.DatabricksTokenFederationProvider._is_same_host") as mock:
+    with patch(
+        "databricks.sql.auth.token_federation.DatabricksTokenFederationProvider._is_same_host"
+    ) as mock:
         yield mock