Add group assumption for OAuth authentication

NEXT_CHANGELOG.md

@@ -4,6 +4,9 @@
 
 ### New Features and Improvements
 
+* Added group assumption through `group_id` / `DATABRICKS_GROUP_ID` for external-browser
+  OAuth, OAuth M2M, and Databricks workload identity federation authentication.
+
 ### Breaking Changes
 
 ### Bug Fixes

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

@@ -151,6 +151,7 @@ private void addOIDCCredentialsProviders(DatabricksConfig config) {
                   namedIdTokenSource.idTokenSource,
                   config.getHttpClient())
               .audience(config.getTokenAudience())
+              .groupId(config.getGroupId())
               .accountId(
                   config.getClientType() == ClientType.ACCOUNT ? config.getAccountId() : null)
               .scopes(config.getScopes())

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ClientCredentials.java

@@ -17,6 +17,7 @@ public class ClientCredentials implements TokenSource {
   public static class Builder {
     private String clientId;
     private String clientSecret;
+    private String groupId;
     private String tokenUrl;
     private HttpClient hc = new CommonsHttpClient.Builder().withTimeoutSeconds(30).build();
 
@@ -37,6 +38,11 @@ public Builder withClientSecret(String clientSecret) {
       return this;
     }
 
+    public Builder withGroupId(String groupId) {
+      this.groupId = groupId;
+      return this;
+    }
+
     public Builder withEndpointParametersSupplier(
         Supplier<Map<String, String>> endpointParamsSupplier) {
       this.endpointParamsSupplier = endpointParamsSupplier;
@@ -67,13 +73,14 @@ public ClientCredentials build() {
       Objects.requireNonNull(this.clientId, "clientId must be specified");
       Objects.requireNonNull(this.tokenUrl, "tokenUrl must be specified");
       return new ClientCredentials(
-          hc, clientId, clientSecret, tokenUrl, endpointParamsSupplier, scopes, position);
+          hc, clientId, clientSecret, groupId, tokenUrl, endpointParamsSupplier, scopes, position);
     }
   }
 
   private HttpClient hc;
   private String clientId;
   private String clientSecret;
+  private String groupId;
   private String tokenUrl;
   private List<String> scopes;
   private AuthParameterPosition position;
@@ -83,13 +90,15 @@ private ClientCredentials(
       HttpClient hc,
       String clientId,
       String clientSecret,
+      String groupId,
       String tokenUrl,
       Supplier<Map<String, String>> endpointParamsSupplier,
       List<String> scopes,
       AuthParameterPosition position) {
     this.hc = hc;
     this.clientId = clientId;
     this.clientSecret = clientSecret;
+    this.groupId = groupId;
     this.tokenUrl = tokenUrl;
     this.endpointParamsSupplier = endpointParamsSupplier;
     this.scopes = scopes;
@@ -103,6 +112,9 @@ public Token getToken() {
     if (scopes != null) {
       params.put("scope", String.join(" ", scopes));
     }
+    if (groupId != null && !groupId.isEmpty()) {
+      params.put("assume_group", groupId);
+    }
     if (endpointParamsSupplier != null) {
       params.putAll(endpointParamsSupplier.get());
     }

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java

@@ -43,13 +43,17 @@ public class DatabricksOAuthTokenSource implements TokenSource {
   /** Scopes to request during token exchange. */
   private final List<String> scopes;
 
+  /** Group the exchanged token assumes. */
+  private final String groupId;
+
   private static final String GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange";
   private static final String SUBJECT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:jwt";
   private static final String GRANT_TYPE_PARAM = "grant_type";
   private static final String SUBJECT_TOKEN_PARAM = "subject_token";
   private static final String SUBJECT_TOKEN_TYPE_PARAM = "subject_token_type";
   private static final String SCOPE_PARAM = "scope";
   private static final String CLIENT_ID_PARAM = "client_id";
+  private static final String ASSUME_GROUP_PARAM = "assume_group";
 
   private DatabricksOAuthTokenSource(Builder builder) {
     this.clientId = builder.clientId;
@@ -60,6 +64,7 @@ private DatabricksOAuthTokenSource(Builder builder) {
     this.idTokenSource = builder.idTokenSource;
     this.httpClient = builder.httpClient;
     this.scopes = builder.scopes == null ? Arrays.asList() : builder.scopes;
+    this.groupId = builder.groupId;
   }
 
   /**
@@ -75,6 +80,7 @@ public static class Builder {
     private String accountId;
     private String audience;
     private List<String> scopes;
+    private String groupId;
 
     /**
      * Creates a new Builder with required parameters.
@@ -133,6 +139,12 @@ public Builder scopes(List<String> scopes) {
       return this;
     }
 
+    /** Sets the Databricks group the exchanged token should assume. */
+    public Builder groupId(String groupId) {
+      this.groupId = groupId;
+      return this;
+    }
+
     /**
      * Builds a new DatabricksOAuthTokenSource instance.
      *
@@ -178,6 +190,9 @@ public Token getToken() {
     if (!Strings.isNullOrEmpty(clientId)) {
       params.put(CLIENT_ID_PARAM, clientId);
     }
+    if (!Strings.isNullOrEmpty(groupId)) {
+      params.put(ASSUME_GROUP_PARAM, groupId);
+    }
 
     OAuthResponse response;
     try {

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java

@@ -70,7 +70,8 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
       if (tokenCache == null) {
         // Create a default FileTokenCache based on config
         Path cachePath =
-            TokenCacheUtils.getCacheFilePath(config.getHost(), clientId, config.getScopes());
+            TokenCacheUtils.getCacheFilePath(
+                config.getHost(), clientId, config.getScopes(), config.getGroupId());
         tokenCache = new FileTokenCache(cachePath);
       }
 
@@ -147,6 +148,7 @@ CachedTokenSource performBrowserAuth(
             .withClientSecret(clientSecret)
             .withHost(config.getHost())
             .withAccountId(config.getAccountId())
+            .withGroupId(config.getGroupId())
             .withRedirectUrl(config.getEffectiveOAuthRedirectUrl())
             .withBrowserTimeout(config.getOAuthBrowserAuthTimeout())
             .withScopes(getScopes(config, oidcEndpoints))

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java

@@ -44,6 +44,7 @@ public static class Builder {
     private String clientSecret;
     private HttpClient hc;
     private String accountId;
+    private String groupId;
     private Optional<Duration> browserTimeout = Optional.empty();
     private OpenIDConnectEndpoints openIDConnectEndpoints;
 
@@ -93,6 +94,11 @@ public Builder withAccountId(String accountId) {
       return this;
     }
 
+    public Builder withGroupId(String groupId) {
+      this.groupId = groupId;
+      return this;
+    }
+
     public Builder withBrowserTimeout(Duration browserTimeout) {
       this.browserTimeout = Optional.of(browserTimeout);
       return this;
@@ -112,6 +118,7 @@ public Builder withBrowserTimeout(Duration browserTimeout) {
   private final boolean isAzure;
   private final OpenIDConnectEndpoints openIDConnectEndpoints;
   private final Optional<Duration> browserTimeout;
+  private final String groupId;
 
   private OAuthClient(Builder b) throws IOException {
     this.clientId = Objects.requireNonNull(b.clientId);
@@ -121,7 +128,11 @@ private OAuthClient(Builder b) throws IOException {
     this.hc = b.hc;
 
     DatabricksConfig config =
-        new DatabricksConfig().setHost(b.host).setAccountId(b.accountId).resolve();
+        new DatabricksConfig()
+            .setHost(b.host)
+            .setAccountId(b.accountId)
+            .setHttpClient(b.hc)
+            .resolve();
     openIDConnectEndpoints = b.openIDConnectEndpoints;
     if (openIDConnectEndpoints == null) {
       throw new DatabricksException(b.host + " does not support OAuth");
@@ -133,6 +144,7 @@ private OAuthClient(Builder b) throws IOException {
     this.authUrl = openIDConnectEndpoints.getAuthorizationEndpoint();
     this.browserTimeout = b.browserTimeout;
     this.scopes = b.scopes;
+    this.groupId = b.groupId;
   }
 
   public String getHost() {
@@ -235,6 +247,9 @@ public Consent initiateConsent() throws MalformedURLException {
     params.put("state", state);
     params.put("code_challenge", challenge);
     params.put("code_challenge_method", "S256");
+    if (groupId != null && !groupId.isEmpty()) {
+      params.put("assume_group", groupId);
+    }
 
     String url = urlEncode(authUrl, params);
 

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthM2MServicePrincipalCredentialsProvider.java

@@ -33,6 +33,7 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
               .withHttpClient(config.getHttpClient())
               .withClientId(config.getClientId())
               .withClientSecret(config.getClientSecret())
+              .withGroupId(config.getGroupId())
               .withTokenUrl(jsonResponse.getTokenEndpoint())
               .withScopes(config.getScopes())
               .withAuthParameterPosition(AuthParameterPosition.HEADER)

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenCacheUtils.java

@@ -24,16 +24,40 @@ public class TokenCacheUtils {
    * @return The path to the token cache file
    */
   public static Path getCacheFilePath(String host, String clientId, List<String> scopes) {
+    return getCacheFilePath(host, clientId, scopes, null);
+  }
+
+  /**
+   * Returns the cache path for an OAuth configuration and its fixed assumed group. The empty-group
+   * path intentionally remains byte-for-byte compatible with older SDK versions.
+   */
+  public static Path getCacheFilePath(
+      String host, String clientId, List<String> scopes, String groupId) {
     try {
       // Create SHA-256 hash of host, client_id, and scopes
       MessageDigest hash = MessageDigest.getInstance("SHA-256");
       for (String chunk : new String[] {host, clientId, String.join(",", scopes)}) {
         hash.update(chunk.getBytes(StandardCharsets.UTF_8));
       }
 
+      // Finalize the legacy cache key before including the group. Keeping this digest unchanged is
+      // important because users without a group may already have tokens stored at the path created
+      // by older SDK versions.
+      byte[] cacheKeyDigest = hash.digest();
+
+      if (groupId != null && !groupId.isEmpty()) {
+        // A token issued for one assumed group must never be reused for another group or for a
+        // non-group session. Hash the fixed-length legacy digest together with the group ID to
+        // create a separate cache namespace.
+        hash.reset();
+        hash.update(cacheKeyDigest);
+        hash.update(groupId.getBytes(StandardCharsets.UTF_8));
+        cacheKeyDigest = hash.digest();
+      }
+
       // Convert hash bytes to hexadecimal string
       StringBuilder hexString = new StringBuilder();
-      for (byte b : hash.digest()) {
+      for (byte b : cacheKeyDigest) {
         String hex = Integer.toHexString(0xff & b);
         if (hex.length() == 1) {
           hexString.append('0');