address comments

Author: madhav-db

src/databricks/sql/auth/token_federation.py

@@ -56,6 +56,20 @@ def __init__(
 
     def add_headers(self, request_headers: Dict[str, str]):
         """Add authentication headers to the request."""
+
+        if self._cached_token and not self._cached_token.is_expired():
+            request_headers["Authorization"] = f"{self._cached_token.token_type} {self._cached_token.access_token}"
+            return
+
+        # Get the external headers first to check if we need token federation
+        self._external_headers = {}
+        self.external_provider.add_headers(self._external_headers)
+
+        # If no Authorization header from external provider, pass through all headers
+        if "Authorization" not in self._external_headers:
+            request_headers.update(self._external_headers)
+            return
+
         token = self._get_token()
         request_headers["Authorization"] = f"{token.token_type} {token.access_token}"
 
@@ -65,11 +79,7 @@ def _get_token(self) -> Token:
         if self._cached_token and not self._cached_token.is_expired():
             return self._cached_token
 
-        # Get the external token
-        self._external_headers = {}
-        self.external_provider.add_headers(self._external_headers)
-
-        # Extract token from Authorization header
+        # Extract token from already-fetched headers
         auth_header = self._external_headers.get("Authorization", "")
         token_type, access_token = self._extract_token_from_header(auth_header)
 

tests/unit/test_auth.py

@@ -164,7 +164,9 @@ def __call__(self, *args, **kwargs) -> HeaderFactory:
         kwargs = {"credentials_provider": MyProvider()}
         mock_http_client = MagicMock()
         auth_provider = get_python_sql_connector_auth_provider(hostname, mock_http_client, **kwargs)
-        self.assertTrue(type(auth_provider).__name__, "ExternalAuthProvider")
+
+        self.assertEqual(type(auth_provider).__name__, "TokenFederationProvider")
+        self.assertEqual(type(auth_provider.external_provider).__name__, "ExternalAuthProvider")
 
         headers = {}
         auth_provider.add_headers(headers)
@@ -199,8 +201,11 @@ def test_get_python_sql_connector_default_auth(self, mock__initial_get_token):
         hostname = "foo.cloud.databricks.com"
         mock_http_client = MagicMock()
         auth_provider = get_python_sql_connector_auth_provider(hostname, mock_http_client)
-        self.assertTrue(type(auth_provider).__name__, "DatabricksOAuthProvider")
-        self.assertTrue(auth_provider._client_id, PYSQL_OAUTH_CLIENT_ID)
+
+        self.assertEqual(type(auth_provider).__name__, "TokenFederationProvider")
+        self.assertEqual(type(auth_provider.external_provider).__name__, "DatabricksOAuthProvider")
+
+        self.assertEqual(auth_provider.external_provider._client_id, PYSQL_OAUTH_CLIENT_ID)
 
 
 class TestClientCredentialsTokenSource: