Refactor token exchange parameters to be instance-specific in DatabricksTokenFederationProvider

Author: madhav-db

src/databricks/sql/auth/token_federation.py

@@ -15,14 +15,6 @@
 
 logger = logging.getLogger(__name__)
 
-# Token exchange constants
-TOKEN_EXCHANGE_PARAMS = {
-    "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
-    "scope": "sql",
-    "subject_token_type": "urn:ietf:params:oauth:token-type:jwt",
-    "return_original_token_if_authenticated": "true",
-}
-
 
 class DatabricksTokenFederationProvider(CredentialsProvider):
     """
@@ -40,6 +32,14 @@ class DatabricksTokenFederationProvider(CredentialsProvider):
         "Content-Type": "application/x-www-form-urlencoded",
     }
 
+    # Token exchange parameters
+    TOKEN_EXCHANGE_PARAMS = {
+        "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
+        "scope": "sql",
+        "subject_token_type": "urn:ietf:params:oauth:token-type:jwt",
+        "return_original_token_if_authenticated": "true",
+    }
+
     def __init__(
         self,
         credentials_provider: CredentialsProvider,
@@ -317,7 +317,7 @@ def _exchange_token(self, access_token: str) -> Token:
             ValueError: If token exchange fails
         """
         # Prepare the request data
-        token_exchange_data = dict(TOKEN_EXCHANGE_PARAMS)
+        token_exchange_data = dict(self.TOKEN_EXCHANGE_PARAMS)
         token_exchange_data["subject_token"] = access_token
 
         # Add client_id if provided