fixes / improvements

Author: tejaskochar-db

databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java

@@ -3,8 +3,10 @@
 import com.databricks.sdk.core.oauth.CachedTokenSource;
 import com.databricks.sdk.core.oauth.OAuthHeaderFactory;
 import com.databricks.sdk.core.oauth.Token;
+import com.databricks.sdk.core.oauth.TokenSource;
 import com.databricks.sdk.core.utils.OSUtils;
 import com.databricks.sdk.support.InternalApi;
+import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import java.nio.charset.StandardCharsets;
 import java.util.*;
@@ -20,6 +22,13 @@ public class DatabricksCliCredentialsProvider implements CredentialsProvider {
 
   private static final ObjectMapper MAPPER = new ObjectMapper();
 
+  /** Thrown when the cached CLI token's scopes don't match the SDK's configured scopes. */
+  static class ScopeMismatchException extends DatabricksException {
+    ScopeMismatchException(String message) {
+      super(message);
+    }
+  }
+
   /**
    * offline_access controls whether the IdP issues a refresh token. It does not grant any API
    * permissions, so its presence or absence should not cause a scope mismatch error.
@@ -104,18 +113,38 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
         return null;
       }
 
+      // Wrap the token source with scope validation so that every token — both the
+      // initial fetch and subsequent refreshes — is checked against the configured scopes.
+      TokenSource effectiveSource;
+      if (config.isScopesExplicitlySet()) {
+        List<String> scopes = config.getScopes();
+        effectiveSource =
+            () -> {
+              Token t = tokenSource.getToken();
+              validateTokenScopes(t, scopes, host);
+              return t;
+            };
+      } else {
+        effectiveSource = tokenSource;
+      }
+
       CachedTokenSource cachedTokenSource =
-          new CachedTokenSource.Builder(tokenSource)
+          new CachedTokenSource.Builder(effectiveSource)
               .setAsyncDisabled(config.getDisableAsyncTokenRefresh())
               .build();
-      Token token =
-          cachedTokenSource.getToken(); // We need this for checking if databricks CLI is installed.
-
-      if (config.isScopesExplicitlySet()) {
-        validateTokenScopes(token, config.getScopes(), config.getHost());
-      }
+      cachedTokenSource.getToken(); // We need this for checking if databricks CLI is installed.
 
       return OAuthHeaderFactory.fromTokenSource(cachedTokenSource);
+    } catch (ScopeMismatchException e) {
+      // Scope validation failed. When the user explicitly selected databricks-cli auth,
+      // surface the mismatch immediately so they get an actionable error. When we're being
+      // tried as part of the default credential chain, step aside so other providers get
+      // a chance.
+      if (DATABRICKS_CLI.equals(config.getAuthType())) {
+        throw e;
+      }
+      LOG.warn("Databricks CLI token scope mismatch, skipping: {}", e.getMessage());
+      return null;
     } catch (DatabricksException e) {
       String stderr = e.getMessage();
       if (stderr.contains("not found")) {
@@ -126,17 +155,6 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
         LOG.info("OAuth not configured or not available");
         return null;
       }
-      // Scope validation failed. When the user explicitly selected databricks-cli auth,
-      // surface the mismatch immediately so they get an actionable error. When we're being
-      // tried as part of the default credential chain, step aside so other providers get
-      // a chance.
-      if (stderr.contains("do not match the configured scopes")) {
-        if (DATABRICKS_CLI.equals(config.getAuthType())) {
-          throw e;
-        }
-        LOG.warn("Databricks CLI token scope mismatch, skipping: {}", e.getMessage());
-        return null;
-      }
       throw e;
     }
   }
@@ -179,15 +197,13 @@ static void validateTokenScopes(Token token, List<String> requestedScopes, Strin
       List<String> sortedRequested = new ArrayList<>(requested);
       Collections.sort(sortedRequested);
 
-      // Build a re-auth command hint with scopes (excluding offline_access)
-      String scopesArg = String.join(",", sortedRequested);
-
-      throw new DatabricksException(
+      throw new ScopeMismatchException(
           String.format(
               "Token issued by Databricks CLI has scopes %s which do not match "
-                  + "the configured scopes %s. Please re-authenticate with the desired scopes "
-                  + "by running `databricks auth login --host %s --scopes %s`.",
-              sortedTokenScopes, sortedRequested, host, scopesArg));
+                  + "the configured scopes %s. Please re-authenticate "
+                  + "with the desired scopes by running `databricks auth login` with the --scopes flag."
+                  + "Scopes default to all-apis.",
+              sortedTokenScopes, sortedRequested));
     }
   }
 
@@ -196,18 +212,18 @@ static void validateTokenScopes(Token token, List<String> requestedScopes, Strin
    * valid JWT.
    */
   private static Map<String, Object> getJwtClaims(String accessToken) {
+    String[] parts = accessToken.split("\\.");
+    if (parts.length != 3) {
+      LOG.debug("Tried to decode access token as JWT, but failed: {} components", parts.length);
+      return null;
+    }
     try {
-      String[] parts = accessToken.split("\\.");
-      if (parts.length != 3) {
-        LOG.debug("Tried to decode access token as JWT, but failed: {} components", parts.length);
-        return null;
-      }
       byte[] payloadBytes = Base64.getUrlDecoder().decode(parts[1]);
       String payloadJson = new String(payloadBytes, StandardCharsets.UTF_8);
       @SuppressWarnings("unchecked")
       Map<String, Object> claims = MAPPER.readValue(payloadJson, Map.class);
       return claims;
-    } catch (Exception e) {
+    } catch (IllegalArgumentException | JsonProcessingException e) {
       LOG.debug("Failed to decode JWT claims: {}", e.getMessage());
       return null;
     }

databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java

@@ -435,11 +435,11 @@ public DatabricksConfig setScopes(List<String> scopes) {
   }
 
   /**
-   * Returns true if scopes were explicitly configured (either directly in code or loaded from a CLI
-   * profile/config file). When scopes are not set, getScopes() defaults to ["all-apis"], which
+   * Returns true if scopes were explicitly configured (either directly in code or loaded from a
+   * config file). When scopes are not set, getScopes() defaults to ["all-apis"], which
    * would cause false-positive mismatches during scope validation.
    */
-  public boolean isScopesExplicitlySet() {
+  boolean isScopesExplicitlySet() {
     return scopes != null && !scopes.isEmpty();
   }
 

databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliScopeValidationTest.java

@@ -85,7 +85,7 @@ void testScopeValidation(
 
     if (expectError) {
       assertThrows(
-          DatabricksException.class,
+          DatabricksCliCredentialsProvider.ScopeMismatchException.class,
           () ->
               DatabricksCliCredentialsProvider.validateTokenScopes(token, configuredScopes, HOST));
     } else {
@@ -116,14 +116,17 @@ void testNonJwtTokenSkipsValidation() {
   @Test
   void testErrorMessageContainsReauthCommand() {
     Token token = makeToken(Collections.singletonMap("scope", "all-apis"));
-    DatabricksException e =
+    DatabricksCliCredentialsProvider.ScopeMismatchException e =
         assertThrows(
-            DatabricksException.class,
+            DatabricksCliCredentialsProvider.ScopeMismatchException.class,
             () ->
                 DatabricksCliCredentialsProvider.validateTokenScopes(
                     token, Arrays.asList("sql", "offline_access"), HOST));
     assertTrue(
-        e.getMessage().contains("databricks auth login --host " + HOST + " --scopes sql"),
+        e.getMessage().contains("databricks auth login"),
         "Expected re-auth command in error message, got: " + e.getMessage());
+    assertTrue(
+        e.getMessage().contains("do not match the configured scopes"),
+        "Expected scope mismatch details in error message, got: " + e.getMessage());
   }
 }