Skip to content

Commit 22adeed

Browse files
sreekanth-dbsreekanth-db
committed
azure utils throwing exception
Signed-off-by: Sreekanth Vadigi <[email protected]>
1 parent 9f9f7a1 commit 22adeed

File tree

3 files changed

+107
-46
lines changed

3 files changed

+107
-46
lines changed

‎databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java‎

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,16 @@
55
import com.fasterxml.jackson.databind.ObjectMapper;
66
import java.util.HashMap;
77
import java.util.Map;
8+
import org.slf4j.Logger;
9+
import org.slf4j.LoggerFactory;
810

911
/**
1012
* Adds refreshed Azure Active Directory (AAD) Service Principal OAuth tokens to every request,
1113
* while automatically resolving different Azure environment endpoints.
1214
*/
1315
public class AzureServicePrincipalCredentialsProvider implements CredentialsProvider {
16+
private static final Logger logger =
17+
LoggerFactory.getLogger(AzureServicePrincipalCredentialsProvider.class);
1418
private final ObjectMapper mapper = new ObjectMapper();
1519
private String tenantId;
1620

@@ -27,11 +31,13 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
2731
return null;
2832
}
2933

30-
this.tenantId =
31-
config.getAzureTenantId() != null
32-
? config.getAzureTenantId()
33-
: AzureUtils.inferTenantId(config);
34-
if (this.tenantId == null) {
34+
try {
35+
this.tenantId =
36+
config.getAzureTenantId() != null
37+
? config.getAzureTenantId()
38+
: AzureUtils.inferTenantId(config);
39+
} catch (Exception e) {
40+
logger.warn("Failed to infer Azure tenant ID: {}", e.getMessage());
3541
return null;
3642
}
3743

‎databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/AzureUtils.java‎

Lines changed: 20 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -108,42 +108,33 @@ public static Optional<String> getAzureWorkspaceResourceId(Workspace workspace)
108108
* Infers the Azure tenant ID from the Databricks workspace login page.
109109
*
110110
* @param config The DatabricksConfig instance
111-
* @return the discovered tenant ID, or null if discovery fails
111+
* @return the discovered tenant ID
112+
* @throws DatabricksException if tenant ID discovery fails
112113
*/
113-
public static String inferTenantId(DatabricksConfig config) {
114+
public static String inferTenantId(DatabricksConfig config) throws DatabricksException {
115+
114116
if (config.getAzureTenantId() != null) {
115117
return config.getAzureTenantId();
116118
}
117119

118-
if (!config.isAzure() || config.getHost() == null) {
119-
logger.warn("Cannot infer tenant ID: workspace is not Azure or host is missing");
120-
return null;
120+
if (config.getHost() == null) {
121+
throw new DatabricksException("Cannot infer tenant ID: host is missing");
122+
}
123+
124+
if (!config.isAzure()) {
125+
throw new DatabricksException("Cannot infer tenant ID: workspace is not Azure");
121126
}
122127

123128
String loginUrl = config.getHost() + AZURE_AUTH_ENDPOINT;
124129

125130
try {
126131
String redirectLocation = getRedirectLocation(config, loginUrl);
127-
if (redirectLocation == null) {
128-
logger.warn("Failed to get redirect location from Azure auth endpoint: {}", loginUrl);
129-
return null;
130-
}
131-
132132
String extractedTenantId = extractTenantIdFromUrl(redirectLocation);
133-
if (extractedTenantId == null) {
134-
logger.warn("Failed to extract tenant ID from redirect URL: {}", redirectLocation);
135-
return null;
136-
}
137-
138133
logger.info("Successfully discovered Azure tenant ID: {}", extractedTenantId);
139134
return extractedTenantId;
140135

141136
} catch (Exception e) {
142-
logger.warn(
143-
"Exception occurred while inferring Azure tenant ID from {}: {}",
144-
loginUrl,
145-
e.getMessage());
146-
return null;
137+
throw new DatabricksException("Failed to infer Azure tenant ID from " + loginUrl, e);
147138
}
148139
}
149140

@@ -154,36 +145,34 @@ private static String getRedirectLocation(DatabricksConfig config, String loginU
154145
Response response = config.getHttpClient().execute(request);
155146

156147
if (response.getStatusCode() != 302) {
157-
logger.warn(
158-
"Expected redirect (302) from {}, got status code: {}",
159-
loginUrl,
160-
response.getStatusCode());
161-
return null;
148+
throw new DatabricksException(
149+
"Expected redirect (302) from "
150+
+ loginUrl
151+
+ ", got status code: "
152+
+ response.getStatusCode());
162153
}
163154

164155
String location = response.getFirstHeader("Location");
165156
if (location == null) {
166-
logger.warn("No Location header in redirect response from {}", loginUrl);
157+
throw new DatabricksException("No Location header in redirect response from " + loginUrl);
167158
}
168159

169160
return location;
170161
}
171162

172-
private static String extractTenantIdFromUrl(String redirectUrl) {
163+
private static String extractTenantIdFromUrl(String redirectUrl) throws DatabricksException {
173164
try {
174165
// Parse: https://login.microsoftonline.com/<tenant-id>/oauth2/authorize?...
175166
URL entraIdUrl = new URL(redirectUrl);
176167
String[] pathSegments = entraIdUrl.getPath().split("/");
177168

178169
if (pathSegments.length < 2) {
179-
logger.warn("Invalid path in Location header: {}", entraIdUrl.getPath());
180-
return null;
170+
throw new DatabricksException("Invalid path in Location header: " + entraIdUrl.getPath());
181171
}
182172

183173
return pathSegments[1];
184174
} catch (Exception e) {
185-
logger.warn("Failed to parse tenant ID from URL {}: {}", redirectUrl, e.getMessage());
186-
return null;
175+
throw new DatabricksException("Failed to parse tenant ID from URL " + redirectUrl, e);
187176
}
188177
}
189178
}

‎databricks-sdk-java/src/test/java/com/databricks/sdk/core/utils/AzureUtilsTest.java‎

Lines changed: 76 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import static org.junit.jupiter.api.Assertions.*;
44

55
import com.databricks.sdk.core.DatabricksConfig;
6+
import com.databricks.sdk.core.DatabricksException;
67
import com.databricks.sdk.core.FixtureServer;
78
import com.databricks.sdk.core.commons.CommonsHttpClient;
89
import java.io.IOException;
@@ -15,9 +16,27 @@ public void testInferTenantId404() throws IOException {
1516
try (FixtureServer server = new FixtureServer().with("GET", "/aad/auth", "", 404)) {
1617
DatabricksConfig config = new DatabricksConfig();
1718
config.setHost(server.getUrl());
19+
config.setAzureWorkspaceResourceId(
20+
"/subscriptions/123/resourceGroups/rg/providers/Microsoft.Databricks/workspaces/ws");
1821
config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build());
19-
String result = AzureUtils.inferTenantId(config);
20-
assertNull(result);
22+
23+
DatabricksException exception =
24+
assertThrows(
25+
DatabricksException.class,
26+
() -> {
27+
AzureUtils.inferTenantId(config);
28+
});
29+
assertEquals(
30+
"Failed to infer Azure tenant ID from " + server.getUrl() + "/aad/auth",
31+
exception.getMessage());
32+
33+
assertNotNull(exception.getCause());
34+
assertInstanceOf(DatabricksException.class, exception.getCause());
35+
DatabricksException cause = (DatabricksException) exception.getCause();
36+
assertEquals(
37+
"Expected redirect (302) from " + server.getUrl() + "/aad/auth, got status code: 404",
38+
cause.getMessage());
39+
2140
assertNull(config.getAzureTenantId());
2241
}
2342
}
@@ -27,9 +46,27 @@ public void testInferTenantIdNoLocationHeader() throws IOException {
2746
try (FixtureServer server = new FixtureServer().with("GET", "/aad/auth", "", 302)) {
2847
DatabricksConfig config = new DatabricksConfig();
2948
config.setHost(server.getUrl());
49+
config.setAzureWorkspaceResourceId(
50+
"/subscriptions/123/resourceGroups/rg/providers/Microsoft.Databricks/workspaces/ws");
3051
config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build());
31-
String result = AzureUtils.inferTenantId(config);
32-
assertNull(result);
52+
53+
DatabricksException exception =
54+
assertThrows(
55+
DatabricksException.class,
56+
() -> {
57+
AzureUtils.inferTenantId(config);
58+
});
59+
assertEquals(
60+
"Failed to infer Azure tenant ID from " + server.getUrl() + "/aad/auth",
61+
exception.getMessage());
62+
63+
assertNotNull(exception.getCause());
64+
assertInstanceOf(DatabricksException.class, exception.getCause());
65+
DatabricksException cause = (DatabricksException) exception.getCause();
66+
assertEquals(
67+
"No Location header in redirect response from " + server.getUrl() + "/aad/auth",
68+
cause.getMessage());
69+
3370
assertNull(config.getAzureTenantId());
3471
}
3572
}
@@ -46,9 +83,26 @@ public void testInferTenantIdUnparsableLocationHeader() throws IOException {
4683
try (FixtureServer server = new FixtureServer().with(fixture)) {
4784
DatabricksConfig config = new DatabricksConfig();
4885
config.setHost(server.getUrl());
86+
config.setAzureWorkspaceResourceId(
87+
"/subscriptions/123/resourceGroups/rg/providers/Microsoft.Databricks/workspaces/ws");
4988
config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build());
50-
String result = AzureUtils.inferTenantId(config);
51-
assertNull(result);
89+
90+
DatabricksException exception =
91+
assertThrows(
92+
DatabricksException.class,
93+
() -> {
94+
AzureUtils.inferTenantId(config);
95+
});
96+
assertEquals(
97+
"Failed to infer Azure tenant ID from " + server.getUrl() + "/aad/auth",
98+
exception.getMessage());
99+
100+
assertNotNull(exception.getCause());
101+
assertInstanceOf(DatabricksException.class, exception.getCause());
102+
DatabricksException cause = (DatabricksException) exception.getCause();
103+
assertEquals(
104+
"Failed to parse tenant ID from URL https://unexpected-location", cause.getMessage());
105+
52106
assertNull(config.getAzureTenantId());
53107
}
54108
}
@@ -79,8 +133,14 @@ public void testInferTenantIdSkipsWhenNotAzure() {
79133
DatabricksConfig config = new DatabricksConfig();
80134
config.setHost("https://my-workspace.cloud.databricks.com"); // non-azure host
81135
config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build());
82-
String result = AzureUtils.inferTenantId(config);
83-
assertNull(result);
136+
137+
DatabricksException exception =
138+
assertThrows(
139+
DatabricksException.class,
140+
() -> {
141+
AzureUtils.inferTenantId(config);
142+
});
143+
assertEquals("Cannot infer tenant ID: workspace is not Azure", exception.getMessage());
84144
assertNull(config.getAzureTenantId());
85145
}
86146

@@ -99,8 +159,14 @@ public void testInferTenantIdSkipsWhenAlreadySet() {
99159
public void testInferTenantIdSkipsWhenNoHost() {
100160
DatabricksConfig config = new DatabricksConfig();
101161
config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build());
102-
String result = AzureUtils.inferTenantId(config);
103-
assertNull(result);
162+
163+
DatabricksException exception =
164+
assertThrows(
165+
DatabricksException.class,
166+
() -> {
167+
AzureUtils.inferTenantId(config);
168+
});
169+
assertEquals("Cannot infer tenant ID: host is missing", exception.getMessage());
104170
assertNull(config.getAzureTenantId());
105171
}
106172
}

0 commit comments

Comments
 (0)