From 32f18f9f66c1a37aec0ae78cfdbe4d763427c999 Mon Sep 17 00:00:00 2001 From: samikshya-chand_data Date: Wed, 28 May 2025 23:17:59 +0530 Subject: [PATCH 1/4] Fix azure u2m --- .../com/databricks/sdk/core/DatabricksConfig.java | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) 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..95183bdf5 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,14 @@ 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) + */ + private boolean shouldUseAzureOidcEndpoints() { + return Objects.equals(getAuthType(), null) || getAzureClientId() != null; + } } From 7defeec74b98380bc3657743f6a3d432224b3bd3 Mon Sep 17 00:00:00 2001 From: samikshya-chand_data Date: Wed, 28 May 2025 23:36:35 +0530 Subject: [PATCH 2/4] Add unit tests --- .../databricks/sdk/core/DatabricksConfig.java | 17 +++++----- .../sdk/core/DatabricksConfigTest.java | 31 +++++++++++++++++++ 2 files changed, 41 insertions(+), 7 deletions(-) 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 95183bdf5..f64e7ee41 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 @@ -410,13 +410,17 @@ public DatabricksConfig setAzureUseMsi(boolean azureUseMsi) { return this; } - /** @deprecated Use {@link #getAzureUseMsi()} instead. */ + /** + * @deprecated Use {@link #getAzureUseMsi()} instead. + */ @Deprecated() public boolean getAzureUseMSI() { return azureUseMsi; } - /** @deprecated Use {@link #setAzureUseMsi(boolean)} instead. */ + /** + * @deprecated Use {@link #setAzureUseMsi(boolean)} instead. + */ @Deprecated public DatabricksConfig setAzureUseMSI(boolean azureUseMsi) { this.azureUseMsi = azureUseMsi; @@ -745,12 +749,11 @@ public String getEffectiveOAuthRedirectUrl() { } /** - * 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) + * 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) */ - private boolean shouldUseAzureOidcEndpoints() { + 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()); + } } From 969640a6d4e4b412224ca15cdba73b6a8680b7c6 Mon Sep 17 00:00:00 2001 From: samikshya-chand_data Date: Wed, 28 May 2025 23:40:21 +0530 Subject: [PATCH 3/4] format fix --- .../java/com/databricks/sdk/core/DatabricksConfig.java | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) 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 f64e7ee41..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 @@ -410,17 +410,13 @@ public DatabricksConfig setAzureUseMsi(boolean azureUseMsi) { return this; } - /** - * @deprecated Use {@link #getAzureUseMsi()} instead. - */ + /** @deprecated Use {@link #getAzureUseMsi()} instead. */ @Deprecated() public boolean getAzureUseMSI() { return azureUseMsi; } - /** - * @deprecated Use {@link #setAzureUseMsi(boolean)} instead. - */ + /** @deprecated Use {@link #setAzureUseMsi(boolean)} instead. */ @Deprecated public DatabricksConfig setAzureUseMSI(boolean azureUseMsi) { this.azureUseMsi = azureUseMsi; From 9729113966310bb31d33cb94ed763dca0f8f3ddc Mon Sep 17 00:00:00 2001 From: samikshya-chand_data Date: Wed, 28 May 2025 23:44:48 +0530 Subject: [PATCH 4/4] Add next changelog --- NEXT_CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) 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