Skip to content

Commit d443694

Browse files
Divyansh-dbDivyansh-db
committed
add azure devops OIDC authentication
1 parent d476b3a commit d443694

File tree

3 files changed

+251
-0
lines changed

3 files changed

+251
-0
lines changed

‎databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,17 @@ private void addOIDCCredentialsProviders(DatabricksConfig config) {
125125
config.getActionsIdTokenRequestUrl(),
126126
config.getActionsIdTokenRequestToken(),
127127
config.getHttpClient())));
128+
129+
// Try to create Azure DevOps token source - if environment variables are missing,
130+
// skip this provider gracefully.
131+
try {
132+
namedIdTokenSources.add(
133+
new NamedIDTokenSource(
134+
"azure-devops-oidc",
135+
new AzureDevOpsIDTokenSource(config.getHttpClient())));
136+
} catch (DatabricksException e) {
137+
LOG.debug("Azure DevOps OIDC provider not available: {}", e.getMessage());
138+
}
128139
// Add new IDTokenSources and ID providers here. Example:
129140
// namedIdTokenSources.add(new NamedIDTokenSource("custom-oidc", new CustomIDTokenSource(...)));
130141

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

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
1+
package com.databricks.sdk.core.oauth;
2+
3+
import com.databricks.sdk.core.DatabricksException;
4+
import com.databricks.sdk.core.http.HttpClient;
5+
import com.databricks.sdk.core.http.Request;
6+
import com.databricks.sdk.core.http.Response;
7+
import com.fasterxml.jackson.databind.ObjectMapper;
8+
import com.fasterxml.jackson.databind.node.ObjectNode;
9+
import com.google.common.base.Strings;
10+
import java.io.IOException;
11+
12+
/**
13+
* AzureDevOpsIDTokenSource retrieves JWT Tokens from Azure DevOps Pipelines. This class implements
14+
* the IDTokenSource interface and provides a method for obtaining ID tokens specifically from Azure
15+
* DevOps Pipeline environment.
16+
*
17+
* <p>This implementation follows the Azure DevOps OIDC token API as documented at:
18+
* https://learn.microsoft.com/en-us/rest/api/azure/devops/distributedtask/oidctoken/create
19+
*/
20+
public class AzureDevOpsIDTokenSource implements IDTokenSource {
21+
/* Access token for authenticating with Azure DevOps API */
22+
private final String azureDevOpsAccessToken;
23+
/* Team Foundation Collection URI (e.g., https://dev.azure.com/organization) */
24+
private final String azureDevOpsTeamFoundationCollectionUri;
25+
/* Plan ID for the current pipeline run */
26+
private final String azureDevOpsPlanId;
27+
/* Job ID for the current pipeline job */
28+
private final String azureDevOpsJobId;
29+
/* Team Project ID where the pipeline is running */
30+
private final String azureDevOpsTeamProjectId;
31+
/* Host type (e.g., "build", "release") */
32+
private final String azureDevOpsHostType;
33+
/* HTTP client for making requests to Azure DevOps */
34+
private final HttpClient httpClient;
35+
/* JSON mapper for parsing response data */
36+
private static final ObjectMapper mapper = new ObjectMapper();
37+
38+
/**
39+
* Constructs a new AzureDevOpsIDTokenSource by reading environment variables.
40+
* This constructor implements fail-early validation - if any required environment
41+
* variables are missing, it will throw a DatabricksException immediately.
42+
*
43+
* @param httpClient The HTTP client to use for making requests
44+
* @throws DatabricksException if any required environment variables are missing
45+
*/
46+
public AzureDevOpsIDTokenSource(HttpClient httpClient) {
47+
if (httpClient == null) {
48+
throw new DatabricksException("HttpClient cannot be null");
49+
}
50+
this.httpClient = httpClient;
51+
52+
// Fail early: validate all required environment variables
53+
this.azureDevOpsAccessToken = validateEnvironmentVariable("SYSTEM_ACCESSTOKEN");
54+
this.azureDevOpsTeamFoundationCollectionUri = validateEnvironmentVariable("SYSTEM_TEAMFOUNDATIONCOLLECTIONURI");
55+
this.azureDevOpsPlanId = validateEnvironmentVariable("SYSTEM_PLANID");
56+
this.azureDevOpsJobId = validateEnvironmentVariable("SYSTEM_JOBID");
57+
this.azureDevOpsTeamProjectId = validateEnvironmentVariable("SYSTEM_TEAMPROJECTID");
58+
this.azureDevOpsHostType = validateEnvironmentVariable("SYSTEM_HOSTTYPE");
59+
}
60+
61+
/**
62+
* Validates that an environment variable is present and not empty.
63+
*
64+
* @param varName The environment variable name
65+
* @return The environment variable value
66+
* @throws DatabricksException if the environment variable is missing or empty
67+
*/
68+
private String validateEnvironmentVariable(String varName) {
69+
String value = System.getenv(varName);
70+
if (Strings.isNullOrEmpty(value)) {
71+
throw new DatabricksException(
72+
String.format("Missing %s, likely not calling from Azure DevOps Pipeline", varName));
73+
}
74+
return value;
75+
}
76+
77+
/**
78+
* Retrieves an ID token from Azure DevOps Pipelines. This method makes an authenticated request
79+
* to Azure DevOps to obtain a JWT token that can later be exchanged for a Databricks access token.
80+
*
81+
*
82+
* <p>Note: The audience parameter is ignored for Azure DevOps OIDC tokens as they have a
83+
* hardcoded audience for Azure AD integration.
84+
*
85+
* @param audience Ignored for Azure DevOps OIDC tokens
86+
* @return An IDToken object containing the JWT token value
87+
* @throws DatabricksException if the token request fails
88+
*/
89+
@Override
90+
public IDToken getIDToken(String audience) {
91+
92+
// Build Azure DevOps OIDC endpoint URL
93+
// Format: {collectionUri}/{teamProjectId}/_apis/distributedtask/hubs/{hostType}/plans/{planId}/jobs/{jobId}/oidctoken?api-version=7.2-preview.1
94+
String requestUrl = String.format(
95+
"%s/%s/_apis/distributedtask/hubs/%s/plans/%s/jobs/%s/oidctoken?api-version=7.2-preview.1",
96+
azureDevOpsTeamFoundationCollectionUri,
97+
azureDevOpsTeamProjectId,
98+
azureDevOpsHostType,
99+
azureDevOpsPlanId,
100+
azureDevOpsJobId);
101+
102+
Request req = new Request("POST", requestUrl)
103+
.withHeader("Authorization", "Bearer " + azureDevOpsAccessToken)
104+
.withHeader("Content-Type", "application/json");
105+
106+
Response resp;
107+
try {
108+
resp = httpClient.execute(req);
109+
} catch (IOException e) {
110+
throw new DatabricksException(
111+
"Failed to request ID token from Azure DevOps at " + requestUrl + ": " + e.getMessage(), e);
112+
}
113+
114+
if (resp.getStatusCode() != 200) {
115+
throw new DatabricksException(
116+
"Failed to request ID token from Azure DevOps: status code "
117+
+ resp.getStatusCode()
118+
+ ", response body: "
119+
+ resp.getBody().toString());
120+
}
121+
122+
// Parse the JSON response
123+
// Azure DevOps returns {"oidcToken":"***"} format, not {"value":"***"} like GitHub Actions
124+
ObjectNode jsonResp;
125+
try {
126+
jsonResp = mapper.readValue(resp.getBody(), ObjectNode.class);
127+
} catch (IOException e) {
128+
throw new DatabricksException(
129+
"Failed to parse Azure DevOps OIDC token response: " + e.getMessage(), e);
130+
}
131+
132+
// Validate response structure and token value
133+
if (!jsonResp.has("oidcToken")) {
134+
throw new DatabricksException("Azure DevOps OIDC token response missing 'oidcToken' field");
135+
}
136+
137+
try {
138+
String tokenValue = jsonResp.get("oidcToken").textValue();
139+
if (Strings.isNullOrEmpty(tokenValue)) {
140+
throw new DatabricksException("Received empty OIDC token from Azure DevOps");
141+
}
142+
return new IDToken(tokenValue);
143+
} catch (IllegalArgumentException e) {
144+
throw new DatabricksException("Received invalid OIDC token from Azure DevOps: " + e.getMessage(), e);
145+
}
146+
}
147+
}

