diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index ac1055a9c..fcd30486d 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -5,6 +5,7 @@ ### New Features and Improvements ### Bug Fixes +- Fix Azure OIDC endpoint selection to support both U2M and M2M authentication flows ([#453](https://github.com/databricks/databricks-sdk-java/pull/454)). ### Documentation 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 de6548982..c2105c959 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 @@ -628,7 +628,8 @@ private OpenIDConnectEndpoints fetchDefaultOidcEndpoints() throws IOException { if (getHost() == null) { return null; } - if (isAzure() && getAzureClientId() != null) { + + if (isAzure() && shouldUseAzureOidcEndpoints()) { Request request = new Request("GET", getHost() + "/oidc/oauth2/v2.0/authorize"); request.setRedirectionBehavior(false); Response resp = getHttpClient().execute(request); @@ -742,4 +743,13 @@ public DatabricksConfig newWithWorkspaceHost(String host) { public String getEffectiveOAuthRedirectUrl() { return redirectUrl != null ? redirectUrl : "http://localhost:8080/callback"; } + + /** + * Determines if Azure-specific OIDC endpoints should be used. This is true in two cases: 1. When + * auth type is not specified (this is only in case of external browser auth) 2. When Azure client + * ID is present (service principal auth) + */ + boolean shouldUseAzureOidcEndpoints() { + return Objects.equals(getAuthType(), null) || getAzureClientId() != null; + } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java index 38b6fcd9c..28bd76510 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java @@ -251,4 +251,35 @@ public void testGetTokenSourceWithOAuth() { assertFalse(tokenSource instanceof ErrorTokenSource); assertEquals(tokenSource.getToken().getAccessToken(), "test-token"); } + + @Test + public void testShouldUseAzureOidcEndpointsForExternalBrowserAuth() { + DatabricksConfig config = + new DatabricksConfig() + .setHost("https://adb-1234567890.0.azuredatabricks.net/") + .setAuthType(null); + assertTrue(config.shouldUseAzureOidcEndpoints()); + } + + @Test + public void testShouldUseAzureOidcEndpointsForServicePrincipal() { + DatabricksConfig config = + new DatabricksConfig() + .setHost("https://adb-1234567890.0.azuredatabricks.net/") + .setAzureClientId("test-client-id") + .setAzureClientSecret("test-client-secret") + .setAzureTenantId("test-tenant-id"); + assertTrue(config.shouldUseAzureOidcEndpoints()); + } + + @Test + public void testShouldNotUseAzureOidcEndpointsForAzureM2M() { + DatabricksConfig config = + new DatabricksConfig() + .setHost("https://adb-1234567890.0.azuredatabricks.net/") + .setAuthType("oauth-m2m") + .setClientId("test-client-id") + .setClientSecret("test-client-secret"); + assertFalse(config.shouldUseAzureOidcEndpoints()); + } }