From 3afdfa8e3d1f087a887e7514607f876fea06ebb6 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Mon, 19 May 2025 12:10:22 +0000 Subject: [PATCH 1/6] Implement direct dataplane access --- .../sdk/core/oauth/DataPlaneTokenSource.java | 102 ++++++++++ .../sdk/core/oauth/EndpointTokenSource.java | 92 +++++++++ .../sdk/core/oauth/TokenEndpointClient.java | 91 +++++++++ .../core/oauth/DataPlaneTokenSourceTest.java | 180 +++++++++++++++++ .../core/oauth/EndpointTokenSourceTest.java | 191 ++++++++++++++++++ .../core/oauth/TokenEndpointClientTest.java | 171 ++++++++++++++++ 6 files changed, 827 insertions(+) create mode 100644 databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DataPlaneTokenSource.java create mode 100644 databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/EndpointTokenSource.java create mode 100644 databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenEndpointClient.java create mode 100644 databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DataPlaneTokenSourceTest.java create mode 100644 databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/EndpointTokenSourceTest.java create mode 100644 databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenEndpointClientTest.java diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DataPlaneTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DataPlaneTokenSource.java new file mode 100644 index 000000000..b12a92dd2 --- /dev/null +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DataPlaneTokenSource.java @@ -0,0 +1,102 @@ +package com.databricks.sdk.core.oauth; + +import com.databricks.sdk.core.http.HttpClient; +import java.util.Objects; +import java.util.concurrent.ConcurrentHashMap; + +/** + * Manages and provides Databricks data plane tokens. This class is responsible for acquiring and + * caching OAuth tokens that are specific to a particular Databricks data plane service endpoint and + * a set of authorization details. It utilizes a {@link DatabricksOAuthTokenSource} for obtaining + * control plane tokens, which may then be exchanged or used to authorize requests for data plane + * tokens. Cached {@link EndpointTokenSource} instances are used to efficiently reuse tokens for + * repeated requests to the same endpoint with the same authorization context. + */ +public class DataPlaneTokenSource { + private final HttpClient httpClient; + private final DatabricksOAuthTokenSource cpTokenSource; + private final ConcurrentHashMap sourcesCache; + + /** + * Caching key for {@link EndpointTokenSource}, based on endpoint and authorization details. This + * is a value object that uniquely identifies a token source configuration. + */ + private static final class TokenSourceKey { + /** The target service endpoint URL. */ + private final String endpoint; + + /** Specific authorization details for the endpoint. */ + private final String authDetails; + + /** + * Constructs a TokenSourceKey. + * + * @param endpoint The target service endpoint URL. + * @param authDetails Specific authorization details. + */ + public TokenSourceKey(String endpoint, String authDetails) { + this.endpoint = endpoint; + this.authDetails = authDetails; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + TokenSourceKey that = (TokenSourceKey) o; + return Objects.equals(endpoint, that.endpoint) + && Objects.equals(authDetails, that.authDetails); + } + + @Override + public int hashCode() { + return Objects.hash(endpoint, authDetails); + } + } + + /** + * Constructs a DataPlaneTokenSource. + * + * @param httpClient The {@link HttpClient} for token requests. + * @param cpTokenSource The {@link DatabricksOAuthTokenSource} for control plane tokens. + * @throws NullPointerException if either parameter is null + */ + public DataPlaneTokenSource(HttpClient httpClient, DatabricksOAuthTokenSource cpTokenSource) { + this.httpClient = Objects.requireNonNull(httpClient, "HTTP client cannot be null"); + this.cpTokenSource = + Objects.requireNonNull(cpTokenSource, "Control plane token source cannot be null"); + this.sourcesCache = new ConcurrentHashMap<>(); + } + + /** + * Retrieves a token for the specified endpoint and authorization details. It uses a cached {@link + * EndpointTokenSource} if available, otherwise creates and caches a new one. + * + * @param endpoint The target data plane service endpoint. + * @param authDetails Authorization details for the endpoint. + * @return The dataplane {@link Token}. + * @throws NullPointerException if either parameter is null + * @throws IllegalArgumentException if either parameter is empty + */ + public Token getToken(String endpoint, String authDetails) { + Objects.requireNonNull(endpoint, "Data plane endpoint URL cannot be null"); + Objects.requireNonNull(authDetails, "Authorization details cannot be null"); + if (endpoint.isEmpty()) { + throw new IllegalArgumentException("Data plane endpoint URL cannot be empty"); + } + if (authDetails.isEmpty()) { + throw new IllegalArgumentException("Authorization details cannot be empty"); + } + TokenSourceKey key = new TokenSourceKey(endpoint, authDetails); + + EndpointTokenSource specificSource = + sourcesCache.computeIfAbsent( + key, k -> new EndpointTokenSource(this.cpTokenSource, k.authDetails, this.httpClient)); + + return specificSource.getToken(); + } +} diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/EndpointTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/EndpointTokenSource.java new file mode 100644 index 000000000..c54e7f6c0 --- /dev/null +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/EndpointTokenSource.java @@ -0,0 +1,92 @@ +package com.databricks.sdk.core.oauth; + +import com.databricks.sdk.core.DatabricksException; +import com.databricks.sdk.core.http.HttpClient; +import java.time.LocalDateTime; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Represents a token source that exchanges a control plane token for an endpoint-specific dataplane + * token. It utilizes an underlying {@link DatabricksOAuthTokenSource} to obtain the initial control + * plane token. + */ +public class EndpointTokenSource extends RefreshableTokenSource { + private static final Logger LOG = LoggerFactory.getLogger(EndpointTokenSource.class); + private static final String JWT_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer"; + private static final String GRANT_TYPE_PARAM = "grant_type"; + private static final String AUTHORIZATION_DETAILS_PARAM = "authorization_details"; + private static final String ASSERTION_PARAM = "assertion"; + private static final String TOKEN_ENDPOINT = "/oidc/v1/token"; + + private final DatabricksOAuthTokenSource cpTokenSource; + private final String authDetails; + private final HttpClient httpClient; + + /** + * Constructs a new EndpointTokenSource. + * + * @param cpTokenSource The {@link DatabricksOAuthTokenSource} used to obtain the control plane + * token. + * @param authDetails The authorization details required for the token exchange. + * @param httpClient The {@link HttpClient} used to make the token exchange request. + * @throws IllegalArgumentException if authDetails is empty. + * @throws NullPointerException if any of the parameters are null. + */ + public EndpointTokenSource( + DatabricksOAuthTokenSource cpTokenSource, String authDetails, HttpClient httpClient) { + this.cpTokenSource = + Objects.requireNonNull(cpTokenSource, "Control plane token source cannot be null"); + this.authDetails = Objects.requireNonNull(authDetails, "Authorization details cannot be null"); + if (authDetails.isEmpty()) { + throw new IllegalArgumentException("Authorization details cannot be empty"); + } + this.httpClient = Objects.requireNonNull(httpClient, "HTTP client cannot be null"); + } + + /** + * Fetches an endpoint-specific dataplane token by exchanging a control plane token. + * + *

This method first obtains a control plane token from the configured {@code cpTokenSource}. + * It then uses this token as an assertion along with the provided {@code authDetails} to request + * a new, more scoped dataplane token from the Databricks OAuth token endpoint ({@value + * #TOKEN_ENDPOINT}). + * + * @return A new {@link Token} containing the exchanged dataplane access token, its type, any + * accompanying refresh token, and its expiry time. + * @throws DatabricksException if the token exchange with the OAuth endpoint fails. + * @throws IllegalArgumentException if the token endpoint url is empty. + * @throws NullPointerException if any of the parameters are null. + */ + @Override + protected Token refresh() { + Token cpToken = cpTokenSource.getToken(); + + Map params = new HashMap<>(); + params.put(GRANT_TYPE_PARAM, JWT_GRANT_TYPE); + params.put(AUTHORIZATION_DETAILS_PARAM, authDetails); + params.put(ASSERTION_PARAM, cpToken.getAccessToken()); + + OAuthResponse oauthResponse; + try { + oauthResponse = TokenEndpointClient.requestToken(this.httpClient, TOKEN_ENDPOINT, params); + } catch (DatabricksException | IllegalArgumentException | NullPointerException e) { + LOG.error( + "Failed to exchange control plane token for dataplane token at endpoint {}: {}", + TOKEN_ENDPOINT, + e.getMessage(), + e); + throw e; + } + + LocalDateTime expiry = LocalDateTime.now().plusSeconds(oauthResponse.getExpiresIn()); + return new Token( + oauthResponse.getAccessToken(), + oauthResponse.getTokenType(), + oauthResponse.getRefreshToken(), + expiry); + } +} diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenEndpointClient.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenEndpointClient.java new file mode 100644 index 000000000..69883dd24 --- /dev/null +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenEndpointClient.java @@ -0,0 +1,91 @@ +package com.databricks.sdk.core.oauth; + +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.Response; +import com.fasterxml.jackson.databind.ObjectMapper; +import java.io.IOException; +import java.util.Map; +import java.util.Objects; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Client for interacting with an OAuth token endpoint. + * + *

This class provides a method to request an OAuth token from a specified token endpoint URL + * using the provided HTTP client and request parameters. It handles the HTTP request and parses the + * JSON response into an {@link OAuthResponse} object. + */ +public final class TokenEndpointClient { + private static final Logger LOG = LoggerFactory.getLogger(TokenEndpointClient.class); + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); + + private TokenEndpointClient() {} + + /** + * Requests an OAuth token from the specified token endpoint. + * + * @param httpClient The {@link HttpClient} to use for making the request. + * @param tokenEndpointUrl The URL of the token endpoint. + * @param params A map of parameters to include in the token request. + * @return An {@link OAuthResponse} containing the token information. + * @throws DatabricksException if an error occurs during the token request or response parsing. + * @throws IllegalArgumentException if the token endpoint URL is empty. + * @throws NullPointerException if any of the parameters are null. + */ + public static OAuthResponse requestToken( + HttpClient httpClient, String tokenEndpointUrl, Map params) + throws DatabricksException { + Objects.requireNonNull(httpClient, "HttpClient cannot be null"); + Objects.requireNonNull(params, "Request parameters map cannot be null"); + Objects.requireNonNull(tokenEndpointUrl, "Token endpoint URL cannot be null"); + + if (tokenEndpointUrl.isEmpty()) { + throw new IllegalArgumentException("Token endpoint URL cannot be empty"); + } + + Response rawResponse; + try { + LOG.debug("Requesting token from endpoint: {}", tokenEndpointUrl); + rawResponse = httpClient.execute(new FormRequest(tokenEndpointUrl, params)); + } catch (IOException e) { + LOG.error("Failed to request token from {}: {}", tokenEndpointUrl, e.getMessage(), e); + throw new DatabricksException( + String.format("Failed to request token from %s: %s", tokenEndpointUrl, e.getMessage()), + e); + } + + OAuthResponse response; + try { + response = OBJECT_MAPPER.readValue(rawResponse.getBody(), OAuthResponse.class); + } catch (IOException e) { + LOG.error( + "Failed to parse OAuth response from token endpoint {}: {}", + tokenEndpointUrl, + e.getMessage(), + e); + throw new DatabricksException( + String.format( + "Failed to parse OAuth response from token endpoint %s: %s", + tokenEndpointUrl, e.getMessage()), + e); + } + + if (response.getErrorCode() != null) { + String errorSummary = + response.getErrorSummary() != null ? response.getErrorSummary() : "No summary provided."; + LOG.error( + "Token request to {} failed with error: {} - {}", + tokenEndpointUrl, + response.getErrorCode(), + errorSummary); + throw new DatabricksException( + String.format( + "Token request failed with error: %s - %s", response.getErrorCode(), errorSummary)); + } + LOG.debug("Successfully obtained token response from {}", tokenEndpointUrl); + return response; + } +} diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DataPlaneTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DataPlaneTokenSourceTest.java new file mode 100644 index 000000000..91418798e --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DataPlaneTokenSourceTest.java @@ -0,0 +1,180 @@ +package com.databricks.sdk.core.oauth; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.*; + +import com.databricks.sdk.core.http.HttpClient; +import com.databricks.sdk.core.http.Response; +import java.io.IOException; +import java.net.URL; +import java.time.LocalDateTime; +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; + +public class DataPlaneTokenSourceTest { + private static final String TEST_ENDPOINT_1 = "https://endpoint1.databricks.com/"; + private static final String TEST_ENDPOINT_2 = "https://endpoint2.databricks.com/"; + private static final String TEST_AUTH_DETAILS_1 = "{\"aud\":\"aud1\"}"; + private static final String TEST_AUTH_DETAILS_2 = "{\"aud\":\"aud2\"}"; + private static final String TEST_CP_TOKEN = "cp-access-token"; + private static final String TEST_TOKEN_TYPE = "Bearer"; + private static final String TEST_REFRESH_TOKEN = "refresh-token"; + private static final int TEST_EXPIRES_IN = 3600; + + private static Stream provideDataPlaneTokenScenarios() throws Exception { + // Mock DatabricksOAuthTokenSource for control plane token + Token cpToken = + new Token(TEST_CP_TOKEN, TEST_TOKEN_TYPE, null, LocalDateTime.now().plusSeconds(600)); + DatabricksOAuthTokenSource mockCpTokenSource = mock(DatabricksOAuthTokenSource.class); + when(mockCpTokenSource.getToken()).thenReturn(cpToken); + + // --- Mock HttpClient for different scenarios --- + // Success JSON for endpoint1/auth1 + String successJson1 = + "{" + + "\"access_token\":\"dp-access-token1\"," + + "\"token_type\":\"Bearer\"," + + "\"refresh_token\":\"refresh-token\"," + + "\"expires_in\":3600" + + "}"; + HttpClient mockSuccessClient1 = mock(HttpClient.class); + when(mockSuccessClient1.execute(any())) + .thenReturn(new Response(successJson1, 200, "OK", new URL(TEST_ENDPOINT_1))); + + // Success JSON for endpoint2/auth2 + String successJson2 = + "{" + + "\"access_token\":\"dp-access-token2\"," + + "\"token_type\":\"Bearer\"," + + "\"refresh_token\":\"refresh-token\"," + + "\"expires_in\":3600" + + "}"; + HttpClient mockSuccessClient2 = mock(HttpClient.class); + when(mockSuccessClient2.execute(any())) + .thenReturn(new Response(successJson2, 200, "OK", new URL(TEST_ENDPOINT_2))); + + // Error response JSON + String errorJson = + "{" + "\"error\":\"invalid_request\"," + "\"error_description\":\"Bad request\"" + "}"; + HttpClient mockErrorClient = mock(HttpClient.class); + when(mockErrorClient.execute(any())) + .thenReturn(new Response(errorJson, 400, "Bad Request", new URL(TEST_ENDPOINT_1))); + + // IOException scenario + HttpClient mockIOExceptionClient = mock(HttpClient.class); + when(mockIOExceptionClient.execute(any())).thenThrow(new IOException("Network error")); + + // For null cpTokenSource + DatabricksOAuthTokenSource nullCpTokenSource = null; + + // For null httpClient + HttpClient nullHttpClient = null; + + // For null/empty endpoint or authDetails + return Stream.of( + Arguments.of( + "Success: endpoint1/auth1", + TEST_ENDPOINT_1, + TEST_AUTH_DETAILS_1, + mockSuccessClient1, + mockCpTokenSource, + new Token( + "dp-access-token1", + TEST_TOKEN_TYPE, + TEST_REFRESH_TOKEN, + LocalDateTime.now().plusSeconds(TEST_EXPIRES_IN)), + null // No exception + ), + Arguments.of( + "Success: endpoint2/auth2 (different cache key)", + TEST_ENDPOINT_2, + TEST_AUTH_DETAILS_2, + mockSuccessClient2, + mockCpTokenSource, + new Token( + "dp-access-token2", + TEST_TOKEN_TYPE, + TEST_REFRESH_TOKEN, + LocalDateTime.now().plusSeconds(TEST_EXPIRES_IN)), + null), + Arguments.of( + "Error response from endpoint", + TEST_ENDPOINT_1, + TEST_AUTH_DETAILS_1, + mockErrorClient, + mockCpTokenSource, + null, + com.databricks.sdk.core.DatabricksException.class), + Arguments.of( + "IOException from HttpClient", + TEST_ENDPOINT_1, + TEST_AUTH_DETAILS_1, + mockIOExceptionClient, + mockCpTokenSource, + null, + com.databricks.sdk.core.DatabricksException.class), + Arguments.of( + "Null cpTokenSource", + TEST_ENDPOINT_1, + TEST_AUTH_DETAILS_1, + mockSuccessClient1, + nullCpTokenSource, + null, + NullPointerException.class), + Arguments.of( + "Null httpClient", + TEST_ENDPOINT_1, + TEST_AUTH_DETAILS_1, + nullHttpClient, + mockCpTokenSource, + null, + NullPointerException.class), + Arguments.of( + "Null endpoint", + null, + TEST_AUTH_DETAILS_1, + mockSuccessClient1, + mockCpTokenSource, + null, + NullPointerException.class), + Arguments.of( + "Null authDetails", + TEST_ENDPOINT_1, + null, + mockSuccessClient1, + mockCpTokenSource, + null, + NullPointerException.class)); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("provideDataPlaneTokenScenarios") + void testDataPlaneTokenSource( + String testName, + String endpoint, + String authDetails, + HttpClient httpClient, + DatabricksOAuthTokenSource cpTokenSource, + Token expectedToken, + Class expectedException) { + if (expectedException != null) { + assertThrows( + expectedException, + () -> { + DataPlaneTokenSource source = new DataPlaneTokenSource(httpClient, cpTokenSource); + source.getToken(endpoint, authDetails); + }); + } else { + DataPlaneTokenSource source = new DataPlaneTokenSource(httpClient, cpTokenSource); + Token token = source.getToken(endpoint, authDetails); + assertNotNull(token); + assertEquals(expectedToken.getAccessToken(), token.getAccessToken()); + assertEquals(expectedToken.getTokenType(), token.getTokenType()); + assertEquals(expectedToken.getRefreshToken(), token.getRefreshToken()); + assertTrue(token.isValid()); + } + } +} diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/EndpointTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/EndpointTokenSourceTest.java new file mode 100644 index 000000000..549077690 --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/EndpointTokenSourceTest.java @@ -0,0 +1,191 @@ +package com.databricks.sdk.core.oauth; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.*; + +import com.databricks.sdk.core.DatabricksException; +import com.databricks.sdk.core.http.HttpClient; +import com.databricks.sdk.core.http.Response; +import java.io.IOException; +import java.net.URL; +import java.time.LocalDateTime; +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 EndpointTokenSourceTest { + private static final String TEST_AUTH_DETAILS = "{\"aud\":\"test-audience\"}"; + private static final String TEST_CP_TOKEN = "cp-access-token"; + private static final String TEST_DP_TOKEN = "dp-access-token"; + private static final String TEST_TOKEN_TYPE = "Bearer"; + private static final String TEST_REFRESH_TOKEN = "refresh-token"; + private static final int TEST_EXPIRES_IN = 3600; + + private static Stream provideEndpointTokenScenarios() throws Exception { + // Success response JSON + String successJson = + "{" + + "\"access_token\":\"" + + TEST_DP_TOKEN + + "\"," + + "\"token_type\":\"" + + TEST_TOKEN_TYPE + + "\"," + + "\"expires_in\":" + + TEST_EXPIRES_IN + + "," + + "\"refresh_token\":\"" + + TEST_REFRESH_TOKEN + + "\"}"; + // Error response JSON + String errorJson = + "{" + + "\"error\":\"invalid_client\"," + + "\"error_description\":\"Client authentication failed\"}"; + // Malformed JSON + String malformedJson = "{not valid json}"; + + // Mock DatabricksOAuthTokenSource for control plane token + Token cpToken = new Token(TEST_CP_TOKEN, TEST_TOKEN_TYPE, LocalDateTime.now().plusMinutes(10)); + DatabricksOAuthTokenSource mockCpTokenSource = mock(DatabricksOAuthTokenSource.class); + when(mockCpTokenSource.getToken()).thenReturn(cpToken); + + // Mock HttpClient for success + HttpClient mockSuccessClient = mock(HttpClient.class); + when(mockSuccessClient.execute(any())) + .thenReturn(new Response(successJson, 200, "OK", new URL("https://test.databricks.com/"))); + + // Mock HttpClient for error response + HttpClient mockErrorClient = mock(HttpClient.class); + when(mockErrorClient.execute(any())) + .thenReturn( + new Response(errorJson, 400, "Bad Request", new URL("https://test.databricks.com/"))); + + // Mock HttpClient for malformed JSON + HttpClient mockMalformedClient = mock(HttpClient.class); + when(mockMalformedClient.execute(any())) + .thenReturn( + new Response(malformedJson, 200, "OK", new URL("https://test.databricks.com/"))); + + // Mock HttpClient for IOException + HttpClient mockIOExceptionClient = mock(HttpClient.class); + when(mockIOExceptionClient.execute(any())).thenThrow(new IOException("Network error")); + + return Stream.of( + Arguments.of( + "Success response", + mockCpTokenSource, + TEST_AUTH_DETAILS, + mockSuccessClient, + null, // No exception expected + TEST_DP_TOKEN, + TEST_TOKEN_TYPE, + TEST_REFRESH_TOKEN, + TEST_EXPIRES_IN), + Arguments.of( + "OAuth error response", + mockCpTokenSource, + TEST_AUTH_DETAILS, + mockErrorClient, + DatabricksException.class, + null, + null, + null, + 0), + Arguments.of( + "Malformed JSON response", + mockCpTokenSource, + TEST_AUTH_DETAILS, + mockMalformedClient, + DatabricksException.class, + null, + null, + null, + 0), + Arguments.of( + "IOException from HttpClient", + mockCpTokenSource, + TEST_AUTH_DETAILS, + mockIOExceptionClient, + DatabricksException.class, + null, + null, + null, + 0), + Arguments.of( + "Null cpTokenSource", + null, + TEST_AUTH_DETAILS, + mockSuccessClient, + NullPointerException.class, + null, + null, + null, + 0), + Arguments.of( + "Null authDetails", + mockCpTokenSource, + null, + mockSuccessClient, + NullPointerException.class, + null, + null, + null, + 0), + Arguments.of( + "Empty authDetails", + mockCpTokenSource, + "", + mockSuccessClient, + IllegalArgumentException.class, + null, + null, + null, + 0), + Arguments.of( + "Null httpClient", + mockCpTokenSource, + TEST_AUTH_DETAILS, + null, + NullPointerException.class, + null, + null, + null, + 0)); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("provideEndpointTokenScenarios") + void testEndpointTokenSource( + String testName, + DatabricksOAuthTokenSource cpTokenSource, + String authDetails, + HttpClient httpClient, + Class expectedException, + String expectedAccessToken, + String expectedTokenType, + String expectedRefreshToken, + int expectedExpiresIn) { + if (expectedException != null) { + assertThrows( + expectedException, + () -> { + EndpointTokenSource source = + new EndpointTokenSource(cpTokenSource, authDetails, httpClient); + source.getToken(); + }); + } else { + EndpointTokenSource source = new EndpointTokenSource(cpTokenSource, authDetails, httpClient); + Token token = source.getToken(); + assertNotNull(token); + assertEquals(expectedAccessToken, token.getAccessToken()); + assertEquals(expectedTokenType, token.getTokenType()); + assertEquals(expectedRefreshToken, token.getRefreshToken()); + // Allow a few seconds of clock skew for expiry + assertTrue(token.isValid()); + assertTrue(token.getAccessToken().length() > 0); + } + } +} diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenEndpointClientTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenEndpointClientTest.java new file mode 100644 index 000000000..581c90143 --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenEndpointClientTest.java @@ -0,0 +1,171 @@ +package com.databricks.sdk.core.oauth; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.*; + +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.Response; +import java.io.IOException; +import java.net.URL; +import java.util.HashMap; +import java.util.Map; +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 TokenEndpointClientTest { + private static final String TOKEN_ENDPOINT_URL = "https://test.databricks.com/oauth/token"; + private static final Map PARAMS = new HashMap<>(); + + private static Stream provideTokenScenarios() throws Exception { + // Success response JSON + String successJson = + "{" + + "\"access_token\":\"test-access-token\"," + + "\"token_type\":\"Bearer\"," + + "\"expires_in\":3600," + + "\"refresh_token\":\"test-refresh-token\"}"; + // Error response JSON + String errorJson = + "{" + + "\"error\":\"invalid_client\"," + + "\"error_description\":\"Client authentication failed\"}"; + // Malformed JSON + String malformedJson = "{not valid json}"; + + // Mock HttpClient for success + HttpClient mockSuccessClient = mock(HttpClient.class); + when(mockSuccessClient.execute(any(FormRequest.class))) + .thenReturn(new Response(successJson, 200, "OK", new URL("https://test.databricks.com/"))); + + // Mock HttpClient for error response + HttpClient mockErrorClient = mock(HttpClient.class); + when(mockErrorClient.execute(any(FormRequest.class))) + .thenReturn( + new Response(errorJson, 400, "Bad Request", new URL("https://test.databricks.com/"))); + + // Mock HttpClient for malformed JSON + HttpClient mockMalformedClient = mock(HttpClient.class); + when(mockMalformedClient.execute(any(FormRequest.class))) + .thenReturn( + new Response(malformedJson, 200, "OK", new URL("https://test.databricks.com/"))); + + // Mock HttpClient for IOException + HttpClient mockIOExceptionClient = mock(HttpClient.class); + when(mockIOExceptionClient.execute(any(FormRequest.class))) + .thenThrow(new IOException("Network error")); + + return Stream.of( + Arguments.of( + "Success response", + mockSuccessClient, + TOKEN_ENDPOINT_URL, + PARAMS, + null, // No exception expected + "test-access-token", + "Bearer", + 3600, + "test-refresh-token"), + Arguments.of( + "OAuth error response", + mockErrorClient, + TOKEN_ENDPOINT_URL, + PARAMS, + DatabricksException.class, + null, + null, + 0, + null), + Arguments.of( + "Malformed JSON response", + mockMalformedClient, + TOKEN_ENDPOINT_URL, + PARAMS, + DatabricksException.class, + null, + null, + 0, + null), + Arguments.of( + "IOException from HttpClient", + mockIOExceptionClient, + TOKEN_ENDPOINT_URL, + PARAMS, + DatabricksException.class, + null, + null, + 0, + null), + Arguments.of( + "Null HttpClient", + null, + TOKEN_ENDPOINT_URL, + PARAMS, + NullPointerException.class, + null, + null, + 0, + null), + Arguments.of( + "Null tokenEndpointUrl", + mockSuccessClient, + null, + PARAMS, + NullPointerException.class, + null, + null, + 0, + null), + Arguments.of( + "Empty tokenEndpointUrl", + mockSuccessClient, + "", + PARAMS, + IllegalArgumentException.class, + null, + null, + 0, + null), + Arguments.of( + "Null params", + mockSuccessClient, + TOKEN_ENDPOINT_URL, + null, + NullPointerException.class, + null, + null, + 0, + null)); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("provideTokenScenarios") + void testRequestToken( + String testName, + HttpClient httpClient, + String tokenEndpointUrl, + Map params, + Class expectedException, + String expectedAccessToken, + String expectedTokenType, + int expectedExpiresIn, + String expectedRefreshToken) { + if (expectedException != null) { + assertThrows( + expectedException, + () -> TokenEndpointClient.requestToken(httpClient, tokenEndpointUrl, params)); + } else { + OAuthResponse response = + TokenEndpointClient.requestToken(httpClient, tokenEndpointUrl, params); + assertNotNull(response); + assertEquals(expectedAccessToken, response.getAccessToken()); + assertEquals(expectedTokenType, response.getTokenType()); + assertEquals(expectedExpiresIn, response.getExpiresIn()); + assertEquals(expectedRefreshToken, response.getRefreshToken()); + } + } +} From 79e360da4a46ac89b40371185ee32b989d32030b Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Fri, 23 May 2025 16:00:54 +0000 Subject: [PATCH 2/6] Add host as a field to token sources --- .../sdk/core/oauth/DataPlaneTokenSource.java | 21 +++++++-- .../sdk/core/oauth/EndpointTokenSource.java | 25 ++++++---- .../core/oauth/DataPlaneTokenSourceTest.java | 46 +++++++++++++------ .../core/oauth/EndpointTokenSourceTest.java | 45 ++++++++++++++---- 4 files changed, 101 insertions(+), 36 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DataPlaneTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DataPlaneTokenSource.java index b12a92dd2..4504d9ca3 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DataPlaneTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DataPlaneTokenSource.java @@ -14,9 +14,9 @@ */ public class DataPlaneTokenSource { private final HttpClient httpClient; - private final DatabricksOAuthTokenSource cpTokenSource; + private final TokenSource cpTokenSource; + private final String host; private final ConcurrentHashMap sourcesCache; - /** * Caching key for {@link EndpointTokenSource}, based on endpoint and authorization details. This * is a value object that uniquely identifies a token source configuration. @@ -62,13 +62,19 @@ public int hashCode() { * Constructs a DataPlaneTokenSource. * * @param httpClient The {@link HttpClient} for token requests. - * @param cpTokenSource The {@link DatabricksOAuthTokenSource} for control plane tokens. + * @param cpTokenSource The {@link TokenSource} for control plane tokens. + * @param host The host for the token exchange request. * @throws NullPointerException if either parameter is null */ - public DataPlaneTokenSource(HttpClient httpClient, DatabricksOAuthTokenSource cpTokenSource) { + public DataPlaneTokenSource(HttpClient httpClient, TokenSource cpTokenSource, String host) { this.httpClient = Objects.requireNonNull(httpClient, "HTTP client cannot be null"); this.cpTokenSource = Objects.requireNonNull(cpTokenSource, "Control plane token source cannot be null"); + this.host = Objects.requireNonNull(host, "Host cannot be null"); + + if (host.isEmpty()) { + throw new IllegalArgumentException("Host cannot be empty"); + } this.sourcesCache = new ConcurrentHashMap<>(); } @@ -85,17 +91,22 @@ public DataPlaneTokenSource(HttpClient httpClient, DatabricksOAuthTokenSource cp public Token getToken(String endpoint, String authDetails) { Objects.requireNonNull(endpoint, "Data plane endpoint URL cannot be null"); Objects.requireNonNull(authDetails, "Authorization details cannot be null"); + if (endpoint.isEmpty()) { throw new IllegalArgumentException("Data plane endpoint URL cannot be empty"); } if (authDetails.isEmpty()) { throw new IllegalArgumentException("Authorization details cannot be empty"); } + TokenSourceKey key = new TokenSourceKey(endpoint, authDetails); EndpointTokenSource specificSource = sourcesCache.computeIfAbsent( - key, k -> new EndpointTokenSource(this.cpTokenSource, k.authDetails, this.httpClient)); + key, + k -> + new EndpointTokenSource( + this.cpTokenSource, k.authDetails, this.httpClient, this.host)); return specificSource.getToken(); } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/EndpointTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/EndpointTokenSource.java index c54e7f6c0..3ca75c441 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/EndpointTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/EndpointTokenSource.java @@ -11,8 +11,7 @@ /** * Represents a token source that exchanges a control plane token for an endpoint-specific dataplane - * token. It utilizes an underlying {@link DatabricksOAuthTokenSource} to obtain the initial control - * plane token. + * token. It utilizes an underlying {@link TokenSource} to obtain the initial control plane token. */ public class EndpointTokenSource extends RefreshableTokenSource { private static final Logger LOG = LoggerFactory.getLogger(EndpointTokenSource.class); @@ -22,29 +21,35 @@ public class EndpointTokenSource extends RefreshableTokenSource { private static final String ASSERTION_PARAM = "assertion"; private static final String TOKEN_ENDPOINT = "/oidc/v1/token"; - private final DatabricksOAuthTokenSource cpTokenSource; + private final TokenSource cpTokenSource; private final String authDetails; private final HttpClient httpClient; + private final String host; /** * Constructs a new EndpointTokenSource. * - * @param cpTokenSource The {@link DatabricksOAuthTokenSource} used to obtain the control plane - * token. + * @param cpTokenSource The {@link TokenSource} used to obtain the control plane token. * @param authDetails The authorization details required for the token exchange. * @param httpClient The {@link HttpClient} used to make the token exchange request. - * @throws IllegalArgumentException if authDetails is empty. + * @param host The host for the token exchange request. + * @throws IllegalArgumentException if authDetails is empty or host is empty. * @throws NullPointerException if any of the parameters are null. */ public EndpointTokenSource( - DatabricksOAuthTokenSource cpTokenSource, String authDetails, HttpClient httpClient) { + TokenSource cpTokenSource, String authDetails, HttpClient httpClient, String host) { this.cpTokenSource = Objects.requireNonNull(cpTokenSource, "Control plane token source cannot be null"); this.authDetails = Objects.requireNonNull(authDetails, "Authorization details cannot be null"); + this.httpClient = Objects.requireNonNull(httpClient, "HTTP client cannot be null"); + this.host = Objects.requireNonNull(host, "Host cannot be null"); + if (authDetails.isEmpty()) { throw new IllegalArgumentException("Authorization details cannot be empty"); } - this.httpClient = Objects.requireNonNull(httpClient, "HTTP client cannot be null"); + if (host.isEmpty()) { + throw new IllegalArgumentException("Host cannot be empty"); + } } /** @@ -64,7 +69,6 @@ public EndpointTokenSource( @Override protected Token refresh() { Token cpToken = cpTokenSource.getToken(); - Map params = new HashMap<>(); params.put(GRANT_TYPE_PARAM, JWT_GRANT_TYPE); params.put(AUTHORIZATION_DETAILS_PARAM, authDetails); @@ -72,7 +76,8 @@ protected Token refresh() { OAuthResponse oauthResponse; try { - oauthResponse = TokenEndpointClient.requestToken(this.httpClient, TOKEN_ENDPOINT, params); + oauthResponse = + TokenEndpointClient.requestToken(this.httpClient, this.host + TOKEN_ENDPOINT, params); } catch (DatabricksException | IllegalArgumentException | NullPointerException e) { LOG.error( "Failed to exchange control plane token for dataplane token at endpoint {}: {}", diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DataPlaneTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DataPlaneTokenSourceTest.java index 91418798e..35b3586d5 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DataPlaneTokenSourceTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DataPlaneTokenSourceTest.java @@ -23,6 +23,7 @@ public class DataPlaneTokenSourceTest { private static final String TEST_TOKEN_TYPE = "Bearer"; private static final String TEST_REFRESH_TOKEN = "refresh-token"; private static final int TEST_EXPIRES_IN = 3600; + private static final String TEST_HOST = "https://test.databricks.com"; private static Stream provideDataPlaneTokenScenarios() throws Exception { // Mock DatabricksOAuthTokenSource for control plane token @@ -31,7 +32,6 @@ private static Stream provideDataPlaneTokenScenarios() throws Excepti DatabricksOAuthTokenSource mockCpTokenSource = mock(DatabricksOAuthTokenSource.class); when(mockCpTokenSource.getToken()).thenReturn(cpToken); - // --- Mock HttpClient for different scenarios --- // Success JSON for endpoint1/auth1 String successJson1 = "{" @@ -56,7 +56,6 @@ private static Stream provideDataPlaneTokenScenarios() throws Excepti when(mockSuccessClient2.execute(any())) .thenReturn(new Response(successJson2, 200, "OK", new URL(TEST_ENDPOINT_2))); - // Error response JSON String errorJson = "{" + "\"error\":\"invalid_request\"," + "\"error_description\":\"Bad request\"" + "}"; HttpClient mockErrorClient = mock(HttpClient.class); @@ -67,12 +66,6 @@ private static Stream provideDataPlaneTokenScenarios() throws Excepti HttpClient mockIOExceptionClient = mock(HttpClient.class); when(mockIOExceptionClient.execute(any())).thenThrow(new IOException("Network error")); - // For null cpTokenSource - DatabricksOAuthTokenSource nullCpTokenSource = null; - - // For null httpClient - HttpClient nullHttpClient = null; - // For null/empty endpoint or authDetails return Stream.of( Arguments.of( @@ -81,6 +74,7 @@ private static Stream provideDataPlaneTokenScenarios() throws Excepti TEST_AUTH_DETAILS_1, mockSuccessClient1, mockCpTokenSource, + TEST_HOST, new Token( "dp-access-token1", TEST_TOKEN_TYPE, @@ -94,6 +88,7 @@ private static Stream provideDataPlaneTokenScenarios() throws Excepti TEST_AUTH_DETAILS_2, mockSuccessClient2, mockCpTokenSource, + TEST_HOST, new Token( "dp-access-token2", TEST_TOKEN_TYPE, @@ -106,6 +101,7 @@ private static Stream provideDataPlaneTokenScenarios() throws Excepti TEST_AUTH_DETAILS_1, mockErrorClient, mockCpTokenSource, + TEST_HOST, null, com.databricks.sdk.core.DatabricksException.class), Arguments.of( @@ -114,6 +110,7 @@ private static Stream provideDataPlaneTokenScenarios() throws Excepti TEST_AUTH_DETAILS_1, mockIOExceptionClient, mockCpTokenSource, + TEST_HOST, null, com.databricks.sdk.core.DatabricksException.class), Arguments.of( @@ -121,15 +118,17 @@ private static Stream provideDataPlaneTokenScenarios() throws Excepti TEST_ENDPOINT_1, TEST_AUTH_DETAILS_1, mockSuccessClient1, - nullCpTokenSource, + null, + TEST_HOST, null, NullPointerException.class), Arguments.of( "Null httpClient", TEST_ENDPOINT_1, TEST_AUTH_DETAILS_1, - nullHttpClient, + null, mockCpTokenSource, + TEST_HOST, null, NullPointerException.class), Arguments.of( @@ -138,6 +137,7 @@ private static Stream provideDataPlaneTokenScenarios() throws Excepti TEST_AUTH_DETAILS_1, mockSuccessClient1, mockCpTokenSource, + TEST_HOST, null, NullPointerException.class), Arguments.of( @@ -146,8 +146,27 @@ private static Stream provideDataPlaneTokenScenarios() throws Excepti null, mockSuccessClient1, mockCpTokenSource, + TEST_HOST, + null, + NullPointerException.class), + Arguments.of( + "Null host", + TEST_ENDPOINT_1, + TEST_AUTH_DETAILS_1, + mockSuccessClient1, + mockCpTokenSource, + null, + null, + NullPointerException.class), + Arguments.of( + "Empty host", + TEST_ENDPOINT_1, + TEST_AUTH_DETAILS_1, + mockSuccessClient1, + mockCpTokenSource, + "", null, - NullPointerException.class)); + IllegalArgumentException.class)); } @ParameterizedTest(name = "{0}") @@ -158,17 +177,18 @@ void testDataPlaneTokenSource( String authDetails, HttpClient httpClient, DatabricksOAuthTokenSource cpTokenSource, + String host, Token expectedToken, Class expectedException) { if (expectedException != null) { assertThrows( expectedException, () -> { - DataPlaneTokenSource source = new DataPlaneTokenSource(httpClient, cpTokenSource); + DataPlaneTokenSource source = new DataPlaneTokenSource(httpClient, cpTokenSource, host); source.getToken(endpoint, authDetails); }); } else { - DataPlaneTokenSource source = new DataPlaneTokenSource(httpClient, cpTokenSource); + DataPlaneTokenSource source = new DataPlaneTokenSource(httpClient, cpTokenSource, host); Token token = source.getToken(endpoint, authDetails); assertNotNull(token); assertEquals(expectedToken.getAccessToken(), token.getAccessToken()); diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/EndpointTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/EndpointTokenSourceTest.java index 549077690..a3af2254f 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/EndpointTokenSourceTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/EndpointTokenSourceTest.java @@ -22,9 +22,9 @@ class EndpointTokenSourceTest { private static final String TEST_TOKEN_TYPE = "Bearer"; private static final String TEST_REFRESH_TOKEN = "refresh-token"; private static final int TEST_EXPIRES_IN = 3600; + private static final String TEST_HOST = "https://test.databricks.com"; private static Stream provideEndpointTokenScenarios() throws Exception { - // Success response JSON String successJson = "{" + "\"access_token\":\"" @@ -39,12 +39,12 @@ private static Stream provideEndpointTokenScenarios() throws Exceptio + "\"refresh_token\":\"" + TEST_REFRESH_TOKEN + "\"}"; - // Error response JSON + String errorJson = "{" + "\"error\":\"invalid_client\"," + "\"error_description\":\"Client authentication failed\"}"; - // Malformed JSON + String malformedJson = "{not valid json}"; // Mock DatabricksOAuthTokenSource for control plane token @@ -79,6 +79,7 @@ private static Stream provideEndpointTokenScenarios() throws Exceptio mockCpTokenSource, TEST_AUTH_DETAILS, mockSuccessClient, + TEST_HOST, null, // No exception expected TEST_DP_TOKEN, TEST_TOKEN_TYPE, @@ -89,6 +90,7 @@ private static Stream provideEndpointTokenScenarios() throws Exceptio mockCpTokenSource, TEST_AUTH_DETAILS, mockErrorClient, + TEST_HOST, DatabricksException.class, null, null, @@ -99,6 +101,7 @@ private static Stream provideEndpointTokenScenarios() throws Exceptio mockCpTokenSource, TEST_AUTH_DETAILS, mockMalformedClient, + TEST_HOST, DatabricksException.class, null, null, @@ -109,6 +112,7 @@ private static Stream provideEndpointTokenScenarios() throws Exceptio mockCpTokenSource, TEST_AUTH_DETAILS, mockIOExceptionClient, + TEST_HOST, DatabricksException.class, null, null, @@ -119,6 +123,7 @@ private static Stream provideEndpointTokenScenarios() throws Exceptio null, TEST_AUTH_DETAILS, mockSuccessClient, + TEST_HOST, NullPointerException.class, null, null, @@ -129,6 +134,7 @@ private static Stream provideEndpointTokenScenarios() throws Exceptio mockCpTokenSource, null, mockSuccessClient, + TEST_HOST, NullPointerException.class, null, null, @@ -139,6 +145,7 @@ private static Stream provideEndpointTokenScenarios() throws Exceptio mockCpTokenSource, "", mockSuccessClient, + TEST_HOST, IllegalArgumentException.class, null, null, @@ -149,10 +156,33 @@ private static Stream provideEndpointTokenScenarios() throws Exceptio mockCpTokenSource, TEST_AUTH_DETAILS, null, + TEST_HOST, + NullPointerException.class, + null, + null, + null, + 0), + Arguments.of( + "Null host", + mockCpTokenSource, + TEST_AUTH_DETAILS, + mockSuccessClient, + null, NullPointerException.class, null, null, null, + 0), + Arguments.of( + "Empty host", + mockCpTokenSource, + TEST_AUTH_DETAILS, + mockSuccessClient, + "", + IllegalArgumentException.class, + null, + null, + null, 0)); } @@ -163,6 +193,7 @@ void testEndpointTokenSource( DatabricksOAuthTokenSource cpTokenSource, String authDetails, HttpClient httpClient, + String host, Class expectedException, String expectedAccessToken, String expectedTokenType, @@ -173,19 +204,17 @@ void testEndpointTokenSource( expectedException, () -> { EndpointTokenSource source = - new EndpointTokenSource(cpTokenSource, authDetails, httpClient); + new EndpointTokenSource(cpTokenSource, authDetails, httpClient, host); source.getToken(); }); } else { - EndpointTokenSource source = new EndpointTokenSource(cpTokenSource, authDetails, httpClient); + EndpointTokenSource source = + new EndpointTokenSource(cpTokenSource, authDetails, httpClient, host); Token token = source.getToken(); assertNotNull(token); assertEquals(expectedAccessToken, token.getAccessToken()); assertEquals(expectedTokenType, token.getTokenType()); assertEquals(expectedRefreshToken, token.getRefreshToken()); - // Allow a few seconds of clock skew for expiry - assertTrue(token.isValid()); - assertTrue(token.getAccessToken().length() > 0); } } } From b60ed0e79a6023fc6fe52104ef11f1643fa66eba Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Fri, 23 May 2025 16:47:27 +0000 Subject: [PATCH 3/6] Updated tests --- .../core/oauth/DataPlaneTokenSourceTest.java | 36 +++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DataPlaneTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DataPlaneTokenSourceTest.java index 35b3586d5..5eb08ac93 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DataPlaneTokenSourceTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DataPlaneTokenSourceTest.java @@ -10,9 +10,11 @@ import java.net.URL; import java.time.LocalDateTime; 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; +import org.mockito.MockedConstruction; public class DataPlaneTokenSourceTest { private static final String TEST_ENDPOINT_1 = "https://endpoint1.databricks.com/"; @@ -197,4 +199,38 @@ void testDataPlaneTokenSource( assertTrue(token.isValid()); } } + + @Test + void testEndpointTokenSourceConstructionCount() throws Exception { + Token cpToken = new Token(TEST_CP_TOKEN, TEST_TOKEN_TYPE, null, LocalDateTime.now().plusSeconds(3600)); + DatabricksOAuthTokenSource mockCpTokenSource = mock(DatabricksOAuthTokenSource.class); + when(mockCpTokenSource.getToken()).thenReturn(cpToken); + + String successJson = "{\"access_token\":\"dp-access-token\",\"token_type\":\"Bearer\",\"refresh_token\":\"refresh-token\",\"expires_in\":3600}"; + HttpClient mockHttpClient = mock(HttpClient.class); + when(mockHttpClient.execute(any())).thenReturn(new Response(successJson, 200, "OK", new URL(TEST_ENDPOINT_1))); + + try (MockedConstruction mockedConstruction = mockConstruction(EndpointTokenSource.class)) { + DataPlaneTokenSource source = new DataPlaneTokenSource(mockHttpClient, mockCpTokenSource, TEST_HOST); + + // First call - should create new EndpointTokenSource + source.getToken(TEST_ENDPOINT_1, TEST_AUTH_DETAILS_1); + assertEquals(1, mockedConstruction.constructed().size(), "First call should create one EndpointTokenSource"); + + // Second call with same endpoint and auth details - should reuse existing EndpointTokenSource + source.getToken(TEST_ENDPOINT_1, TEST_AUTH_DETAILS_1); + assertEquals(1, mockedConstruction.constructed().size(), "This call should reuse the existing EndpointTokenSource"); + + // Call with different endpoint - should create new EndpointTokenSource + source.getToken(TEST_ENDPOINT_2, TEST_AUTH_DETAILS_2); + assertEquals(2, mockedConstruction.constructed().size(), "Different endpoint should create new EndpointTokenSource"); + + // Call with different auth details - should create new EndpointTokenSource + source.getToken(TEST_ENDPOINT_1, TEST_AUTH_DETAILS_2); + assertEquals(3, mockedConstruction.constructed().size(), "Different auth details should create new EndpointTokenSource"); + + source.getToken(TEST_ENDPOINT_2, TEST_AUTH_DETAILS_2); + assertEquals(3, mockedConstruction.constructed().size(), "This call should reuse the existing EndpointTokenSource"); + } + } } From b56c4d597e8dca41e1bdfbe6131ce33141add9d5 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Fri, 23 May 2025 16:48:47 +0000 Subject: [PATCH 4/6] Updated tests --- .../com/databricks/sdk/core/oauth/DataPlaneTokenSourceTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DataPlaneTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DataPlaneTokenSourceTest.java index 5eb08ac93..a3724ff3d 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DataPlaneTokenSourceTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DataPlaneTokenSourceTest.java @@ -201,7 +201,7 @@ void testDataPlaneTokenSource( } @Test - void testEndpointTokenSourceConstructionCount() throws Exception { + void testEndpointTokenSourceCaching() throws Exception { Token cpToken = new Token(TEST_CP_TOKEN, TEST_TOKEN_TYPE, null, LocalDateTime.now().plusSeconds(3600)); DatabricksOAuthTokenSource mockCpTokenSource = mock(DatabricksOAuthTokenSource.class); when(mockCpTokenSource.getToken()).thenReturn(cpToken); From fd72018c1b8652a22245dc215bf5a83a56d47205 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Sun, 25 May 2025 20:53:53 +0000 Subject: [PATCH 5/6] Fix javadoc --- .../databricks/sdk/core/oauth/DataPlaneTokenSource.java | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DataPlaneTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DataPlaneTokenSource.java index 4504d9ca3..8d48c2dff 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DataPlaneTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DataPlaneTokenSource.java @@ -64,7 +64,8 @@ public int hashCode() { * @param httpClient The {@link HttpClient} for token requests. * @param cpTokenSource The {@link TokenSource} for control plane tokens. * @param host The host for the token exchange request. - * @throws NullPointerException if either parameter is null + * @throws NullPointerException if any parameter is null. + * @throws IllegalArgumentException if the host is empty. */ public DataPlaneTokenSource(HttpClient httpClient, TokenSource cpTokenSource, String host) { this.httpClient = Objects.requireNonNull(httpClient, "HTTP client cannot be null"); @@ -85,8 +86,9 @@ public DataPlaneTokenSource(HttpClient httpClient, TokenSource cpTokenSource, St * @param endpoint The target data plane service endpoint. * @param authDetails Authorization details for the endpoint. * @return The dataplane {@link Token}. - * @throws NullPointerException if either parameter is null - * @throws IllegalArgumentException if either parameter is empty + * @throws NullPointerException if either parameter is null. + * @throws IllegalArgumentException if either parameter is empty. + * @throws DatabricksException if the token request fails. */ public Token getToken(String endpoint, String authDetails) { Objects.requireNonNull(endpoint, "Data plane endpoint URL cannot be null"); From ba40e9492b6d5f07fd20b0692656477d463279fc Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Sun, 25 May 2025 20:54:58 +0000 Subject: [PATCH 6/6] Fix formatting --- .../core/oauth/DataPlaneTokenSourceTest.java | 40 ++++++++++++++----- 1 file changed, 30 insertions(+), 10 deletions(-) diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DataPlaneTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DataPlaneTokenSourceTest.java index a3724ff3d..5887c4ee1 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DataPlaneTokenSourceTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DataPlaneTokenSourceTest.java @@ -202,35 +202,55 @@ void testDataPlaneTokenSource( @Test void testEndpointTokenSourceCaching() throws Exception { - Token cpToken = new Token(TEST_CP_TOKEN, TEST_TOKEN_TYPE, null, LocalDateTime.now().plusSeconds(3600)); + Token cpToken = + new Token(TEST_CP_TOKEN, TEST_TOKEN_TYPE, null, LocalDateTime.now().plusSeconds(3600)); DatabricksOAuthTokenSource mockCpTokenSource = mock(DatabricksOAuthTokenSource.class); when(mockCpTokenSource.getToken()).thenReturn(cpToken); - String successJson = "{\"access_token\":\"dp-access-token\",\"token_type\":\"Bearer\",\"refresh_token\":\"refresh-token\",\"expires_in\":3600}"; + String successJson = + "{\"access_token\":\"dp-access-token\",\"token_type\":\"Bearer\",\"refresh_token\":\"refresh-token\",\"expires_in\":3600}"; HttpClient mockHttpClient = mock(HttpClient.class); - when(mockHttpClient.execute(any())).thenReturn(new Response(successJson, 200, "OK", new URL(TEST_ENDPOINT_1))); + when(mockHttpClient.execute(any())) + .thenReturn(new Response(successJson, 200, "OK", new URL(TEST_ENDPOINT_1))); - try (MockedConstruction mockedConstruction = mockConstruction(EndpointTokenSource.class)) { - DataPlaneTokenSource source = new DataPlaneTokenSource(mockHttpClient, mockCpTokenSource, TEST_HOST); + try (MockedConstruction mockedConstruction = + mockConstruction(EndpointTokenSource.class)) { + DataPlaneTokenSource source = + new DataPlaneTokenSource(mockHttpClient, mockCpTokenSource, TEST_HOST); // First call - should create new EndpointTokenSource source.getToken(TEST_ENDPOINT_1, TEST_AUTH_DETAILS_1); - assertEquals(1, mockedConstruction.constructed().size(), "First call should create one EndpointTokenSource"); + assertEquals( + 1, + mockedConstruction.constructed().size(), + "First call should create one EndpointTokenSource"); // Second call with same endpoint and auth details - should reuse existing EndpointTokenSource source.getToken(TEST_ENDPOINT_1, TEST_AUTH_DETAILS_1); - assertEquals(1, mockedConstruction.constructed().size(), "This call should reuse the existing EndpointTokenSource"); + assertEquals( + 1, + mockedConstruction.constructed().size(), + "This call should reuse the existing EndpointTokenSource"); // Call with different endpoint - should create new EndpointTokenSource source.getToken(TEST_ENDPOINT_2, TEST_AUTH_DETAILS_2); - assertEquals(2, mockedConstruction.constructed().size(), "Different endpoint should create new EndpointTokenSource"); + assertEquals( + 2, + mockedConstruction.constructed().size(), + "Different endpoint should create new EndpointTokenSource"); // Call with different auth details - should create new EndpointTokenSource source.getToken(TEST_ENDPOINT_1, TEST_AUTH_DETAILS_2); - assertEquals(3, mockedConstruction.constructed().size(), "Different auth details should create new EndpointTokenSource"); + assertEquals( + 3, + mockedConstruction.constructed().size(), + "Different auth details should create new EndpointTokenSource"); source.getToken(TEST_ENDPOINT_2, TEST_AUTH_DETAILS_2); - assertEquals(3, mockedConstruction.constructed().size(), "This call should reuse the existing EndpointTokenSource"); + assertEquals( + 3, + mockedConstruction.constructed().size(), + "This call should reuse the existing EndpointTokenSource"); } } }