Enhance token federation refresh to get fresh external tokens

Author: madhav-db

src/databricks/sql/auth/token_federation.py

@@ -283,33 +283,56 @@ def _is_same_host(self, url1: str, url2: str) -> bool:
 
     def _refresh_token(self, access_token: str, token_type: str) -> Dict[str, str]:
         """
-        Attempt to refresh an expired token.
+        Attempt to refresh an expired token by first getting a fresh external token
+        and then exchanging it for a new Databricks token.
 
-        For most OAuth implementations, refreshing involves a new token exchange
-        with the latest external 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
+            access_token: The original external access token (will be replaced)
             token_type: The token type (Bearer, etc.)
 
         Returns:
             The headers with the fresh token
         """
         try:
-            logger.info("Refreshing expired token via new token exchange")
-            # For most federation implementations, refresh is just a new token exchange
-            token_claims = self._parse_jwt_claims(access_token)
+            logger.info("Refreshing expired token by getting a new external token")
+
+            # ENHANCEMENT: Get a fresh token from the underlying credentials provider
+            # instead of reusing the same access_token
+            fresh_headers = self.credentials_provider()()
+
+            # Extract the fresh token from the headers
+            auth_header = fresh_headers.get("Authorization", "")
+            if not auth_header:
+                logger.error("No Authorization header in fresh headers")
+                return self.external_provider_headers
+
+            parts = auth_header.split(" ", 1)
+            if len(parts) != 2:
+                logger.error(f"Invalid Authorization header format: {auth_header}")
+                return self.external_provider_headers
+
+            fresh_token_type = parts[0]
+            fresh_access_token = parts[1]
+
+            logger.debug("Got fresh external token")
+
+            # Now process the fresh token
+            token_claims = self._parse_jwt_claims(fresh_access_token)
             idp_type = self._detect_idp_from_claims(token_claims)
 
-            # Perform a new token exchange
-            refreshed_token = self._exchange_token(access_token, idp_type)
+            # Perform a new token exchange with the fresh token
+            refreshed_token = self._exchange_token(fresh_access_token, idp_type)
 
             # Update the stored token
             self.last_exchanged_token = refreshed_token
-            self.last_external_token = access_token
+            self.last_external_token = fresh_access_token
 
             # Create new headers with the refreshed token
-            headers = dict(self.external_provider_headers)
+            headers = dict(fresh_headers)  # Use the fresh headers as base
             headers[
                 "Authorization"
             ] = f"{refreshed_token.token_type} {refreshed_token.access_token}"

tests/unit/test_token_federation.py

@@ -114,6 +114,63 @@ def test_init_oidc_discovery(self, mock_get_endpoints, mock_requests_get):
         federation_provider.token_endpoint = None
         federation_provider._init_oidc_discovery()
         self.assertEqual(federation_provider.token_endpoint, "https://example.com/oidc/v1/token")
+    
+    @patch('databricks.sql.auth.token_federation.DatabricksTokenFederationProvider._parse_jwt_claims')
+    @patch('databricks.sql.auth.token_federation.DatabricksTokenFederationProvider._exchange_token')
+    @patch('databricks.sql.auth.token_federation.DatabricksTokenFederationProvider._is_same_host')
+    def test_token_refresh(self, mock_is_same_host, mock_exchange_token, mock_parse_jwt):
+        """Test token refresh functionality for approaching expiry."""
+        # Set up mocks
+        mock_parse_jwt.return_value = {"iss": "https://login.microsoftonline.com/tenant"}
+        mock_is_same_host.return_value = False
+        
+        # Create a simple credentials provider that returns a fixed token
+        external_token = "test_token"
+        creds_provider = SimpleCredentialsProvider(external_token)
+        
+        # Set up the token federation provider
+        federation_provider = DatabricksTokenFederationProvider(
+            creds_provider, "example.com", "client_id"
+        )
+        
+        # Mock the token exchange to return a known token
+        future_time = datetime.now(tz=timezone.utc) + timedelta(hours=1)
+        mock_exchange_token.return_value = Token(
+            "exchanged_token_1", "Bearer", expiry=future_time
+        )
+        
+        # First call to get initial headers and token - this should trigger an exchange
+        headers_factory = federation_provider()
+        headers = headers_factory()
+        
+        # Verify the exchange happened
+        mock_exchange_token.assert_called_with(external_token, "azure")
+        self.assertEqual(headers["Authorization"], "Bearer exchanged_token_1")
+        
+        # Reset the mocks to track the next call
+        mock_exchange_token.reset_mock()
+        
+        # Now simulate an approaching expiry
+        near_expiry = datetime.now(tz=timezone.utc) + timedelta(minutes=4)
+        federation_provider.last_exchanged_token = Token(
+            "exchanged_token_1", "Bearer", expiry=near_expiry
+        )
+        federation_provider.last_external_token = external_token
+        
+        # Set up the mock to return a different token for the refresh
+        mock_exchange_token.return_value = Token(
+            "exchanged_token_2", "Bearer", expiry=future_time
+        )
+        
+        # Make a second call which should trigger refresh
+        headers = headers_factory()
+        
+        # Verify the token was exchanged with the SAME external token (current implementation)
+        # This is different from the JDBC driver approach which gets a fresh token
+        mock_exchange_token.assert_called_once_with(external_token, "azure")
+        
+        # Verify the headers contain the new token
+        self.assertEqual(headers["Authorization"], "Bearer exchanged_token_2")
 
 
 class TestTokenFederationFactory(unittest.TestCase):

