From 9c7ae67d55b957b095ad82a7260b8a554c65479a Mon Sep 17 00:00:00 2001 From: Sreekanth Vadigi Date: Fri, 11 Jul 2025 13:57:21 +0530 Subject: [PATCH 1/8] infer azure tenant id Signed-off-by: Sreekanth Vadigi --- .../databricks/sdk/core/DatabricksConfig.java | 86 +++++++++++++++++- ...reServicePrincipalCredentialsProvider.java | 4 +- .../sdk/core/DatabricksConfigTest.java | 87 +++++++++++++++++++ 3 files changed, 174 insertions(+), 3 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 0bc0b868a..064a32194 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 @@ -14,10 +14,15 @@ import java.io.File; import java.io.IOException; import java.lang.reflect.Field; +import java.net.URL; import java.util.*; import org.apache.http.HttpMessage; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; public class DatabricksConfig { + + private static final Logger logger = LoggerFactory.getLogger(DatabricksConfig.class); private CredentialsProvider credentialsProvider = new DefaultCredentialsProvider(); @ConfigAttribute(env = "DATABRICKS_HOST") @@ -726,7 +731,7 @@ private DatabricksConfig clone(Set fieldsToSkip) { } public DatabricksConfig clone() { - return clone(new HashSet<>()); + return clone(new HashSet<>(Collections.singletonList("logger"))); } public DatabricksConfig newWithWorkspaceHost(String host) { @@ -736,6 +741,7 @@ public DatabricksConfig newWithWorkspaceHost(String host) { // The config for WorkspaceClient has a different host and Azure Workspace resource // ID, and also omits // the account ID. + "logger", "host", "accountId", "azureWorkspaceResourceId", @@ -755,4 +761,82 @@ public DatabricksConfig newWithWorkspaceHost(String host) { public String getEffectiveOAuthRedirectUrl() { return redirectUrl != null ? redirectUrl : "http://localhost:8080/callback"; } + + private static final String AZURE_AUTH_ENDPOINT = "/aad/auth"; + + /** + * [Internal] Load the Azure tenant ID from the Azure Databricks login page. If the tenant ID is + * already set, this method does nothing. + */ + public void loadAzureTenantId() { + + if (!isAzure() || azureTenantId != null || host == null) { + return; + } + + String loginUrl = host + AZURE_AUTH_ENDPOINT; + logger.debug("Loading tenant ID from {}", loginUrl); + + try { + String redirectLocation = getRedirectLocation(loginUrl); + if (redirectLocation == null) { + return; + } + + String extractedTenantId = extractTenantIdFromUrl(redirectLocation); + if (extractedTenantId == null) { + return; + } + + this.azureTenantId = extractedTenantId; + logger.debug("Loaded tenant ID: {}", this.azureTenantId); + + } catch (Exception e) { + logger.warn("Failed to load tenant ID: {}", e.getMessage()); + } + } + + private String getRedirectLocation(String loginUrl) throws IOException { + + Request request = new Request("GET", loginUrl); + request.setRedirectionBehavior(false); + Response response = getHttpClient().execute(request); + int statusCode = response.getStatusCode(); + + if (statusCode / 100 != 3) { + logger.warn( + "Failed to get tenant ID from {}: expected status code 3xx, got {}", + loginUrl, + statusCode); + return null; + } + + String location = response.getFirstHeader("Location"); + if (location == null) { + logger.warn("No Location header in response from {}", loginUrl); + } + + return location; + } + + private String extractTenantIdFromUrl(String redirectUrl) { + try { + // The Location header has the following form: + // https://login.microsoftonline.com//oauth2/authorize?... + // The domain may change depending on the Azure cloud (e.g. login.microsoftonline.us for US + // Government cloud). + URL entraIdUrl = new URL(redirectUrl); + String[] pathSegments = entraIdUrl.getPath().split("/"); + + if (pathSegments.length < 2) { + logger.warn("Invalid path in Location header: {}", entraIdUrl.getPath()); + return null; + } + + return pathSegments[1]; + } catch (Exception e) { + logger.warn("Failed to extract tenant ID from URL {}: {}", redirectUrl, e.getMessage()); + return null; + } + } } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java index 6e6df86eb..e5b81e341 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java @@ -22,12 +22,12 @@ public String authType() { public OAuthHeaderFactory configure(DatabricksConfig config) { if (!config.isAzure() || config.getAzureClientId() == null - || config.getAzureClientSecret() == null - || config.getAzureTenantId() == null) { + || config.getAzureClientSecret() == null) { return null; } AzureUtils.ensureHostPresent( config, mapper, AzureServicePrincipalCredentialsProvider::tokenSourceFor); + config.loadAzureTenantId(); CachedTokenSource inner = tokenSourceFor(config, config.getEffectiveAzureLoginAppId()); CachedTokenSource cloud = tokenSourceFor(config, config.getAzureEnvironment().getServiceManagementEndpoint()); 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 b3ac333a3..43544d3df 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 @@ -250,4 +250,91 @@ public void testGetTokenSourceWithOAuth() { assertFalse(tokenSource instanceof ErrorTokenSource); assertEquals(tokenSource.getToken().getAccessToken(), "test-token"); } + + @Test + public void testLoadAzureTenantId404() throws IOException { + try (FixtureServer server = new FixtureServer().with("GET", "/aad/auth", "", 404)) { + DatabricksConfig config = new DatabricksConfig(); + config.setHost(server.getUrl()); + config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); + config.loadAzureTenantId(); + assertNull(config.getAzureTenantId()); + } + } + + @Test + public void testLoadAzureTenantIdNoLocationHeader() throws IOException { + try (FixtureServer server = new FixtureServer().with("GET", "/aad/auth", "", 302)) { + DatabricksConfig config = new DatabricksConfig(); + config.setHost(server.getUrl()); + config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); + config.loadAzureTenantId(); + assertNull(config.getAzureTenantId()); + } + } + + @Test + public void testLoadAzureTenantIdUnparsableLocationHeader() throws IOException { + FixtureServer.FixtureMapping fixture = + new FixtureServer.FixtureMapping.Builder() + .validateMethod("GET") + .validatePath("/aad/auth") + .withRedirect("https://unexpected-location", 302) + .build(); + + try (FixtureServer server = new FixtureServer().with(fixture)) { + DatabricksConfig config = new DatabricksConfig(); + config.setHost(server.getUrl()); + config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); + config.loadAzureTenantId(); + assertNull(config.getAzureTenantId()); + } + } + + @Test + public void testLoadAzureTenantIdHappyPath() throws IOException { + FixtureServer.FixtureMapping fixture = + new FixtureServer.FixtureMapping.Builder() + .validateMethod("GET") + .validatePath("/aad/auth") + .withRedirect("https://login.microsoftonline.com/test-tenant-id/oauth2/authorize", 302) + .build(); + + try (FixtureServer server = new FixtureServer().with(fixture)) { + DatabricksConfig config = new DatabricksConfig(); + config.setHost(server.getUrl()); + config.setAzureWorkspaceResourceId( + "/subscriptions/123/resourceGroups/rg/providers/Microsoft.Databricks/workspaces/ws"); + config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); + config.loadAzureTenantId(); + assertEquals("test-tenant-id", config.getAzureTenantId()); + } + } + + @Test + public void testLoadAzureTenantIdSkipsWhenNotAzure() throws IOException { + DatabricksConfig config = new DatabricksConfig(); + config.setHost("https://my-workspace.cloud.databricks.com"); // non-azure host + config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); + config.loadAzureTenantId(); + assertNull(config.getAzureTenantId()); + } + + @Test + public void testLoadAzureTenantIdSkipsWhenAlreadySet() throws IOException { + DatabricksConfig config = new DatabricksConfig(); + config.setHost("https://adb-123.0.azuredatabricks.net"); + config.setAzureTenantId("existing-tenant-id"); + config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); + config.loadAzureTenantId(); + assertEquals("existing-tenant-id", config.getAzureTenantId()); + } + + @Test + public void testLoadAzureTenantIdSkipsWhenNoHost() throws IOException { + DatabricksConfig config = new DatabricksConfig(); + config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); + config.loadAzureTenantId(); + assertNull(config.getAzureTenantId()); + } } From 56859bfe621053c2f6826b079e4c09e67ac5143f Mon Sep 17 00:00:00 2001 From: Sreekanth Vadigi Date: Fri, 11 Jul 2025 14:21:14 +0530 Subject: [PATCH 2/8] unit test fix Signed-off-by: Sreekanth Vadigi --- .../main/java/com/databricks/sdk/core/DatabricksConfig.java | 5 ++--- 1 file changed, 2 insertions(+), 3 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 064a32194..5b0ce5577 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 @@ -762,8 +762,6 @@ public String getEffectiveOAuthRedirectUrl() { return redirectUrl != null ? redirectUrl : "http://localhost:8080/callback"; } - private static final String AZURE_AUTH_ENDPOINT = "/aad/auth"; - /** * [Internal] Load the Azure tenant ID from the Azure Databricks login page. If the tenant ID is * already set, this method does nothing. @@ -774,7 +772,8 @@ public void loadAzureTenantId() { return; } - String loginUrl = host + AZURE_AUTH_ENDPOINT; + final String azureAuthEndpoint = "/aad/auth"; + String loginUrl = host + azureAuthEndpoint; logger.debug("Loading tenant ID from {}", loginUrl); try { From 312c32a3df029eacd8f6484572bd3e82d9c345d3 Mon Sep 17 00:00:00 2001 From: Sreekanth Vadigi Date: Thu, 31 Jul 2025 18:57:05 +0530 Subject: [PATCH 3/8] changelog Signed-off-by: Sreekanth Vadigi --- NEXT_CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 6bd12e9f3..16174ee42 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -3,6 +3,7 @@ ## Release v0.57.0 ### New Features and Improvements +- Azure Service Principal credential provider can now automatically discover tenant ID when not explicitly provided ### Bug Fixes From f06c006793da697a566b7ffcc6108c47c1038a57 Mon Sep 17 00:00:00 2001 From: Sreekanth Vadigi Date: Sun, 3 Aug 2025 18:29:13 +0530 Subject: [PATCH 4/8] addressing review comments Signed-off-by: Sreekanth Vadigi --- .../databricks/sdk/core/DatabricksConfig.java | 32 +++++++++++++------ ...reServicePrincipalCredentialsProvider.java | 7 +++- .../sdk/core/DatabricksConfigTest.java | 21 ++++++++---- 3 files changed, 42 insertions(+), 18 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 5b0ce5577..2355d628d 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 @@ -23,6 +23,10 @@ public class DatabricksConfig { private static final Logger logger = LoggerFactory.getLogger(DatabricksConfig.class); + + /** Azure authentication endpoint for tenant ID discovery */ + private static final String AZURE_AUTH_ENDPOINT = "/aad/auth"; + private CredentialsProvider credentialsProvider = new DefaultCredentialsProvider(); @ConfigAttribute(env = "DATABRICKS_HOST") @@ -731,7 +735,7 @@ private DatabricksConfig clone(Set fieldsToSkip) { } public DatabricksConfig clone() { - return clone(new HashSet<>(Collections.singletonList("logger"))); + return clone(new HashSet<>(Arrays.asList("logger", "AZURE_AUTH_ENDPOINT"))); } public DatabricksConfig newWithWorkspaceHost(String host) { @@ -742,6 +746,7 @@ public DatabricksConfig newWithWorkspaceHost(String host) { // ID, and also omits // the account ID. "logger", + "AZURE_AUTH_ENDPOINT", "host", "accountId", "azureWorkspaceResourceId", @@ -765,33 +770,40 @@ public String getEffectiveOAuthRedirectUrl() { /** * [Internal] Load the Azure tenant ID from the Azure Databricks login page. If the tenant ID is * already set, this method does nothing. + * + * @return true if tenant ID is available (either was already set or successfully loaded), false otherwise */ - public void loadAzureTenantId() { + public boolean loadAzureTenantId() { - if (!isAzure() || azureTenantId != null || host == null) { - return; + if (azureTenantId != null) { + return true; // Tenant ID already available - success + } + + if (!isAzure() || host == null) { + return false; // Configuration issue - can't perform operation } - final String azureAuthEndpoint = "/aad/auth"; - String loginUrl = host + azureAuthEndpoint; + String loginUrl = host + AZURE_AUTH_ENDPOINT; logger.debug("Loading tenant ID from {}", loginUrl); try { String redirectLocation = getRedirectLocation(loginUrl); if (redirectLocation == null) { - return; + return false; // Failed to get redirect location } String extractedTenantId = extractTenantIdFromUrl(redirectLocation); if (extractedTenantId == null) { - return; + return false; // Failed to extract tenant ID } this.azureTenantId = extractedTenantId; logger.debug("Loaded tenant ID: {}", this.azureTenantId); + return true; // Successfully loaded } catch (Exception e) { logger.warn("Failed to load tenant ID: {}", e.getMessage()); + return false; } } @@ -802,9 +814,9 @@ private String getRedirectLocation(String loginUrl) throws IOException { Response response = getHttpClient().execute(request); int statusCode = response.getStatusCode(); - if (statusCode / 100 != 3) { + if (statusCode != 302) { logger.warn( - "Failed to get tenant ID from {}: expected status code 3xx, got {}", + "Failed to get tenant ID from {}: expected status code 302, got {}", loginUrl, statusCode); return null; diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java index e5b81e341..fb0654ce6 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java @@ -27,7 +27,12 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { } AzureUtils.ensureHostPresent( config, mapper, AzureServicePrincipalCredentialsProvider::tokenSourceFor); - config.loadAzureTenantId(); + + boolean tenantIdLoaded = config.loadAzureTenantId(); + if (!tenantIdLoaded) { + return null; + } + CachedTokenSource inner = tokenSourceFor(config, config.getEffectiveAzureLoginAppId()); CachedTokenSource cloud = tokenSourceFor(config, config.getAzureEnvironment().getServiceManagementEndpoint()); 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 43544d3df..568de664e 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 @@ -257,7 +257,8 @@ public void testLoadAzureTenantId404() throws IOException { DatabricksConfig config = new DatabricksConfig(); config.setHost(server.getUrl()); config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); - config.loadAzureTenantId(); + boolean result = config.loadAzureTenantId(); + assertFalse(result); assertNull(config.getAzureTenantId()); } } @@ -268,7 +269,8 @@ public void testLoadAzureTenantIdNoLocationHeader() throws IOException { DatabricksConfig config = new DatabricksConfig(); config.setHost(server.getUrl()); config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); - config.loadAzureTenantId(); + boolean result = config.loadAzureTenantId(); + assertFalse(result); assertNull(config.getAzureTenantId()); } } @@ -286,7 +288,8 @@ public void testLoadAzureTenantIdUnparsableLocationHeader() throws IOException { DatabricksConfig config = new DatabricksConfig(); config.setHost(server.getUrl()); config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); - config.loadAzureTenantId(); + boolean result = config.loadAzureTenantId(); + assertFalse(result); assertNull(config.getAzureTenantId()); } } @@ -306,7 +309,8 @@ public void testLoadAzureTenantIdHappyPath() throws IOException { config.setAzureWorkspaceResourceId( "/subscriptions/123/resourceGroups/rg/providers/Microsoft.Databricks/workspaces/ws"); config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); - config.loadAzureTenantId(); + boolean result = config.loadAzureTenantId(); + assertTrue(result); assertEquals("test-tenant-id", config.getAzureTenantId()); } } @@ -316,7 +320,8 @@ public void testLoadAzureTenantIdSkipsWhenNotAzure() throws IOException { DatabricksConfig config = new DatabricksConfig(); config.setHost("https://my-workspace.cloud.databricks.com"); // non-azure host config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); - config.loadAzureTenantId(); + boolean result = config.loadAzureTenantId(); + assertFalse(result); assertNull(config.getAzureTenantId()); } @@ -326,7 +331,8 @@ public void testLoadAzureTenantIdSkipsWhenAlreadySet() throws IOException { config.setHost("https://adb-123.0.azuredatabricks.net"); config.setAzureTenantId("existing-tenant-id"); config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); - config.loadAzureTenantId(); + boolean result = config.loadAzureTenantId(); + assertTrue(result); assertEquals("existing-tenant-id", config.getAzureTenantId()); } @@ -334,7 +340,8 @@ public void testLoadAzureTenantIdSkipsWhenAlreadySet() throws IOException { public void testLoadAzureTenantIdSkipsWhenNoHost() throws IOException { DatabricksConfig config = new DatabricksConfig(); config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); - config.loadAzureTenantId(); + boolean result = config.loadAzureTenantId(); + assertFalse(result); assertNull(config.getAzureTenantId()); } } From cf281fd9cf5b65b0b380aba70c3dfcd7798a02e8 Mon Sep 17 00:00:00 2001 From: Sreekanth Vadigi Date: Mon, 11 Aug 2025 18:39:41 +0530 Subject: [PATCH 5/8] moved inferring logic to azure utils Signed-off-by: Sreekanth Vadigi --- .../databricks/sdk/core/DatabricksConfig.java | 96 +--------------- ...reServicePrincipalCredentialsProvider.java | 18 +-- .../databricks/sdk/core/utils/AzureUtils.java | 84 ++++++++++++++ .../sdk/core/DatabricksConfigTest.java | 93 --------------- .../sdk/core/utils/AzureUtilsTest.java | 106 ++++++++++++++++++ 5 files changed, 201 insertions(+), 196 deletions(-) create mode 100644 databricks-sdk-java/src/test/java/com/databricks/sdk/core/utils/AzureUtilsTest.java 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 2355d628d..5c7118698 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 @@ -14,19 +14,11 @@ import java.io.File; import java.io.IOException; import java.lang.reflect.Field; -import java.net.URL; import java.util.*; import org.apache.http.HttpMessage; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; public class DatabricksConfig { - private static final Logger logger = LoggerFactory.getLogger(DatabricksConfig.class); - - /** Azure authentication endpoint for tenant ID discovery */ - private static final String AZURE_AUTH_ENDPOINT = "/aad/auth"; - private CredentialsProvider credentialsProvider = new DefaultCredentialsProvider(); @ConfigAttribute(env = "DATABRICKS_HOST") @@ -735,7 +727,7 @@ private DatabricksConfig clone(Set fieldsToSkip) { } public DatabricksConfig clone() { - return clone(new HashSet<>(Arrays.asList("logger", "AZURE_AUTH_ENDPOINT"))); + return clone(new HashSet<>()); } public DatabricksConfig newWithWorkspaceHost(String host) { @@ -745,8 +737,6 @@ public DatabricksConfig newWithWorkspaceHost(String host) { // The config for WorkspaceClient has a different host and Azure Workspace resource // ID, and also omits // the account ID. - "logger", - "AZURE_AUTH_ENDPOINT", "host", "accountId", "azureWorkspaceResourceId", @@ -766,88 +756,4 @@ public DatabricksConfig newWithWorkspaceHost(String host) { public String getEffectiveOAuthRedirectUrl() { return redirectUrl != null ? redirectUrl : "http://localhost:8080/callback"; } - - /** - * [Internal] Load the Azure tenant ID from the Azure Databricks login page. If the tenant ID is - * already set, this method does nothing. - * - * @return true if tenant ID is available (either was already set or successfully loaded), false otherwise - */ - public boolean loadAzureTenantId() { - - if (azureTenantId != null) { - return true; // Tenant ID already available - success - } - - if (!isAzure() || host == null) { - return false; // Configuration issue - can't perform operation - } - - String loginUrl = host + AZURE_AUTH_ENDPOINT; - logger.debug("Loading tenant ID from {}", loginUrl); - - try { - String redirectLocation = getRedirectLocation(loginUrl); - if (redirectLocation == null) { - return false; // Failed to get redirect location - } - - String extractedTenantId = extractTenantIdFromUrl(redirectLocation); - if (extractedTenantId == null) { - return false; // Failed to extract tenant ID - } - - this.azureTenantId = extractedTenantId; - logger.debug("Loaded tenant ID: {}", this.azureTenantId); - return true; // Successfully loaded - - } catch (Exception e) { - logger.warn("Failed to load tenant ID: {}", e.getMessage()); - return false; - } - } - - private String getRedirectLocation(String loginUrl) throws IOException { - - Request request = new Request("GET", loginUrl); - request.setRedirectionBehavior(false); - Response response = getHttpClient().execute(request); - int statusCode = response.getStatusCode(); - - if (statusCode != 302) { - logger.warn( - "Failed to get tenant ID from {}: expected status code 302, got {}", - loginUrl, - statusCode); - return null; - } - - String location = response.getFirstHeader("Location"); - if (location == null) { - logger.warn("No Location header in response from {}", loginUrl); - } - - return location; - } - - private String extractTenantIdFromUrl(String redirectUrl) { - try { - // The Location header has the following form: - // https://login.microsoftonline.com//oauth2/authorize?... - // The domain may change depending on the Azure cloud (e.g. login.microsoftonline.us for US - // Government cloud). - URL entraIdUrl = new URL(redirectUrl); - String[] pathSegments = entraIdUrl.getPath().split("/"); - - if (pathSegments.length < 2) { - logger.warn("Invalid path in Location header: {}", entraIdUrl.getPath()); - return null; - } - - return pathSegments[1]; - } catch (Exception e) { - logger.warn("Failed to extract tenant ID from URL {}: {}", redirectUrl, e.getMessage()); - return null; - } - } } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java index fb0654ce6..0555fed46 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java @@ -12,6 +12,7 @@ */ public class AzureServicePrincipalCredentialsProvider implements CredentialsProvider { private final ObjectMapper mapper = new ObjectMapper(); + private String tenantId; @Override public String authType() { @@ -25,14 +26,15 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { || config.getAzureClientSecret() == null) { return null; } - AzureUtils.ensureHostPresent( - config, mapper, AzureServicePrincipalCredentialsProvider::tokenSourceFor); - - boolean tenantIdLoaded = config.loadAzureTenantId(); - if (!tenantIdLoaded) { + + this.tenantId = config.getAzureTenantId() != null ? config.getAzureTenantId() : AzureUtils.inferTenantId(config); + if (this.tenantId == null) { return null; } - + + AzureUtils.ensureHostPresent( + config, mapper, this::tokenSourceFor); + CachedTokenSource inner = tokenSourceFor(config, config.getEffectiveAzureLoginAppId()); CachedTokenSource cloud = tokenSourceFor(config, config.getAzureEnvironment().getServiceManagementEndpoint()); @@ -60,9 +62,9 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { * @return A CachedTokenSource instance capable of fetching OAuth tokens for the specified Azure * resource. */ - private static CachedTokenSource tokenSourceFor(DatabricksConfig config, String resource) { + private CachedTokenSource tokenSourceFor(DatabricksConfig config, String resource) { String aadEndpoint = config.getAzureEnvironment().getActiveDirectoryEndpoint(); - String tokenUrl = aadEndpoint + config.getAzureTenantId() + "/oauth2/token"; + String tokenUrl = aadEndpoint + this.tenantId + "/oauth2/token"; Map endpointParams = new HashMap<>(); endpointParams.put("resource", resource); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/AzureUtils.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/AzureUtils.java index 09cea6e86..0f1368766 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/AzureUtils.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/AzureUtils.java @@ -10,12 +10,20 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.node.ObjectNode; import java.io.IOException; +import java.net.URL; import java.util.Map; import java.util.Optional; import java.util.function.BiFunction; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; public class AzureUtils { + private static final Logger logger = LoggerFactory.getLogger(AzureUtils.class); + + /** Azure authentication endpoint for tenant ID discovery */ + private static final String AZURE_AUTH_ENDPOINT = "/aad/auth"; + public static String getWorkspaceFromJsonResponse(ObjectNode jsonResponse) throws IOException { JsonNode properties = jsonResponse.get("properties"); if (properties == null) { @@ -95,4 +103,80 @@ public static Optional getAzureWorkspaceResourceId(Workspace workspace) workspace.getWorkspaceName()); return Optional.of(resourceId); } + + /** + * Infers the Azure tenant ID from the Databricks workspace login page. + * + * @param config The DatabricksConfig instance + * @return the discovered tenant ID, or null if discovery fails + */ + public static String inferTenantId(DatabricksConfig config) { + if (config.getAzureTenantId() != null) { + return config.getAzureTenantId(); + } + + if (!config.isAzure() || config.getHost() == null) { + logger.warn("Cannot infer tenant ID: workspace is not Azure or host is missing"); + return null; + } + + String loginUrl = config.getHost() + AZURE_AUTH_ENDPOINT; + + try { + String redirectLocation = getRedirectLocation(config, loginUrl); + if (redirectLocation == null) { + logger.warn("Failed to get redirect location from Azure auth endpoint: {}", loginUrl); + return null; + } + + String extractedTenantId = extractTenantIdFromUrl(redirectLocation); + if (extractedTenantId == null) { + logger.warn("Failed to extract tenant ID from redirect URL: {}", redirectLocation); + return null; + } + + logger.info("Successfully discovered Azure tenant ID: {}", extractedTenantId); + return extractedTenantId; + + } catch (Exception e) { + logger.warn("Exception occurred while inferring Azure tenant ID from {}: {}", loginUrl, e.getMessage()); + return null; + } + } + + private static String getRedirectLocation(DatabricksConfig config, String loginUrl) throws IOException { + Request request = new Request("GET", loginUrl); + request.setRedirectionBehavior(false); + Response response = config.getHttpClient().execute(request); + + if (response.getStatusCode() != 302) { + logger.warn("Expected redirect (302) from {}, got status code: {}", loginUrl, response.getStatusCode()); + return null; + } + + String location = response.getFirstHeader("Location"); + if (location == null) { + logger.warn("No Location header in redirect response from {}", loginUrl); + } + + return location; + } + + private static String extractTenantIdFromUrl(String redirectUrl) { + try { + // Parse: https://login.microsoftonline.com//oauth2/authorize?... + URL entraIdUrl = new URL(redirectUrl); + String[] pathSegments = entraIdUrl.getPath().split("/"); + + if (pathSegments.length < 2) { + logger.warn("Invalid path in Location header: {}", entraIdUrl.getPath()); + return null; + } + + return pathSegments[1]; + } catch (Exception e) { + logger.warn("Failed to parse tenant ID from URL {}: {}", redirectUrl, e.getMessage()); + return 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 568de664e..a025bcaa6 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,97 +251,4 @@ public void testGetTokenSourceWithOAuth() { assertEquals(tokenSource.getToken().getAccessToken(), "test-token"); } - @Test - public void testLoadAzureTenantId404() throws IOException { - try (FixtureServer server = new FixtureServer().with("GET", "/aad/auth", "", 404)) { - DatabricksConfig config = new DatabricksConfig(); - config.setHost(server.getUrl()); - config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); - boolean result = config.loadAzureTenantId(); - assertFalse(result); - assertNull(config.getAzureTenantId()); - } - } - - @Test - public void testLoadAzureTenantIdNoLocationHeader() throws IOException { - try (FixtureServer server = new FixtureServer().with("GET", "/aad/auth", "", 302)) { - DatabricksConfig config = new DatabricksConfig(); - config.setHost(server.getUrl()); - config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); - boolean result = config.loadAzureTenantId(); - assertFalse(result); - assertNull(config.getAzureTenantId()); - } - } - - @Test - public void testLoadAzureTenantIdUnparsableLocationHeader() throws IOException { - FixtureServer.FixtureMapping fixture = - new FixtureServer.FixtureMapping.Builder() - .validateMethod("GET") - .validatePath("/aad/auth") - .withRedirect("https://unexpected-location", 302) - .build(); - - try (FixtureServer server = new FixtureServer().with(fixture)) { - DatabricksConfig config = new DatabricksConfig(); - config.setHost(server.getUrl()); - config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); - boolean result = config.loadAzureTenantId(); - assertFalse(result); - assertNull(config.getAzureTenantId()); - } - } - - @Test - public void testLoadAzureTenantIdHappyPath() throws IOException { - FixtureServer.FixtureMapping fixture = - new FixtureServer.FixtureMapping.Builder() - .validateMethod("GET") - .validatePath("/aad/auth") - .withRedirect("https://login.microsoftonline.com/test-tenant-id/oauth2/authorize", 302) - .build(); - - try (FixtureServer server = new FixtureServer().with(fixture)) { - DatabricksConfig config = new DatabricksConfig(); - config.setHost(server.getUrl()); - config.setAzureWorkspaceResourceId( - "/subscriptions/123/resourceGroups/rg/providers/Microsoft.Databricks/workspaces/ws"); - config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); - boolean result = config.loadAzureTenantId(); - assertTrue(result); - assertEquals("test-tenant-id", config.getAzureTenantId()); - } - } - - @Test - public void testLoadAzureTenantIdSkipsWhenNotAzure() throws IOException { - DatabricksConfig config = new DatabricksConfig(); - config.setHost("https://my-workspace.cloud.databricks.com"); // non-azure host - config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); - boolean result = config.loadAzureTenantId(); - assertFalse(result); - assertNull(config.getAzureTenantId()); - } - - @Test - public void testLoadAzureTenantIdSkipsWhenAlreadySet() throws IOException { - DatabricksConfig config = new DatabricksConfig(); - config.setHost("https://adb-123.0.azuredatabricks.net"); - config.setAzureTenantId("existing-tenant-id"); - config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); - boolean result = config.loadAzureTenantId(); - assertTrue(result); - assertEquals("existing-tenant-id", config.getAzureTenantId()); - } - - @Test - public void testLoadAzureTenantIdSkipsWhenNoHost() throws IOException { - DatabricksConfig config = new DatabricksConfig(); - config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); - boolean result = config.loadAzureTenantId(); - assertFalse(result); - assertNull(config.getAzureTenantId()); - } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/utils/AzureUtilsTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/utils/AzureUtilsTest.java new file mode 100644 index 000000000..fea9e36ba --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/utils/AzureUtilsTest.java @@ -0,0 +1,106 @@ +package com.databricks.sdk.core.utils; + +import static org.junit.jupiter.api.Assertions.*; + +import com.databricks.sdk.core.DatabricksConfig; +import com.databricks.sdk.core.commons.CommonsHttpClient; +import com.databricks.sdk.core.FixtureServer; +import java.io.IOException; +import org.junit.jupiter.api.Test; + +public class AzureUtilsTest { + + @Test + public void testInferTenantId404() throws IOException { + try (FixtureServer server = new FixtureServer().with("GET", "/aad/auth", "", 404)) { + DatabricksConfig config = new DatabricksConfig(); + config.setHost(server.getUrl()); + config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); + String result = AzureUtils.inferTenantId(config); + assertNull(result); + assertNull(config.getAzureTenantId()); + } + } + + @Test + public void testInferTenantIdNoLocationHeader() throws IOException { + try (FixtureServer server = new FixtureServer().with("GET", "/aad/auth", "", 302)) { + DatabricksConfig config = new DatabricksConfig(); + config.setHost(server.getUrl()); + config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); + String result = AzureUtils.inferTenantId(config); + assertNull(result); + assertNull(config.getAzureTenantId()); + } + } + + @Test + public void testInferTenantIdUnparsableLocationHeader() throws IOException { + FixtureServer.FixtureMapping fixture = + new FixtureServer.FixtureMapping.Builder() + .validateMethod("GET") + .validatePath("/aad/auth") + .withRedirect("https://unexpected-location", 302) + .build(); + + try (FixtureServer server = new FixtureServer().with(fixture)) { + DatabricksConfig config = new DatabricksConfig(); + config.setHost(server.getUrl()); + config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); + String result = AzureUtils.inferTenantId(config); + assertNull(result); + assertNull(config.getAzureTenantId()); + } + } + + @Test + public void testInferTenantIdHappyPath() throws IOException { + FixtureServer.FixtureMapping fixture = + new FixtureServer.FixtureMapping.Builder() + .validateMethod("GET") + .validatePath("/aad/auth") + .withRedirect("https://login.microsoftonline.com/test-tenant-id/oauth2/authorize", 302) + .build(); + + try (FixtureServer server = new FixtureServer().with(fixture)) { + DatabricksConfig config = new DatabricksConfig(); + config.setHost(server.getUrl()); + config.setAzureWorkspaceResourceId( + "/subscriptions/123/resourceGroups/rg/providers/Microsoft.Databricks/workspaces/ws"); + config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); + String result = AzureUtils.inferTenantId(config); + assertEquals("test-tenant-id", result); + assertNull(config.getAzureTenantId()); // Config should remain unchanged + } + } + + @Test + public void testInferTenantIdSkipsWhenNotAzure() { + DatabricksConfig config = new DatabricksConfig(); + config.setHost("https://my-workspace.cloud.databricks.com"); // non-azure host + config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); + String result = AzureUtils.inferTenantId(config); + assertNull(result); + assertNull(config.getAzureTenantId()); + } + + @Test + public void testInferTenantIdSkipsWhenAlreadySet() { + DatabricksConfig config = new DatabricksConfig(); + config.setHost("https://adb-123.0.azuredatabricks.net"); + config.setAzureTenantId("existing-tenant-id"); + config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); + String result = AzureUtils.inferTenantId(config); + assertEquals("existing-tenant-id", result); + assertEquals("existing-tenant-id", config.getAzureTenantId()); // Config should remain unchanged + } + + @Test + public void testInferTenantIdSkipsWhenNoHost() { + DatabricksConfig config = new DatabricksConfig(); + config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); + String result = AzureUtils.inferTenantId(config); + assertNull(result); + assertNull(config.getAzureTenantId()); + } +} From 9f9f7a12f1334e87553b1e6fe841f94e0235c52a Mon Sep 17 00:00:00 2001 From: Sreekanth Vadigi Date: Mon, 11 Aug 2025 19:04:38 +0530 Subject: [PATCH 6/8] code formatting Signed-off-by: Sreekanth Vadigi --- .../databricks/sdk/core/DatabricksConfig.java | 1 - ...reServicePrincipalCredentialsProvider.java | 14 ++++--- .../databricks/sdk/core/utils/AzureUtils.java | 39 +++++++++++-------- .../sdk/core/DatabricksConfigTest.java | 1 - .../sdk/core/utils/AzureUtilsTest.java | 2 +- 5 files changed, 32 insertions(+), 25 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 5c7118698..0bc0b868a 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 @@ -18,7 +18,6 @@ import org.apache.http.HttpMessage; public class DatabricksConfig { - private CredentialsProvider credentialsProvider = new DefaultCredentialsProvider(); @ConfigAttribute(env = "DATABRICKS_HOST") diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java index 0555fed46..94a5acecf 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java @@ -26,15 +26,17 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { || config.getAzureClientSecret() == null) { return null; } - - this.tenantId = config.getAzureTenantId() != null ? config.getAzureTenantId() : AzureUtils.inferTenantId(config); + + this.tenantId = + config.getAzureTenantId() != null + ? config.getAzureTenantId() + : AzureUtils.inferTenantId(config); if (this.tenantId == null) { return null; } - - AzureUtils.ensureHostPresent( - config, mapper, this::tokenSourceFor); - + + AzureUtils.ensureHostPresent(config, mapper, this::tokenSourceFor); + CachedTokenSource inner = tokenSourceFor(config, config.getEffectiveAzureLoginAppId()); CachedTokenSource cloud = tokenSourceFor(config, config.getAzureEnvironment().getServiceManagementEndpoint()); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/AzureUtils.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/AzureUtils.java index 0f1368766..8e7abd0ce 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/AzureUtils.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/AzureUtils.java @@ -106,7 +106,7 @@ public static Optional getAzureWorkspaceResourceId(Workspace workspace) /** * Infers the Azure tenant ID from the Databricks workspace login page. - * + * * @param config The DatabricksConfig instance * @return the discovered tenant ID, or null if discovery fails */ @@ -114,12 +114,12 @@ public static String inferTenantId(DatabricksConfig config) { if (config.getAzureTenantId() != null) { return config.getAzureTenantId(); } - + if (!config.isAzure() || config.getHost() == null) { logger.warn("Cannot infer tenant ID: workspace is not Azure or host is missing"); return null; } - + String loginUrl = config.getHost() + AZURE_AUTH_ENDPOINT; try { @@ -128,52 +128,59 @@ public static String inferTenantId(DatabricksConfig config) { logger.warn("Failed to get redirect location from Azure auth endpoint: {}", loginUrl); return null; } - + String extractedTenantId = extractTenantIdFromUrl(redirectLocation); if (extractedTenantId == null) { logger.warn("Failed to extract tenant ID from redirect URL: {}", redirectLocation); return null; } - + logger.info("Successfully discovered Azure tenant ID: {}", extractedTenantId); return extractedTenantId; - + } catch (Exception e) { - logger.warn("Exception occurred while inferring Azure tenant ID from {}: {}", loginUrl, e.getMessage()); + logger.warn( + "Exception occurred while inferring Azure tenant ID from {}: {}", + loginUrl, + e.getMessage()); return null; } } - - private static String getRedirectLocation(DatabricksConfig config, String loginUrl) throws IOException { + + private static String getRedirectLocation(DatabricksConfig config, String loginUrl) + throws IOException { Request request = new Request("GET", loginUrl); request.setRedirectionBehavior(false); Response response = config.getHttpClient().execute(request); - + if (response.getStatusCode() != 302) { - logger.warn("Expected redirect (302) from {}, got status code: {}", loginUrl, response.getStatusCode()); + logger.warn( + "Expected redirect (302) from {}, got status code: {}", + loginUrl, + response.getStatusCode()); return null; } - + String location = response.getFirstHeader("Location"); if (location == null) { logger.warn("No Location header in redirect response from {}", loginUrl); } - + return location; } - + private static String extractTenantIdFromUrl(String redirectUrl) { try { // Parse: https://login.microsoftonline.com//oauth2/authorize?... URL entraIdUrl = new URL(redirectUrl); String[] pathSegments = entraIdUrl.getPath().split("/"); - + if (pathSegments.length < 2) { logger.warn("Invalid path in Location header: {}", entraIdUrl.getPath()); return null; } - return pathSegments[1]; + return pathSegments[1]; } catch (Exception e) { logger.warn("Failed to parse tenant ID from URL {}: {}", redirectUrl, e.getMessage()); return 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 a025bcaa6..b3ac333a3 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 @@ -250,5 +250,4 @@ public void testGetTokenSourceWithOAuth() { assertFalse(tokenSource instanceof ErrorTokenSource); assertEquals(tokenSource.getToken().getAccessToken(), "test-token"); } - } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/utils/AzureUtilsTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/utils/AzureUtilsTest.java index fea9e36ba..057d41878 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/utils/AzureUtilsTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/utils/AzureUtilsTest.java @@ -3,8 +3,8 @@ import static org.junit.jupiter.api.Assertions.*; import com.databricks.sdk.core.DatabricksConfig; -import com.databricks.sdk.core.commons.CommonsHttpClient; import com.databricks.sdk.core.FixtureServer; +import com.databricks.sdk.core.commons.CommonsHttpClient; import java.io.IOException; import org.junit.jupiter.api.Test; From 22adeedfedcc804e489fdc804d12be280030a33f Mon Sep 17 00:00:00 2001 From: Sreekanth Vadigi Date: Wed, 13 Aug 2025 15:52:30 +0530 Subject: [PATCH 7/8] azure utils throwing exception Signed-off-by: Sreekanth Vadigi --- ...reServicePrincipalCredentialsProvider.java | 16 ++-- .../databricks/sdk/core/utils/AzureUtils.java | 51 +++++------ .../sdk/core/utils/AzureUtilsTest.java | 86 ++++++++++++++++--- 3 files changed, 107 insertions(+), 46 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java index 94a5acecf..640244b9d 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java @@ -5,12 +5,16 @@ import com.fasterxml.jackson.databind.ObjectMapper; import java.util.HashMap; import java.util.Map; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * Adds refreshed Azure Active Directory (AAD) Service Principal OAuth tokens to every request, * while automatically resolving different Azure environment endpoints. */ public class AzureServicePrincipalCredentialsProvider implements CredentialsProvider { + private static final Logger logger = + LoggerFactory.getLogger(AzureServicePrincipalCredentialsProvider.class); private final ObjectMapper mapper = new ObjectMapper(); private String tenantId; @@ -27,11 +31,13 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { return null; } - this.tenantId = - config.getAzureTenantId() != null - ? config.getAzureTenantId() - : AzureUtils.inferTenantId(config); - if (this.tenantId == null) { + try { + this.tenantId = + config.getAzureTenantId() != null + ? config.getAzureTenantId() + : AzureUtils.inferTenantId(config); + } catch (Exception e) { + logger.warn("Failed to infer Azure tenant ID: {}", e.getMessage()); return null; } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/AzureUtils.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/AzureUtils.java index 8e7abd0ce..2b6ea21fe 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/AzureUtils.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/AzureUtils.java @@ -108,42 +108,33 @@ public static Optional getAzureWorkspaceResourceId(Workspace workspace) * Infers the Azure tenant ID from the Databricks workspace login page. * * @param config The DatabricksConfig instance - * @return the discovered tenant ID, or null if discovery fails + * @return the discovered tenant ID + * @throws DatabricksException if tenant ID discovery fails */ - public static String inferTenantId(DatabricksConfig config) { + public static String inferTenantId(DatabricksConfig config) throws DatabricksException { + if (config.getAzureTenantId() != null) { return config.getAzureTenantId(); } - if (!config.isAzure() || config.getHost() == null) { - logger.warn("Cannot infer tenant ID: workspace is not Azure or host is missing"); - return null; + if (config.getHost() == null) { + throw new DatabricksException("Cannot infer tenant ID: host is missing"); + } + + if (!config.isAzure()) { + throw new DatabricksException("Cannot infer tenant ID: workspace is not Azure"); } String loginUrl = config.getHost() + AZURE_AUTH_ENDPOINT; try { String redirectLocation = getRedirectLocation(config, loginUrl); - if (redirectLocation == null) { - logger.warn("Failed to get redirect location from Azure auth endpoint: {}", loginUrl); - return null; - } - String extractedTenantId = extractTenantIdFromUrl(redirectLocation); - if (extractedTenantId == null) { - logger.warn("Failed to extract tenant ID from redirect URL: {}", redirectLocation); - return null; - } - logger.info("Successfully discovered Azure tenant ID: {}", extractedTenantId); return extractedTenantId; } catch (Exception e) { - logger.warn( - "Exception occurred while inferring Azure tenant ID from {}: {}", - loginUrl, - e.getMessage()); - return null; + throw new DatabricksException("Failed to infer Azure tenant ID from " + loginUrl, e); } } @@ -154,36 +145,34 @@ private static String getRedirectLocation(DatabricksConfig config, String loginU Response response = config.getHttpClient().execute(request); if (response.getStatusCode() != 302) { - logger.warn( - "Expected redirect (302) from {}, got status code: {}", - loginUrl, - response.getStatusCode()); - return null; + throw new DatabricksException( + "Expected redirect (302) from " + + loginUrl + + ", got status code: " + + response.getStatusCode()); } String location = response.getFirstHeader("Location"); if (location == null) { - logger.warn("No Location header in redirect response from {}", loginUrl); + throw new DatabricksException("No Location header in redirect response from " + loginUrl); } return location; } - private static String extractTenantIdFromUrl(String redirectUrl) { + private static String extractTenantIdFromUrl(String redirectUrl) throws DatabricksException { try { // Parse: https://login.microsoftonline.com//oauth2/authorize?... URL entraIdUrl = new URL(redirectUrl); String[] pathSegments = entraIdUrl.getPath().split("/"); if (pathSegments.length < 2) { - logger.warn("Invalid path in Location header: {}", entraIdUrl.getPath()); - return null; + throw new DatabricksException("Invalid path in Location header: " + entraIdUrl.getPath()); } return pathSegments[1]; } catch (Exception e) { - logger.warn("Failed to parse tenant ID from URL {}: {}", redirectUrl, e.getMessage()); - return null; + throw new DatabricksException("Failed to parse tenant ID from URL " + redirectUrl, e); } } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/utils/AzureUtilsTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/utils/AzureUtilsTest.java index 057d41878..435471a32 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/utils/AzureUtilsTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/utils/AzureUtilsTest.java @@ -3,6 +3,7 @@ import static org.junit.jupiter.api.Assertions.*; import com.databricks.sdk.core.DatabricksConfig; +import com.databricks.sdk.core.DatabricksException; import com.databricks.sdk.core.FixtureServer; import com.databricks.sdk.core.commons.CommonsHttpClient; import java.io.IOException; @@ -15,9 +16,27 @@ public void testInferTenantId404() throws IOException { try (FixtureServer server = new FixtureServer().with("GET", "/aad/auth", "", 404)) { DatabricksConfig config = new DatabricksConfig(); config.setHost(server.getUrl()); + config.setAzureWorkspaceResourceId( + "/subscriptions/123/resourceGroups/rg/providers/Microsoft.Databricks/workspaces/ws"); config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); - String result = AzureUtils.inferTenantId(config); - assertNull(result); + + DatabricksException exception = + assertThrows( + DatabricksException.class, + () -> { + AzureUtils.inferTenantId(config); + }); + assertEquals( + "Failed to infer Azure tenant ID from " + server.getUrl() + "/aad/auth", + exception.getMessage()); + + assertNotNull(exception.getCause()); + assertInstanceOf(DatabricksException.class, exception.getCause()); + DatabricksException cause = (DatabricksException) exception.getCause(); + assertEquals( + "Expected redirect (302) from " + server.getUrl() + "/aad/auth, got status code: 404", + cause.getMessage()); + assertNull(config.getAzureTenantId()); } } @@ -27,9 +46,27 @@ public void testInferTenantIdNoLocationHeader() throws IOException { try (FixtureServer server = new FixtureServer().with("GET", "/aad/auth", "", 302)) { DatabricksConfig config = new DatabricksConfig(); config.setHost(server.getUrl()); + config.setAzureWorkspaceResourceId( + "/subscriptions/123/resourceGroups/rg/providers/Microsoft.Databricks/workspaces/ws"); config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); - String result = AzureUtils.inferTenantId(config); - assertNull(result); + + DatabricksException exception = + assertThrows( + DatabricksException.class, + () -> { + AzureUtils.inferTenantId(config); + }); + assertEquals( + "Failed to infer Azure tenant ID from " + server.getUrl() + "/aad/auth", + exception.getMessage()); + + assertNotNull(exception.getCause()); + assertInstanceOf(DatabricksException.class, exception.getCause()); + DatabricksException cause = (DatabricksException) exception.getCause(); + assertEquals( + "No Location header in redirect response from " + server.getUrl() + "/aad/auth", + cause.getMessage()); + assertNull(config.getAzureTenantId()); } } @@ -46,9 +83,26 @@ public void testInferTenantIdUnparsableLocationHeader() throws IOException { try (FixtureServer server = new FixtureServer().with(fixture)) { DatabricksConfig config = new DatabricksConfig(); config.setHost(server.getUrl()); + config.setAzureWorkspaceResourceId( + "/subscriptions/123/resourceGroups/rg/providers/Microsoft.Databricks/workspaces/ws"); config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); - String result = AzureUtils.inferTenantId(config); - assertNull(result); + + DatabricksException exception = + assertThrows( + DatabricksException.class, + () -> { + AzureUtils.inferTenantId(config); + }); + assertEquals( + "Failed to infer Azure tenant ID from " + server.getUrl() + "/aad/auth", + exception.getMessage()); + + assertNotNull(exception.getCause()); + assertInstanceOf(DatabricksException.class, exception.getCause()); + DatabricksException cause = (DatabricksException) exception.getCause(); + assertEquals( + "Failed to parse tenant ID from URL https://unexpected-location", cause.getMessage()); + assertNull(config.getAzureTenantId()); } } @@ -79,8 +133,14 @@ public void testInferTenantIdSkipsWhenNotAzure() { DatabricksConfig config = new DatabricksConfig(); config.setHost("https://my-workspace.cloud.databricks.com"); // non-azure host config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); - String result = AzureUtils.inferTenantId(config); - assertNull(result); + + DatabricksException exception = + assertThrows( + DatabricksException.class, + () -> { + AzureUtils.inferTenantId(config); + }); + assertEquals("Cannot infer tenant ID: workspace is not Azure", exception.getMessage()); assertNull(config.getAzureTenantId()); } @@ -99,8 +159,14 @@ public void testInferTenantIdSkipsWhenAlreadySet() { public void testInferTenantIdSkipsWhenNoHost() { DatabricksConfig config = new DatabricksConfig(); config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build()); - String result = AzureUtils.inferTenantId(config); - assertNull(result); + + DatabricksException exception = + assertThrows( + DatabricksException.class, + () -> { + AzureUtils.inferTenantId(config); + }); + assertEquals("Cannot infer tenant ID: host is missing", exception.getMessage()); assertNull(config.getAzureTenantId()); } } From 0bb96a29d558d49d18c9f9ae25a41ff7f5b52a4f Mon Sep 17 00:00:00 2001 From: Sreekanth Vadigi Date: Wed, 13 Aug 2025 16:57:49 +0530 Subject: [PATCH 8/8] removing logger from azureutils Signed-off-by: Sreekanth Vadigi --- .../java/com/databricks/sdk/core/utils/AzureUtils.java | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/AzureUtils.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/AzureUtils.java index 2b6ea21fe..c5b0aaaa0 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/AzureUtils.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/AzureUtils.java @@ -14,13 +14,9 @@ import java.util.Map; import java.util.Optional; import java.util.function.BiFunction; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; public class AzureUtils { - private static final Logger logger = LoggerFactory.getLogger(AzureUtils.class); - /** Azure authentication endpoint for tenant ID discovery */ private static final String AZURE_AUTH_ENDPOINT = "/aad/auth"; @@ -129,9 +125,7 @@ public static String inferTenantId(DatabricksConfig config) throws DatabricksExc try { String redirectLocation = getRedirectLocation(config, loginUrl); - String extractedTenantId = extractTenantIdFromUrl(redirectLocation); - logger.info("Successfully discovered Azure tenant ID: {}", extractedTenantId); - return extractedTenantId; + return extractTenantIdFromUrl(redirectLocation); } catch (Exception e) { throw new DatabricksException("Failed to infer Azure tenant ID from " + loginUrl, e);