‎databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/AzureDevOpsIDTokenSourceTest.java‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
package com.databricks.sdk.core.oauth;
2+
3+
import static org.junit.jupiter.api.Assertions.*;
4+
import static org.mockito.ArgumentMatchers.any;
5+
import static org.mockito.Mockito.*;
6+
7+
import com.databricks.sdk.core.DatabricksException;
8+
import com.databricks.sdk.core.http.HttpClient;
9+
import com.databricks.sdk.core.http.Request;
10+
import com.databricks.sdk.core.http.Response;
11+
import java.io.ByteArrayInputStream;
12+
import java.io.IOException;
13+
import java.nio.charset.StandardCharsets;
14+
import org.junit.jupiter.api.BeforeEach;
15+
import org.junit.jupiter.api.Test;
16+
import org.mockito.ArgumentCaptor;
17+
import org.mockito.Mock;
18+
import org.mockito.MockitoAnnotations;
19+
20+
/**
21+
* Unit tests for AzureDevOpsIDTokenSource.
22+
*
23+
* Note: These tests focus on the core functionality. Environment variable validation
24+
* tests are limited since the class now reads directly from System.getenv().
25+
* Integration tests should be used to test the full environment variable behavior.
26+
*/
27+
public class AzureDevOpsIDTokenSourceTest {
28+
29+
@Mock private HttpClient httpClient;
30+
@Mock private Response response;
31+
32+
@BeforeEach
33+
void setUp() {
34+
MockitoAnnotations.openMocks(this);
35+
}
36+
37+
@Test
38+
void testNullHttpClient() {
39+
// Act & Assert
40+
DatabricksException exception = assertThrows(DatabricksException.class,
41+
() -> new AzureDevOpsIDTokenSource(null));
42+
assertTrue(exception.getMessage().contains("HttpClient cannot be null"));
43+
}
44+
45+
/**
46+
* Test that audience parameter is ignored (Azure DevOps has hardcoded audience).
47+
* This test verifies that the URL construction doesn't include audience parameter.
48+
*/
49+
@Test
50+
void testAudienceParameterIgnored() throws IOException {
51+
// This test can only run if environment variables are set (e.g., in Azure DevOps)
52+
// Skip if not in Azure DevOps environment
53+
if (System.getenv("SYSTEM_ACCESSTOKEN") == null) {
54+
return; // Skip test if not in Azure DevOps environment
55+
}
56+
57+
// Arrange
58+
String audience = "https://databricks.com";
59+
String expectedToken = "eyJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiJ9...";
60+
String responseBody = String.format("{\"oidcToken\":\"%s\"}", expectedToken);
61+
62+
when(response.getStatusCode()).thenReturn(200);
63+
when(response.getBody()).thenReturn(new ByteArrayInputStream(responseBody.getBytes(StandardCharsets.UTF_8)));
64+
when(httpClient.execute(any(Request.class))).thenReturn(response);
65+
66+
AzureDevOpsIDTokenSource tokenSource = new AzureDevOpsIDTokenSource(httpClient);
67+
68+
// Act
69+
IDToken result = tokenSource.getIDToken(audience);
70+
71+
// Assert
72+
assertNotNull(result);
73+
assertEquals(expectedToken, result.getValue());
74+
75+
// Verify the request URL does NOT include the audience parameter (it's ignored)
76+
ArgumentCaptor<Request> requestCaptor = ArgumentCaptor.forClass(Request.class);
77+
verify(httpClient).execute(requestCaptor.capture());
78+
79+
Request capturedRequest = requestCaptor.getValue();
80+
String requestUri = capturedRequest.getUri().toString();
81+
assertFalse(requestUri.contains("audience="),
82+
"Audience parameter should be ignored for Azure DevOps OIDC tokens");
83+
assertTrue(requestUri.contains("api-version=7.2-preview.1"));
84+
}
85+
86+
/**
87+
* Note: Most environment variable validation tests are not included here
88+
* since the class now reads directly from System.getenv(). These should be
89+
* tested in integration tests where the environment can be controlled.
90+
*
91+
* The tests below focus on the core HTTP functionality that can be unit tested.
92+
*/
93+
}

0 commit comments

Comments
 (0)