tests/unit/test_token_federation_jdbc.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python3
+
+"""
+Unit tests for the JDBC-style token refresh in Databricks SQL connector.
+
+This test verifies that the token federation implementation follows the JDBC driver's approach
+of getting a fresh external token before exchanging it for a Databricks token during refresh.
+"""
+
+import unittest
+from unittest.mock import patch, MagicMock
+from datetime import datetime, timezone, timedelta
+
+from databricks.sql.auth.token_federation import (
+    DatabricksTokenFederationProvider,
+    Token
+)
+
+
+class RefreshingCredentialsProvider:
+    """
+    A credentials provider that returns different tokens on each call.
+    This simulates providers like Azure AD that can refresh their tokens.
+    """
+    
+    def __init__(self):
+        self.call_count = 0
+    
+    def auth_type(self):
+        return "bearer"
+    
+    def __call__(self, *args, **kwargs):
+        def get_headers():
+            self.call_count += 1
+            # Return a different token each time to simulate fresh tokens
+            return {"Authorization": f"Bearer fresh_token_{self.call_count}"}
+        return get_headers
+
+
+class TestJdbcStyleTokenRefresh(unittest.TestCase):
+    """Tests for the JDBC-style token refresh implementation."""
+    
+    @patch('databricks.sql.auth.token_federation.DatabricksTokenFederationProvider._parse_jwt_claims')
+    @patch('databricks.sql.auth.token_federation.DatabricksTokenFederationProvider._exchange_token')
+    @patch('databricks.sql.auth.token_federation.DatabricksTokenFederationProvider._is_same_host')
+    def test_refresh_gets_fresh_token(self, mock_is_same_host, mock_exchange_token, mock_parse_jwt):
+        """Test that token refresh first gets a fresh external token."""
+        # Set up mocks
+        mock_parse_jwt.return_value = {"iss": "https://login.microsoftonline.com/tenant"}
+        mock_is_same_host.return_value = False
+        
+        # Create a credentials provider that returns different tokens on each call
+        refreshing_provider = RefreshingCredentialsProvider()
+        
+        # Set up the token federation provider
+        federation_provider = DatabricksTokenFederationProvider(
+            refreshing_provider, "example.com", "client_id"
+        )
+        
+        # Set up mock for token exchange
+        future_time = datetime.now(tz=timezone.utc) + timedelta(hours=1)
+        mock_exchange_token.return_value = Token(
+            "exchanged_token_1", "Bearer", expiry=future_time
+        )
+        
+        # First call to get initial headers and token
+        headers_factory = federation_provider()
+        headers = headers_factory()
+        
+        # Verify the first exchange happened
+        mock_exchange_token.assert_called_with("fresh_token_1", "azure")
+        self.assertEqual(headers["Authorization"], "Bearer exchanged_token_1")
+        self.assertEqual(refreshing_provider.call_count, 1)
+        
+        # Reset the mock to track the next call
+        mock_exchange_token.reset_mock()
+        
+        # Now simulate an approaching expiry
+        near_expiry = datetime.now(tz=timezone.utc) + timedelta(minutes=4)
+        federation_provider.last_exchanged_token = Token(
+            "exchanged_token_1", "Bearer", expiry=near_expiry
+        )
+        federation_provider.last_external_token = "fresh_token_1"
+        
+        # Set up the mock to return a different token for the refresh
+        mock_exchange_token.return_value = Token(
+            "exchanged_token_2", "Bearer", expiry=future_time
+        )
+        
+        # Make a second call which should trigger refresh
+        headers = headers_factory()
+        
+        # With JDBC-style implementation:
+        # 1. Should call credentials provider again to get fresh token
+        self.assertEqual(refreshing_provider.call_count, 2)
+        
+        # 2. Should exchange the FRESH token (fresh_token_2), not the stored one (fresh_token_1)
+        mock_exchange_token.assert_called_once_with("fresh_token_2", "azure")
+        
+        # 3. Should return headers with the new Databricks token
+        self.assertEqual(headers["Authorization"], "Bearer exchanged_token_2")
+
+
+if __name__ == "__main__":
+    unittest.main()