From bb922453d67f3a00c8f6a1e87441ad02f17d4a11 Mon Sep 17 00:00:00 2001 From: Mihai Mitrea Date: Wed, 26 Aug 2026 08:08:45 +0000 Subject: [PATCH] Add group configuration and auth validation --- .../sdk/core/AzureCliCredentialsProvider.java | 10 ++ .../sdk/core/AzureMsiCredentialsProvider.java | 10 ++ .../sdk/core/BasicCredentialsProvider.java | 11 ++ .../DatabricksCliCredentialsProvider.java | 10 ++ .../databricks/sdk/core/DatabricksConfig.java | 16 +++ .../GoogleCredentialsCredentialsProvider.java | 10 ++ .../sdk/core/GoogleIdCredentialsProvider.java | 10 ++ .../databricks/sdk/core/GroupAssumption.java | 22 ++++ .../NotebookNativeCredentialsProvider.java | 10 ++ .../sdk/core/PatCredentialsProvider.java | 11 ++ .../AzureGithubOidcCredentialsProvider.java | 10 ++ ...reServicePrincipalCredentialsProvider.java | 10 ++ .../sdk/core/DatabricksConfigTest.java | 60 +++++++++++ ...supportedGroupCredentialsProviderTest.java | 101 ++++++++++++++++++ .../test/resources/testdata/.databrickscfg | 4 + 15 files changed, 305 insertions(+) create mode 100644 databricks-sdk-java/src/main/java/com/databricks/sdk/core/GroupAssumption.java create mode 100644 databricks-sdk-java/src/test/java/com/databricks/sdk/core/UnsupportedGroupCredentialsProviderTest.java diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/AzureCliCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/AzureCliCredentialsProvider.java index 77f9c9fdb..da0ddeb20 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/AzureCliCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/AzureCliCredentialsProvider.java @@ -81,6 +81,16 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { return null; } + if (GroupAssumption.isRequested(config)) { + // Return null during automatic discovery so the chain can continue. If the user explicitly + // requested this provider, throw an actionable error explaining why it cannot be used. + if (authType().equals(config.getAuthType())) { + throw GroupAssumption.unsupportedAuth(authType()); + } + + return null; + } + try { AzureUtils.ensureHostPresent(config, mapper, this::tokenSourceFor); String resource = config.getEffectiveAzureLoginAppId(); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/AzureMsiCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/AzureMsiCredentialsProvider.java index 4b7b86522..b86dbce8c 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/AzureMsiCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/AzureMsiCredentialsProvider.java @@ -42,6 +42,16 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { return null; } + if (GroupAssumption.isRequested(config)) { + // Return null during automatic discovery so the chain can continue. If the user explicitly + // requested this provider, throw an actionable error explaining why it cannot be used. + if (authType().equals(config.getAuthType())) { + throw GroupAssumption.unsupportedAuth(authType()); + } + + return null; + } + LOG.debug("Generating AAD token via Azure MSI"); AzureUtils.ensureHostPresent(config, mapper, this::tokenSourceFor); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/BasicCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/BasicCredentialsProvider.java index bf396ed54..74339a0c5 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/BasicCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/BasicCredentialsProvider.java @@ -22,6 +22,17 @@ public HeaderFactory configure(DatabricksConfig config) { if (username == null || password == null || host == null) { return null; } + + if (GroupAssumption.isRequested(config)) { + // Return null during automatic discovery so the chain can continue. If the user explicitly + // requested this provider, throw an actionable error explaining why it cannot be used. + if (authType().equals(config.getAuthType())) { + throw GroupAssumption.unsupportedAuth(authType()); + } + + return null; + } + byte[] bytes = String.format("%s:%s", config.getUsername(), config.getPassword()).getBytes(); String base64 = Base64.getEncoder().encodeToString(bytes); Map headers = new HashMap<>(); 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 index d2f7631ae..598e3713d 100644 --- 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 @@ -84,6 +84,16 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { return null; } + if (GroupAssumption.isRequested(config)) { + // Return null during automatic discovery so the chain can continue. If the user explicitly + // requested this provider, throw an actionable error explaining why it cannot be used. + if (authType().equals(config.getAuthType())) { + throw GroupAssumption.unsupportedAuth(authType()); + } + + return null; + } + try { CliTokenSource tokenSource = getDatabricksCliTokenSource(config); if (tokenSource == 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 index ed45b1d1b..7a50eed72 100644 --- 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 @@ -41,6 +41,13 @@ public class DatabricksConfig { @ConfigAttribute(env = "DATABRICKS_WORKSPACE_ID") private String workspaceId; + /** + * ID of the Databricks group to assume through role-based access control (RBAC). A client uses + * this group for its full lifetime. + */ + @ConfigAttribute(env = "DATABRICKS_GROUP_ID") + private String groupId; + @ConfigAttribute(env = "DATABRICKS_TOKEN", auth = "pat", sensitive = true) private String token; @@ -339,6 +346,15 @@ public DatabricksConfig setWorkspaceId(String workspaceId) { return this; } + public String getGroupId() { + return groupId; + } + + public DatabricksConfig setGroupId(String groupId) { + this.groupId = groupId; + return this; + } + public String getDatabricksCliPath() { return this.databricksCliPath; } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleCredentialsCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleCredentialsCredentialsProvider.java index 1852355f2..85b0f303c 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleCredentialsCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleCredentialsCredentialsProvider.java @@ -34,6 +34,16 @@ public HeaderFactory configure(DatabricksConfig config) { return null; } + if (GroupAssumption.isRequested(config)) { + // Return null during automatic discovery so the chain can continue. If the user explicitly + // requested this provider, throw an actionable error explaining why it cannot be used. + if (authType().equals(config.getAuthType())) { + throw GroupAssumption.unsupportedAuth(authType()); + } + + return null; + } + ServiceAccountCredentials serviceAccountCredentials; try { diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleIdCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleIdCredentialsProvider.java index 675f95b69..9bcbafca2 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleIdCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleIdCredentialsProvider.java @@ -31,6 +31,16 @@ public HeaderFactory configure(DatabricksConfig config) { return null; } + if (GroupAssumption.isRequested(config)) { + // Return null during automatic discovery so the chain can continue. If the user explicitly + // requested this provider, throw an actionable error explaining why it cannot be used. + if (authType().equals(config.getAuthType())) { + throw GroupAssumption.unsupportedAuth(authType()); + } + + return null; + } + GoogleCredentials googleCredentials; try { googleCredentials = GoogleCredentials.getApplicationDefault(); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GroupAssumption.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GroupAssumption.java new file mode 100644 index 000000000..f19bd5066 --- /dev/null +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GroupAssumption.java @@ -0,0 +1,22 @@ +package com.databricks.sdk.core; + +import com.databricks.sdk.support.InternalApi; + +/** Shared utilities for group assumption. */ +@InternalApi +public final class GroupAssumption { + private GroupAssumption() {} + + public static boolean isRequested(DatabricksConfig config) { + return config.getGroupId() != null && !config.getGroupId().isEmpty(); + } + + /** Creates an actionable error for an explicitly configured unsupported authentication type. */ + public static DatabricksException unsupportedAuth(String authType) { + String message = + "auth type \"%s\" does not support group assumption. " + + "Use Databricks OAuth or workload identity federation authentication"; + + return new DatabricksException(String.format(message, authType)); + } +} diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/NotebookNativeCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/NotebookNativeCredentialsProvider.java index ec3259ae3..544bb1f7c 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/NotebookNativeCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/NotebookNativeCredentialsProvider.java @@ -27,6 +27,16 @@ public String authType() { @Override public HeaderFactory configure(DatabricksConfig config) { + if (GroupAssumption.isRequested(config)) { + // Return null during automatic discovery so the chain can continue. If the user explicitly + // requested this provider, throw an actionable error explaining why it cannot be used. + if (authType().equals(config.getAuthType())) { + throw GroupAssumption.unsupportedAuth(authType()); + } + + return null; + } + if (System.getenv("DATABRICKS_RUNTIME_VERSION") == null) { LOG.debug("DBR not detected, skipping runtime auth"); return null; diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/PatCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/PatCredentialsProvider.java index 710678396..da7b1e5b0 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/PatCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/PatCredentialsProvider.java @@ -20,6 +20,17 @@ public HeaderFactory configure(DatabricksConfig config) { if (token == null || host == null) { return null; } + + if (GroupAssumption.isRequested(config)) { + // Return null during automatic discovery so the chain can continue. If the user explicitly + // requested this provider, throw an actionable error explaining why it cannot be used. + if (authType().equals(config.getAuthType())) { + throw GroupAssumption.unsupportedAuth(authType()); + } + + return null; + } + Map headers = new HashMap<>(); headers.put("Authorization", String.format("Bearer %s", token)); return () -> headers; diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureGithubOidcCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureGithubOidcCredentialsProvider.java index 6a5a8c278..f3986c12f 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureGithubOidcCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureGithubOidcCredentialsProvider.java @@ -33,6 +33,16 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { return null; } + if (GroupAssumption.isRequested(config)) { + // Return null during automatic discovery so the chain can continue. If the user explicitly + // requested this provider, throw an actionable error explaining why it cannot be used. + if (authType().equals(config.getAuthType())) { + throw GroupAssumption.unsupportedAuth(authType()); + } + + return null; + } + Optional idToken = requestIdToken(config); if (!idToken.isPresent()) { return null; diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java index 3ee6679b4..8c19b0415 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java @@ -33,6 +33,16 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { return null; } + if (GroupAssumption.isRequested(config)) { + // Return null during automatic discovery so the chain can continue. If the user explicitly + // requested this provider, throw an actionable error explaining why it cannot be used. + if (authType().equals(config.getAuthType())) { + throw GroupAssumption.unsupportedAuth(authType()); + } + + return null; + } + try { this.tenantId = config.getAzureTenantId() != null diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java index f1c885065..d43ade2dd 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java @@ -180,6 +180,7 @@ public void testNewWithWorkspaceHost() { .setAuthType("oauth-m2m") .setClientId("my-client-id") .setClientSecret("my-client-secret") + .setGroupId("group-123") .setAccountId("account-id") .setHost("https://account.cloud.databricks.com"); String workspaceHost = "https://workspace.cloud.databricks.com"; @@ -190,6 +191,7 @@ public void testNewWithWorkspaceHost() { assert newWorkspaceConfig.getAuthType().equals("oauth-m2m"); assert newWorkspaceConfig.getClientId().equals("my-client-id"); assert newWorkspaceConfig.getClientSecret().equals("my-client-secret"); + assert newWorkspaceConfig.getGroupId().equals("group-123"); } @Test @@ -199,6 +201,7 @@ public void testClone() { .setAuthType("oauth-m2m") .setClientId("my-client-id") .setClientSecret("my-client-secret") + .setGroupId("group-123") .setAccountId("account-id") .setHost("https://account.cloud.databricks.com"); @@ -208,6 +211,7 @@ public void testClone() { assert newWorkspaceConfig.getAuthType().equals("oauth-m2m"); assert newWorkspaceConfig.getClientId().equals("my-client-id"); assert newWorkspaceConfig.getClientSecret().equals("my-client-secret"); + assert newWorkspaceConfig.getGroupId().equals("group-123"); } @Test @@ -281,6 +285,7 @@ public void testEnvironmentVariableLoading() { env.put("DATABRICKS_OAUTH_BROWSER_AUTH_TIMEOUT", "30"); env.put("DATABRICKS_DEBUG_TRUNCATE_BYTES", "100"); env.put("DATABRICKS_RATE_LIMIT", "50"); + env.put("DATABRICKS_GROUP_ID", "environment-group"); DatabricksConfig config = new DatabricksConfig(); config.resolve(new Environment(env, new ArrayList<>(), System.getProperty("os.name"))); @@ -288,6 +293,61 @@ public void testEnvironmentVariableLoading() { assertEquals(Duration.ofSeconds(30), config.getOAuthBrowserAuthTimeout()); assertEquals(Integer.valueOf(100), config.getDebugTruncateBytes()); assertEquals(Integer.valueOf(50), config.getRateLimit()); + assertEquals("environment-group", config.getGroupId()); + } + + // Verifies that group ID follows the standard configuration precedence: code, environment, then + // profile. + @Test + public void testGroupIdCodeEnvironmentProfilePrecedence() { + Map env = new HashMap<>(); + env.put("HOME", TestOSUtils.resource("/testdata")); + Environment profileEnvironment = + new Environment(env, new ArrayList<>(), System.getProperty("os.name")); + + DatabricksConfig fromProfile = + new DatabricksConfig() + .setProfile("group") + .setHttpClient( + request -> { + throw new IOException("offline test"); + }); + fromProfile.resolve(profileEnvironment); + assertEquals("profile-group", fromProfile.getGroupId()); + + DatabricksConfig withoutGroup = + new DatabricksConfig() + .setProfile("scope-empty") + .setHttpClient( + request -> { + throw new IOException("offline test"); + }); + withoutGroup.resolve(profileEnvironment); + assertNull(withoutGroup.getGroupId()); + + env.put("DATABRICKS_GROUP_ID", "environment-group"); + Environment environmentOverride = + new Environment(env, new ArrayList<>(), System.getProperty("os.name")); + DatabricksConfig fromEnvironment = + new DatabricksConfig() + .setProfile("group") + .setHttpClient( + request -> { + throw new IOException("offline test"); + }); + fromEnvironment.resolve(environmentOverride); + assertEquals("environment-group", fromEnvironment.getGroupId()); + + DatabricksConfig fromCode = + new DatabricksConfig() + .setProfile("group") + .setGroupId("code-group") + .setHttpClient( + request -> { + throw new IOException("offline test"); + }); + fromCode.resolve(environmentOverride); + assertEquals("code-group", fromCode.getGroupId()); } @Test diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/UnsupportedGroupCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/UnsupportedGroupCredentialsProviderTest.java new file mode 100644 index 000000000..8be3441cd --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/UnsupportedGroupCredentialsProviderTest.java @@ -0,0 +1,101 @@ +package com.databricks.sdk.core; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.databricks.sdk.core.oauth.AzureGithubOidcCredentialsProvider; +import com.databricks.sdk.core.oauth.AzureServicePrincipalCredentialsProvider; +import java.util.stream.Stream; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +class UnsupportedGroupCredentialsProviderTest { + // Verifies that explicitly selecting an applicable provider without group support throws the + // existing SDK error type with an actionable message. + @ParameterizedTest(name = "{0}") + @MethodSource("unsupportedProviders") + void applicableUnsupportedProviderRejectsGroup( + String authType, CredentialsProvider provider, DatabricksConfig config) { + config.setGroupId("group-123").setAuthType(authType); + + DatabricksException error = + assertThrows(DatabricksException.class, () -> provider.configure(config)); + + assertEquals(DatabricksException.class, error.getClass()); + assertTrue(error.getMessage().contains(authType)); + assertTrue(error.getMessage().contains("does not support group assumption")); + } + + // Verifies that an applicable provider without group support returns null during automatic + // discovery, allowing the credential chain to continue to another provider. + @ParameterizedTest(name = "{0}") + @MethodSource("unsupportedProviders") + void automaticChainSkipsApplicableUnsupportedProvider( + String authType, CredentialsProvider provider, DatabricksConfig config) { + config.setGroupId("group-123").setAuthType(null); + + assertNull(provider.configure(config)); + } + + private static Stream unsupportedProviders() { + String azureHost = "https://adb-123.4.azuredatabricks.net"; + String gcpHost = "https://123.4.gcp.databricks.com"; + return Stream.of( + Arguments.of( + "pat", + new PatCredentialsProvider(), + new DatabricksConfig() + .setHost("https://workspace.cloud.databricks.com") + .setToken("token")), + Arguments.of( + "basic", + new BasicCredentialsProvider(), + new DatabricksConfig() + .setHost("https://workspace.cloud.databricks.com") + .setUsername("user") + .setPassword("password")), + Arguments.of( + "databricks-cli", + new DatabricksCliCredentialsProvider(), + new DatabricksConfig() + .setHost("https://workspace.cloud.databricks.com") + .setAuthType("databricks-cli")), + Arguments.of( + "runtime", + new NotebookNativeCredentialsProvider(), + new DatabricksConfig().setAuthType("runtime")), + Arguments.of( + "azure-cli", + new AzureCliCredentialsProvider(), + new DatabricksConfig().setHost(azureHost)), + Arguments.of( + "azure-msi", + new AzureMsiCredentialsProvider(), + new DatabricksConfig().setHost(azureHost).setAzureUseMsi(true)), + Arguments.of( + "azure-client-secret", + new AzureServicePrincipalCredentialsProvider(), + new DatabricksConfig() + .setHost(azureHost) + .setAzureClientId("client") + .setAzureClientSecret("secret")), + Arguments.of( + "github-oidc-azure", + new AzureGithubOidcCredentialsProvider(), + new DatabricksConfig() + .setHost(azureHost) + .setAzureClientId("client") + .setAzureTenantId("tenant")), + Arguments.of( + "google-credentials", + new GoogleCredentialsCredentialsProvider(), + new DatabricksConfig().setHost(gcpHost).setGoogleCredentials("credentials")), + Arguments.of( + "google-id", + new GoogleIdCredentialsProvider(), + new DatabricksConfig().setHost(gcpHost).setGoogleServiceAccount("service-account"))); + } +} diff --git a/databricks-sdk-java/src/test/resources/testdata/.databrickscfg b/databricks-sdk-java/src/test/resources/testdata/.databrickscfg index 3b4fa447d..d3b113e74 100644 --- a/databricks-sdk-java/src/test/resources/testdata/.databrickscfg +++ b/databricks-sdk-java/src/test/resources/testdata/.databrickscfg @@ -50,3 +50,7 @@ scopes = clusters:read [scope-multiple] host = https://example.cloud.databricks.com scopes = clusters, jobs, pipelines, iam:read, files:read, mlflow, model-serving:read + +[group] +host = https://example.cloud.databricks.com +group_id = profile-group