diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java index 23fd990b9..35127219c 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java @@ -796,18 +796,6 @@ private OpenIDConnectEndpoints fetchDefaultOidcEndpoints() throws IOException { if (getHostType() == HostType.UNIFIED) { return getUnifiedOidcEndpoints(getAccountId()); } - - if (isAzure() && getAzureClientId() != null) { - Request request = new Request("GET", getHost() + "/oidc/oauth2/v2.0/authorize"); - request.setRedirectionBehavior(false); - Response resp = getHttpClient().execute(request); - String realAuthUrl = resp.getFirstHeader("location"); - if (realAuthUrl == null) { - return null; - } - return new OpenIDConnectEndpoints( - realAuthUrl.replaceAll("/authorize", "/token"), realAuthUrl); - } if (isAccountClient() && getAccountId() != null) { String prefix = getHost() + "/oidc/accounts/" + getAccountId(); return new OpenIDConnectEndpoints(prefix + "/v1/token", prefix + "/v1/authorize"); 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 dbd1cba9f..5a3b1044d 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 @@ -7,6 +7,7 @@ import com.databricks.sdk.core.DatabricksException; import com.databricks.sdk.core.FixtureServer; import com.databricks.sdk.core.HeaderFactory; +import com.databricks.sdk.core.OpenIDConnectEndpoints; import com.databricks.sdk.core.commons.CommonsHttpClient; import com.databricks.sdk.core.http.HttpClient; import com.databricks.sdk.core.http.Request; @@ -303,7 +304,8 @@ void cacheWithValidRefreshableTokenTest() throws IOException { any(DatabricksConfig.class), any(String.class), any(String.class), - any(TokenCache.class)); + any(TokenCache.class), + any(OpenIDConnectEndpoints.class)); // Verify token was NOT saved back to cache (we're using the cached one as-is). Mockito.verify(mockTokenCache, Mockito.never()).save(any(Token.class)); @@ -363,7 +365,7 @@ void cacheWithValidNonRefreshableTokenTest() throws IOException { // Verify performBrowserAuth was NOT called. Mockito.verify(provider, Mockito.never()) - .performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class)); + .performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class), any(OpenIDConnectEndpoints.class)); // Verify no token was saved (we're using the cached one as-is). Mockito.verify(mockTokenCache, Mockito.never()).save(any(Token.class)); @@ -430,7 +432,9 @@ void cacheWithInvalidAccessTokenValidRefreshTest() throws IOException { any(DatabricksConfig.class), any(String.class), any(String.class), - any(TokenCache.class)); + any(TokenCache.class), + any(OpenIDConnectEndpoints.class) + ); // Verify token was saved back to cache Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); @@ -508,7 +512,7 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { Mockito.spy(new ExternalBrowserCredentialsProvider(mockTokenCache)); Mockito.doReturn(cachedTokenSource) .when(provider) - .performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class)); + .performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class), any(OpenIDConnectEndpoints.class)); // Spy on the config to inject the endpoints DatabricksConfig spyConfig = Mockito.spy(config); @@ -527,7 +531,7 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { // Verify performBrowserAuth was called since refresh failed Mockito.verify(provider, Mockito.times(1)) - .performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class)); + .performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class), any(OpenIDConnectEndpoints.class)); // Verify token was saved after browser auth (for the new token) Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); @@ -572,17 +576,26 @@ void cacheWithInvalidTokensTest() throws IOException { new DatabricksConfig() .setAuthType("external-browser") .setHost("https://test.databricks.com") - .setClientId("test-client-id"); + .setClientId("test-client-id") + .setHttpClient(mockHttpClient); // Create our provider and mock the browser auth method ExternalBrowserCredentialsProvider provider = Mockito.spy(new ExternalBrowserCredentialsProvider(mockTokenCache)); Mockito.doReturn(cachedTokenSource) .when(provider) - .performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class)); + .performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class), any(OpenIDConnectEndpoints.class)); + + // Spy on the config to inject the endpoints + OpenIDConnectEndpoints endpoints = + new OpenIDConnectEndpoints( + "https://test.databricks.com/oidc/v1/token", + "https://test.databricks.com/oidc/v1/authorize"); + DatabricksConfig spyConfig = Mockito.spy(config); + Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints(); // Configure provider - HeaderFactory headerFactory = provider.configure(config); + HeaderFactory headerFactory = provider.configure(spyConfig); assertNotNull(headerFactory); // Verify headers contain the browser auth token (fallback) Map headers = headerFactory.headers(); @@ -593,7 +606,7 @@ void cacheWithInvalidTokensTest() throws IOException { // Verify performBrowserAuth was called since we had an invalid token Mockito.verify(provider, Mockito.times(1)) - .performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class)); + .performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class), any(OpenIDConnectEndpoints.class)); // Verify token was saved after browser auth (for the new token) Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); @@ -609,7 +622,7 @@ void doNotAddOfflineAccessScopeWhenDisableOauthRefreshTokenIsTrue() { .setScopes(Arrays.asList("my-test-scope")); ExternalBrowserCredentialsProvider provider = new ExternalBrowserCredentialsProvider(); - List scopes = provider.getScopes(config); + List scopes = provider.getScopes(config, null); assertEquals(1, scopes.size()); assertTrue(scopes.contains("my-test-scope")); @@ -625,7 +638,7 @@ void doNotRemoveUserProvidedScopesWhenDisableOauthRefreshTokenIsTrue() { .setScopes(Arrays.asList("my-test-scope", "offline_access")); ExternalBrowserCredentialsProvider provider = new ExternalBrowserCredentialsProvider(); - List scopes = provider.getScopes(config); + List scopes = provider.getScopes(config, null); assertEquals(2, scopes.size()); assertTrue(scopes.contains("offline_access")); @@ -641,7 +654,7 @@ void addOfflineAccessScopeWhenDisableOauthRefreshTokenIsFalse() { .setScopes(Arrays.asList("my-test-scope")); ExternalBrowserCredentialsProvider provider = new ExternalBrowserCredentialsProvider(); - List scopes = provider.getScopes(config); + List scopes = provider.getScopes(config, null); assertEquals(2, scopes.size()); assertTrue(scopes.contains("offline_access")); diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/integration/DatabricksOidcAccountIT.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/integration/DatabricksOidcAccountIT.java new file mode 100644 index 000000000..6e7f31e66 --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/integration/DatabricksOidcAccountIT.java @@ -0,0 +1,73 @@ +package com.databricks.sdk.integration; + +import com.databricks.sdk.AccountClient; +import com.databricks.sdk.core.DatabricksConfig; +import com.databricks.sdk.integration.framework.EnvContext; +import com.databricks.sdk.integration.framework.EnvOrSkip; +import com.databricks.sdk.integration.framework.EnvTest; +import com.databricks.sdk.service.iam.ListAccountServicePrincipalsRequest; +import com.databricks.sdk.service.iam.ServicePrincipal; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +@ExtendWith(EnvTest.class) +@EnvContext("account") +public class DatabricksOidcAccountIT { + private static final Logger LOG = LoggerFactory.getLogger(DatabricksOidcAccountIT.class); + + @Test + void testAccountOAuthM2MAuth( + @EnvOrSkip("CLOUD_ENV") String cloudEnv, + @EnvOrSkip("DATABRICKS_HOST") String host, + @EnvOrSkip("DATABRICKS_ACCOUNT_ID") String accountId, + @EnvOrSkip("TEST_DATABRICKS_CLIENT_ID") String clientId, + @EnvOrSkip("TEST_DATABRICKS_CLIENT_SECRET") String clientSecret) { + LOG.info("Cloud environment: {}", cloudEnv); + + // Create account client with OAuth M2M authentication + DatabricksConfig config = + new DatabricksConfig() + .setHost(host) + .setAccountId(accountId) + .setClientId(clientId) + .setClientSecret(clientSecret) + .setAuthType("oauth-m2m"); + + AccountClient ac = new AccountClient(config); + + // List service principals to verify authentication works + Iterable servicePrincipals = + ac.servicePrincipals().list(new ListAccountServicePrincipalsRequest()); + servicePrincipals.iterator().next(); + } + + @Test + void testAccountAzureClientSecretAuth( + @EnvOrSkip("CLOUD_ENV") String cloudEnv, + @EnvOrSkip("DATABRICKS_HOST") String host, + @EnvOrSkip("DATABRICKS_ACCOUNT_ID") String accountId, + @EnvOrSkip("ARM_CLIENT_ID") String azureClientId, + @EnvOrSkip("ARM_CLIENT_SECRET") String azureClientSecret, + @EnvOrSkip("ARM_TENANT_ID") String azureTenantId) { + LOG.info("Cloud environment: {}", cloudEnv); + + // Create account client with Azure client secret authentication + DatabricksConfig config = + new DatabricksConfig() + .setHost(host) + .setAccountId(accountId) + .setAzureClientId(azureClientId) + .setAzureClientSecret(azureClientSecret) + .setAzureTenantId(azureTenantId) + .setAuthType("azure-client-secret"); + + AccountClient ac = new AccountClient(config); + + // List service principals to verify authentication works + Iterable servicePrincipals = + ac.servicePrincipals().list(new ListAccountServicePrincipalsRequest()); + servicePrincipals.iterator().next(); + } +} diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/integration/DatabricksOidcWorkspaceIT.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/integration/DatabricksOidcWorkspaceIT.java new file mode 100644 index 000000000..e81fd90c6 --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/integration/DatabricksOidcWorkspaceIT.java @@ -0,0 +1,72 @@ +package com.databricks.sdk.integration; + +import static org.junit.jupiter.api.Assertions.assertNotNull; + +import com.databricks.sdk.WorkspaceClient; +import com.databricks.sdk.core.DatabricksConfig; +import com.databricks.sdk.integration.framework.EnvContext; +import com.databricks.sdk.integration.framework.EnvOrSkip; +import com.databricks.sdk.integration.framework.EnvTest; +import com.databricks.sdk.service.iam.User; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +@ExtendWith(EnvTest.class) +@EnvContext("ucws") +public class DatabricksOidcWorkspaceIT { + private static final Logger LOG = LoggerFactory.getLogger(DatabricksOidcWorkspaceIT.class); + + @Test + void testWorkspaceOAuthM2MAuth( + @EnvOrSkip("CLOUD_ENV") String cloudEnv, + @EnvOrSkip("DATABRICKS_HOST") String host, + @EnvOrSkip("TEST_DATABRICKS_CLIENT_ID") String clientId, + @EnvOrSkip("TEST_DATABRICKS_CLIENT_SECRET") String clientSecret) { + LOG.info("Cloud environment: {}", cloudEnv); + + // Create workspace client with OAuth M2M authentication + DatabricksConfig config = + new DatabricksConfig() + .setHost(host) + .setClientId(clientId) + .setClientSecret(clientSecret) + .setAuthType("oauth-m2m"); + + WorkspaceClient ws = new WorkspaceClient(config); + + // Call the "me" API + User me = ws.currentUser().me(); + + // Verify we got a valid response + assertNotNull(me.getUserName(), "Expected non-empty UserName"); + } + + @Test + void testWorkspaceAzureClientSecretAuth( + @EnvOrSkip("CLOUD_ENV") String cloudEnv, + @EnvOrSkip("DATABRICKS_HOST") String host, + @EnvOrSkip("ARM_CLIENT_ID") String azureClientId, + @EnvOrSkip("ARM_CLIENT_SECRET") String azureClientSecret, + @EnvOrSkip("ARM_TENANT_ID") String azureTenantId) { + LOG.info("Cloud environment: {}", cloudEnv); + + // Create workspace client with Azure client secret authentication + DatabricksConfig config = + new DatabricksConfig() + .setHost(host) + .setAzureClientId(azureClientId) + .setAzureClientSecret(azureClientSecret) + .setAzureTenantId(azureTenantId) + .setAuthType("azure-client-secret"); + + WorkspaceClient ws = new WorkspaceClient(config); + + // Call the "me" API + User me = ws.currentUser().me(); + + // Verify we got a valid response + assertNotNull(me.getUserName(), "Expected non-empty UserName"); + } +}