Fix regression in sdk (changes in token federation)

NEXT_CHANGELOG.md

@@ -13,5 +13,6 @@
 - Fixed SQL syntax error when LIKE queries contain empty ESCAPE clauses.
 - Fix: driver failing to authenticate on token update in U2M flow.
 - Fix: driver failing to parse complex data types with nullable attributes.
+- Fixed: Resolved SDK token-caching regression causing token refresh on every call. SDK is now configured once to avoid excessive token endpoint hits and rate limiting.
 ---
 *Note: When making changes, please add your change under the appropriate section with a brief description.* 

src/main/java/com/databricks/jdbc/auth/DatabricksTokenFederationProvider.java

@@ -46,6 +46,7 @@ public class DatabricksTokenFederationProvider implements CredentialsProvider, T
   private static final JdbcLogger LOGGER =
       JdbcLoggerFactory.getLogger(DatabricksTokenFederationProvider.class);
   private Token token;
+  private HeaderFactory externalHeaderFactory;
   private static final Map<String, String> TOKEN_EXCHANGE_PARAMS =
       Map.of(
           "grant_type",
@@ -69,6 +70,9 @@ public DatabricksTokenFederationProvider(
     this.credentialsProvider = credentialsProvider;
     this.externalProviderHeaders = new HashMap<>();
     this.hc = DatabricksHttpClientFactory.getInstance().getClient(connectionContext);
+    // Initialize a minimal config; real config will be provided via configure(databricksConfig)
+    this.config = null;
+    this.externalHeaderFactory = null;
     this.token =
         new Token(
             DatabricksJdbcConstants.EMPTY_STRING,
@@ -85,6 +89,7 @@ public DatabricksTokenFederationProvider(
     this.connectionContext = connectionContext;
     this.credentialsProvider = credentialsProvider;
     this.config = config;
+    this.externalHeaderFactory = this.credentialsProvider.configure(this.config);
     this.externalProviderHeaders = new HashMap<>();
     this.token =
         new Token(
@@ -113,6 +118,8 @@ public HeaderFactory configure(DatabricksConfig databricksConfig) {
     }
 
     this.config = databricksConfig;
+    // Call the underlying provider's configure ONCE and cache the HeaderFactory
+    this.externalHeaderFactory = this.credentialsProvider.configure(this.config);
     return () -> {
       Token exchangedToken = getToken();
       Map<String, String> headers = new HashMap<>(this.externalProviderHeaders);
@@ -124,7 +131,11 @@ public HeaderFactory configure(DatabricksConfig databricksConfig) {
   }
 
   public Token getToken() {
-    this.externalProviderHeaders = this.credentialsProvider.configure(this.config).headers();
+    if (this.externalHeaderFactory == null) {
+      // Lazy-initialize if configure(databricksConfig) was not called yet
+      this.externalHeaderFactory = this.credentialsProvider.configure(this.config);
+    }
+    this.externalProviderHeaders = this.externalHeaderFactory.headers();
     String[] tokenInfo = extractTokenInfoFromHeader(this.externalProviderHeaders);
     String accessTokenType = tokenInfo[0];
     String accessToken = tokenInfo[1];

src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java

@@ -377,18 +377,21 @@ public void setupM2MConfig() throws DatabricksParsingException {
                   connectionContext, new AzureServicePrincipalCredentialsProvider()));
     } else {
       databricksConfig
-          .setAuthType(DatabricksJdbcConstants.M2M_AUTH_TYPE)
           .setClientId(connectionContext.getClientId())
           .setClientSecret(connectionContext.getClientSecret());
       if (connectionContext.useJWTAssertion()) {
-        databricksConfig.setCredentialsProvider(
-            new DatabricksTokenFederationProvider(
-                connectionContext,
-                new PrivateKeyClientCredentialProvider(connectionContext, databricksConfig)));
+        CredentialsProvider jwtProvider =
+            new PrivateKeyClientCredentialProvider(connectionContext, databricksConfig);
+        databricksConfig
+            .setAuthType(jwtProvider.authType())
+            .setCredentialsProvider(
+                new DatabricksTokenFederationProvider(connectionContext, jwtProvider));
       } else {
-        databricksConfig.setCredentialsProvider(
-            new DatabricksTokenFederationProvider(
-                connectionContext, new OAuthM2MServicePrincipalCredentialsProvider()));
+        CredentialsProvider m2mProvider = new OAuthM2MServicePrincipalCredentialsProvider();
+        databricksConfig
+            .setAuthType(DatabricksJdbcConstants.M2M_AUTH_TYPE)
+            .setCredentialsProvider(
+                new DatabricksTokenFederationProvider(connectionContext, m2mProvider));
       }
     }
   }

src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java

@@ -170,7 +170,7 @@ void testM2MWithJWT() throws DatabricksSQLException {
     assertEquals("https://sample-host.18.azuredatabricks.net", config.getHost());
     assertEquals("test-client", config.getClientId());
     assertEquals("custom-oauth-m2m", provider.authType());
-    assertEquals(DatabricksJdbcConstants.M2M_AUTH_TYPE, config.getAuthType());
+    assertEquals(provider.authType(), config.getAuthType());
     assertEquals(
         PrivateKeyClientCredentialProvider.class, provider.getCredentialsProvider().getClass());
   }