Fix regression in sdk (changes in token federation) (#1052)

- The issue :
- An [E2E
test](https://docs.google.com/document/d/1McX4IgD-ZBTtiNXNUrEemsjbVj3oeoQYDwegjjln6UA/edit?ouid=113104269373381935368&tab=t.0#heading=h.k832j8h2svg)
conducted with JDBC, with additional SDK-level logging, confirmed that
the SDK was refreshing tokens on every call—even when successive calls
occurred within seconds. This behavior occurred despite Databricks M2M
tokens being valid for 59 minutes. As a result, the token endpoint was
hit excessively, eventually triggering global rate limits and throttling
for the IP.
- Each [getToken call in the
SDK](https://github.com/databricks/databricks-jdbc/blob/a17b84c1a0418094a8434f56246c764fa235d19b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksHttpTTransport.java#L166)
fetched a new token from the server because the SDK did not cache tokens
correctly. This was traced to a regression in the SDK’s token management
logic.
- This PR ensures that the SDK is configured once in the constructor,
preventing repeated token endpoint calls. We also plan to perform a
broader SDK code audit to identify and address any similar issues going
forward.

- Tested manually using M2M flow : The token retrieval is now performed
only once when M2M creds are used.
- Unit tests

- Internal doc :
https://docs.google.com/document/d/1McX4IgD-ZBTtiNXNUrEemsjbVj3oeoQYDwegjjln6UA/edit?tab=t.g07tag19b223

---------

Co-authored-by: Gopal Lal <[email protected]>

Author: samikshya-db

NEXT_CHANGELOG.md

@@ -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/auth/DatabricksTokenFederationProvider.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/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.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());
   }