diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 6328d4f6d..1fe26a18c 100644 --- a/NEXT_CHANGELOG.md +++ b/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 diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java index 3b1a98026..3fc30e874 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java +++ b/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()) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ClientCredentials.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ClientCredentials.java index 8cee3ef29..e30ea15ec 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ClientCredentials.java +++ b/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> 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 scopes; private AuthParameterPosition position; @@ -83,6 +90,7 @@ private ClientCredentials( HttpClient hc, String clientId, String clientSecret, + String groupId, String tokenUrl, Supplier> endpointParamsSupplier, List scopes, @@ -90,6 +98,7 @@ private ClientCredentials( 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()); } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java index 250abb241..8007b7643 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java @@ -43,6 +43,9 @@ public class DatabricksOAuthTokenSource implements TokenSource { /** Scopes to request during token exchange. */ private final List 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"; @@ -50,6 +53,7 @@ public class DatabricksOAuthTokenSource implements TokenSource { 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 scopes; + private String groupId; /** * Creates a new Builder with required parameters. @@ -133,6 +139,12 @@ public Builder scopes(List 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 { diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java index 7c8d0fe4a..a9cf56cf5 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java +++ b/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)) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java index 9af534703..8ed6dba95 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java +++ b/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 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 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); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthM2MServicePrincipalCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthM2MServicePrincipalCredentialsProvider.java index c702d5707..72441a627 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthM2MServicePrincipalCredentialsProvider.java +++ b/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) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenCacheUtils.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenCacheUtils.java index fa7ff8706..7c39af100 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenCacheUtils.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenCacheUtils.java @@ -24,6 +24,15 @@ public class TokenCacheUtils { * @return The path to the token cache file */ public static Path getCacheFilePath(String host, String clientId, List 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 scopes, String groupId) { try { // Create SHA-256 hash of host, client_id, and scopes MessageDigest hash = MessageDigest.getInstance("SHA-256"); @@ -31,9 +40,24 @@ public static Path getCacheFilePath(String host, String clientId, List s 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'); diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/AuthProfilesTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/AuthProfilesTest.java index 332d4c5be..8cfa66dc3 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/AuthProfilesTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/AuthProfilesTest.java @@ -1,13 +1,15 @@ package com.databricks.sdk.core; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; -import com.databricks.sdk.core.http.HttpClient; import com.databricks.sdk.core.http.Request; -import com.databricks.sdk.core.http.Response; +import com.databricks.sdk.core.oauth.Consent; +import com.databricks.sdk.core.oauth.OAuthClient; +import com.databricks.sdk.core.oauth.OpenIDConnectEndpoints; import com.databricks.sdk.core.utils.Environment; import java.io.IOException; -import java.net.URL; import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.nio.file.Path; @@ -317,6 +319,27 @@ void oauthM2M(HostProfile p) { assertEquals("oauth-m2m", config.getAuthType()); } + // Verifies that M2M sends assume_group in the token form for workspace, account, and unified + // hosts without changing the discovered token endpoint. + @ParameterizedTest(name = "{0}") + @MethodSource("allProfiles") + void oauthM2MForwardsGroupOnExistingEndpoint(HostProfile p) { + DatabricksConfig config = + profileConfig(p) + .setClientId("test-client") + .setClientSecret("test-secret") + .setGroupId("group-123") + .setAuthType("oauth-m2m"); + MappingHttpClient client = withOidc(httpClientFor(p), p); + config.setHttpClient(client); + + assertEquals("Bearer test-token", resolveAndAuthenticate(config, emptyEnvironment())); + + Request tokenRequest = client.singleRequest("POST", p.tokenPath()); + assertTrue(tokenRequest.getBodyString().contains("assume_group=group-123")); + assertFalse(tokenRequest.getUrl().contains("?o=")); + } + // ---- GitHub OIDC ----------------------------------------------------------- @ParameterizedTest(name = "{0}") @@ -352,6 +375,74 @@ void envOIDC(HostProfile p) { assertEquals("env-oidc", config.getAuthType()); } + // Verifies that workload identity federation sends assume_group in the token-exchange form for + // every host type without adding the historical o= routing query parameter. + @ParameterizedTest(name = "{0}") + @MethodSource("allProfiles") + void envOIDCForwardsGroupOnExistingEndpoint(HostProfile p) { + DatabricksConfig config = + profileConfig(p).setClientId("test-client").setGroupId("group-123").setAuthType("env-oidc"); + MappingHttpClient client = withOidc(httpClientFor(p), p); + config.setHttpClient(client); + + assertEquals( + "Bearer test-token", + resolveAndAuthenticate( + config, environmentWith("DATABRICKS_OIDC_TOKEN", "test-oidc-token"))); + + Request tokenRequest = client.singleRequest("POST", p.tokenPath()); + assertTrue(tokenRequest.getBodyString().contains("assume_group=group-123")); + assertFalse(tokenRequest.getUrl().contains("?o=")); + } + + // Verifies that browser OAuth on workspace, account, and unified hosts adds assume_group only to + // the authorization request. The discovered endpoints remain unchanged, and the authorization + // code exchange neither receives assume_group nor adds the historical o= routing + // query parameter. + @ParameterizedTest(name = "{0}") + @MethodSource("allProfiles") + void externalBrowserForwardsGroupOnlyOnAuthorization(HostProfile p) throws IOException { + DatabricksConfig config = + profileConfig(p) + .setClientId("test-client") + .setClientSecret("test-secret") + .setGroupId("group-123") + .setAuthType("external-browser"); + MappingHttpClient client = withOidc(httpClientFor(p), p); + config.setHttpClient(client); + config.resolve(emptyEnvironment()); + + OpenIDConnectEndpoints endpoints = config.getDatabricksOidcEndpoints(); + OAuthClient oauthClient = + new OAuthClient.Builder() + .withHttpClient(client) + .withClientId(config.getClientId()) + .withClientSecret(config.getClientSecret()) + .withHost(config.getHost()) + .withAccountId(config.getAccountId()) + .withGroupId(config.getGroupId()) + .withRedirectUrl(config.getEffectiveOAuthRedirectUrl()) + .withScopes(config.getScopes()) + .withOpenIDConnectEndpoints(endpoints) + .build(); + + Consent consent = oauthClient.initiateConsent(); + assertEquals(p.authorizationEndpoint(), consent.getAuthUrl().split("\\?", 2)[0]); + assertTrue(consent.getAuthUrl().contains("assume_group=group-123")); + assertFalse(consent.getAuthUrl().contains("?o=")); + assertEquals(p.tokenEndpoint(), consent.getTokenUrl()); + + Map callback = new HashMap<>(); + callback.put("code", "authorization-code"); + callback.put("state", consent.getState()); + consent.exchangeCallbackParameters(callback); + + Request tokenRequest = client.singleRequest("POST", p.tokenPath()); + assertTrue(tokenRequest.getBodyString().contains("grant_type=authorization_code")); + assertFalse(tokenRequest.getBodyString().contains("assume_group")); + assertFalse(tokenRequest.getUrl().contains("?o=")); + } + // ---- File OIDC ------------------------------------------------------------- @ParameterizedTest(name = "{0}") @@ -440,33 +531,4 @@ void hostMetadataResolutionPopulatesDiscoveryUrl(HostProfile p) { assertEquals(TEST_WORKSPACE_ID, config.getWorkspaceId()); } } - - // ---- Minimal HttpClient fixture ------------------------------------------- - - /** - * Matches requests on {@code "METHOD path"} and returns a stubbed JSON body with HTTP 200. Every - * test must register a mapping for {@code GET /.well-known/databricks-config} (see {@link - * #httpClientFor(HostProfile)}); unmapped requests fail loudly with {@link IOException} so a - * missing fixture cannot silently fall through. - */ - private static class MappingHttpClient implements HttpClient { - private final Map mappings = new HashMap<>(); - - MappingHttpClient put(String key, String jsonBody) { - mappings.put(key, jsonBody); - return this; - } - - @Override - public Response execute(Request in) throws IOException { - String rawUrl = in.getUrl(); - URL url = new URL(rawUrl); - String key = in.getMethod() + " " + url.getPath(); - String body = mappings.get(key); - if (body == null) { - throw new IOException("No mock for " + key + " (url=" + rawUrl + ")"); - } - return new Response(body, 200, "OK", url); - } - } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/GroupDefaultCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/GroupDefaultCredentialsProviderTest.java new file mode 100644 index 000000000..3f8858ed9 --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/GroupDefaultCredentialsProviderTest.java @@ -0,0 +1,45 @@ +package com.databricks.sdk.core; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.databricks.sdk.core.http.Request; +import org.junit.jupiter.api.Test; + +class GroupDefaultCredentialsProviderTest { + // Verifies that the default chain skips an unsupported provider and reaches a group-aware M2M + // provider instead of using credentials for the unassumed identity. + @Test + void defaultChainContinuesAfterUnsupportedProvider() { + MappingHttpClient httpClient = + new MappingHttpClient() + .put("GET /.well-known/databricks-config", "{}") + .put( + "GET /oidc/.well-known/oauth-authorization-server", + "{\"token_endpoint\":\"https://workspace.example/oidc/v1/token\"}") + .put( + "POST /oidc/v1/token", + "{\"token_type\":\"Bearer\",\"access_token\":\"role-token\"," + + "\"expires_in\":3600}"); + + DatabricksConfig config = + new DatabricksConfig() + .setHost("https://workspace.example") + .setDiscoveryUrl( + "https://workspace.example/oidc/.well-known/oauth-authorization-server") + .setToken("normal-pat-must-not-be-used") + .setClientId("test-client") + .setClientSecret("test-secret") + .setGroupId("group-123"); + config.setHttpClient(httpClient); + + DefaultCredentialsProvider provider = new DefaultCredentialsProvider(); + HeaderFactory headers = provider.configure(config); + + assertEquals("Bearer role-token", headers.headers().get("Authorization")); + assertEquals("oauth-m2m", provider.authType()); + + Request tokenRequest = httpClient.singleRequest("POST", "/oidc/v1/token"); + assertTrue(tokenRequest.getBodyString().contains("assume_group=group-123")); + } +} diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/MappingHttpClient.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/MappingHttpClient.java new file mode 100644 index 000000000..3dab09115 --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/MappingHttpClient.java @@ -0,0 +1,68 @@ +package com.databricks.sdk.core; + +import com.databricks.sdk.core.http.HttpClient; +import com.databricks.sdk.core.http.Request; +import com.databricks.sdk.core.http.Response; +import java.io.IOException; +import java.net.URL; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * Minimal HTTP fixture that matches requests on {@code "METHOD path"} and returns a stubbed JSON + * body with HTTP 200. Unmapped requests fail loudly so a missing fixture cannot silently fall + * through. + */ +class MappingHttpClient implements HttpClient { + private final Map mappings = new HashMap<>(); + private final List requests = new ArrayList<>(); + + MappingHttpClient put(String key, String jsonBody) { + mappings.put(key, jsonBody); + return this; + } + + /** + * Returns the only recorded request with the given HTTP method and URL path. The host and query + * parameters are intentionally ignored so parameterized host-profile tests can inspect the + * request body independently of the endpoint host. Fails when no request or multiple requests + * match because either result would make the test assertion ambiguous. + */ + Request singleRequest(String method, String path) { + Request match = null; + for (Request request : requests) { + try { + if (method.equals(request.getMethod()) + && path.equals(new URL(request.getUrl()).getPath())) { + if (match != null) { + throw new AssertionError("Multiple requests matched " + method + " " + path); + } + match = request; + } + } catch (IOException e) { + throw new AssertionError("Invalid request URL", e); + } + } + if (match == null) { + throw new AssertionError("No request matched " + method + " " + path); + } + return match; + } + + @Override + public Response execute(Request request) throws IOException { + requests.add(request); + + String rawUrl = request.getUrl(); + URL url = new URL(rawUrl); + String key = request.getMethod() + " " + url.getPath(); + String body = mappings.get(key); + if (body == null) { + throw new IOException("No mock for " + key + " (url=" + rawUrl + ")"); + } + + return new Response(body, 200, "OK", url); + } +} diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSourceTest.java index 93c0fc816..b412814e4 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSourceTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSourceTest.java @@ -4,17 +4,23 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.*; +import com.databricks.sdk.core.DatabricksConfig; import com.databricks.sdk.core.DatabricksException; import com.databricks.sdk.core.http.FormRequest; import com.databricks.sdk.core.http.HttpClient; +import com.databricks.sdk.core.http.Request; import com.databricks.sdk.core.http.Response; import com.fasterxml.jackson.databind.ObjectMapper; import java.io.IOException; import java.net.MalformedURLException; import java.net.URL; +import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; +import java.util.List; import java.util.Map; import java.util.stream.Stream; +import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.MethodSource; import org.mockito.Mockito; @@ -286,13 +292,12 @@ private static Stream provideTestCases() throws MalformedURLException } } - private static HttpClient createMockHttpClient( - FormRequest expectedRequest, int statusCode, String responseBody) { + private static HttpClient createMockHttpClient(FormRequest request, int status, String body) { try { HttpClient mockHttpClient = Mockito.mock(HttpClient.class); - String statusMessage = statusCode == 200 ? "OK" : "Bad Request"; - when(mockHttpClient.execute(expectedRequest)) - .thenReturn(new Response(responseBody, statusCode, statusMessage, new URL(TEST_HOST))); + String statusMessage = status == 200 ? "OK" : "Bad Request"; + when(mockHttpClient.execute(request)) + .thenReturn(new Response(body, status, statusMessage, new URL(TEST_HOST))); return mockHttpClient; } catch (IOException e) { throw new RuntimeException("Failed to create mock HTTP client", e); @@ -342,4 +347,138 @@ void testTokenSource(TestCase testCase) { verify(testCase.idTokenSource, atLeastOnce()).getIDToken(testCase.expectedAudience); } } + + // Verifies that every workload identity token exchange sends the configured group to the token + // endpoint, rather than only the first exchange. + @Test + void everyTokenExchangeIncludesTheGroup() throws Exception { + RecordingClient httpClient = new RecordingClient(ignored -> tokenResponse("token")); + + DatabricksOAuthTokenSource source = + builder(httpClient, "group-123").scopes(Collections.singletonList("all-apis")).build(); + + assertEquals("token", source.getToken().getAccessToken()); + assertEquals("token", source.getToken().getAccessToken()); + assertEquals(2, httpClient.requests.size()); + assertAllRequestsContain(httpClient, "assume_group=group-123"); + } + + // Verifies that separate WIF clients for normal access and different groups each cache only + // their own access token. + @Test + void groupCachesAreIsolatedByClient() throws Exception { + RecordingClient httpClient = + new RecordingClient( + request -> { + String requestBody = request.getBodyString(); + if (requestBody.contains("assume_group=group-a")) { + return tokenResponse("token-group-a"); + } + if (requestBody.contains("assume_group=group-b")) { + return tokenResponse("token-group-b"); + } + return tokenResponse("token-normal"); + }); + + assertCachedToken(httpClient, null, "Bearer token-normal"); + assertCachedToken(httpClient, "group-a", "Bearer token-group-a"); + assertCachedToken(httpClient, "group-b", "Bearer token-group-b"); + + assertEquals(3, httpClient.requests.size()); + } + + // Verifies that adding a group does not wrap or otherwise change errors returned by the token + // endpoint. + @Test + void groupServerFailureIsReturnedNormally() throws Exception { + RecordingClient httpClient = + new RecordingClient( + ignored -> + new Response( + "{\"error\":\"invalid_request\",\"error_description\":\"assume_group is not" + + " supported at the account level\"}", + 400, + "Bad Request", + new URL(TEST_HOST))); + + DatabricksOAuthTokenSource source = builder(httpClient, "group-123").build(); + + DatabricksException error = assertThrows(DatabricksException.class, source::getToken); + assertEquals(DatabricksException.class, error.getClass()); + assertEquals( + "Token request failed with error: invalid_request - " + + "assume_group is not supported at the account level", + error.getMessage()); + assertNull(error.getCause()); + assertEquals(1, httpClient.requests.size()); + } + + /** Verifies that every request recorded by the client contains the expected form field. */ + private static void assertAllRequestsContain(RecordingClient client, String field) { + for (Request request : client.requests) { + assertTrue(request.getBodyString().contains(field)); + } + } + + /** Creates a token source builder with the fixed WIF configuration used by these tests. */ + private static DatabricksOAuthTokenSource.Builder builder(HttpClient client, String groupId) { + try { + OpenIDConnectEndpoints endpoints = + new OpenIDConnectEndpoints(TEST_TOKEN_ENDPOINT, TEST_AUTHORIZATION_ENDPOINT); + IDTokenSource idTokenSource = mock(IDTokenSource.class); + when(idTokenSource.getIDToken(any())).thenReturn(new IDToken(TEST_ID_TOKEN)); + + return new DatabricksOAuthTokenSource.Builder( + TEST_CLIENT_ID, TEST_HOST, endpoints, idTokenSource, client) + .groupId(groupId); + } catch (MalformedURLException e) { + throw new RuntimeException("Failed to create test OIDC endpoints", e); + } + } + + /** Creates a successful OAuth response containing the requested access token. */ + private static Response tokenResponse(String token) throws MalformedURLException { + return new Response( + String.format( + "{\"access_token\":\"%s\",\"token_type\":\"Bearer\",\"expires_in\":3600}", token), + 200, + "OK", + new URL(TEST_HOST)); + } + + /** Verifies that one client fetches its expected token once and then reuses the cached token. */ + private static void assertCachedToken(RecordingClient client, String group, String header) { + DatabricksOAuthTokenSource source = builder(client, group).build(); + TokenSourceCredentialsProvider provider = + new TokenSourceCredentialsProvider(source, "test-oidc"); + OAuthHeaderFactory headers = + provider.configure(new DatabricksConfig().setDisableAsyncTokenRefresh(true)); + + assertEquals(header, headers.headers().get("Authorization")); + assertEquals(header, headers.headers().get("Authorization")); + } + + /** Produces an HTTP response for a recorded token request. */ + @FunctionalInterface + private interface TokenResponseFactory { + Response create(Request request) throws IOException; + } + + /** Records token requests and delegates response creation to a test-specific factory. */ + private static class RecordingClient implements HttpClient { + private final List requests = new ArrayList<>(); + private final TokenResponseFactory responseFactory; + + /** Creates a recording client that uses the supplied factory for every response. */ + private RecordingClient(TokenResponseFactory responseFactory) { + this.responseFactory = responseFactory; + } + + /** Records the request before returning the response selected by the test. */ + @Override + public Response execute(Request request) throws IOException { + requests.add(request); + return responseFactory.create(request); + } + } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java index 284727ebb..04476b5b1 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java @@ -41,6 +41,7 @@ void clientAndConsentTest() throws IOException { .setAuthType("external-browser") .setHost(fixtures.getUrl()) .setClientId("test-client-id") + .setGroupId("group-123") .setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); config.resolve(); @@ -53,6 +54,7 @@ void clientAndConsentTest() throws IOException { .withClientId(config.getClientId()) .withClientSecret(config.getClientSecret()) .withHost(config.getHost()) + .withGroupId(config.getGroupId()) .withOpenIDConnectEndpoints(config.getDatabricksOidcEndpoints()) .withRedirectUrl(config.getEffectiveOAuthRedirectUrl()) .withScopes(config.getScopes()) @@ -68,6 +70,8 @@ void clientAndConsentTest() throws IOException { assertTrue(authUrl.contains("client_id=test-client-id")); assertTrue(authUrl.contains("redirect_uri=http%3A%2F%2Flocalhost%3A8080%2Fcallback")); assertTrue(authUrl.contains("scope=all-apis")); + assertTrue(authUrl.contains("assume_group=group-123")); + assertFalse(testConsent.getTokenUrl().contains("assume_group")); } } @@ -117,6 +121,9 @@ void clientAndConsentTestWithCustomRedirectUrl() throws IOException { assertTrue(authUrl.contains("client_id=test-client-id")); assertTrue(authUrl.contains("redirect_uri=http%3A%2F%2Flocalhost%3A8010")); assertTrue(authUrl.contains("scope=sql")); + // Verifies that existing browser authorization requests remain unchanged when no group is + // configured. + assertFalse(authUrl.contains("assume_group")); } } @@ -420,6 +427,7 @@ void cacheWithInvalidAccessTokenValidRefreshTest() throws IOException { .setAuthType("external-browser") .setHost("https://test.databricks.com") .setClientId("test-client-id") + .setGroupId("group-123") .setHttpClient(mockHttpClient); // We need to provide OIDC endpoints for token refresh @@ -447,7 +455,10 @@ void cacheWithInvalidAccessTokenValidRefreshTest() throws IOException { Mockito.verify(mockTokenCache, Mockito.times(1)).load(); // Verify HTTP call was made to refresh the token - Mockito.verify(mockHttpClient, Mockito.times(1)).execute(any(Request.class)); + ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(Request.class); + Mockito.verify(mockHttpClient, Mockito.times(1)).execute(requestCaptor.capture()); + assertTrue(requestCaptor.getValue().getBodyString().contains("grant_type=refresh_token")); + assertFalse(requestCaptor.getValue().getBodyString().contains("assume_group")); // Verify performBrowserAuth was NOT called since refresh succeeded Mockito.verify(provider, Mockito.never()) diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileTokenCacheTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileTokenCacheTest.java index 710fbf1b8..73f961ead 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileTokenCacheTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileTokenCacheTest.java @@ -66,6 +66,26 @@ void testNullPathRejection() { "Should throw NullPointerException for null path"); } + // Verifies that existing non-group cache paths remain compatible while each group receives an + // isolated cache entry. + @Test + void testGroupCacheIsolationAndNoGroupCompatibility() { + Path legacyPath = TokenCacheUtils.getCacheFilePath(TEST_HOST, TEST_CLIENT_ID, TEST_SCOPES); + Path nullGroupPath = + TokenCacheUtils.getCacheFilePath(TEST_HOST, TEST_CLIENT_ID, TEST_SCOPES, null); + Path emptyGroupPath = + TokenCacheUtils.getCacheFilePath(TEST_HOST, TEST_CLIENT_ID, TEST_SCOPES, ""); + Path groupAPath = + TokenCacheUtils.getCacheFilePath(TEST_HOST, TEST_CLIENT_ID, TEST_SCOPES, "group-a"); + Path groupBPath = + TokenCacheUtils.getCacheFilePath(TEST_HOST, TEST_CLIENT_ID, TEST_SCOPES, "group-b"); + + assertEquals(legacyPath, nullGroupPath); + assertEquals(legacyPath, emptyGroupPath); + assertNotEquals(legacyPath, groupAPath); + assertNotEquals(groupAPath, groupBPath); + } + @Test void testOverwriteToken() { // Given two tokens saved in sequence diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/OAuthM2MServicePrincipalCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/OAuthM2MServicePrincipalCredentialsProviderTest.java new file mode 100644 index 000000000..46681d83c --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/OAuthM2MServicePrincipalCredentialsProviderTest.java @@ -0,0 +1,167 @@ +package com.databricks.sdk.core.oauth; + +import static org.junit.jupiter.api.Assertions.*; + +import com.databricks.sdk.core.DatabricksConfig; +import com.databricks.sdk.core.DatabricksException; +import com.databricks.sdk.core.http.HttpClient; +import com.databricks.sdk.core.http.Request; +import com.databricks.sdk.core.http.Response; +import java.io.IOException; +import java.net.URL; +import java.util.ArrayList; +import java.util.List; +import org.junit.jupiter.api.Test; + +class OAuthM2MServicePrincipalCredentialsProviderTest { + // Verifies that M2M includes assume_group every time an expired client-credentials token is + // minted, rather than only on the initial request. + @Test + void everyTokenMintIncludesTheGroup() { + TokenHttpClient httpClient = + new TokenHttpClient( + (request, requestNumber) -> + new Response( + String.format( + "{\"access_token\":\"token-%d\",\"token_type\":\"Bearer\",\"expires_in\":0}", + requestNumber), + 200, + "OK", + new URL(request.getUrl()))); + + DatabricksConfig config = + new DatabricksConfig() + .setHost("https://accounts.cloud.databricks.com") + .setAccountId("account-123") + .setDiscoveryUrl("https://accounts.cloud.databricks.com/discovery") + .setClientId("client-id") + .setClientSecret("client-secret") + .setGroupId("group-123") + .setHttpClient(httpClient) + .setDisableAsyncTokenRefresh(true); + + OAuthHeaderFactory headers = + new OAuthM2MServicePrincipalCredentialsProvider().configure(config); + + assertEquals("Bearer token-1", headers.headers().get("Authorization")); + assertEquals("Bearer token-2", headers.headers().get("Authorization")); + assertEquals(2, httpClient.tokenBodies.size()); + for (String body : httpClient.tokenBodies) { + assertTrue(body.contains("grant_type=client_credentials")); + assertTrue(body.contains("assume_group=group-123")); + assertFalse(body.contains("refresh_token")); + } + } + + // Verifies that a server rejection is returned after one grouped M2M request and is never + // retried without assume_group. + @Test + void groupServerRejectionDoesNotFallback() { + TokenHttpClient httpClient = + new TokenHttpClient( + (request, ignoredRequestNumber) -> + new Response( + "{\"error\":\"invalid_target\"}", + 400, + "Bad Request", + new URL(request.getUrl()))); + + DatabricksConfig config = config(httpClient, "group-123"); + + OAuthHeaderFactory headers = + new OAuthM2MServicePrincipalCredentialsProvider().configure(config); + + assertThrows(DatabricksException.class, headers::headers); + assertEquals(1, httpClient.tokenBodies.size()); + assertTrue(httpClient.tokenBodies.get(0).contains("assume_group=group-123")); + } + + // Verifies that separate M2M clients for normal access and different groups each cache only + // their own access token. + @Test + void cachesAreIsolatedByClient() { + TokenHttpClient httpClient = + new TokenHttpClient( + (request, ignoredRequestNumber) -> { + String body = request.getBodyString(); + String token = "token-normal"; + + if (body.contains("assume_group=group-a")) { + token = "token-group-a"; + } else if (body.contains("assume_group=group-b")) { + token = "token-group-b"; + } + + return new Response( + String.format( + "{\"access_token\":\"%s\",\"token_type\":\"Bearer\",\"expires_in\":3600}", + token), + 200, + "OK", + new URL(request.getUrl())); + }); + + String[][] testCases = { + {null, "Bearer token-normal"}, + {"group-a", "Bearer token-group-a"}, + {"group-b", "Bearer token-group-b"} + }; + + for (String[] testCase : testCases) { + OAuthHeaderFactory headers = + new OAuthM2MServicePrincipalCredentialsProvider() + .configure(config(httpClient, testCase[0])); + + assertEquals(testCase[1], headers.headers().get("Authorization")); + assertEquals(testCase[1], headers.headers().get("Authorization")); + } + + assertEquals(testCases.length, httpClient.tokenBodies.size()); + } + + /** Creates a fixed M2M configuration using the supplied HTTP client and optional group. */ + private static DatabricksConfig config(HttpClient httpClient, String groupId) { + return new DatabricksConfig() + .setHost("https://accounts.cloud.databricks.com") + .setAccountId("account-123") + .setDiscoveryUrl("https://accounts.cloud.databricks.com/discovery") + .setClientId("client-id") + .setClientSecret("client-secret") + .setGroupId(groupId) + .setHttpClient(httpClient) + .setDisableAsyncTokenRefresh(true); + } + + /** Produces the token endpoint response for a recorded token request. */ + @FunctionalInterface + private interface TokenResponder { + Response respond(Request request, int requestNumber) throws IOException; + } + + /** Handles discovery uniformly while allowing each test to define token endpoint behavior. */ + private static class TokenHttpClient implements HttpClient { + final List tokenBodies = new ArrayList<>(); + private final TokenResponder tokenResponder; + + /** Creates a client using the supplied behavior for token endpoint requests. */ + TokenHttpClient(TokenResponder tokenResponder) { + this.tokenResponder = tokenResponder; + } + + /** Returns OIDC discovery metadata for GET requests and records every token request body. */ + @Override + public Response execute(Request request) throws IOException { + if (request.getMethod().equals("GET")) { + return new Response( + "{\"token_endpoint\":\"https://accounts.cloud.databricks.com/oidc/accounts/account-123/v1/token\"," + + "\"authorization_endpoint\":\"https://accounts.cloud.databricks.com/oidc/accounts/account-123/v1/authorize\"}", + 200, + "OK", + new URL(request.getUrl())); + } + + tokenBodies.add(request.getBodyString()); + return tokenResponder.respond(request, tokenBodies.size()); + } + } +} diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenSourceCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenSourceCredentialsProviderTest.java index 8d2d68fd4..b2da9411e 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenSourceCredentialsProviderTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenSourceCredentialsProviderTest.java @@ -9,6 +9,7 @@ import java.time.Instant; import java.util.Map; import java.util.stream.Stream; +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; @@ -45,6 +46,19 @@ void testTokenScenarios( assertEquals(TEST_AUTH_TYPE, provider.authType()); } + // Verifies that configuring a group does not change how an ordinary token-source failure is + // handled by the credential chain. + @Test + void groupDoesNotChangeTokenFailureHandling() { + TokenSource tokenSource = mock(TokenSource.class); + when(tokenSource.getToken()).thenThrow(new DatabricksException("Token retrieval failed")); + + provider = new TokenSourceCredentialsProvider(tokenSource, TEST_AUTH_TYPE); + DatabricksConfig config = new DatabricksConfig().setGroupId("group-123"); + + assertNull(provider.configure(config)); + } + /** Provides test scenarios */ private static Stream provideTokenScenarios() { // Mock behaviour of successful token retrieval