Validate Databricks CLI token scopes against SDK configuration

tejaskochar-db · claude · tejaskochar-db · commit 2a1667e80834 · 2026-03-05T06:23:38.000Z
Detects when a cached Databricks CLI token was issued with different
OAuth scopes than what the SDK configuration requires. Surfaces an
actionable error telling the user how to re-authenticate instead of
silently making requests with the wrong scopes.

Co-Authored-By: Claude Opus 4.6 &lt;noreply@anthropic.com&gt;
diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java
@@ -2,8 +2,12 @@
 
 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.utils.OSUtils;
 import com.databricks.sdk.support.InternalApi;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import java.nio.charset.StandardCharsets;
 import java.util.*;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -15,9 +19,14 @@ public class DatabricksCliCredentialsProvider implements CredentialsProvider {
 
   public static final String DATABRICKS_CLI = "databricks-cli";
 
-  static final String ERR_CUSTOM_SCOPES_NOT_SUPPORTED =
-      "custom scopes are not supported with databricks-cli auth; "
-          + "scopes are determined by what was last used when logging in with `databricks auth login`";
+  private static final ObjectMapper MAPPER = new ObjectMapper();
+
+  /**
+   * 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.
+   */
+  private static final Set<String> SCOPES_IGNORED_FOR_COMPARISON =
+      Collections.singleton("offline_access");
 
   @Override
   public String authType() {
@@ -96,15 +105,16 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
         return null;
       }
 
-      if (config.isScopesExplicitlySet()) {
-        throw new DatabricksException(ERR_CUSTOM_SCOPES_NOT_SUPPORTED);
-      }
-
       CachedTokenSource cachedTokenSource =
           new CachedTokenSource.Builder(tokenSource)
               .setAsyncDisabled(config.getDisableAsyncTokenRefresh())
               .build();
-      cachedTokenSource.getToken(); // We need this for checking if databricks CLI is installed.
+      Token token =
+          cachedTokenSource.getToken(); // We need this for checking if databricks CLI is installed.
+
+      if (config.isScopesExplicitlySet()) {
+        validateTokenScopes(token, config.getScopes(), config.getHost());
+      }
 
       return OAuthHeaderFactory.fromTokenSource(cachedTokenSource);
     } catch (DatabricksException e) {
@@ -117,7 +127,108 @@ 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;
     }
   }
+
+  /**
+   * Validate that the token's scopes match the requested scopes from the config.
+   *
+   * <p>The {@code databricks auth token} command does not accept scopes yet. It returns whatever
+   * token was cached from the last {@code databricks auth login}. If a user configures specific
+   * scopes in the SDK config but their cached CLI token was issued with different scopes, requests
+   * will silently use the wrong scopes. This check surfaces that mismatch early with an actionable
+   * error telling the user how to re-authenticate with the correct scopes.
+   */
+  static void validateTokenScopes(Token token, List<String> requestedScopes, String host) {
+    Map<String, Object> claims = getJwtClaims(token.getAccessToken());
+    if (claims == null) {
+      LOG.debug("Could not decode token as JWT to validate scopes");
+      return;
+    }
+
+    Object tokenScopesRaw = claims.get("scope");
+    if (tokenScopesRaw == null) {
+      LOG.debug("Token does not contain 'scope' claim, skipping scope validation");
+      return;
+    }
+
+    Set<String> tokenScopes = parseScopeClaim(tokenScopesRaw);
+    if (tokenScopes == null) {
+      LOG.debug("Unexpected 'scope' claim type: {}", tokenScopesRaw.getClass());
+      return;
+    }
+
+    tokenScopes.removeAll(SCOPES_IGNORED_FOR_COMPARISON);
+    Set<String> requested = new HashSet<>(requestedScopes);
+    requested.removeAll(SCOPES_IGNORED_FOR_COMPARISON);
+
+    if (!tokenScopes.equals(requested)) {
+      List<String> sortedTokenScopes = new ArrayList<>(tokenScopes);
+      Collections.sort(sortedTokenScopes);
+      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(
+          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));
+    }
+  }
+
+  /**
+   * Decode a JWT access token and return its payload claims. Returns null if the token is not a
+   * valid JWT.
+   */
+  private static Map<String, Object> getJwtClaims(String accessToken) {
+    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) {
+      LOG.debug("Failed to decode JWT claims: {}", e.getMessage());
+      return null;
+    }
+  }
+
+  /**
+   * Parse the JWT "scope" claim, which can be either a space-delimited string or a JSON array.
+   * Returns null if the type is unexpected.
+   */
+  private static Set<String> parseScopeClaim(Object scopeClaim) {
+    if (scopeClaim instanceof String) {
+      return new HashSet<>(Arrays.asList(((String) scopeClaim).split("\\s+")));
+    } else if (scopeClaim instanceof List) {
+      Set<String> scopes = new HashSet<>();
+      for (Object s : (List<?>) scopeClaim) {
+        scopes.add(String.valueOf(s));
+      }
+      return scopes;
+    }
+    return null;
+  }
 }
diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java
@@ -51,12 +51,6 @@ public class DatabricksConfig {
   @ConfigAttribute(auth = "oauth")
   private List<String> scopes;
 
-  // Temporary field to track if scopes were explicitly set by the user.
-  // This is used to ensure users don't set explicit scopes when using 
-  // `databricks-cli` auth, as it does not respect the scopes.
-  // TODO: Remove this field once the `auth token` command supports scopes.
-  private boolean scopesExplicitlySet = false;
-
   @ConfigAttribute(env = "DATABRICKS_REDIRECT_URL", auth = "oauth")
   private String redirectUrl;
 
@@ -437,12 +431,16 @@ public List<String> getScopes() {
 
   public DatabricksConfig setScopes(List<String> scopes) {
     this.scopes = scopes;
-    this.scopesExplicitlySet = true;
     return this;
   }
 
+  /**
+   * 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
+   * would cause false-positive mismatches during scope validation.
+   */
   public boolean isScopesExplicitlySet() {
-    return scopesExplicitlySet;
+    return scopes != null && !scopes.isEmpty();
   }
 
   public String getProfile() {
diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java
@@ -1,6 +1,5 @@
 package com.databricks.sdk.core;
 
-import static com.databricks.sdk.core.DatabricksCliCredentialsProvider.ERR_CUSTOM_SCOPES_NOT_SUPPORTED;
 import static org.junit.jupiter.api.Assertions.*;
 
 import java.util.Arrays;
@@ -141,43 +140,6 @@ void testBuildHostArgs_UnifiedHostFalse_WithAccountHost() {
         cmd);
   }
 
-  @Test
-  void testConfigure_ErrorsWhenScopesExplicitlySet() {
-    DatabricksConfig config =
-        new DatabricksConfig()
-            .setHost(HOST)
-            .setDatabricksCliPath(CLI_PATH)
-            .setScopes(Arrays.asList("sql"));
-
-    DatabricksException e =
-        assertThrows(DatabricksException.class, () -> provider.configure(config));
-    assertEquals(ERR_CUSTOM_SCOPES_NOT_SUPPORTED, e.getMessage());
-  }
-
-  @Test
-  void testConfigure_SkipsWhenCliNotFoundEvenWithScopes() {
-    // When CLI is not available, the provider should return null (skip)
-    // rather than throwing an error about scopes.
-    DatabricksConfig config =
-        new DatabricksConfig()
-            .setHost(HOST)
-            .setScopes(Arrays.asList("sql"));
-
-    assertNull(provider.configure(config));
-  }
-
-  @Test
-  void testConfigure_NoErrorWhenNoScopes() {
-    DatabricksConfig config = new DatabricksConfig().setHost(HOST);
-
-    try {
-      provider.configure(config);
-    } catch (Exception e) {
-      // May fail for other reasons (CLI not found, env not set), but must not be the scope error
-      assertNotEquals(ERR_CUSTOM_SCOPES_NOT_SUPPORTED, e.getMessage());
-    }
-  }
-
   @Test
   void testScopesExplicitlySetFlag() {
     DatabricksConfig config = new DatabricksConfig();
diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliScopeValidationTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliScopeValidationTest.java
@@ -0,0 +1,133 @@
+package com.databricks.sdk.core;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import com.databricks.sdk.core.oauth.Token;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import java.nio.charset.StandardCharsets;
+import java.time.Instant;
+import java.util.*;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+
+class DatabricksCliScopeValidationTest {
+
+  private static final String HOST = "https://my-workspace.cloud.databricks.com";
+  private static final ObjectMapper MAPPER = new ObjectMapper();
+
+  /** Builds a fake JWT (header.payload.signature) with the given claims. */
+  private static String makeJwt(Map<String, Object> claims) {
+    try {
+      String header =
+          Base64.getUrlEncoder()
+              .withoutPadding()
+              .encodeToString("{\"alg\":\"none\"}".getBytes(StandardCharsets.UTF_8));
+      String payload =
+          Base64.getUrlEncoder()
+              .withoutPadding()
+              .encodeToString(MAPPER.writeValueAsBytes(claims));
+      return header + "." + payload + ".sig";
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  private static Token makeToken(Map<String, Object> claims) {
+    return new Token(makeJwt(claims), "Bearer", Instant.now().plusSeconds(3600));
+  }
+
+  static List<Arguments> scopeValidationCases() {
+    return Arrays.asList(
+        // Exact match (offline_access filtered out).
+        Arguments.of(
+            Collections.singletonMap("scope", "sql offline_access"),
+            Collections.singletonList("sql"),
+            false,
+            "match"),
+        // Mismatch throws.
+        Arguments.of(
+            Collections.singletonMap("scope", "all-apis offline_access"),
+            Collections.singletonList("sql"),
+            true,
+            "mismatch"),
+        // offline_access on token only — still equivalent.
+        Arguments.of(
+            Collections.singletonMap("scope", "all-apis offline_access"),
+            Collections.singletonList("all-apis"),
+            false,
+            "offline_access_on_token_only"),
+        // offline_access in config only — still equivalent.
+        Arguments.of(
+            Collections.singletonMap("scope", "all-apis"),
+            Arrays.asList("all-apis", "offline_access"),
+            false,
+            "offline_access_in_config_only"),
+        // Scope claim as list instead of string.
+        Arguments.of(
+            new HashMap<String, Object>() {
+              {
+                put("scope", Arrays.asList("sql", "offline_access"));
+              }
+            },
+            Collections.singletonList("sql"),
+            false,
+            "scope_as_list"));
+  }
+
+  @ParameterizedTest(name = "{3}")
+  @MethodSource("scopeValidationCases")
+  void testScopeValidation(
+      Map<String, Object> tokenClaims,
+      List<String> configuredScopes,
+      boolean expectError,
+      String testName) {
+    Token token = makeToken(tokenClaims);
+
+    if (expectError) {
+      assertThrows(
+          DatabricksException.class,
+          () ->
+              DatabricksCliCredentialsProvider.validateTokenScopes(
+                  token, configuredScopes, HOST));
+    } else {
+      assertDoesNotThrow(
+          () ->
+              DatabricksCliCredentialsProvider.validateTokenScopes(
+                  token, configuredScopes, HOST));
+    }
+  }
+
+  @Test
+  void testNoScopeClaimSkipsValidation() {
+    Token token = makeToken(Collections.singletonMap("sub", "user@example.com"));
+    assertDoesNotThrow(
+        () ->
+            DatabricksCliCredentialsProvider.validateTokenScopes(
+                token, Collections.singletonList("sql"), HOST));
+  }
+
+  @Test
+  void testNonJwtTokenSkipsValidation() {
+    Token token = new Token("opaque-token-string", "Bearer", Instant.now().plusSeconds(3600));
+    assertDoesNotThrow(
+        () ->
+            DatabricksCliCredentialsProvider.validateTokenScopes(
+                token, Collections.singletonList("sql"), HOST));
+  }
+
+  @Test
+  void testErrorMessageContainsReauthCommand() {
+    Token token = makeToken(Collections.singletonMap("scope", "all-apis"));
+    DatabricksException e =
+        assertThrows(
+            DatabricksException.class,
+            () ->
+                DatabricksCliCredentialsProvider.validateTokenScopes(
+                    token, Arrays.asList("sql", "offline_access"), HOST));
+    assertTrue(
+        e.getMessage().contains("databricks auth login --host " + HOST + " --scopes sql"),
+        "Expected re-auth command in error message, got: " + e.getMessage());
+  }
+}
